PyTorch中带状矩阵与稠密矩阵乘法的高效实现问询

你的判断是对的：用diag_embed重构出完整的三对角矩阵后再做乘法，PyTorch完全不会识别它的特殊结构，只会把它当成普通的稠密矩阵处理，计算量还是O(N²K)，和原始的A@B没区别，甚至因为重构矩阵的额外内存开销和计算步骤，速度反而更慢。

要真正利用带状矩阵的稀疏性实现O(NK)级别的计算效率，你需要跳过矩阵重构，直接用对角向量和B做元素级的并行计算。以下是具体实现方案：

三对角矩阵乘B的高效实现 #

直接提取三个对角向量，分别计算它们对结果的贡献，再合并：

import torch

def tridiag_matrix_mult(main_diag, upper_diag, lower_diag, B):
    # 输入维度说明：
    # main_diag: [..., N]，主对角线向量
    # upper_diag: [..., N-1]，上对角线向量（偏移+1）
    # lower_diag: [..., N-1]，下对角线向量（偏移-1）
    # B: [..., N, K]，待乘的N×K矩阵
    
    # 主对角线的贡献：每个元素乘以B的对应行
    result = main_diag.unsqueeze(-1) * B
    
    # 上对角线的贡献：元素i乘以B的第i行，加到结果的第i+1行
    result[..., 1:, :] += upper_diag.unsqueeze(-1) * B[..., :-1, :]
    
    # 下对角线的贡献：元素i乘以B的第i+1行，加到结果的第i行
    result[..., :-1, :] += lower_diag.unsqueeze(-1) * B[..., 1:, :]
    
    return result

调用时直接用你已经提取的三个对角向量，不用重构矩阵：

# 假设A是[d1, d2, N, N]，B是[d1, d2, N, K]
main_diag = torch.diagonal(A, dim1=-2, dim2=-1)
upper_diag = torch.diagonal(A, offset=1, dim1=-2, dim2=-1)
lower_diag = torch.diagonal(A, offset=-1, dim1=-2, dim2=-1)

# 高效计算三对角矩阵乘B
fast_result = tridiag_matrix_mult(main_diag, upper_diag, lower_diag, B)

扩展到任意带宽的带状矩阵 #

如果需要支持width参数（比如width=2对应五对角矩阵），可以遍历所有偏移量，提取对应对角线后累加贡献：

def band_matrix_mult(A, width, B):
    # A: [..., N, N]，原始矩阵
    # width: 带宽，比如width=1对应三对角
    # B: [..., N, K]，待乘矩阵
    result = torch.zeros_like(B)
    N = A.size(-1)
    
    # 遍历从-width到width的所有偏移量
    for offset in range(-width, width + 1):
        diag = torch.diagonal(A, offset=offset, dim1=-2, dim2=-1)
        diag_expand = diag.unsqueeze(-1)
        
        if offset == 0:
            # 主对角线，直接乘B的所有行
            result += diag_expand * B
        elif offset > 0:
            # 上偏移：diag元素对应A的(i, i+offset)，乘B的第i行，加到结果的第i+offset行
            result[..., offset:, :] += diag_expand * B[..., :-offset, :]
        else:
            # 下偏移：offset为负，取绝对值
            abs_offset = -offset
            # diag元素对应A的(i+abs_offset, i)，乘B的第i+abs_offset行，加到结果的第i行
            result[..., :-abs_offset, :] += diag_expand * B[..., abs_offset:, :]
    
    return result

性能优势说明 #

• 计算量：普通稠密矩阵乘是O(N²K)，而带状矩阵乘是O((2width+1)*NK)，当width远小于N时（比如width=1，N=10000），计算量直接降到原来的1/5000，速度提升非常明显。
• GPU优化：PyTorch的元素级操作、广播和切片都是高度优化的GPU操作，能充分利用CUDA核心的并行能力，不会有额外的开销。
• 内存节省：不需要存储N×N的稠密带状矩阵，只需要存储(2width+1)个长度为N的向量，内存占用从O(N²)降到O(N)。

性能验证示例 #

可以用torch.utils.benchmark测试对比：

from torch.utils.benchmark import Timer

N = 10000
K = 100
A = torch.randn(N, N, device="cuda")
B = torch.randn(N, K, device="cuda")

# 方法1：原始A@B
t1 = Timer(stmt="A @ B", globals={"A": A, "B": B})
print("原始矩阵乘：", t1.timeit(10))

# 方法2：重构三对角矩阵再乘
main_diag = torch.diagonal(A)
upper_diag = torch.diagonal(A, offset=1)
lower_diag = torch.diagonal(A, offset=-1)
tridiag = torch.diag_embed(main_diag) + torch.diag_embed(upper_diag, offset=1) + torch.diag_embed(lower_diag, offset=-1)
t2 = Timer(stmt="tridiag @ B", globals={"tridiag": tridiag, "B": B})
print("重构三对角矩阵乘：", t2.timeit(10))

# 方法3：手动实现的高效乘
t3 = Timer(stmt="tridiag_matrix_mult(main_diag, upper_diag, lower_diag, B)", 
           globals={"tridiag_matrix_mult": tridiag_matrix_mult, 
                    "main_diag": main_diag, "upper_diag": upper_diag, "lower_diag": lower_diag, "B": B})
print("高效三对角矩阵乘：", t3.timeit(10))

测试结果会显示，方法3的速度比前两者快几个数量级。

内容的提问来源于stack exchange，提问作者VIVID