Transformer自注意力机制计算复杂度深度解析：从O(l²d)到优化实践

1次阅读

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

传统 RNN 处理长度为 $l$ 的序列时，计算复杂度为 $O(l \cdot d^2)$（$d$ 为特征维度），而 Transformer 的自注意力机制需要计算所有位置对之间的关系，复杂度飙升至 $O(l^2 \cdot d)$。具体来看：

计算过程拆解 ：
Query-Key 乘积：$(l \times d) \times (d \times l) \rightarrow O(l^2 d)$
Softmax 归一化：$O(l^2)$
Value 加权求和：$(l \times l) \times (l \times d) \rightarrow O(l^2 d)$
对比实验数据 ：
当 $l=1024, d=512$ 时：
- RNN：~2.6G FLOPs
- Self-Attention：~1.1T FLOPs（相差 400 倍！）

多头机制将计算拆分为 $h$ 个独立头，实现显存和计算资源的并行利用：

数学表达 ：
$\text{MultiHead}(Q,K,V) = \text{Concat}(head_1,…,head_h)W^O$
其中 $head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$
并行优势 ：
计算量不变，但 GPU 可同时处理多个头
显存占用从 $O(l^2d)$ 降至 $O(h \cdot (l^2 \cdot \frac{d}{h}))$

通过限制每个位置只能关注特定区域，将稠密矩阵变为稀疏矩阵：

固定模式 （如 Longformer 的滑动窗口）：
每个 token 只关注前后 $w$ 个位置
复杂度从 $O(l^2d)$ 降至 $O(l \cdot w \cdot d)$
动态模式 （如 Reformer 的 LSH 注意力）：
使用哈希函数将相似 token 分到同一桶
复杂度 $O(l \log l \cdot d)$

import torch
from einops import rearrange

def sparse_attention(Q, K, V, mask):
    """
    Q/K/V: [batch, heads, seq_len, dim]  # 输入张量形状
    mask: [seq_len, seq_len]  # 稀疏掩码矩阵
    """
    # Step 1: 计算原始注意力分数
    attn_scores = torch.einsum('bhid,bhjd->bhij', Q, K)  # [b,h,l,l]

    # Step 2: 应用稀疏掩码
    attn_scores = attn_scores.masked_fill(mask == 0, -1e9)

    # Step 3: 标准化与输出
    attn_weights = torch.softmax(attn_scores, dim=-1)
    return torch.einsum('bhij,bhjd->bhid', attn_weights, V)