如何不使用中间3D矩阵实现或近似目标矩阵计算？

公式化简 #

我们可以先对原式的内层求和做拆解，避免生成3D中间张量：
原式为：
$$C_{ij} = \max_k \bigg{ \sum_l \bigg( W_{ik}W_{kl}W_{lj} - W_{ik}W_{kj} \bigg) \bigg}$$
将求和项拆分，把和$l$无关的项提出求和外：
$$
\begin{align*}
\sum_l \bigg( W_{ik}W_{kl}W_{lj} - W_{ik}W_{kj} \bigg) &= W_{ik}\sum_l W_{kl}W_{lj} - W_{ik}W_{kj}\sum_l 1 \
&= W_{ik} \cdot (W^2){kj} - N \cdot W{ik} W_{kj} \
&= W_{ik} \cdot \left( (W^2){kj} - N W{kj} \right)
\end{align*}
$$
其中$W^2$是$W$和自身的矩阵乘积，$N$是矩阵边长。
令中间矩阵$B = W^2 - N \cdot W$，则原式简化为：
$$C_{ij} = \max_{k} \left( W_{ik} \cdot B_{kj} \right)$$
这个运算本质是Max-Product半环矩阵乘法，不需要构造3D张量即可高效实现。

精确实现方案（PyTorch） #

方案1：低显存循环实现

仅需$O(N^2)$显存，循环逻辑高度可并行，速度远快于逐元素循环：

import torch

def compute_C(W: torch.Tensor) -> torch.Tensor:
    N = W.shape[0]
    # 计算中间矩阵B
    W2 = W @ W
    B = W2 - N * W
    # 初始化C为负无穷
    C = torch.full((N, N), -torch.inf, device=W.device, dtype=W.dtype)
    # 循环k，每次计算外积后取逐元素最大值
    for k in range(N):
        temp = W[:, k, None] @ B[None, k, :]
        C = torch.max(C, temp)
    return C

如果想要进一步提速，可以将k按块拆分，减少循环次数：

def compute_C_block(W: torch.Tensor, block_size: int = 32) -> torch.Tensor:
    N = W.shape[0]
    W2 = W @ W
    B = W2 - N * W
    C = torch.full((N, N), -torch.inf, device=W.device, dtype=W.dtype)
    # 按块遍历k
    for start in range(0, N, block_size):
        end = min(start + block_size, N)
        W_block = W[:, start:end]
        B_block = B[start:end, :]
        temp = W_block.unsqueeze(-1) * B_block.unsqueeze(0)
        temp_max = temp.max(dim=1).values
        C = torch.max(C, temp_max)
    return C

方案2：超大规模矩阵分块实现

如果$N$过大导致$N \times N$的$C$矩阵都无法放到显存中，可以对$i,j$维度同时分块，按需计算和存储结果，显存占用仅为$O(block_size^2)$。

近似实现方案 #

当$N$超过$10^{4$量级，精确计算的$O(N}3)$复杂度不可接受时，可以用以下近似方法：

• Top-k采样近似：对每个行向量$W[i]$仅保留数值最大的前K个元素（K通常取10~100），仅在这部分k值上计算最大值，复杂度降至$O(KN^2)$，如果$W$是非负稀疏矩阵，精度损失可以忽略。
• 低秩分解近似：如果$W$满足低秩假设，可以先对$W$做SVD分解得到$W \approx U V^T$（$U,V$为$N \times r$矩阵，$r \ll N$），再基于低秩矩阵计算Max-Product乘积，复杂度降至$O(r^2 N)$。

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