4.2 Attention Mechanism
Attention 是 content-addressed memory。给定 query，模型在一组 key-value memory 中检索相关 value，并用可微 softmax 完成加权求和。
Running Example: Pronoun Resolution
看句子：
The animal did not cross the street because it was tired.当模型处理 it 时，它需要判断 it 更像指代 animal 还是 street。attention 的直觉就是：当前位置产生一个 query，历史 token 产生 keys 和 values。若 it 的 query 与 animal 的 key 更相似，那么 animal 的 value 就会被更多读入当前表示。
用极简数值例子表示，假设 query 为
\[ q_{\text{it}}=[1,0], \]
两个 key 为
\[ k_{\text{animal}}=[2,0], \qquad k_{\text{street}}=[0,1]. \]
点积分数是
\[ q^\top k_{\text{animal}}=2, \qquad q^\top k_{\text{street}}=0. \]
softmax 后：
\[ \alpha_{\text{animal}} = \frac{e^2}{e^2+e^0}\approx0.88, \qquad \alpha_{\text{street}}\approx0.12. \]
因此当前 token 的表示主要读取 animal 的 value。这当然是玩具例子，但它展示了 attention 的基本机械过程：similarity score \(\rightarrow\) probability \(\rightarrow\) weighted sum。
Q, K, V
Given \(Q\in\mathbb{R}^{n_q\times d_k}\), \(K\in\mathbb{R}^{n_k\times d_k}\), and \(V\in\mathbb{R}^{n_k\times d_v}\), \[ \operatorname{Attention}(Q,K,V) = \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right)V. \]
\(QK^\top\) 是相似度矩阵。第 \(i\) 行表示第 \(i\) 个 query 对所有 keys 的检索分数；softmax 后得到概率分布；乘 \(V\) 后得到 weighted memory readout。
QK as a Learned Bilinear Compatibility
在 self-attention 中，\(Q,K,V\) 都来自同一个 hidden sequence \(X\)：
\[ Q=XW^Q,\qquad K=XW^K,\qquad V=XW^V. \]
于是第 \(i\) 个 token 对第 \(j\) 个 token 的 raw score 是
\[ s_{ij} = \frac{(x_iW^Q)(x_jW^K)^\top}{\sqrt{d_h}} = \frac{x_i W^Q(W^K)^\top x_j^\top}{\sqrt{d_h}}. \]
令
\[ A=W^Q(W^K)^\top, \]
则 \(s_{ij}\) 是由矩阵 \(A\) 定义的 bilinear compatibility。也就是说，attention 不只是“点积相似度”，而是先把 token 表示投影到 query/key 两个子空间，再用一个低秩双线性形式判断兼容性。
Attention compatibility is the learned scalar score that decides how strongly a query position addresses a key position before softmax normalization.
这个视角解释了三个现象：
- \(W^Q\) 决定“当前位置想找什么”；
- \(W^K\) 决定“每个历史位置以什么方式被索引”；
- \(W^V\) 决定“被读到之后提供什么内容”。
由于
\[ \operatorname{rank}(A) = \operatorname{rank}(W^Q(W^K)^\top) \leq d_h, \]
单个 head 的 score 几何受 head dimension 限制。多头 attention 不是简单重复同一个检索器，而是在多个低秩 compatibility 子空间里并行检索，再通过 \(W^O\) 混合回 residual stream。
For one attention head without mask bias, the raw score matrix \[ S=\frac{QK^\top}{\sqrt{d_h}} \] has rank at most \(d_h\).
矩阵 \(Q\in\mathbb{R}^{T\times d_h}\)，\(K\in\mathbb{R}^{T\times d_h}\)。由矩阵秩不等式：
\[ \operatorname{rank}(QK^\top) \leq \min(\operatorname{rank}(Q),\operatorname{rank}(K)) \leq d_h. \]
除以常数 \(\sqrt{d_h}\) 不改变 rank。注意 softmax、mask bias、多个 heads 和后续非线性会改变最终函数族；这个定理只描述 raw score matrix 的线性代数结构。
Why Divide by \(\sqrt{d_k}\)
若 \(q_i,k_i\) 独立、均值 \(0\)、方差 \(1\)，则
\[ q^\top k = \sum_{i=1}^{d_k} q_i k_i, \qquad \operatorname{Var}(q^\top k)=d_k. \]
不缩放时，维度越大 logits 方差越大，softmax 越容易饱和，梯度越小。除以 \(\sqrt{d_k}\) 后 logits 方差回到约 \(1\)。
假设 \(q_i,k_i\) 独立，且
\[ \mathbb{E}[q_i]=\mathbb{E}[k_i]=0, \qquad \operatorname{Var}(q_i)=\operatorname{Var}(k_i)=1. \]
由于 \(q_i,k_i\) 独立，
\[ \mathbb{E}[q_i k_i]=0, \]
且
\[ \operatorname{Var}(q_i k_i) = \mathbb{E}[q_i^2k_i^2] - \mathbb{E}[q_ik_i]^2 = \mathbb{E}[q_i^2]\mathbb{E}[k_i^2] =1. \]
点积是
\[ q^\top k=\sum_{i=1}^{d_k}q_ik_i. \]
独立项方差相加：
\[ \operatorname{Var}(q^\top k)=d_k. \]
因此
\[ \operatorname{Var}\left(\frac{q^\top k}{\sqrt{d_k}}\right)=1. \]
这不是说真实网络里 \(q,k\) 一定独立同分布，而是说明 scaling 的初始化尺度动机：让 attention logits 不随 head dimension 系统性变大。
Softmax as Differentiable Retrieval
对于单个 query，attention weights 是
\[ \alpha_j = \frac{\exp(s_j)} {\sum_r\exp(s_r)}, \qquad s_j=\frac{q^\top k_j}{\sqrt{d_k}}. \]
输出是
\[ o=\sum_j\alpha_jv_j. \]
这像在 memory table 里检索，但不是硬检索一个 value，而是对所有 value 做 soft read。softmax 的 Jacobian 为
\[ \frac{\partial \alpha_i}{\partial s_j} = \alpha_i(\mathbf{1}_{i=j}-\alpha_j). \]
对
\[ \alpha_i=\frac{e^{s_i}}{Z}, \qquad Z=\sum_r e^{s_r}, \]
求导。若 \(i=j\)：
\[ \frac{\partial \alpha_i}{\partial s_i} = \frac{e^{s_i}Z-e^{s_i}e^{s_i}}{Z^2} = \alpha_i(1-\alpha_i). \]
若 \(i\ne j\)：
\[ \frac{\partial \alpha_i}{\partial s_j} = -\frac{e^{s_i}e^{s_j}}{Z^2} = -\alpha_i\alpha_j. \]
合并即
\[ \frac{\partial \alpha_i}{\partial s_j} = \alpha_i(\mathbf{1}_{i=j}-\alpha_j). \]
这个式子解释了两个训练现象：
- 如果 softmax 极端尖锐，最大项 \(\alpha_i\approx1\)，很多梯度接近 0；
- attention 的梯度不是每个 key 独立变化，而是所有 keys 通过归一化互相竞争。
Raising one attention logit lowers the normalized probability mass available to other visible tokens. Masking, temperature, and scaling all change this competition.
Temperature, Entropy, and Sharpness
有时 attention 会显式或隐式改变 temperature：
\[ \alpha_j(\tau) = \frac{\exp(s_j/\tau)} {\sum_r\exp(s_r/\tau)}. \]
\(\tau\) 小时分布更尖锐，接近 hard retrieval；\(\tau\) 大时分布更平，接近平均读取。标准 scaled dot-product attention 相当于把 score 先除以 \(\sqrt{d_h}\)，让初始化时的 effective temperature 不随 head dimension 漂移。
可以把 entropy 写成：
\[ H(\alpha) = -\sum_j \alpha_j\log\alpha_j. \]
令 \(\beta=1/\tau\)，\(\alpha_j=\operatorname{softmax}(\beta s)_j\)。有一个很有用的结论：
\[ \frac{dH}{d\tau} = \frac{\operatorname{Var}_{j\sim\alpha}(s_j)}{\tau^3} \geq0. \]
也就是说，temperature 增大时 attention entropy 单调不降。它不是一个只影响“随机性”的采样超参，而是在训练和推理里都会改变 value mixing 的平滑程度。
写
\[ \log\alpha_j=\beta s_j-\log Z, \qquad Z=\sum_r e^{\beta s_r}. \]
Entropy 为
\[ H = -\sum_j\alpha_j\log\alpha_j = -\beta\mathbb{E}_\alpha[s]+\log Z. \]
对 \(\beta\) 求导。因为
\[ \frac{d}{d\beta}\log Z = \mathbb{E}_\alpha[s], \]
并且
\[ \frac{d}{d\beta}\mathbb{E}_\alpha[s] = \operatorname{Var}_\alpha(s), \]
所以
\[ \frac{dH}{d\beta} = -\mathbb{E}_\alpha[s] -\beta\operatorname{Var}_\alpha(s) + \mathbb{E}_\alpha[s] = -\beta\operatorname{Var}_\alpha(s). \]
又 \(\beta=1/\tau\)，\(d\beta/d\tau=-1/\tau^2\)，因此
\[ \frac{dH}{d\tau} = \frac{\operatorname{Var}_\alpha(s)}{\tau^3} \geq0. \]
When attention becomes nearly one-hot, the softmax Jacobian has very small entries for most keys. The model may still copy through the selected value, but gradients that would teach alternative key matches become weak.
Attention as Kernel Regression
Attention 也可以看成一种可学习的 kernel smoother。给定 query \(q\) 和 keys \(k_j\)，定义 kernel:
\[ K(q,k_j)=\exp\left(\frac{q^\top k_j}{\sqrt{d_k}}\right). \]
则
\[ o(q) = \sum_j \frac{K(q,k_j)} {\sum_r K(q,k_r)} v_j. \]
这与 Nadaraya-Watson kernel regression 形式相同：相似的 keys 给出更大权重，对 values 做加权平均。不同点在于 Transformer 里的 \(q,k,v\) 都是由神经网络学出来的，不是手工定义的特征。
这个视角很有用：attention 的关键不是“点积”本身，而是三步：
- learn a query space；
- learn a key space for addressing；
- learn a value space for content。
同一个 token 的 key 和 value 可以承担不同功能：key 决定“会不会被读到”，value 决定“被读到后提供什么信息”。
Values Matter as Much as Weights
只看 attention map 容易误解。输出是
\[ o_i=\sum_j\alpha_{ij}v_j. \]
如果 \(\alpha_{ij}\) 很大但 \(v_j\) 在当前方向上信息很小，贡献仍可能有限；如果某个 value 经过 \(W^V\) 投影后很强，即使 attention 权重中等，也可能对 residual stream 产生明显影响。
更进一步，多层网络中第 \(\ell\) 层某个 head 的输出还会经过
\[ O_{\text{head}}W^O \]
写回 residual stream，再被后续 MLP 和 attention 改写。因此 attention weights 是诊断入口，不是完整解释。
Attention Output as a Convex Combination
在没有 output projection 之前，单个 query 的 attention 输出
\[ o_i=\sum_j\alpha_{ij}v_j \]
满足
\[ \alpha_{ij}\geq0, \qquad \sum_j\alpha_{ij}=1. \]
所以 \(o_i\) 位于 visible value vectors 的 convex hull 中。这一点很重要：单个 attention read 本身是在“从已有 values 中混合”，不是凭空创造新方向。新的方向来自多头拼接、output projection、residual addition、MLP 非线性和层间组合。
For a fixed query and visible key set, the pre-projection attention output lies in the convex hull of the corresponding value vectors.
这也解释了为什么 value vectors 不能被忽略。如果所有 \(v_j\) 在某个语义方向上都没有信息，再漂亮的 attention map 也读不出那个信息；如果 \(W^V\) 把某个 token 的特征变成强 value，较小的 attention weight 也可能产生大贡献。
Multi-Head Attention
单个 attention head 只能在一个子空间里检索。Multi-head attention 并行做多组投影：
\[ \operatorname{head}_i = \operatorname{Attention}(XW_i^Q,XW_i^K,XW_i^V), \]
\[ \operatorname{MHA}(X) = [\operatorname{head}_1;\ldots;\operatorname{head}_H]W^O. \]
不同 head 可以学习语法、位置、实体指代、复制、局部窗口、分隔符等不同模式。
Head Dimensions and Parameter Count
设 hidden size 为 \(d\)，head 数为 \(H\)，每个 head 维度为
\[ d_h=\frac{d}{H}. \]
标准 MHA 通常用四个线性投影：
\[ W^Q,W^K,W^V,W^O\in\mathbb{R}^{d\times d}. \]
所以 attention block 的主要参数量约为
\[ 4d^2. \]
改变 head 数 \(H\) 时，如果 \(d\) 固定，参数量不变，但每个 head 的维度 \(d_h\) 变小。更多 heads 给模型更多并行检索子空间，但每个 head 的表示容量更小；这就是为什么实际模型通常让 \(d_h\) 保持在 64、128 这类比较稳定的范围。
更细地看，head 数改变的是三件事：
| Quantity | Formula when \(d\) fixed | Effect |
|---|---|---|
| per-head dimension | \(d_h=d/H\) | smaller compatibility subspace |
| raw score rank per head | \(\leq d_h\) | lower-rank score geometry |
| number of score matrices | \(H\) | more parallel retrieval patterns |
因此 “more heads” 不是单调更强。若 \(d_h\) 太小，每个 head 的 Q/K compatibility 过窄；若 heads 太少，不同关系模式会被迫挤在同一个 score matrix 里。现代 LLM 常把 \(d_h\) 稳定在一个经验区间，再通过层数、宽度、GQA/MQA 等系统约束调节。
With fixed hidden size, increasing the number of heads reduces each head’s dimension. It changes the factorization of attention capacity rather than simply adding parameters.
Cross-Attention
Self-attention 中 \(Q,K,V\) 都来自同一个序列 \(X\)。Cross-attention 中 query 来自 target sequence，key/value 来自 source sequence：
\[ Q=Y W^Q, \qquad K=X W^K, \qquad V=X W^V. \]
这用于 encoder-decoder 模型，例如翻译：
source encoder states -> K,V
target decoder states -> Qdecoder 在生成目标语言 token 时，用 target prefix 产生 query，再到 source sentence 的 memory 中读取相关信息。Decoder-only LLM 通常没有显式 cross-attention，而是把所有条件都写进同一个 prompt，通过 self-attention 读取。
MHA, MQA, and GQA
现代 LLM serving 中，KV cache 很贵，所以 attention 结构也会围绕 KV cache 改。
| Variant | Query heads | KV heads | KV cache |
|---|---|---|---|
| MHA | \(H_q\) | \(H_q\) | largest |
| MQA | \(H_q\) | \(1\) | smallest |
| GQA | \(H_q\) | \(H_{kv}<H_q\) | middle |
Grouped-query attention 的计算可以理解为：多个 query heads 共享同一组 key/value heads。若 \(H_q=32,H_{kv}=8\)，则每 4 个 query heads 共享一个 KV head。这样模型仍有 32 个 query 子空间，但 cache 只存 8 组 K/V。
KV head sharing reduces the number of key/value projections and cache heads while retaining more query heads. It trades some attention expressivity for lower memory bandwidth during decoding.
如果 batch 为 \(B\)，decode 到长度 \(T\)，head dimension 为 \(d_h\)，KV dtype bytes 为 \(s\)，单层 KV cache 大小约为：
\[ M_{\text{KV}} = 2BT H_{kv}d_hs. \]
这里的 \(2\) 来自 key 和 value。MHA、MQA、GQA 的主要系统差异可以直接从这个式子看出来：
\[ \frac{M_{\text{GQA}}}{M_{\text{MHA}}} = \frac{H_{kv}}{H_q}, \qquad \frac{M_{\text{MQA}}}{M_{\text{MHA}}} = \frac{1}{H_q}. \]
所以 GQA 不是“小技巧”，而是 decode-time memory bandwidth 设计。大模型服务时，每生成一个 token 都要读取历史 KV；减少 \(H_{kv}\) 会直接降低每步读取量。
实现 GQA 时常见做法是把 KV heads repeat 到 query heads：
def repeat_kv(x, n_rep):
# x: [B, H_kv, T, D]
if n_rep == 1:
return x
b, h, t, d = x.shape
x = x[:, :, None, :, :].expand(b, h, n_rep, t, d)
return x.reshape(b, h * n_rep, t, d)这个 expand 不复制 KV，本质上是 stride trick；后续 attention kernel 可能会用专门的 grouped layout 避免真的 materialize repeat。
Attention as Message Passing
Self-attention 可以看成完全图上的 message passing：
\[ h_i' = \sum_{j} \alpha_{ij}W^V h_j, \qquad \alpha_{ij} = \operatorname{softmax}_j \left( \frac{(W^Qh_i)^\top(W^Kh_j)}{\sqrt{d_k}} \right). \]
这与 GNN 的区别是：GNN 通常给定稀疏 graph，Transformer 动态学习 dense graph。
Permutation Equivariance Without Position
如果没有 positional encoding，也没有 causal/padding mask，self-attention 对 token 顺序是 permutation equivariant。也就是说，如果把输入 token 顺序整体打乱，输出会以同样方式被打乱，而不会知道“第几个 token”。
Let \(P\) be a permutation matrix and define self-attention without positional features or masks as \[ \operatorname{Attn}(X)=\operatorname{softmax}\left(\frac{XW^Q(W^K)^\top X^\top}{\sqrt{d_h}}\right)XW^V. \] Then \[ \operatorname{Attn}(PX)=P\operatorname{Attn}(X). \]
令
\[ Q=XW^Q,\qquad K=XW^K,\qquad V=XW^V. \]
输入变成 \(PX\) 后：
\[ Q'=PQ,\qquad K'=PK,\qquad V'=PV. \]
score matrix 为
\[ S' = \frac{Q'(K')^\top}{\sqrt{d_h}} = \frac{P Q K^\top P^\top}{\sqrt{d_h}} = P S P^\top. \]
row-wise softmax 对 permutation 保持等变：
\[ \operatorname{softmax}(P S P^\top) = P\operatorname{softmax}(S)P^\top. \]
因此
\[ \operatorname{Attn}(PX) = P\operatorname{softmax}(S)P^\top P V = P\operatorname{softmax}(S)V = P\operatorname{Attn}(X). \]
这个定理说明 position encoding 和 mask 不是可有可无的装饰。没有它们，self-attention 更像集合函数，对顺序没有内在感知；加入 causal mask 后，模型不仅知道位置，还被限制为一个有向时间图；加入 RoPE/ALiBi/absolute position 后，score 里又混入了相对或绝对位置信息。
A sequence function \(F\) is permutation equivariant if permuting the input sequence permutes the output sequence in the same way: \(F(PX)=PF(X)\).
Complexity
标准 attention 的时间和显存复杂度是
\[ O(T^2d). \]
这使长上下文成为系统瓶颈。后续的 FlashAttention、sliding window、linear attention、state-space models，都是在不同方向上缓解这个 \(T^2\)。
Attention weights show where information is read from in a layer/head, but model computation also includes value vectors, MLPs, residual paths, and later layers. They are useful diagnostics, not complete causal explanations.
更细地看，self-attention 主要有三块开销：
| Operation | Shape intuition | Cost |
|---|---|---|
| QKV projection | \([BT,d]\times[d,3d]\) | \(O(BTd^2)\) |
| score matrix | \(QK^\top\) | \(O(BH T^2 d_h)\) |
| value mixing | \(PV\) | \(O(BH T^2 d_h)\) |
当 \(T\) 小、\(d\) 大时，projection/MLP 可能主导；当 \(T\) 很长时，\(T^2\) attention 主导。很多“attention 优化”只优化第二/三项，不会减少 QKV projection 和 MLP 的成本。
Masked Stable Attention
真实 attention 的计算顺序通常是：
project Q,K,V
-> apply RoPE / position operation
-> compute scores = QK^T / sqrt(d_h)
-> add mask bias
-> stable softmax
-> attention dropout during training
-> multiply by V
-> output projectionmask 一般以 additive bias 进入 logits：
\[ S_{ij} = \frac{q_i^\top k_j}{\sqrt{d_h}} +M_{ij}, \]
其中可见位置 \(M_{ij}=0\)，不可见位置 \(M_{ij}=-\infty\)。softmax 后不可见位置概率为 0。
数值稳定 softmax：
\[ \operatorname{softmax}(s)_j = \frac{\exp(s_j-m)} {\sum_r\exp(s_r-m)}, \qquad m=\max_r s_r. \]
If every logit in a row is masked to \(-\infty\), stable softmax still has an undefined denominator. Kernels and masks must ensure each query has at least one visible key or explicitly handle empty rows.
最小 PyTorch 实现：
import math
import torch
import torch.nn.functional as F
def attention(q, k, v, mask=None, dropout_p=0.0, training=False):
# q,k,v: [B, H, T, Dh]
scores = q @ k.transpose(-2, -1)
scores = scores / math.sqrt(q.shape[-1])
if mask is not None:
# mask should broadcast to [B, H, Tq, Tk], True means visible.
scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
probs = torch.softmax(scores.float(), dim=-1).to(q.dtype)
probs = F.dropout(probs, p=dropout_p, training=training)
out = probs @ v
return out两个细节值得注意：
- softmax 常转 FP32 做归一化，再转回低精度；
- mask 的 shape 通常要能 broadcast 到
[B,H,T_q,T_k]。
Gradient Flow Through Attention
attention 的输出为
\[ O=PV, \qquad P=\operatorname{softmax}(S), \qquad S=\frac{QK^\top}{\sqrt{d_h}}. \]
给定上游梯度 \(G=\partial\mathcal{L}/\partial O\)，对 value 的梯度：
\[ \frac{\partial\mathcal{L}}{\partial V} = P^\top G. \]
对 attention probability：
\[ \frac{\partial\mathcal{L}}{\partial P} = GV^\top. \]
softmax backward 对每一行独立：
\[ \frac{\partial\mathcal{L}}{\partial S_i} = P_i\odot \left( \frac{\partial\mathcal{L}}{\partial P_i} - \left\langle \frac{\partial\mathcal{L}}{\partial P_i},P_i \right\rangle \mathbf{1} \right). \]
最后：
\[ \frac{\partial\mathcal{L}}{\partial Q} = \frac{1}{\sqrt{d_h}} \frac{\partial\mathcal{L}}{\partial S}K, \qquad \frac{\partial\mathcal{L}}{\partial K} = \frac{1}{\sqrt{d_h}} \left( \frac{\partial\mathcal{L}}{\partial S} \right)^\top Q. \]
这说明 attention 训练不只是更新 value content。若某个 head 读错位置，loss 会通过 softmax competition 回到 Q/K 空间，改变“应该匹配谁”；同时通过 \(P^\top G\) 改变被读 token 的 value 内容。
KV Cache and Incremental Decoding
训练时对完整序列并行计算：
Q,K,V: [B, H, T, Dh]
scores: [B, H, T, T]decode 时每步只新增一个 query：
q_t: [B, Hq, 1, Dh]
K_cache: [B, Hkv, t, Dh]
V_cache: [B, Hkv, t, Dh]
scores: [B, Hq, 1, t]伪代码：
def decode_step(x_t, cache):
q, k_new, v_new = project_qkv(x_t)
k_new, v_new = apply_rope_with_offset(k_new, v_new, cache.length)
cache.append(k_new, v_new)
k, v = cache.k, cache.v
out = attention(q, repeat_kv(k, n_rep), repeat_kv(v, n_rep))
return out, cache这里 position offset 很关键。若 RoPE/position id 每步都从 0 开始，模型会把所有 decode token 当作同一位置，生成会迅速崩坏。
Prefill computes all queries against all keys in parallel. Decode computes one or a few new queries against a growing KV cache. Bugs often come from reusing training-time mask or position code in decode without offsets.
Implementation Checklist
写 attention 时至少验证：
- Q/K/V projection 后 shape 是否是
[B,H,T,Dh]； - scale 是否用
sqrt(d_h)而不是sqrt(d_model)； - mask 是否 broadcast 到 score matrix，并且没有 all-masked rows；
- causal mask 是否禁止未来 token；
- padding mask 是否同时影响 attention 和 loss；
- softmax 是否在稳定 dtype 中计算；
- dropout 是否只在 training mode 对 attention probabilities 生效；
- GQA/MQA 的 KV repeat 是否没有误复制或错配 heads；
- KV cache append 是否使用正确 position offset；
- attention weights、value vectors、output projection 是否一起看，而不是只解释 heatmap。