Attention 的复杂度

Self-Attention 的核心是 $n \times n$ 的权重矩阵——序列长度 $n$ 每翻一倍，attention 的内存与计算大约翻四倍。这就是 $O(n^2)$ 复杂度。BERT 上限 512、GPT-3 2048、现代 LLM 推 128K 上下文，背后全是算法与工程在跟这个平方律搏斗。

直觉：全班两两握手的次数

$n$ 个人，每人要和其他 $n$ 人配对寒暄，大约 $n^2$ 次握手。128 个 token 要算 $128^2 = 16{,}384$ 对相似度；1024 个 token 要 $1{,}048{,}576$ 对——长度线性增，配对次数平方增。

定义

对单头 attention，设序列长 $n$，$d_k = d_{\text{model}}/h$。

时间复杂度（乘加各计 1 FLOP）：

$h$ 头合计，$h \cdot d_k = d_{\text{model}}$，故 attention 矩阵部分为 $O(n^2 d_{\text{model}})$。

空间复杂度（存 attention 权重）：$O(n^2)$ 每头；$h$ 头则 $O(h n^2)$。

相对 FFN 的 $O(n d_{\text{model}} d_{\text{ff}})$，当 $n$ 很大时 $n^2$ 项主导。

例子

例 1：BERT-base，$d=768, h=12, d_k=64$

例 2：内存（fp32，单层单 batch，12 头）

训练时还需存梯度与 optimizer state，显存压力更大。

例 3：与线性层对比（$n=512$）

• QKV+FFN：$\approx O(n d^2) \sim 10^9$ FLOPs
• Attention 矩阵：$\approx O(n^2 d) \sim 10^9$ FLOPs

$n=512, d=768$ 时二者同量级；$n=8192$ 时 $n^2$ 项碾压 $nd^2$。

自己拖一下

三张表分别列了 FLOPs、内存、对比，但「谁在什么时候压倒谁」是一条连续曲线，拼不出来。拖 $n$，看紫色那两段（$n^2$ 项）怎么从一条缝长到吞掉整条：

这里有个反直觉的点值得记住：交叉点在 $n = 2d$。LLaMA-7B（$d=4096$）要到 $n=8192$ 才轮到 attention 项超过投影/FFN——也就是说，8K 以下的上下文，大头一直是那些稠密矩阵乘，不是 attention。「attention 是平方所以最贵」这句流行的话，只在长上下文下才对。真正一直可怕的是那块 $O(n^2)$ 的显存（拖到 128K 看那个数字），这才是 FlashAttention 要解决的。

# 粗算 BERT-base 单层 attention FLOPs
def attention_flops(n, d_model=768, n_heads=12):
    d_k = d_model // n_heads
    qkv = 6 * n * d_model ** 2
    scores = 2 * n_heads * n ** 2 * d_k
    av = 2 * n_heads * n ** 2 * d_k
    out_proj = 2 * n * d_model ** 2
    return qkv + scores + av + out_proj

for n in [128, 512, 2048, 8192]:
    print(n, attention_flops(n) / 1e9, "GFLOPs")

量化：缓解 $O(n^2)$ 的方向

常见误区

1. $O(n^2)$ 只指 Self-Attention：Cross-Attention 是 $O(n_t n_s)$，源句长 $n_s$、目标句长 $n_t$。
2. Decoder 推理仍是 $O(n^2)$ 每步：无 cache 时生成第 $t$ token 要算 $t \times t$；有 KV cache 后每步约 $O(t)$，总生成 $O(n^2)$。
3. 加头数不影响 $n^2$ 项：$h$ 与 $d_k$ 成反比，$h n^2 d_k = n^2 d_{\text{model}}$。

小结

Attention 对序列长度 $n$ 是 $O(n^2)$ 时间与空间（权重矩阵）。BERT-base 在 $n=512$ 尚可；128K 上下文必须靠 FlashAttention、稀疏、GQA 等组合优化。下一篇看位置编码如何随长度扩展而演进。

下一篇：位置编码的变体