Attention Free Transformer：如何突破传统注意力机制的计算瓶颈

1次阅读

共计 1435 个字符，预计需要花费 4 分钟才能阅读完成。

Transformer 模型的核心组件——自注意力机制（Self-Attention）虽然功能强大，但其计算复杂度随着序列长度呈平方级增长（O(n²)）。这意味着处理 1024 个 token 的序列时，需要计算约 100 万次点积操作。更具体地说，给定序列长度 n，标准注意力矩阵的空间复杂度为 O(n²)，这使得处理长文本（如 4000+token 的文档）或高分辨率图像时，显存消耗迅速成为瓶颈。

标准 Transformer：计算复杂度为 O(n²d)，其中 d 为特征维度
$$\text{Attention}(Q,K,V) = \text{softmax}(\frac{QK^T}{\sqrt{d}})V$$
Linformer：通过低秩近似降为 O(ndk)，k 为投影维度
$$\hat{K} = K\cdot W_k,\ \hat{V} = V\cdot W_v$$
AFT 系列 ：实现线性复杂度 O(nd)
$$Y = \sigma(Q) \odot \frac{\sum_{i=1}^n \exp(K_i + w_{i,j}) \odot V_i}{\sum_{i=1}^n \exp(K_i + w_{i,j})}$$
（其中⊙表示逐元素乘，w 为可学习的位置偏置）

AFT-full：全局建模能力，保留完整的位置偏置矩阵 w∈ℝ^{n×n}
AFT-local：仅考虑窗口内位置（如半径 128），w∈ℝ^{n×(2r+1)}
AFT-simple：完全移除位置偏置，仅保留元素级交互

# PyTorch 实现核心计算流程
def aft_forward(Q, K, V, w=None):
    """
    Q: [batch, n, d]
    K: [batch, n, d]
    V: [batch, n, d]
    w: [n,n] (AFT-full) 或 [n,2r+1] (AFT-local)
    """
    Q_sig = torch.sigmoid(Q)  # 门控机制
    K_exp = torch.exp(K)

    if w is not None:
        # 位置偏置处理
        K_exp = K_exp.unsqueeze(1) * torch.exp(w).unsqueeze(-1)

    KV = (K_exp * V).sum(dim=1, keepdim=True)
    Z = K_exp.sum(dim=1, keepdim=True)
    return Q_sig * (KV / Z)

采用分解式设计：
$$w_{i,j} = \phi(i)^T \theta(j)$$
其中 ϕ,θ 为低维投影（如 8 维），将空间复杂度从 O(n²) 降至 O(n)