详细分析大语言模型attention的计算复杂度，从数学角度分析

大语言模型（LLM）中 Attention 的计算复杂度：系统与数学视角

下面以**单层多头自注意力（Multi-Head Self-Attention, MHA）**为基准，分别给出时间/空间复杂度、精确到常数的 FLOPs 估算、训练与推理（含 KV Cache）阶段的差异，以及若干改进/近似注意力的复杂度对比与直觉化解释。

1) 标准（密集）自注意力的精确计算量

设序列长度为 $n$，模型维度为 $d$，头数为 $h$，每头维度 $d_h=d/h$。输入矩阵 $X\in {\mathbb{R}}^{n\times d}$。

1.1 线性投影

通常用三组权重将 $X$ 投影为 $Q,K,V$：

$$Q=X{W}_Q,K=X{W}_K,V=X{W}_V,W_Q,W_K,W_V\in {\mathbb{R}}^{d\times d}.$$

• 每次矩阵乘法（$n\times d$ 与 $d\times d$）的乘加 FLOPs $\approx 2n{d}^2$。
• 三个投影合计 FLOPs $\approx 3\times 2n{d}^2=6n{d}^2$。

输出拼接后还要过一次输出投影 $W_O\in {\mathbb{R}}^{d\times d}$：再加 $2n{d}^2$ FLOPs。

线性投影总 FLOPs（与序列无关项）：

$$\boxed{8n{d}^2}$$

1.2 注意力核心（按头分块计算）

对每个头 $i$:

$$A_i=\mathrm{softmax}\left(\frac{Q_i{K}_i^{\top }}{\sqrt{d_h}}\right),O_i=A_i{V}_i.$$

• $Q_i{K}_i^{\top }$: $(n\times d_h)\cdot (d_h\times n)\Rightarrow 2n^2{d}_h$ FLOPs。
• softmax：约 $O(n^2)$（相对前两项常数级，可忽略在主项里）。
• $A_i{V}_i$: $(n\times n)\cdot (n\times d_h)\Rightarrow 2n^2{d}_h$ FLOPs。

对 $h$ 个头求和：

$$\text{Attention 核心 FLOPs}\approx h\cdot (2n^2{d}_h+2n^2{d}_h)=4n^2(h{d}_h)=\boxed{4n^{2}d}.$$

1.3 单层总 FLOPs（前向）

$$\boxed{\text{FLOPs}\approx 8n{d}^2+4n^{2}d}$$

这条式子非常关键：当 $n\ll d$ 时，$8n{d}^2$（投影与 MLP 类似量级）主导；当 $n\gg d$ 时，$4n^{2}d$（注意力矩阵）主导。
临界处在 $8n{d}^2\approx 4n^{2}d\Rightarrow n\approx 2d$。

注：实际 Transformer 还含有前馈网络（FFN/MLP），其 FLOPs 约为常见扩张倍数 $m$（如 4）下的 $\approx 2nd\cdot md+2nmd\cdot d\approx 4mn{d}^2$（前后两次线性层合计，忽略激活常数），常见 $m=4$ 时约 $\sim 16n{d}^2$，常常与上面的 $8n{d}^2$ 同量级甚至更大。因此在短上下文下，MLP 往往比注意力更贵；在超长上下文下，注意力 $n^2$ 项会快速成为主导。

2) 空间复杂度（内存/显存）

• 保存 $Q,K,V$：$O(nd)$。
• 朴素实现会显式构造注意力矩阵 $A\in {\mathbb{R}}^{n\times n}$：$O(n^2)$。
• 训练反向需要缓存中间量与梯度：常见为 $O(n^2)$ 级别的额外显存；可用**激活重计算（checkpointing）**以 $ \sim 2\times$ 计算换显存，降到 $\tilde{O}(nd)$ 级别。
• FlashAttention 通过分块（tiling）与在线 softmax，将峰值显存由 $O(n^2)$ 降到 $\boxed{O(nd)}$，时间复杂度仍是 $O(n^{2}d)$ 但常数显著下降（IO 友好）。

3) 自回归推理与 KV Cache

设提示长度 $L$，需生成 $G$ 个新 token，总长度 $T=L+G$。

3.1 预填充（prefill，一次性算完前 $L$ 个位置）

• 复杂度与训练前向同型：$\boxed{8L{d}^2+4L^{2}d}$。

3.2 增量生成（decoding，步进式生成）

开启 KV Cache 时，每步只需：

• 为当前步做投影（$\sim 8d^2$ 常数级 w.r.t. $n$）。
• 与历史 $K,V$ 做注意力：对第 $t$ 步，注意力代价 $\sim O(td)$（读出并点积到长度 $t$ 的缓存）。

因此从 $t=L$ 到 $L+G-1$ 的总注意力代价为：

$$\sum _{t=L}^{L+G-1}O(td)=O(d\cdot \frac{(L+G-1+L)G}{2})=\boxed{O(G(L+G)d)}.$$

直觉：每步与越来越长的缓存相乘，步均线性涨，因此整体是二次和。

3.3 KV Cache 显存大小

每层、每 token 需要缓存 $K$ 与 $V$ 各 $d$ 维（更精确是 $h\cdot d_h=d$），共 $2d$ 个元素。
若用 FP16（2 字节/元素），每层每 token 占：

$$\boxed{2d\times 2\text{ bytes}=4d\text{ bytes}}.$$

• 举例：$d=4096$ 时，每层每 token $=4\times 4096=16384$ 字节 $\Rightarrow 16\text{KB}$。

• 若层数 $N=32$，每 token 的 KV 缓存为 $16\text{KB}\times 32=512\text{KB}$。

• 因此总 KV 显存 $\approx 512\text{KB}\times T$。

  • $T=8,192$（8K）时：$512\times 8192\text{KB}=4,194,304\text{KB}\approx 4\mathbf{GB}$。
  • $T=32,768$（32K）时：$\approx 16GB$。
  • $T=131,072$（128K）时：$\approx 64GB$。

推理瓶颈直觉：解码阶段往往内存带宽受限（每步要从显存连续读大量 $K,V$），而不是纯算力受限。Flash-Decoding、PagedAttention、张量并行/流水并行的 IO 优化都在缓解这点。

3.4 MQA/GQA 对 KV 显存与吞吐的影响

• MQA（Multi-Query Attention）：多个头共享同一组 $K,V$（即 $h_K=h_V=1$），KV 缓存从 $O(h)$ 降到 $O(1)$，解码阶段读带宽显著下降，内存占用近似按头数缩小。
• GQA（Grouped-Query Attention）：每组头共享一组 $K,V$。若分成 $g$ 组，KV 显存从 $O(h)$ 降到 $O(g)$。

4) 交叉注意力（Cross-Attention）

目标序列长度 $n_q$，源序列长度 $n_k$（如编码器-解码器结构）：

• $Q{K}^{\top }$：$2n_q{n}_k{d}_h$ 每头；合计 $\approx 2n_q{n}_{k}d$。
• $AV$：同量级，再来 $2n_q{n}_{k}d$。
• 合计注意力核 FLOPs $\approx \boxed{4n_q{n}_{k}d}$；再加四次 $n{d}^2$ 级的投影项。

5) 近似/稀疏/线性注意力的复杂度对比

目标都是降低 $n^2$ 到接近线性或 $n\log n$ 级别，同时尽量控制常数与误差。

• 局部/滑窗注意力（window size $w$）：
  时间 $\boxed{O(nwd)}$，显存 $\boxed{O(nw)}$。适合长文本但侧重局部依赖；可配少量全局 token。
• Block/稀疏模式（如 BigBird/Longformer）：
  时间 $\boxed{O(n(w+g+r)d)}$，其中 $w$=滑窗宽，$g$=全局 token 数，$r$=随机块连边数。
• LSH/Reformer：期望时间 $\boxed{O(n\log n\cdot d)}$，但实现与常数较复杂。
• Nyströmformer（秩 $r$）：$\boxed{O(nrd+r^{2}d)}$；当 $r\ll n$ 时接近线性。
• 线性注意力（核化/FAVOR+ 等）：
  典型推导将
  $\mathrm{softmax}(Q{K}^{\top })$ 近似为 $\varphi (Q)\varphi (K{)}^{\top }$，可先算
  $S=\varphi (K{)}^{\top }V\in {\mathbb{R}}^{d_{\varphi }\times d}$（代价 $O(n{d}_{\varphi }d)$），再算
  $\varphi (Q)S$（代价 $O(n{d}_{\varphi }d)$），总 $\boxed{O(n{d}_{\varphi }d)}$。
  若 $d_{\varphi }\sim d$，则为 $O(n{d}^2)$——对极长序列更划算；但当 $n$ 不大时常数未必占优。
• FlashAttention：仍是 $O(n^{2}d)$ 时间，但将显存峰值压到 $O(nd)$，并显著降低 IO，实际速度常常大幅提升。

结论：当 $n$ 远大于 $d$ 时，上述方法能把主项从 $n^2$ 降到近线性；当 $n$ 与 $d$ 同量级或 $n$ 较小，近似法的收益变小甚至不如标准注意力（常数/误差/实现复杂度）。

6) 数学角度的要点与直觉

1. 主导项来自两个 $n\times n$ 的乘法：
   $⟨Q_i,K_i⟩$ 形成 $n\times n$ 的打分矩阵，以及将其与 $V_i$ 相乘。二者分别贡献 $2n^2{d}_h$ FLOPs/头，合计 $4n^{2}d$。
2. 归一化与稳定性：
   除以 $\sqrt{d_h}$ 保持点积分布方差稳定，避免 softmax 过陡；softmax 的计算量是 $O(n^2)$ 次指数/加法/除法，通常不是主导项。
3. 与快速矩阵乘法的理论界：
   若使用 Strassen/以后算法，正方形矩阵乘法可达 $O(n^{\omega })$, $\omega \approx 2.37$。但注意力的矩阵形状是 $(n\times d_h)\cdot (d_h\times n)$，且在 GPU 上高度优化的 GEMM（经典 $O(n^3)$ 常数小）更实用；工业界并不使用快速矩阵乘法来降注意力的幂指数。
4. 计算-IO（算强度）视角：
   解码阶段每步需要从显存顺序读取 $O(td)$ 的 $K,V$，但计算仅 $O(td)$ 级点积/加和，算强度低，内存带宽成为瓶颈。Flash 系列方法通过分块/重排提升局部性与重用，效果显著。

7) 速查表（单层、前向）

• 总 FLOPs：$\boxed{8n{d}^2+4n^{2}d}$

  • $8n{d}^2$：4 次 $d\times d$ 线性（$Q,K,V,O$）
  • $4n^{2}d$：注意力两次 $n\times n$ 乘法（$Q{K}^{\top }$ 与 $AV$）

• 显存峰值（朴素）：$\boxed{O(n^2)+O(nd)}$；FlashAttention：$\boxed{O(nd)}$。

• 自回归解码（含 KV Cache）总注意力成本：$\boxed{O(G(L+G)d)}$；KV 显存：每层每 token $\boxed{4d\text{ bytes}}$（FP16）。

8) 实用建议

• 短上下文（$n\lesssim d$）：投影与 MLP 往往主导；优化点在张量并行、低精度、算子融合。
• 长上下文（$n\gg d$）：注意力 $n^2$ 主导；优先考虑 FlashAttention、稀疏/滑窗/近似注意力、分块推理。
• 推理吞吐：使用 MQA/GQA 显著降低 KV 显存与带宽；配合 Paged KV、分块解码提升大上下文的可扩展性。
• 训练显存：激活检查点 + FlashAttention 能将注意力显存从 $O(n^2)$ 压至近 $O(nd)$。