线性注意力

为什么标准注意力是瓶颈？ Transformer的核心——自注意力机制，其计算复杂度是序列长度的二次方 O(n²)。这个n²使得处理长序列时的计算和内存开销急剧增长。 标准注意力的计算过程 import torch import torch.nn.functional as F import math def standard_attention(Q, K, V): """ 标准自注意力 Q, K, V: (batch, num_heads, seq_len, d_head) """ d = Q.shape[-1] # 注意力分数: (batch, heads, seq, seq) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d) # Softmax归一化 attn = F.softmax(scores, dim=-1) # 加权求和 output = torch.matmul(attn, V) # (batch, heads, seq, d_head) return output # 计算复杂度分析 # scores矩阵大小: seq_len × seq_len # 当 seq_len = 8192 时, 每个head的scores矩阵 = 8192² = 67M 个元素 # 当 seq_len = 32768 时, 每个head的scores矩阵 = 32768² = 1B 个元素 # 显存随序列长度平方增长！ 不同序列长度的开销对比 序列长度 注意力矩阵大小 显存(FP16, 32头) 计算量(FLOPs) 4K 16M 1GB 0.13T 8K 64M 4GB 0.52T 32K 1B 64GB 8.4T 128K 16B 1TB 134T 1M 1T 64TB 8.2PT 这个表格清楚地展示了为什么标准注意力无法扩展到超长序列。 ...

O(n²)到O(n)的吸引力 标准注意力的O(n²)复杂度是长序列处理的核心障碍。线性注意力旨在将复杂度降至O(n)，使百万级token序列成为可能。但这个目标在保持注意力核心功能的同时实现极其困难——注意力之所以是O(n²)，正是因为它需要建模所有Query-Key对之间的关系。 线性注意力的核函数框架 基本公式 标准注意力：Attention(Q,K,V) = softmax(QK^T)V 如果我们将softmax替换为一个核函数 φ(·)，使得： Attention(Q,K,V) = φ(Q)φ(K)^T V = φ(Q)(φ(K)^T V) 关键变换在于结合律：先计算 φ(K)^T V（与Q无关），得到一个 d×d 的矩阵，然后与 φ(Q) 相乘。这使得复杂度从 O(n²d) 降为 O(nd²)。当 n » d 时，这是显著的速度提升。 Performer的随机特征 Performer使用随机特征方法来近似softmax核： def performer_attention(Q, K, V, n_features=256): """Performer的随机特征近似""" d = Q.shape[-1] # 生成随机投影矩阵 omega = torch.randn(d, n_features) / math.sqrt(d) # 非负映射函数 def phi(x): return torch.exp(x @ omega - 0.5 * (x ** 2).sum(-1, keepdim=True)) / math.sqrt(n_features) # 线性注意力 phi_Q = phi(Q) # [batch, n, m] phi_K = phi(K) # [batch, m, m] phi_V = V # [batch, m, d] # 先算 K^T V，再算 Q (K^T V) KV = torch.einsum('bmd,bme->bde', phi_K, phi_V) # [batch, d, d] output = torch.einsum('bnd,bde->bne', phi_Q, KV) # [batch, n, d] # 归一化 normalizer = torch.einsum('bnd,bmd->bn', phi_Q, phi_K).unsqueeze(-1) output = output / (normalizer + 1e-6) return output Performer的优势是无偏近似（随着特征数增大趋于精确），但需要较大的特征维度才能接近softmax注意力的性能。 ...