变换器与自注意力 Transformer & Self-Attention

所属：01-fundamentals / 01-llm-core
面试权重：★★★
来源：通用领域知识；未做一手考证

一句话#

Transformer 用自注意力（Self-Attention）取代 RNN 的循环结构，让序列中任意两个位置直接交互、且整层计算可写成一次大的矩阵乘；代价是注意力本身对序列长度是 O(n²) 的计算与显存，并且它天生对顺序不敏感，必须靠位置编码把顺序信息注入进去。

为什么重要（面试角度）#

这是所有 LLM 问题的地基，面试官用它做筛人：答不出 softmax(QKᵀ/√d)V、说不清为什么除以 √d_k、不知道 Pre-LN 与 Post-LN 的区别，后面的 KV Cache / MoE / 长上下文就不用谈了。

更值得注意的是，面试官很少满足于背公式，通常会顺着追问三层：

为什么这么设计 —— 除以 √d_k 的方差解释、mask 的作用、多头为什么不是简单堆参数。
代价在哪 —— O(n²) 从哪来、prefill 与 decode 的瓶颈为什么完全不同。
工程上的变体 —— FlashAttention 改变了什么、GQA/MLA 在改哪一部分（→ 见 kv-cache.md）。

核心内容#

注意力公式#

\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V

其中 $Q = XW_Q,\ K = XW_K,\ V = XW_V$ ， $X \in \mathbb{R}^{n \times d}$ ， $M$ 是加性 mask（因果掩码时，未来位置置为 $-\infty$ ，softmax 后权重为 0）。

为什么除以 $\sqrt{d_k}$ ：若 $q,k$ 各维近似独立、零均值单位方差，则点积 $q\cdot k$ 的方差是 $d_k$ 。不缩放时， $d_k$ 越大 logits 越极端，softmax 会逼近 one-hot，梯度趋近于 0（饱和），训练不稳定。除以 $\sqrt{d_k}$ 把方差拉回 1。

softmax 的数值实现：先减去每行最大值再取指数，避免 exp 上溢。这个"减最大值"在 FlashAttention 里对应 online softmax 的在线重标定。

为什么是 O(n²)#

不缓存时，每层的注意力矩阵是 $n \times n$ ：

部分	每层 FLOPs 量级
$QK^\top$ 与 $AV$	$4n^2 d$
Q/K/V/O 四个投影	$8nd^2$
FFN（ $d_{ff}=4d$ ）	$16nd^2$

所以当 $n \gtrsim 4d$ 时，注意力的 $n^2$ 项开始主导；短序列时反而是 FFN 更贵。这解释了两件事：长上下文的成本主要来自注意力，而短序列推理的算力主要花在 FFN 与权重读取上。

多头注意力（Multi-Head Attention）#

把 $d$ 维切成 $h$ 份，每份 $d_h = d/h$ ，各自算一遍注意力，再拼回来过 $W_O$ ：

\text{head}_i = \text{Attention}(XW_Q^{(i)}, XW_K^{(i)}, XW_V^{(i)}), \quad \text{MHA} = \text{Concat}(\text{head}_1,\dots,\text{head}_h)W_O

关键点：多头不是为了增加参数量，而是增加"表示子空间"的数量。总参数量与单头 $d \times d$ 的投影基本同阶，但不同头可以各自关注不同模式（局部/全局、句法/指代、以及训练中自发出现的 induction head 等）。这也意味着头是冗余的——推理时剪掉部分头通常掉点很少，这正是后来 MQA/GQA 能成立的实验基础。

位置信息必须外挂#

自注意力对输入序列是**置换等变（permutation-equivariant）**的：把 token 打乱，输出只是跟着打乱，模型内部没有任何顺序概念。"猫追狗"和"狗追猫"在没有位置编码时对它是同一个集合。因此必须加位置编码（→ 见 position-encoding.md）。

前馈网络与残差结构#

每层 = Attention + FFN，各自包一层残差与归一化：

x = x + Attention(Norm(x))     # Pre-LN
x = x + FFN(Norm(x))

FFN：两层 MLP，中间维度 $d_{ff}$ 通常约 $4d$ 。现代模型多用门控激活 SwiGLU/GeGLU： $\text{FFN}(x) = (xW_{gate} \odot \text{swish}(xW_{up}))W_{down}$ 。注意 SwiGLU 有三个矩阵，所以为了让参数量对齐 $4d^2$ 的两层 MLP， $d_{ff}$ 要取约 $\frac{8}{3}d$ （例如 $d=4096$ 时取 11008 而非常见的 16384）。
Pre-LN vs Post-LN：原始 Transformer 是 Post-LN（Norm(x + Sublayer(x))），训练需要 warmup 且深层不稳；现代 LLM 几乎都是 Pre-LN，梯度流更顺、对 warmup 依赖显著降低，代价是同等深度下表达力略弱（可用 final norm 缓解）。
RMSNorm vs LayerNorm：RMSNorm 只做 $\gamma \cdot x / \sqrt{\text{mean}(x^2)+\epsilon}$ ，不减均值、无 bias，更省算力、效果相当，Llama 系普遍采用。
无 bias：多数现代 LLM（PaLM/Llama 系）去掉了线性层的 bias，理由是收益太小、模型更大时收益更小，还有利于推理融合。

训练 vs 推理：两种完全不同的瓶颈#

阶段	计算形态	瓶颈
Prefill（吃 prompt）	一次前向算完所有位置	计算受限，注意力 $O(n^2)$ 主导
Decode（逐 token 生成）	每步只算 1 个新 query	显存带宽受限，瓶颈是读权重和读 KV Cache

所以 decode 阶段的算力利用率极低，这直接催生了 KV Cache、GQA/MLA、PagedAttention、连续批处理这一整条推理优化路线（→ 见 kv-cache.md、06-agent-infra/）。

FlashAttention 改变了什么#

FlashAttention 是精确注意力，不是近似。它做两件事：

分块（tiling）：把 Q/K/V 切块搬进 SRAM，用 online softmax 边算边重标定，从不把 $n \times n$ 矩阵写回 HBM。显存从 $O(n^2)$ 降到 $O(n)$ 。
重计算：反向传播时不存注意力矩阵，而是重新算一遍（算力换带宽）。

本质是 IO-aware 优化：在 GPU 上，HBM 读写比算术更贵。这也解释了它的一个反直觉结果——它通常同时更快也更省显存，因为它减少了内存搬运。

常见面试问法#

Q：手写注意力公式，并解释每个操作的作用。 A： $\text{softmax}(QK^\top/\sqrt{d_k} + M)V$ 。 $QK^\top$ 算相似度； $1/\sqrt{d_k}$ 控制 logits 方差防 softmax 饱和； $M$ 是加性 mask（因果掩码把未来位置置 $-\infty$ ）；softmax 归一化成权重；乘 $V$ 做加权求和。

Q：为什么多头注意力比单头好？参数量不是差不多吗？ A：参数量同阶，但多头把表示空间切成多个子空间并行做注意力，不同头能学到不同类型的关系（局部/长程、句法/语义），表达更多样的对齐模式。头之间高度冗余，所以可以剪枝，也可以让多个 query 头共享 KV（GQA）。

Q：Transformer 的复杂度是多少？和 RNN 比优劣势？ A：注意力 $O(n^2d)$ 计算、 $O(n^2)$ 中间激活（不优化时），FFN $O(nd^2)$ 。相对 RNN 的 $O(nd^2)$ ：短序列时 RNN 更省，长序列时 Transformer 的 $n^2$ 更贵；但 RNN 无法并行（时间步串行），Transformer 训练时可整序列并行，这才是决定性优势。后续线性注意力/SSM（Mamba 等）试图把 $n^2$ 打回 $n$ ，代价是表达能力与工程成熟度。

Q：Pre-LN 和 Post-LN 的区别？为什么现代模型都用 Pre-LN？ A：归一化位置不同。Post-LN 在残差之后归一化，深层梯度不稳、强依赖 warmup；Pre-LN 在子层输入前归一化，残差是恒等通路，梯度直通，训练稳定、可用更大学习率。代价是最终输出未归一化、深层有效表达略弱，一般加一个 final norm 弥补。

Q：FlashAttention 是近似算法吗？它为什么能又快又省显存？ A：不是近似，输出与标准注意力在数值上等价（差异仅来自浮点累加顺序）。它通过分块 + online softmax 避免把 $n \times n$ 矩阵落到 HBM，并在反向用重计算换显存。加速来自减少 HBM 读写，而不是减少计算量——这也是"IO-aware"这个词的含义。

易错点#

❌ 忘记 $1/\sqrt{d_k}$ ，或说不清它是为了"防止梯度爆炸"——准确说法是控制 logits 方差、避免 softmax 饱和导致梯度消失。
❌ 认为位置信息是 Transformer 自带的。没有位置编码，自注意力完全无序。
❌ 把"多头"说成"增加参数量/增加深度"。
❌ 认为 FlashAttention 是稀疏/近似注意力（那是 Longformer/BigBird 那一类）。
❌ 说"Transformer 复杂度是 O(n²)"却不说清这是注意力那部分、以及 decode 阶段有 KV Cache 后每步是 O(n) 而非 O(n²)。
❌ 混淆 Pre-LN 与 Post-LN 的方向，或说 RMSNorm 比 LayerNorm "更准"——它主要是更省，效果相当。
❌ 漏掉 SwiGLU 的参数配平问题，说 FFN 中间维度一律是 $4d$ 。

自测#

1. 写出注意力公式并解释 $1/\sqrt{d_k}$ 的作用。如果去掉它会怎样？

答案

$\text{Attention}(Q,K,V)=\text{softmax}(QK^\top/\sqrt{d_k}+M)V$ 。假设 q、k 各维独立同分布、零均值单位方差，则点积方差为 $d_k$ ，维度越大 logits 幅度越大，softmax 趋于 one-hot，输出几乎只由最大值决定，梯度随之消失（饱和）。除以 $\sqrt{d_k}$ 把方差归一化到 1，保持 softmax 处于梯度良好的区间，也让不同 head_dim 下的训练超参可迁移。

2. 为什么说自注意力本身没有顺序概念？位置信息是怎么进去的？

答案

注意力是"对集合加权求和"：打乱输入顺序，输出只是跟着置换，score 计算只看内容不看位置。所以必须在输入（或在 Q/K 上）注入位置信息：绝对位置编码（正弦、可学习）、相对位置偏置（T5 bias、ALiBi），或 RoPE 这类作用在 Q/K 上的旋转编码。RoPE 作用于 Q/K 而不作用于 V，是因为它要影响的是注意力权重，不是被加权的内容。

3. Prefill 和 decode 的瓶颈为什么不同？这个区别导致了哪些工程优化？

答案

Prefill 一次并行处理 n 个位置，是大矩阵乘，算术强度高，属于计算受限，注意力 $O(n^2)$ 项主导。Decode 每步只产生 1 个 token，计算量极小但要读全部权重和整个 KV Cache，算术强度极低，属于显存带宽受限。由此衍生：KV Cache 避免重算（→ kv-cache.md）、GQA/MLA 压缩 KV 体积、MQA 减少读取、PagedAttention 消除显存碎片、连续批处理提高每步的 batch 以摊薄权重读取、量化降低字节数。

4. 有面试官说"Transformer 的复杂度是 O(n²)，所以长上下文注定昂贵"。这句话哪里不完整？

答案

要区分计算与显存、训练与推理。注意力确实有 $O(n^2 d)$ 计算和 $O(n^2)$ 中间激活（训练时）；但推理 decode 阶段有 KV Cache，每步只需 $O(nd)$ （对新 query 与缓存的 K/V 做注意力）而非 $O(n^2)$ ，KV Cache 的显存占用是 $O(nd)$ 级别的线性增长。另外 FFN 部分是 $O(nd^2)$ ，与 n 线性。真正的成本结论应该是：长上下文下注意力的算力占比上升、KV Cache 显存线性增长，且原始注意力的显存峰值可用 FlashAttention 从 $O(n^2)$ 降到 $O(n)$ 。

5. 如果让你把一个 12 层的标准 Transformer 改成现代 LLM 的形态，你会做哪些改动？

答案

① 归一化改 Pre-LN + RMSNorm，去掉 LayerNorm 的均值与 bias；② FFN 换成 SwiGLU，中间维度按 8/3·d 配平参数量；③ 线性层去 bias；④ 加 RoPE 而不是可学习绝对位置编码；⑤ 因果掩码（decoder-only），去掉 encoder-decoder 的 cross-attention；⑥ 注意力改 GQA（多个 query 头共享 KV 头）以压缩 KV Cache；⑦ 用 FlashAttention 实现注意力；⑧ 词表扩大（多语言需约 10 万+）以缩短序列长度。这些改动大多可以单独消融，其中 RoPE、Pre-LN/RMSNorm、SwiGLU、GQA 是现代开源 LLM 的公约数。