Skip to content

Transformer 手撕与注意力机制

这一页不是给算法研究员做论文复现,而是让测试人员真正建立"模型是怎么从一串 token 走到下一个 token"的机制感。只有把这一层想清楚,很多现象才不会再神秘,比如为什么 prompt 顺序会影响结果、为什么长上下文会失忆、为什么同样主题换一种表达模型轨迹就变了。本章从数学推导、底层原理、源码剖析、实战案例四个维度,把 Transformer 的每一个关键模块拆到你可以手算、可以写代码验证的程度。

教学导读

**定位:**这一章负责回答"模型为什么会受顺序影响、为什么会遗忘、为什么一个词改掉整条输出轨迹就变了"。 **前置依赖:**建议先理解 token、embedding 和基础矩阵运算直觉。 **适用场景:**Prompt 敏感性测试、长上下文稳定性测试、RAG 证据利用分析、推理偏航解释。 **学完产出:**你应该能口头讲清 token 到 next-token 的主链路,手算一个极简 attention 例子,并能阅读 HuggingFace 源码中的注意力实现。

先说结论。

测试不需要把所有训练细节都推到矩阵微分,但至少要能把"为什么这段提示词顺序一换,输出就完全不一样"讲出结构层解释。Transformer 的核心不是记忆力,而是每一层都在做"基于当前上下文,重新决定谁该关注谁"。理解了这一点,你就会明白:模型的"不稳定"不是 bug,而是注意力机制的结构性特征——输入分布变了,权重分配就变了,输出自然不同。 更重要的是,当你看到模型"一本正经地胡说八道"时,你可以不再停留在"模型幻觉"的标签上,而是追问:是哪一层的注意力把错误的证据关联了?是 FFN 层的参数记忆了错误的知识?还是 causal mask 导致模型看不到后面的更正信息?这些结构化的追问,才是一个专业测试人员和普通使用者的差别。

01. 从 token 到 embedding:模型是怎么把文本变成可计算对象的

文本进模型后首先会被 tokenized,再被映射为 embedding 向量。这里的 embedding 不是检索中的句向量,而是模型内部对每一个 token 的独立表示。随后再叠加位置信息,模型才知道"这些 token 不只是一起出现,而且有先后顺序"。

token ids → token embedding → position encoding → Transformer blocks → next-token logits

1.1 为什么位置编码是必须的

如果没有位置编码,模型只会看到一个 token 集合(bag of tokens),不知道"退款流程说明"和"说明退款流程"有什么次序差异。注意力机制本身是排列等变的(permutation equivariant):交换输入顺序,输出也会等价地交换。如果不注入位置信息,"ABC"和"CBA"对模型来说数学上是等价的。对测试来说,这直接意味着:prompt 的内容相同,不代表顺序无关;顺序是有效信号,而位置编码就是让模型感知这个信号的通道。

01.2 Token Embedding + Positional Encoding 的完整矩阵运算推导

1.2.1 Token Embedding 的本质:查表

假设词表大小为 V,embedding 维度为 d_model。模型维护一个 embedding 矩阵 W_E ∈ ℝ^(V × d_model)。对于一个 token id = 42,embedding 操作本质上就是取 W_E 的第 42 行。 用 one-hot 表示来理解这个过程更清晰。如果我们用 one-hot 向量 e_42 ∈ ℝ^V(只有第 42 个位置是 1,其余全 0)来表示这个 token,那么: x_embed = e_42 · W_E = W_E[42, :] 对于一个长度为 n 的输入序列 [t_1, t_2, ..., t_n],我们可以构造 one-hot 矩阵 X_onehot ∈ ℝ^(n × V),然后做一次矩阵乘法得到整个序列的 embedding: X_embed = X_onehot · W_E ∈ ℝ^(n × d_model) 实际工程中不会真的构造 one-hot 矩阵(太稀疏太大),而是用索引直接取行,但数学等价性是一样的。

python
import numpy as np

V = 1000        # 词表大小
d_model = 64    # embedding 维度
seq_len = 4     # 序列长度

W_E = np.random.randn(V, d_model) * 0.02

token_ids = [42, 7, 256, 99]
X_embed = W_E[token_ids]    # shape: (4, 64)

# 等价于 one-hot 矩阵乘法(仅用于理解,不要在工程中这样写)
X_onehot = np.zeros((seq_len, V))
for i, tid in enumerate(token_ids):
    X_onehot[i, tid] = 1.0
X_embed_v2 = X_onehot @ W_E  # shape: (4, 64)

print("两种方式结果一致:", np.allclose(X_embed, X_embed_v2))

1.2.2 绝对位置编码的矩阵形式

原始 Transformer(Vaswani 2017)使用正弦余弦位置编码。对于位置 pos 和维度索引 i: PE(pos, 2i) = sin(pos / 10000^(2i/d_model)) PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model)) 这样每个位置都得到一个 d_model 维的向量。整个序列的位置编码矩阵 PE ∈ ℝ^(n × d_model)。最终送入 Transformer block 的输入是: X_input = X_embed + PE ∈ ℝ^(n × d_model) 这里用的是逐元素加法,不是拼接。直觉上,"token 是什么"和"token 在哪"两种信息被叠加到同一个向量空间中。

python
def sinusoidal_pe(seq_len, d_model):
    PE = np.zeros((seq_len, d_model))
    position = np.arange(seq_len)[:, np.newaxis]          # (n, 1)
    div_term = 10000 ** (np.arange(0, d_model, 2) / d_model)  # (d/2,)
    PE[:, 0::2] = np.sin(position / div_term)
    PE[:, 1::2] = np.cos(position / div_term)
    return PE

PE = sinusoidal_pe(seq_len=4, d_model=64)
X_input = X_embed + PE   # (4, 64) — 这就是进入第一个 Transformer block 的输入

测试启发。

位置编码是加法叠加的,这意味着 token 语义信息和位置信息在同一个向量空间中"混合"。如果模型训练时只见过最长 4096 个位置,那第 4097 个位置的编码就是"没见过的"。这就是为什么超长上下文测试时模型表现会突然下降——不是"记忆力不够",而是位置编码进入了训练分布之外的区域。

01.5 位置编码再讲深一点:RoPE 的旋转矩阵推导

绝对位置编码的直觉是"第 1 个位置就给第 1 个标签,第 2 个位置给第 2 个标签"。RoPE(Rotary Position Embedding)的核心思想完全不同:它不直接加一个位置向量,而是对 Q 和 K 向量施加一个依赖于位置的旋转变换,使得两个 token 之间的点积只取决于它们的相对位置差。

1.5.1 二维旋转的直觉

考虑最简单的情况:embedding 只有 2 维。对位置 m 的向量 [x_1, x_2] 施加角度为 的旋转: R(mθ) · [x_1, x_2]^T = [x_1·cos(mθ) - x_2·sin(mθ), x_1·sin(mθ) + x_2·cos(mθ)] 当位置 m 的 query 和位置 n 的 key 做点积时: (R(mθ)·q)^T · (R(nθ)·k) = q^T · R((n-m)θ) · k 点积的结果只依赖于 (n-m),即相对位置差!这就是 RoPE 的核心性质。

1.5.2 高维推广

对于 d_model 维的向量,RoPE 把它分成 d_model/2 组,每组 2 维。每组使用不同的旋转频率 θ_i = 10000^(-2i/d_model),施加独立旋转。低频组捕捉远距离相对位置关系,高频组捕捉近距离细粒度位置关系。

方式数学形式直觉测试启发
绝对位置编码X + PE"你在第几个位置"超出训练长度后外推困难
RoPER(mθ)·Q, R(nθ)·K"你和别人相差多远"也编码进来理论上有更好的外推性,长上下文模型常用
ALiBiscore - λ·|m-n|距离越远惩罚越大显式衰减远距离注意力,长文推理更保守
python
def apply_rope_2d(x, position, theta=10000.0):
    """对一个 2 维向量施加 RoPE 旋转"""
    d = x.shape[-1]
    freqs = 1.0 / (theta ** (np.arange(0, d, 2) / d))
    angles = position * freqs  # 每组频率乘以位置
    cos_a, sin_a = np.cos(angles), np.sin(angles)
    x_rot = np.empty_like(x)
    x_rot[..., 0::2] = x[..., 0::2] * cos_a - x[..., 1::2] * sin_a
    x_rot[..., 1::2] = x[..., 0::2] * sin_a + x[..., 1::2] * cos_a
    return x_rot

q = np.array([1.0, 0.5, -0.3, 0.8])
q_at_pos3 = apply_rope_2d(q, position=3)
print("原始 q:", q)
print("旋转后 q (pos=3):", q_at_pos3)

02. Q、K、V 是什么:从线性投影到语义角色

最容易记错的是把 Q、K、V 当成三份神秘向量。其实更直白的理解是:

  • Q(Query):当前 token 想找什么信息。"我正在找和退款相关的上下文。"
  • K(Key):每个 token 提供什么索引标签,方便别人判断自己是否相关。"我是关于政策的内容。"
  • V(Value):真正被聚合传递的内容。"如果你认为我相关,拿走这个信息。"

2.1 Q、K、V 的生成过程

Q、K、V 并不是天然存在的——它们是通过线性投影(矩阵乘法)从同一个输入 X 变换而来: Q = X · W_Q ∈ ℝ^(n × d_k) K = X · W_K ∈ ℝ^(n × d_k) V = X · W_V ∈ ℝ^(n × d_v) 其中 W_Q ∈ ℝ^(d_model × d_k)W_K ∈ ℝ^(d_model × d_k)W_V ∈ ℝ^(d_model × d_v)。在标准 Transformer 中,d_k = d_v = d_model / hh 是头数)。 换句话说,attention 不是把别的 token 原封不动搬过来,而是先把所有 token 投影到"提问空间"和"索引空间"中比较相关性,再从"内容空间"中按权重聚合。三个投影矩阵是模型训练时学到的参数——模型通过大量数据学会了"什么样的 token 应该关注什么样的上下文"。

# 完整的 Q, K, V 生成
d_model = 64
d_k = d_v = 16
n = 4

X = np.random.randn(n, d_model)    # 4 个 token 的输入表示

W_Q = np.random.randn(d_model, d_k) * 0.02
W_K = np.random.randn(d_model, d_k) * 0.02
W_V = np.random.randn(d_model, d_v) * 0.02

Q = X @ W_Q    # (4, 16)
K = X @ W_K    # (4, 16)
V = X @ W_V    # (4, 16)

print(f"X shape: {X.shape}")
print(f"Q shape: {Q.shape}, K shape: {K.shape}, V shape: {V.shape}")

03. Self-Attention 的完整数学推导

这是整个 Transformer 最核心的部分。我们一步一步把公式拆开。

3.1 第一步:计算原始分数矩阵

每个 query 和所有 key 做点积,得到一个"相关性分数"矩阵: S = Q · K^T ∈ ℝ^(n × n) 其中 S[i][j] 表示第 i 个 token 的 query 对第 j 个 token 的 key 的相关性。直觉上,这就是一张"谁关注谁"的原始得分表。

3.2 第二步:缩放

S_scaled = S / √d_k 这个缩放操作看起来简单,但背后有严格的数学原因(下一节详细证明)。

3.3 第三步:Mask(可选)

对于 decoder(生成模型),需要施加 causal mask,把"未来"位置的分数设为 -∞: S_masked[i][j] = S_scaled[i][j] if j ≤ i S_masked[i][j] = -∞ if j > i

3.4 第四步:Softmax 归一化

对每一行(即每个 query)独立做 softmax,得到注意力权重: A[i][j] = exp(S_masked[i][j]) / Σ_k exp(S_masked[i][k]) softmax 有两个关键性质:(1) 所有权重非负;(2) 每行之和为 1。这意味着注意力输出是 value 向量的凸组合。被 mask 掉的位置因为分数是 -∞exp(-∞) = 0,权重自然为零。

3.5 第五步:加权聚合

Output = A · V ∈ ℝ^(n × d_v) 每个 token 的输出是所有 value 向量的加权和,权重就是注意力矩阵中对应的那一行。

3.6 完整公式一览

Attention(Q, K, V) = softmax(Q · K^T / √d_k) · V 用具体的数字走一遍。假设有 3 个 token:退款d_k = 2

Q = [[1, 0],    # "我" 的 query
     [1, 1],    # "想" 的 query
     [1, 2]]    # "退款" 的 query

K = [[1, 0],    # "我" 的 key
     [1, 1],    # "想" 的 key
     [0, 2]]    # "退款" 的 key

V = [[0.2, 0.1],
     [0.5, 0.3],
     [0.9, 0.8]]

Step 1: S = Q @ K^T
S = [[1*1+0*0,  1*1+0*1,  1*0+0*2],     [[1, 1, 0],
     [1*1+1*0,  1*1+1*1,  1*0+1*2],  =   [1, 2, 2],
     [1*1+2*0,  1*1+2*1,  1*0+2*2]]      [1, 3, 4]]

Step 2: S_scaled = S / √2 ≈ S / 1.414
S_scaled ≈ [[0.71, 0.71, 0.00],
            [0.71, 1.41, 1.41],
            [0.71, 2.12, 2.83]]

Step 3: Causal Mask(假设用 decoder)
S_masked ≈ [[0.71,  -∞,   -∞ ],
            [0.71, 1.41,  -∞ ],
            [0.71, 2.12, 2.83]]

Step 4: Softmax(每行独立)
A ≈ [[1.00, 0.00, 0.00],
     [0.33, 0.67, 0.00],
     [0.04, 0.26, 0.70]]

Step 5: Output = A @ V
output[0] = 1.00*[0.2,0.1] + 0*[0.5,0.3] + 0*[0.9,0.8] = [0.20, 0.10]
output[1] = 0.33*[0.2,0.1] + 0.67*[0.5,0.3] + 0*[...]   ≈ [0.40, 0.23]
output[2] = 0.04*[0.2,0.1] + 0.26*[0.5,0.3] + 0.70*[0.9,0.8] ≈ [0.77, 0.64]

注意力分数矩阵示意(Causal Mask 后 softmax) query \ key 我 想 退款 我 1.00 0.00 0.00 想 0.33 0.67 0.00 退款 0.04 0.26 0.70

关键洞察。

"退款"这个 token 虽然可以看到所有前面的 token,但它把 70% 的注意力放在自己身上,26% 放在"想"上,只有 4% 放在"我"上。这就是注意力的"选择性"——不是均匀看,而是根据语义相关性有选择地分配。这解释了为什么在长上下文中,远处的指令容易被"忽略"——不是模型看不到,而是注意力权重太低了。

03.2 为什么要除以 √d_k — 方差分析的严格证明

这个问题几乎是所有面试和深度理解 Transformer 的必考点。我们从概率统计的角度严格证明。

3.2.1 问题设定

假设 Q 和 K 的每个元素都是独立同分布的随机变量,均值为 0,方差为 1(标准初始化假设)。即 q_i ~ N(0,1)k_i ~ N(0,1),且彼此独立。

3.2.2 点积的方差推导

一个 query 向量和一个 key 向量的点积为: s = q · k = Σ_{i=1}^{d_k} q_i · k_i 因为 E[q_i] = 0E[k_i] = 0,所以: E[q_i · k_i] = E[q_i] · E[k_i] = 0 Var[q_i · k_i] = E[(q_i · k_i)^2] - (E[q_i · k_i])^2 = E[q_i^2] · E[k_i^2] - 0 = Var[q_i] · Var[k_i] = 1 × 1 = 1 因为 d_k 个乘积项是独立的,点积的方差是各项方差之和: E[s] = 0 Var[s] = d_k × 1 = d_k

3.2.3 为什么方差大了很糟糕

如果 d_k = 512,那么点积的标准差就是 √512 ≈ 22.6。这意味着分数值会在 [-50, +50] 这样的范围里波动。当这些极端值送进 softmax 时: softmax([50, 1, -30]) ≈ [1.0000, 0.0000, 0.0000] softmax 的输出几乎变成了 one-hot!这就是所谓的"梯度消失区"——softmax 在输入值差异很大时,梯度几乎为零,训练会停滞。

3.2.4 除以 √d_k 的效果

s_scaled = s / √d_k Var[s_scaled] = Var[s] / d_k = d_k / d_k = 1 缩放后,点积的方差恢复为 1,分数值在合理范围内波动,softmax 的输出既不会太尖锐(one-hot)也不会太平坦(uniform),梯度保持健康。

# 实验验证:点积方差与 dk 的关系
for dk in [2, 16, 64, 256, 512]:
    q = np.random.randn(10000, dk)
    k = np.random.randn(10000, dk)
    scores = np.sum(q * k, axis=1)
    scores_scaled = scores / np.sqrt(dk)
    print(f"dk={dk:4d}  raw_var={np.var(scores):8.2f}  scaled_var={np.var(scores_scaled):.4f}")

dk= 2 raw_var= 2.01 scaled_var=1.0050

dk= 512 raw_var= 512.37 scaled_var=1.0007 无论 d_k 多大,缩放后方差始终约为 1。这就是 √d_k 的数学功效。

测试启发。

这个缩放机制解释了为什么模型天然会偏向少数高分关联(softmax 的指数放大效应),同时也解释了为什么不同 embedding 维度的模型在同一任务上可能表现不同——维度不仅影响容量,还通过缩放因子影响注意力的"锐度"。

03.5 最小 NumPy 版 Attention:完整实现与验证

python
import numpy as np

def softmax(x, axis=-1):
    """数值稳定的 softmax:减去最大值防止 exp 溢出"""
    x_max = np.max(x, axis=axis, keepdims=True)
    e_x = np.exp(x - x_max)
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: (n, d_k)
    K: (n, d_k)
    V: (n, d_v)
    mask: (n, n) — True 表示该位置被屏蔽
    """
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)       # (n, n)
    if mask is not None:
        scores = np.where(mask, -1e9, scores)
    weights = softmax(scores, axis=-1)     # (n, n)
    output = weights @ V                   # (n, d_v)
    return output, weights

# 演示
Q = np.array([[1., 0.], [1., 1.], [1., 2.]])
K = np.array([[1., 0.], [1., 1.], [0., 2.]])
V = np.array([[0.2, 0.1], [0.5, 0.3], [0.9, 0.8]])

# Causal mask: 上三角为 True(表示屏蔽未来位置)
n = Q.shape[0]
causal_mask = np.triu(np.ones((n, n), dtype=bool), k=1)

output, weights = scaled_dot_product_attention(Q, K, V, mask=causal_mask)
print("注意力权重矩阵:")
print(np.round(weights, 4))
print("\n输出:")
print(np.round(output, 4))

你真正要记住的不是代码本身,而是这三步:先算相关性分数,再做归一化权重,最后加权聚合 value。这就是注意力的核心工作流。上面的 softmax 函数减去了最大值,这是工程中必须做的数值稳定性处理——否则 exp 很容易溢出到 inf

04. Causal Mask:为什么生成模型不能偷看未来

大语言模型在生成第 t 个 token 时,只能看自己之前的 token,不能看后面的 token。这个约束就是 causal mask 实现的。没有它,模型训练时就会"作弊"——直接复制答案而不是学习预测。

4.1 为什么是下三角矩阵

Causal mask 的可见性矩阵是一个下三角矩阵。下三角的含义是:第 i 行(即第 i 个 token 做 query 时),只有 j ≤ i 的列(即位置在它之前或等于它自身的 token)是可见的。

假设序列长度为 5,可见性矩阵(1=可见,0=屏蔽):

     t0  t1  t2  t3  t4
t0 [ 1   0   0   0   0 ]   ← t0 只能看自己
t1 [ 1   1   0   0   0 ]   ← t1 能看 t0 和自己
t2 [ 1   1   1   0   0 ]   ← t2 能看 t0, t1, 自己
t3 [ 1   1   1   1   0 ]   ← t3 能看 t0~t3
t4 [ 1   1   1   1   1 ]   ← t4 能看全部

这就是 np.tril(np.ones((5,5))) — 下三角矩阵。
屏蔽矩阵是它的补集:np.triu(np.ones((5,5)), k=1)

4.2 Mask 在 softmax 前施加

实现上,我们不是把 mask 位置的权重设为 0(在 softmax 之后),而是在 softmax 之前把对应分数设为 -∞(实践中用 -1e9float('-inf'))。这样 exp(-∞) = 0,softmax 之后自然权重为零,数学上更干净。

python
def causal_mask(seq_len):
    """生成 causal mask:True 表示需要被屏蔽的位置"""
    return np.triu(np.ones((seq_len, seq_len), dtype=bool), k=1)

# 为什么用 triu(k=1) 而不是 1 - tril?
# triu(k=1) 得到严格上三角(对角线上方),
# 即 mask[i][j] = True 当 j > i,表示"未来"位置
mask = causal_mask(5)
print("Mask (True=屏蔽):")
print(mask.astype(int))

# 在 attention 计算中使用
scores = np.random.randn(5, 5)
scores_masked = np.where(mask, -1e9, scores)
print("\nMasked scores (未来位置变成 -1e9):")
print(np.round(scores_masked, 2))

4.3 Padding Mask vs Causal Mask

注意区分两种 mask:

  • Causal Mask:防止看到未来 token,是下三角矩阵,所有 decoder 生成模型都用。
  • Padding Mask:防止关注到 padding 位置,是不规则的,取决于每条样本的实际长度。batch 训练时需要把短序列用 [PAD] 补齐,这些位置不应参与注意力计算。

两种 mask 可以叠加使用:combined_mask = causal_mask | padding_mask

测试启发。

生成模型不是先想好全文再一次性输出,而是每一步都基于"到当前为止"的上下文继续往下走。于是前面某个 token 的细微变化,可能改变后续整条轨迹。这就是为什么很多 case 看起来"只是改了一个词",最终却完全变样。理解了 causal mask,你就理解了自回归生成的本质——每一步都是一个不可逆的决策点。

05. Multi-Head Attention 深度剖析

5.1 为什么需要多头

单头注意力只能学到一种"关注模式"——但语言中的依赖关系是多维度的。一个 token 可能同时需要:(1) 关注句法上的主语,(2) 关注语义上的上下文,(3) 关注格式上的模板。多头注意力让模型能并行地在不同的子空间里学习不同类型的关系。

5.2 维度变换详解

假设 d_model = 512,头数 h = 8,则每个头的维度 d_k = d_v = d_model / h = 64。 对每个头 i ∈ {1, ..., h}: Q_i = X · W_Q^i W_Q^i ∈ ℝ^(512 × 64) K_i = X · W_K^i W_K^i ∈ ℝ^(512 × 64) V_i = X · W_V^i W_V^i ∈ ℝ^(512 × 64)

head_i = Attention(Q_i, K_i, V_i) ∈ ℝ^(n × 64)

拼接: MultiHead = Concat(head_1, ..., head_8) ∈ ℝ^(n × 512)

输出投影: Output = MultiHead · W_O W_O ∈ ℝ^(512 × 512)

5.3 参数量精确计算

参数矩阵形状参数量说明
W_Q(所有头合并)512 × 512262,144等价于 8 个 512×64 拼在一起
W_K(所有头合并)512 × 512262,144同上
W_V(所有头合并)512 × 512262,144同上
W_O(输出投影)512 × 512262,144把拼接后的多头映射回 d_model
合计(不含 bias)1,048,576约 100 万参数 / 每个注意力层

在工程实现中,8 个头的 W_Q^i 通常不是分开存储的,而是合并成一个大矩阵 W_Q ∈ ℝ^(512 × 512),然后通过 reshape 和 transpose 拆分到各个头。这就是你在 PyTorch 源码中看到的 view(batch, seq_len, n_heads, head_dim).transpose(1, 2) 的含义。

python
def multi_head_attention(X, W_Q, W_K, W_V, W_O, n_heads, mask=None):
    """
    X: (n, d_model)
    W_Q, W_K, W_V: (d_model, d_model)
    W_O: (d_model, d_model)
    """
    n, d_model = X.shape
    d_k = d_model // n_heads

    Q = X @ W_Q   # (n, d_model)
    K = X @ W_K
    V = X @ W_V

    # 拆分多头: (n, d_model) -> (n_heads, n, d_k)
    Q = Q.reshape(n, n_heads, d_k).transpose(1, 0, 2)
    K = K.reshape(n, n_heads, d_k).transpose(1, 0, 2)
    V = V.reshape(n, n_heads, d_k).transpose(1, 0, 2)

    # 每个头独立做 attention
    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k)  # (h, n, n)
    if mask is not None:
        scores = np.where(mask, -1e9, scores)
    weights = softmax(scores, axis=-1)
    heads = weights @ V   # (h, n, d_k)

    # 拼接回来: (h, n, d_k) -> (n, d_model)
    concat = heads.transpose(1, 0, 2).reshape(n, d_model)
    output = concat @ W_O
    return output, weights

05.5 GQA / MQA:参数共享如何降低推理成本

5.5.1 问题背景

标准 Multi-Head Attention (MHA) 中,每个头都有独立的 Q、K、V 投影。对于 h=32 头的模型,KV Cache 需要存储 32 组 K 和 V——在长序列推理时,这是巨大的显存开销。

5.5.2 三种注意力变体对比

MHA(标准)

32 个 Q 头,32 个 K 头,32 个 V 头。每个头独立。KV Cache 最大。

MQA(Multi-Query)

32 个 Q 头,但只有 1 个 K 头和 1 个 V 头被所有 Q 共享。KV Cache 缩小 32 倍。质量略有损失。

GQA(Grouped-Query)

32 个 Q 头分成 8 组,每组共享 1 个 K 头和 1 个 V 头。KV Cache 缩小 4 倍。质量接近 MHA。Llama 2/3 采用此方案。 MHA: Q_heads=32, K_heads=32, V_heads=32 → KV 参数量 = 2 × 32 × d_k × d_model MQA: Q_heads=32, K_heads=1, V_heads=1 → KV 参数量 = 2 × 1 × d_k × d_model GQA: Q_heads=32, K_heads=8, V_heads=8 → KV 参数量 = 2 × 8 × d_k × d_model GQA 的实现很简单:计算 K 和 V 时用较少的头数,然后在 attention 计算前把每个 KV 头复制给它负责的那几个 Q 头。

# GQA 的核心逻辑
n_q_heads = 32
n_kv_heads = 8   # 每 4 个 Q 头共享 1 个 KV 头
n_groups = n_q_heads // n_kv_heads  # = 4

# K shape: (n_kv_heads, seq_len, d_k) = (8, n, 64)
# 扩展为: (n_q_heads, seq_len, d_k) = (32, n, 64)
# K_expanded = K.repeat(n_groups, axis=0)  — 每个 KV 头复制 4 次

测试启发。

如果你在测试中发现某个模型在长序列上推理速度特别快但质量略降,很可能是用了 MQA 或 GQA。这不是 bug,是设计权衡。了解这个背景可以帮你更准确地评估"质量下降多少是合理的"。

06. LayerNorm vs RMSNorm:归一化的数学推导

6.1 Layer Normalization

对于一个 d 维向量 x,LayerNorm 的计算如下: μ = (1/d) Σ_{i=1}^{d} x_i (均值) σ² = (1/d) Σ_{i=1}^{d} (x_i - μ)² (方差) LayerNorm(x)_i = γ_i · (x_i - μ) / √(σ² + ε) + β_i (归一化 + 仿射变换) 其中 γ(gain)和 β(bias)是可学习参数,ε 是防止除零的小常数(通常 1e-5)。

6.2 RMSNorm

RMSNorm(Root Mean Square Normalization)去掉了减均值的步骤,只用均方根做归一化: RMS(x) = √( (1/d) Σ_{i=1}^{d} x_i² ) RMSNorm(x)_i = γ_i · x_i / RMS(x) 为什么可以去掉均值?研究发现 LayerNorm 的效果主要来自缩放不变性(re-scaling invariance),而不是平移不变性(re-centering invariance)。RMSNorm 保留了缩放不变性,计算量更少(少了一次均值计算和减法),在 Llama、Gemma 等现代模型中被广泛使用。

6.3 代码实现对比

python
def layer_norm(x, gamma, beta, eps=1e-5):
    mu = np.mean(x, axis=-1, keepdims=True)
    sigma2 = np.var(x, axis=-1, keepdims=True)
    x_norm = (x - mu) / np.sqrt(sigma2 + eps)
    return gamma * x_norm + beta

def rms_norm(x, gamma, eps=1e-6):
    rms = np.sqrt(np.mean(x ** 2, axis=-1, keepdims=True) + eps)
    return gamma * (x / rms)

x = np.array([1.0, -2.0, 3.0, 0.5])
gamma = np.ones_like(x)
beta = np.zeros_like(x)

print("LayerNorm:", layer_norm(x, gamma, beta))
print("RMSNorm:  ", rms_norm(x, gamma))

6.4 Pre-Norm vs Post-Norm

原始 Transformer 用 Post-Norm(先做 attention,再 norm),但现代模型几乎都用 Pre-Norm(先 norm,再做 attention)。 Post-Norm: x + Norm(Attention(x)) — 训练不稳定,需要 warmup Pre-Norm: x + Attention(Norm(x)) — 训练更稳定,Llama 等均采用

测试启发。

Pre-Norm 让 Norm 前的残差路径保持"原始信号"不被归一化修改,这增强了深层网络的梯度流动。但也意味着最后一层的输出没有经过归一化——一些模型会在最后额外加一个 RMSNorm。测试时如果发现模型输出的 logits 数值范围异常,可以检查是否和 norm 位置有关。

06.5 FFN 与激活函数:SwiGLU / GeGLU 的推导

6.5.1 标准 FFN

Transformer 中每个 block 除了 attention,还有一个 position-wise feed-forward network: FFN(x) = W_2 · σ(W_1 · x + b_1) + b_2 其中 W_1 ∈ ℝ^(d_model × d_ff),W_2 ∈ ℝ^(d_ff × d_model) 通常 d_ff = 4 × d_model 原始 Transformer 用 ReLU 作为 σ。但 ReLU 的问题在于:它把所有负值直接截为 0("死神经元"),信息损失较大。

6.5.2 GLU 家族

GLU(Gated Linear Unit)引入了门控机制:输出不只是简单的非线性变换,而是由一个"门"来控制信息通过的比例。 GLU(x) = (W_1 · x) ⊙ σ(W_gate · x) σ = sigmoid SwiGLU(x) = (W_1 · x) ⊙ SiLU(W_gate · x) SiLU(z) = z · sigmoid(z) GeGLU(x) = (W_1 · x) ⊙ GELU(W_gate · x) GELU(z) ≈ z · Φ(z) 其中 是逐元素乘法。注意 GLU 家族需要两个权重矩阵(W_1W_gate),所以参数量增加了。为了保持总参数量不变,实际使用中会把 d_ff 调小到 2/3 × 4 × d_model

6.5.3 SwiGLU 的直觉

SiLU(也叫 Swish)函数 z · sigmoid(z) 的特点是在负值区域不会完全截断——它允许少量负值通过,这比 ReLU 保留了更多信息。门控机制则让模型自适应地决定"这个维度的信息要通过多少"。两者结合就是 SwiGLU,被 Llama、PaLM 等模型证明效果优于 ReLU。

python
def silu(x):
    """SiLU / Swish 激活函数"""
    return x * (1 / (1 + np.exp(-x)))

def swiglu_ffn(x, W1, W_gate, W2):
    """SwiGLU FFN 前向传播"""
    return (silu(x @ W_gate) * (x @ W1)) @ W2

# 参数量对比
d_model, d_ff = 512, 2048
print(f"标准 FFN 参数量:  {2 * d_model * d_ff:,}")      # W1 + W2
d_ff_glu = int(2/3 * d_ff * 2)  # 调整后的 d_ff
print(f"SwiGLU FFN 参数量: {3 * d_model * (d_ff_glu//2):,}")  # W1 + W_gate + W2

06.8 残差连接与深度稳定性

残差连接帮助信息和梯度在深层网络里更稳定地传递。其数学形式极其简单: output = x + SubLayer(x) 为什么这一步如此关键?考虑反向传播时梯度的传递: ∂output/∂x = I + ∂SubLayer(x)/∂x 即使 ∂SubLayer(x)/∂x 很小甚至接近零,梯度至少还有一个恒等矩阵 I 可以传递。这就是残差连接防止梯度消失的根本原因。

没有 Residual

每一层都必须从头学好完整表示。梯度经过 L 层后衰减为 O(α^L),深层几乎学不到东西。

有 Residual

每一层只需学习少量修正(delta)。梯度可以通过"高速公路"直达底层。这解释了为什么现代模型可以堆到 80+ 层。

一个直观理解:

没有 residual(串行):
output = f_L(f_{L-1}(...f_2(f_1(x))...))
→ 梯度需要连乘 L 次,极易消失或爆炸

有 residual(跳跃连接):
output = x + f_1(x) + f_2(x+f_1(x)) + ...
→ 信息可以在"高速公路"上直通,不经过任何变换层

测试上的价值:
这解释了为什么模型会表现出"整体保留原意,但局部被新上下文带偏"的现象——
残差连接保留了原始信息,而 attention+FFN 层在做局部修正。

07. KV Cache 原理:推理时为什么可以缓存 K 和 V

7.1 自回归生成的重复计算问题

在自回归生成中,模型每一步只生成一个新 token,但需要对"之前所有 token"重新计算 attention。如果不做优化,生成 n 个 token 的计算量是 O(n³)(每一步都重新计算前面所有 token 的 Q/K/V 和 attention)。

7.2 核心洞察

观察 attention 公式的一个关键性质:当生成第 t+1 个 token 时,前 t 个 token 的 K 和 V 不会改变(因为它们的输入没变,投影权重也没变)。唯一新增的是第 t+1 个 token 自己的 K 和 V。 第 t 步:K_cache = [K_1, K_2, ..., K_t] 第 t+1 步:只需计算 K_{t+1},然后拼接 K_cache = [K_1, K_2, ..., K_t, K_{t+1}]

同理 V_cache 也只需追加,不需要重算前面的。 而 Q 只需要计算当前位置的 Q_{t+1}。

7.3 为什么 Q 不缓存

因为 causal mask 的约束,第 t+1 个位置的 Q 只会和 K_1, ..., K_{t+1} 做 attention——后续位置的 Q 和前面的 K 无关。所以每一步只需要当前位置的 Q,不需要历史 Q。

7.4 显存开销计算

每一层的 KV Cache 大小 = 2 × seq_len × n_kv_heads × d_k × sizeof(dtype)

示例:Llama 2 70B

  • 80 层,8 个 KV 头(GQA),d_k = 128,FP16
  • 4096 长度的 KV Cache = 2 × 4096 × 8 × 128 × 2 bytes × 80 层 = 2 × 4096 × 8 × 128 × 2 × 80 = ~10.7 GB
# KV Cache 的简化模拟
class KVCache:
    def __init__(self):
        self.k_cache = None  # (accumulated_len, d_k)
        self.v_cache = None

    def update(self, new_k, new_v):
        """追加新 token 的 K 和 V"""
        if self.k_cache is None:
            self.k_cache = new_k
            self.v_cache = new_v
        else:
            self.k_cache = np.concatenate([self.k_cache, new_k], axis=0)
            self.v_cache = np.concatenate([self.v_cache, new_v], axis=0)
        return self.k_cache, self.v_cache

cache = KVCache()
for step in range(5):
    new_k = np.random.randn(1, 64)   # 新 token 的 K
    new_v = np.random.randn(1, 64)   # 新 token 的 V
    k_all, v_all = cache.update(new_k, new_v)
    print(f"Step {step}: K cache shape = {k_all.shape}")

测试启发。

KV Cache 是推理服务性能的关键瓶颈。如果你的测试涉及长上下文或高并发场景,需要理解:(1) KV Cache 的显存消耗随序列长度线性增长;(2) 在 batch 推理中,不同请求的 KV Cache 长度不同,会导致显存碎片;(3) 很多推理框架用 PagedAttention(vLLM 的核心技术)来管理 KV Cache 内存,类似操作系统的虚拟内存分页。

07.5 FlashAttention 核心思想

7.5.1 问题:标准 attention 的 IO 瓶颈

标准 attention 需要在 GPU 的高带宽内存(HBM)中存储完整的 n×n 注意力矩阵。对于 n=8192,这个矩阵有 6700 万个元素。问题不在于计算(GPU 算力充足),而在于 IO——在 HBM 和 SRAM(on-chip)之间反复搬运数据的带宽成为瓶颈。

7.5.2 三个核心创新

IO-Aware

不再以"减少浮点运算"为优化目标,而是以"减少 HBM 读写次数"为目标。用额外的计算(重算而不是存储中间结果)来换取更少的内存访问。

Tiling(分块计算)

把 Q、K、V 分成小块(tile),每次只把一小块加载到 SRAM 中计算。每个 tile 大小选择为刚好能放进 SRAM,避免了 HBM 的反复读写。

Online Softmax

标准 softmax 需要先扫一遍所有分数找最大值、再扫一遍算指数和。Online softmax 可以在一次遍历中完成——通过维护一个 running max 和 running sum,边遍历边更新。

7.5.3 Online Softmax 的核心算法

标准 softmax 需要两轮遍历(first pass 找 max,second pass 算 exp 和 sum)。Online softmax 只需一轮:

python
def online_softmax(scores):
    """一轮遍历完成 softmax(概念演示)"""
    n = len(scores)
    m = -np.inf   # running max
    d = 0.0       # running denominator
    for j in range(n):
        m_new = max(m, scores[j])
        d = d * np.exp(m - m_new) + np.exp(scores[j] - m_new)
        m = m_new
    # 最终 softmax(scores[j]) = exp(scores[j] - m) / d
    return np.array([np.exp(s - m) / d for s in scores])

scores = np.array([2.0, 1.0, 0.5, 3.0])
print("标准 softmax:", softmax(scores))
print("Online softmax:", online_softmax(scores))

测试启发。

FlashAttention 是一个纯工程优化——它的数学结果和标准 attention 完全一致(bit-exact 或接近 bit-exact),不改变模型行为。如果你在测试中发现开启/关闭 FlashAttention 后模型输出不同,差异应该只来自浮点精度(FP16 vs FP32),而不是逻辑差异。如果差异很大,那可能是实现 bug。

08. 完整最小 Transformer Decoder(纯 NumPy)

以下是一个约 100 行的最小 decoder-only Transformer,包含 embedding、RMSNorm、causal self-attention、SwiGLU FFN、残差连接。它可以完整地做一次 forward pass(不含训练逻辑)。

python
import numpy as np

# 基础组件 ──────────────────────────────────────────

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

def silu(x):
    return x * (1 / (1 + np.exp(-x)))

def rms_norm(x, gamma, eps=1e-6):
    rms = np.sqrt(np.mean(x ** 2, axis=-1, keepdims=True) + eps)
    return gamma * (x / rms)

# Attention ─────────────────────────────────────────

def causal_attention(x, wq, wk, wv, wo, n_heads):
    n, d = x.shape
    dk = d // n_heads
    Q, K, V = x @ wq, x @ wk, x @ wv

    # reshape to (n_heads, n, dk)
    Q = Q.reshape(n, n_heads, dk).transpose(1, 0, 2)
    K = K.reshape(n, n_heads, dk).transpose(1, 0, 2)
    V = V.reshape(n, n_heads, dk).transpose(1, 0, 2)

    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(dk)
    mask = np.triu(np.ones((n, n), dtype=bool), k=1)
    scores = np.where(mask, -1e9, scores)
    attn = softmax(scores, axis=-1)
    out = (attn @ V).transpose(1, 0, 2).reshape(n, d)
    return out @ wo

# FFN (SwiGLU) ──────────────────────────────────────

def swiglu_ffn(x, w1, w_gate, w2):
    return (silu(x @ w_gate) * (x @ w1)) @ w2

# Transformer Block ────────────────────────────────

def transformer_block(x, params):
    # Pre-Norm + Attention + Residual
    h = rms_norm(x, params['norm1_g'])
    h = causal_attention(h, params['wq'], params['wk'],
                         params['wv'], params['wo'], params['n_heads'])
    x = x + h

    # Pre-Norm + FFN + Residual
    h = rms_norm(x, params['norm2_g'])
    h = swiglu_ffn(h, params['w1'], params['w_gate'], params['w2'])
    x = x + h
    return x

# 完整 Decoder ─────────────────────────────────────

def mini_decoder_forward(token_ids, model):
    # Token Embedding
    x = model['embed'][token_ids]          # (n, d_model)

    # 简化的位置编码(用正弦编码)
    n, d = x.shape
    pos = np.arange(n)[:, None]
    div = 10000 ** (np.arange(0, d, 2) / d)
    pe = np.zeros((n, d))
    pe[:, 0::2] = np.sin(pos / div)
    pe[:, 1::2] = np.cos(pos / div)
    x = x + pe

    # 逐层处理
    for block_params in model['blocks']:
        x = transformer_block(x, block_params)

    # 最终 RMSNorm + 线性头 → logits
    x = rms_norm(x, model['final_norm_g'])
    logits = x @ model['embed'].T          # 权重共享(weight tying)
    return logits

# 初始化参数并运行 ─────────────────────────────────

np.random.seed(42)
V, d, n_heads, d_ff, n_layers = 256, 64, 4, 128, 2
scale = 0.02

def make_block():
    dk = d
    d_ff_glu = int(d_ff * 2 / 3)
    return {
        'norm1_g': np.ones(d), 'norm2_g': np.ones(d),
        'wq': np.random.randn(d, dk)*scale,
        'wk': np.random.randn(d, dk)*scale,
        'wv': np.random.randn(d, dk)*scale,
        'wo': np.random.randn(dk, d)*scale,
        'w1': np.random.randn(d, d_ff_glu)*scale,
        'w_gate': np.random.randn(d, d_ff_glu)*scale,
        'w2': np.random.randn(d_ff_glu, d)*scale,
        'n_heads': n_heads,
    }

model = {
    'embed': np.random.randn(V, d) * scale,
    'blocks': [make_block() for _ in range(n_layers)],
    'final_norm_g': np.ones(d),
}

token_ids = np.array([10, 42, 7, 128, 200])
logits = mini_decoder_forward(token_ids, model)
print(f"Input tokens: {token_ids}")
print(f"Logits shape: {logits.shape}")    # (5, 256)
print(f"Next-token prediction (argmax): {np.argmax(logits[-1])}")
print(f"Top-5 logits at last position: {np.argsort(logits[-1])[-5:][::-1]}")

代码要点。

  • 权重共享(Weight Tying):输出层的权重直接复用 embedding 矩阵的转置,这在 GPT-2、Llama 等模型中很常见,能减少约 V × d_model 个参数。
  • Pre-Norm 结构:先 RMSNorm 再做 attention/FFN,和 Llama 保持一致。
  • SwiGLU FFN:用门控激活函数而不是简单 ReLU。
  • 这个最小实现虽然参数是随机的(未训练),但其前向传播逻辑和真实 Llama 模型完全一致。

09. HuggingFace LlamaAttention 源码剖析

下面是 HuggingFace transformers 库中 LlamaAttention 类的核心逻辑剖析。理解这段代码,你就理解了工业级 Transformer 的注意力实现。

9.1 投影层定义

# 来自 transformers/models/llama/modeling_llama.py(简化版)
class LlamaAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.hidden_size = config.hidden_size           # 4096
        self.num_heads = config.num_attention_heads      # 32
        self.num_kv_heads = config.num_key_value_heads   # 8(GQA)
        self.head_dim = self.hidden_size // self.num_heads  # 128

        # Q 投影:输出 32 个头 × 128 维 = 4096
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        # K 投影:输出 8 个头 × 128 维 = 1024(GQA 所以头数少)
        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        # V 投影:同 K
        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        # 输出投影
        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)

注意 bias=False——Llama 的所有 Linear 层都不使用偏置项,这是刻意的设计选择,可以减少参数量并配合 RMSNorm 工作。

9.2 前向传播核心逻辑

python
def forward(self, hidden_states, attention_mask=None, position_ids=None,
            past_key_value=None):
    bsz, q_len, _ = hidden_states.size()

    # Step 1: 线性投影
    query_states = self.q_proj(hidden_states)    # (bsz, q_len, 4096)
    key_states = self.k_proj(hidden_states)      # (bsz, q_len, 1024)
    value_states = self.v_proj(hidden_states)    # (bsz, q_len, 1024)

    # Step 2: reshape 成多头形式
    # Q: (bsz, q_len, 32, 128) -> (bsz, 32, q_len, 128)
    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
    # K: (bsz, q_len, 8, 128) -> (bsz, 8, q_len, 128)
    key_states = key_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
    value_states = value_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

    # Step 3: 应用 RoPE 旋转位置编码(只对 Q 和 K 施加)
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

    # Step 4: KV Cache 更新
    if past_key_value is not None:
        key_states = torch.cat([past_key_value[0], key_states], dim=2)
        value_states = torch.cat([past_key_value[1], value_states], dim=2)
    past_key_value = (key_states, value_states)

    # Step 5: GQA — 把 8 个 KV 头扩展到 32 个 Q 头
    # repeat_kv: (bsz, 8, kv_len, 128) -> (bsz, 32, kv_len, 128)
    key_states = repeat_kv(key_states, self.num_heads // self.num_kv_heads)
    value_states = repeat_kv(value_states, self.num_heads // self.num_kv_heads)

    # Step 6: 计算注意力
    attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
    attn_weights = attn_weights + attention_mask    # causal mask
    attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
    attn_output = torch.matmul(attn_weights, value_states)

    # Step 7: 拼接多头 + 输出投影
    attn_output = attn_output.transpose(1, 2).contiguous().reshape(bsz, q_len, self.hidden_size)
    attn_output = self.o_proj(attn_output)

    return attn_output, attn_weights, past_key_value

9.3 源码关键点解读

代码行对应概念为什么这样做
view(...).transpose(1,2)多头拆分把 d_model 维度拆成 (n_heads, head_dim),并把 head 维度提前以便并行计算
apply_rotary_pos_embRoPE对 Q 和 K 施加旋转位置编码,V 不施加(V 不参与相关性计算)
torch.cat([past_kv, new_kv])KV Cache把历史 KV 和当前步骤新计算的 KV 拼接
repeat_kv(..., 4)GQA 扩展8 个 KV 头各复制 4 次变成 32 个,对齐 Q 头数
softmax(..., dtype=float32)精度保护即使模型是 FP16/BF16,softmax 也在 FP32 下计算以避免精度损失
.contiguous().reshape内存布局transpose 后内存不连续,需要 contiguous() 才能 reshape

测试启发。

从源码中可以看到几个测试关注点:(1) softmax 强制使用 FP32——如果某个推理框架跳过了这一步,在极端输入下可能产生 NaN;(2) KV Cache 的拼接操作意味着长序列推理时会产生大量内存分配和拷贝,这是性能瓶颈的来源;(3) GQA 的 repeat_kv 操作在某些硬件上可能有效率问题。

10. 从注意力矩阵看模型行为:可视化与测试启发

10.1 注意力热力图能告诉我们什么

将注意力权重矩阵可视化为热力图,可以直观地看到每个 token "关注"了哪些其他 token。不同层、不同头的注意力模式差异很大:

常见注意力模式

  • 局部注意力:只关注相邻 token(通常在浅层出现)
  • 全局注意力:关注特定的"锚点" token,如句号、[CLS]
  • 对角线模式:每个 token 主要关注自己(identity-like)
  • 垂直条纹:所有 token 都关注某个特定位置——通常是 system prompt 的关键词

可视化代码

python
import matplotlib.pyplot as plt

def plot_attention(weights, tokens):
    fig, ax = plt.subplots(figsize=(8, 6))
    im = ax.imshow(weights, cmap='Blues')
    ax.set_xticks(range(len(tokens)))
    ax.set_yticks(range(len(tokens)))
    ax.set_xticklabels(tokens, rotation=45)
    ax.set_yticklabels(tokens)
    ax.set_xlabel('Key')
    ax.set_ylabel('Query')
    plt.colorbar(im)
    plt.title('Attention Weights')
    plt.tight_layout()
    plt.savefig('attention_vis.png')
    plt.show()

10.2 注意力模式与模型行为的关联

观察到的模式可能含义测试行动
System prompt 对应列权重很低模型"忽略"了系统指令检查是否是长上下文稀释导致,尝试在 prompt 末尾重复关键指令
某个 token 的行权重集中在自己该位置缺乏上下文信息可能是 tokenizer 将关键词切碎了,检查 tokenization 结果
注意力均匀分布(flat)模型对上下文"困惑",没有明确的信息源prompt 可能太模糊,需要增加具体性
最后几个 token 的注意力集中在最近的 token近因效应(recency bias)将关键信息放在 prompt 末尾而非开头

10.3 使用 BertViz 工具

对于基于 HuggingFace 的模型,可以用 bertviz 库快速可视化:

python
from transformers import AutoTokenizer, AutoModel
from bertviz import model_view

tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
model = AutoModel.from_pretrained("bert-base-chinese", output_attentions=True)

inputs = tokenizer("用户想要退款但是不知道流程", return_tensors="pt")
outputs = model(**inputs)
attention = outputs.attentions   # tuple of (batch, heads, seq, seq)

tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
model_view(attention, tokens)    # 在 Jupyter 中显示交互式可视化

测试启发。

注意力可视化不是"调模型"的工具,而是"理解模型行为"的诊断工具。当你遇到一个难以复现的异常输出时,查看注意力矩阵可以帮你区分:是模型"没看到"关键信息(注意力权重低),还是"看到了但理解错了"(注意力权重高但输出错误)。前者可能是 prompt 设计问题,后者可能是模型能力边界问题。

11. 长上下文为什么更容易失忆和偏航

理论上 attention 可以看到整段上下文,但"看得到"不等于"有效利用"。上下文越长,注意力分布越稀、噪声越多、前面指令被后续内容稀释的概率越高。于是你会看到:

  • 多轮对话越长,早期约束越容易失效。
  • RAG 片段一多,模型更容易拼错证据。
  • 相同任务下,短 prompt 更稳定,超长 prompt 更飘。

11.1 "Lost in the Middle" 现象

2023 年 Stanford 的研究(Liu et al.)发现一个重要规律:当相关信息放在长上下文的开头或末尾时,模型利用率最高;放在中间时利用率显著下降。这被称为"Lost in the Middle"效应。其根本原因与注意力分布的形态有关: 在长序列中,softmax 权重趋势:

  • 开头位置:因为所有 token 的 Q 都和第一个位置的 K 有交互机会,累积关注较高
  • 末尾位置:因为位置近(近因效应),在 causal 设置下权重偏高
  • 中间位置:两头不靠,注意力权重被稀释

直觉:想象一排 100 人坐在长桌上,你在最右端(当前 token)。 你自然更容易注意到最左边(开头)和右边刚坐过来的人(末尾), 中间的人很容易被忽略。

11.2 为什么 attention 不是"完美的全局检索"

很多人误解 attention = 完美的信息检索。但 softmax 的归一化性质意味着注意力是一个零和博弈: Σ_j A[i][j] = 1

如果序列长度从 100 增加到 10000,同一个"关键证据" token 分到的权重 从 ~0.1 稀释到 ~0.001。模型需要更好地"学会"识别关键信息才能保持性能。

现象结构层解释测试动作
早期 system 指令失效长上下文中被后续 token 稀释做多轮记忆衰减回归
RAG 命中却未使用检索片段虽在窗口内,但注意力利用不足检查 groundedness 与引用率
一个词改动后全局结果变化自回归生成轨迹被前面分支放大对关键提示词做敏感性测试
中间位置信息丢失Lost in the Middle 效应测试不同位置放置关键信息的效果差异
多文档问答时张冠李戴多个文档的 K/V 在注意力中竞争测试文档顺序对答案归属的影响

12. 从测试视角:Transformer 原理驱动的测试策略总览

12.1 基于 Transformer 结构的测试维度

注意力层测试

  • Prompt 顺序敏感性
  • 关键信息位置效应
  • 长上下文记忆衰减
  • 多文档归属准确性

FFN / 参数测试

  • 知识截止日期验证
  • 事实冲突测试
  • 领域知识边界探测
  • 量化精度损失评估

推理优化测试

  • KV Cache 正确性
  • FlashAttention 精度一致性
  • GQA 质量退化评估
  • 长序列推理性能基准

12.2 可执行的测试清单

  1. 顺序敏感性测试:对 prompt 中的指令、证据段落做排列组合,观察输出变化幅度。如果同一组信息换个顺序结果就变了,说明注意力分配对位置敏感。
  2. 最小改写对照:只改一个关键词(如"退款"→"换货"),对比输出轨迹差异。差异越大,说明这个位置的 token 对后续注意力的影响越大。
  3. 记忆衰减曲线:在第 1、5、10、20 轮对话中重复同一个问题,看答案质量是否逐步下降。下降曲线的形状反映了注意力在长上下文中的稀释速度。
  4. 证据利用率:在 RAG 场景中,给模型 5 段证据,只有 1 段正确,测试模型是否引用了正确证据。变换正确证据的位置(开头/中间/末尾),检查 Lost in the Middle 效应。
  5. 推理一致性:同一输入在开启/关闭 FlashAttention、不同量化级别(FP16/INT8/INT4)下的输出对比。差异应在可接受范围内。
  6. 结构化归因:把"输出错了"拆解为:提示顺序问题(注意力层)、检索证据问题(上下文利用)、知识错误(FFN 参数记忆)、格式错误(decoder 轨迹惯性)。用结构化语言而不是"模型抽风"来描述问题。

一句话总结。

理解 Transformer 不是为了替代算法工程师,而是为了在测试中建立"机制直觉"——当你知道模型内部发生了什么,你就能设计出更精准的测试用例,写出更有洞察力的 bug 报告,做出更合理的质量判断。

13. 练习题

13.1 手算题

  1. 已知某个 query 对 3 个 key 的点积分数分别是 [2, 1, 0],softmax 后谁权重最大?请计算具体数值(保留两位小数)。
  2. 如果第 3 个位置被 causal mask 掉,它还能参与本轮注意力吗?为什么?
  3. 为什么多轮对话里后面的话更容易影响当前输出?
  4. d_k = 256 时,Q·K^T 的点积方差是多少?除以 √d_k 后方差变为多少?
  5. 一个 d_model = 2048h = 16 头的 Multi-Head Attention 层(不含 bias),总参数量是多少?

13.2 代码题

  1. 修改上面的最小 NumPy decoder,添加 GQA 支持(n_kv_heads < n_heads)。
  2. 实现一个简单的 KV Cache,让 decoder 能逐 token 生成而不是每次重算全部。
  3. 写一个函数,接受注意力权重矩阵,输出每个 token 的"注意力熵"——熵越高说明注意力越分散。

13.3 口头复述题

不用看笔记,试着用 5 分钟讲完:token → embedding + position → Q/K/V 投影 → 缩放点积 → causal mask → softmax → 加权聚合 V → 多头拼接 → FFN → 残差 + Norm → logits → 下一个 token。如果你能讲顺,这一章就真正内化了。

13.4 参考答案要点

  • 分数 [2, 1, 0] 经 softmax 后:e^2 / (e^2+e^1+e^0) ≈ 7.39/12.08 ≈ 0.61e^1/12.08 ≈ 0.22e^0/12.08 ≈ 0.17。第一个位置权重最大。
  • 被 causal mask 掉的位置不能参与本轮注意力,因为在 softmax 前其分数被设为 -∞exp(-∞)=0,权重为零。
  • 多轮对话里后面的 token 更容易影响当前输出:(1) 位置更近,RoPE 编码的相对距离更小;(2) 在 causal mask 下,后面的 token 总能看到前面的,形成更强的信息流;(3) 自回归生成中,最近生成的内容直接构成下一步的"上下文"。
  • d_k=256:原始方差 = 256,除以 √256=16 后方差 = 256/256 = 1
  • 参数量:W_Q + W_K + W_V + W_O = 4 × 2048 × 2048 = 16,777,216(约 1677 万参数)。

13.5 自测标准

学完这一页后,你应该能:

  • 手工解释 Q、K、V 各自负责什么,以及它们是如何从输入生成的。
  • 看懂 attention 公式里每一项的作用,并能解释为什么要除以 √d_k
  • 解释为什么提示词顺序和长上下文会改变结果——从注意力分布稀释的角度。
  • 把生成不稳定和自回归轨迹联系起来,而不是停留在"模型随机"。
  • 理解 MHA/GQA/MQA 的区别,并能计算 attention 层的参数量。
  • 说清楚 KV Cache 为什么能加速推理,以及它的显存开销。
  • 能阅读 HuggingFace 的 LlamaAttention 源码并理解每一步。
  • 用注意力可视化来辅助分析模型行为,而不只是盯着最终输出。