Transformer自注意力机制计算复杂度优化：从O(l²)到线性化的实战策略

1次阅读

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

标准自注意力的计算复杂度来源于 QK^T 矩阵乘法。具体来看，当输入序列长度为 $l$，特征维度为 $d$ 时：

计算查询矩阵 $Q$ 和键矩阵 $K$：
$$Q = XW_Q, \quad K = XW_K$$
其中 $X \in \mathbb{R}^{l \times d}$，$W_Q, W_K \in \mathbb{R}^{d \times d_k}$，这两步都是 $O(ld^2)$ 复杂度
计算注意力分数矩阵：
$$A = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})$$
这里的 $QK^T$ 矩阵乘法产生 $O(l^2d)$ 复杂度，当 $l \gg d$ 时成为主要瓶颈
最终输出：
$$O = AV$$
同样需要 $O(l^2d)$ 计算量

方法	核心思想	计算复杂度	适用场景	典型精度损失
Linformer	低秩投影 KV 矩阵	O(kl)	长文档建模	<5%
Reformer	LSH 分桶近似注意力	O(llog l)	相似度高的序列	7-10%
Longformer	滑动窗口 + 全局 token	O(lw)	局部依赖强的数据	3-8%
Performer	随机特征映射	O(ldlog d)	通用场景	5-15%

（w 为窗口大小，k 为投影维度，通常 k≪l）

import torch
from torch import nn

class LinformerAttention(nn.Module):
    def __init__(self, dim, seq_len, k=256, heads=8):
        super().__init__()
        self.proj_k = nn.Parameter(torch.randn(seq_len, k))  # 投影矩阵 E
        self.proj_v = nn.Parameter(torch.randn(seq_len, k))  # 投影矩阵 F

        # 初始化监控变量
        self.max_mem = 0

    def forward(self, q, k, v):
        # q: [b, h, l, d_k], k/v: [b, h, l, d_k]
        b, h, l, d = q.shape

        # 投影降维 [l, k] @ [b, h, l, d_k]^T -> [b, h, k, d_k]
        k = torch.einsum('lk,bhld->bhkd', self.proj_k, k) 
        v = torch.einsum('lk,bhld->bhkd', self.proj_v, v)

        # 计算注意力 [b, h, l, d_k] @ [b, h, d_k, k] -> [b, h, l, k]
        attn = torch.softmax(q @ k.transpose(-2,-1) / (d ** 0.5), 
            dim=-1
        )

        # 输出 [b, h, l, k] @ [b, h, k, d_k] -> [b, h, l, d_k]
        out = attn @ v

        # 记录峰值显存
        self.max_mem = max(
            self.max_mem, 
            torch.cuda.max_memory_allocated() // (1024 ** 2)
        )

        return out

关键点说明：
– 投影矩阵 $E,F \in \mathbb{R}^{l \times k}$ 将 KV 序列长度从 $l$ 压缩到 $k$
– einsum操作实现批量矩阵乘法同时保持维度清晰
– 复杂度从 $O(l^2d)$ 降到 $O(lkd)$

在 IMDB 电影评论分类任务上（序列 padding 到相同长度），测得不同方法的显存占用：