Transformer 架构拆解：注意力机制到底在算什么

点积是一次可微的软检索

把每个 token 表示成 d_model 维向量，一句话就是矩阵 X ∈ R^(n×d)。注意力层用三个可学习投影把它变成三份：

Q = X W_Q    K = X W_K    V = X W_V

几何含义比公式更好用：K 是”我能被什么检索到”的标签，Q 是”我现在要找什么”的探针，V 是”命中之后真正取走的内容”。q_i 与 k_j 的点积衡量第 i 个位置想从第 j 个位置取多少信息，softmax 把所有 j 的竞争归一化成非负、和为 1 的权重，再对 V 加权求和。

这不就是字典查找吗。普通查找用 argmax 取出唯一条目，不可导；注意力用 softmax 取出所有条目的加权混合，处处可导，所以能塞进网络里端到端训练。认准”软检索”这个定位，后面大部分变体都能归到”换一种检索结构”：稀疏注意力是限制候选集，KV 缓存是缓存 K/V 避免重算，线性注意力是换掉相似度核函数。

除以 √d_k 是为了救 softmax 的梯度

设 q、k 的每个分量独立、均值 0、方差 1。点积是 d_k 项之和：

q·k = Σ_{i=1}^{d_k} q_i k_i

每一项方差为 1（两个独立标准正态之积，方差 = E[q²]E[k²] = 1），d_k 项相加后方差就是 d_k，标准差为 √d_k。d_k = 64 时点积的典型绝对值在 8 附近，d_k = 4096 时在 64 附近。

softmax 的麻烦在于它是饱和函数。若某个 logit 比其余大出 10，对应概率已接近 1，其余接近 0，雅可比矩阵几乎全零，反向传播的梯度被掐断。训练初期尤其致命：参数还没学到有意义的对齐模式，logit 的分散纯粹来自维度带来的方差放大，一层层 softmax 早早进入死区。

除以 √d_k 把 logit 的标准差拉回 1，让分布在整个训练过程中保持在 softmax 的敏感区。这不是什么精妙推导，就是一次方差归一化。

import numpy as np

rng = np.random.default_rng(0)
d = 64
q = rng.normal(size=(20000, d))
k = rng.normal(size=(20000, d))
dot = (q * k).sum(-1)
print(dot.std())                # ≈ 8.0  ≈ sqrt(64)
print((dot / np.sqrt(d)).std()) # ≈ 1.0

多头注意力切的是子空间

单头注意力把整个 d_model 维空间压成一组注意力权重，所有语义关系挤在同一个相似度度量里。多头做法是把 d_model 切成 h 份，每份 d_k = d_model / h 维，各自独立算注意力，再拼接后做一次线性投影：

MultiHead(X) = Concat(head_1, ..., head_h) W_O
head_i = Attention(X W_Qi, X W_Ki, X W_Vi)

关键在于参数与算力不变。单头的 W_Q 是 d_model×d_model；多头每份是 d_model×(d_model/h)，h 份拼起来仍是 d_model×d_model。FLOPs 同样守恒。所以多头不是”堆更多计算”，而是”把同样的计算拆到 h 个互不干扰的低维子空间里”，让不同头各自 specialize——经验上确实能观察到有的头盯句法依存、有的头盯指代、有的头只关注相邻位置。

实现上别写循环，用 reshape + transpose 把头维提到 batch 维：

def split_heads(T, n_heads):
    # T: (batch, n, d_model) -> (batch, n_heads, n, d_head)
    b, n, d = T.shape
    return T.reshape(b, n, n_heads, d // n_heads).transpose(0, 2, 1, 3)

没有位置编码，注意力看不见顺序

自注意力对输入是置换等变的。把输入 token 的顺序打乱，输出的注意力权重矩阵只是行列同步置换，每个 token 聚合到的内容完全一样，只是跟着搬家。证明很短：注意力分数只由 q_i·k_j 决定，与 i、j 的绝对下标无关，softmax 逐行归一化也无关。

换句话说，纯自注意力看到的是一个集合，不是序列。”猫追狗”和”狗追猫”在它眼里是同一袋词。位置编码就是把位置信息重新注入回去。

早期用正弦编码 PE(pos, 2i) = sin(pos / 10000^(2i/d))，加在输入 embedding 上。它的好处是对任意长度都有定义，且不同频率的正弦组合使相对位置可由线性变换表示。现在更常见的是 RoPE：不对 embedding 做加法，而是把 Q、K 向量在复平面上按位置旋转角度 pos·θ_i。旋转后 q_i·k_j 只依赖相对位置 (i – j)：

(R_m q)·(R_n k) = q·R_{n-m} k

这个性质让模型天然以相对距离建模，也是长上下文外推相对容易做的原因之一。

O(n²) 是结构性的，绕不过去只能换招

注意力要算 n×n 的分数矩阵，时间和显存都随序列长度平方增长。序列从 4k 拉到 128k，长度 32 倍，注意力部分开销是 1024 倍。这不是常数优化能解决的，必须改计算结构。

要分清”降复杂度”和”降常数”两类。FlashAttention 属于后者：它不改变 n² 这个数量级，但通过把 Q/K/V 分块载入 SRAM、避免把完整 n×n 矩阵写回显存，把显存占用从 O(n²) 压到 O(n)，实际墙钟时间常常快数倍。MQA/GQA 则是针对自回归解码的带宽瓶颈——解码阶段每生成一个 token 都要读一遍全部 KV， Arithmetic intensity 极低，减少 KV 头数是直接有效的。

一个 numpy 版最小实现

import numpy as np

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def attention(X, W_q, W_k, W_v, causal=True):
    # X: (n, d_model)
    n = X.shape[0]
    Q, K, V = X @ W_q, X @ W_k, X @ W_v
    scores = Q @ K.T / np.sqrt(K.shape[-1])      # (n, n)
    if causal:
        mask = np.triu(np.ones((n, n)), k=1).astype(bool)
        scores = np.where(mask, -np.inf, scores)
    return softmax(scores) @ V, softmax(scores)

三十行不到就是全部核心。工业实现的差异全在工程：半精度、kernel 融合、KV 分页、批处理调度，数学上没有任何新东西。

注意力权重不等于”模型在关注什么”

把注意力权重画成热力图当成解释，是很常见也很容易翻车的做法。权重高只说明这个位置的 V 被大量混入输出，并不等价于它对最终预测负有因果责任。两个原因：一是注意力之后还有残差连接和 FFN，被高权重聚合进来的信息可能在后续层被完全改写；二是同一层的多个头会互相抵消，A 头强关注 x、B 头把 x 的贡献减掉，净效果接近零，但两张热力图都很亮。

想做归因， Ablation 或者梯度类方法比看权重靠谱。直接把某个头的输出置零、或者把某条注意力边屏蔽，观察下游指标的变化量，这才是可操作的因果证据。代价是要跑多次前向，长输入下不便宜。

另一个常被忽略的现象是注意力沉淀：序列开头的几个 token 会吸走大量注意力权重，即使它们在语义上毫无作用。这被认为是 softmax 归一化的副产物——模型需要一个”无处可去时倒掉”的接收端。做长上下文实验时如果发现首个 token 权重异常高，先别急着找数据问题。

几个反复踩的坑

把 mask 加在 softmax 之前而不是之后。见过有人先 softmax 再乘 mask，得到的行和不为 1，输出整体被压暗，而且被 mask 的位置变成了绝对零而非”不可见”，数值上还能传梯度进去。

数值稳定那行 x - x.max() 不能省。fp16 下 exp 超过约 11 就溢出成 inf，整行变 NaN，而且 NaN 会顺着残差流一路污染到 loss，排查起来很痛苦。

多头拼接后别忘了 W_O。去掉它等价于强制各头的输出子空间互不混合，实测会掉点，而且掉得不明显——这种静默退化比直接报错危险得多。