主题
注意力机制:Q、K、V 到底在做什么
一句话:注意力是一次可微的查表——每个位置发出一个 Query,和所有位置的 Key 比较,softmax 得到权重,再对 Value 加权求和。它把“按时间一步步传”换成“任意两个位置一步直达”,换来并行和长距离依赖,代价是序列长度的平方级计算。
为什么需要
在注意力之前,序列建模主要靠循环网络:第 t 个位置依赖第 t-1 个位置的隐藏状态。训练时必须按时间顺序跑,序列越长墙钟时间越长,GPU 并行度用不上;第 1 个 token 的信息要经过 t-1 次状态更新才能到达第 t 个位置,路径长、梯度容易衰减,远处细节留不住。
卷积网络可以并行,但每个位置只看局部窗口,要扩大感受野就得堆很多层,任意两个位置之间的路径仍然随距离增长。
注意力的做法很直接:让每个位置直接和所有位置建立连接。第 1 个 token 和第 1000 个 token 之间只隔一次注意力运算,路径长度是常数;所有位置的计算可以写成矩阵乘法,训练时完全并行。代价也很直接:n 个位置两两比较,计算量是 O(n²),这就是长上下文昂贵的根源。
核心机制
一、三个角色,同一份输入
Query、Key、Value 不是三种输入,而是同一个 token 向量经过三个线性变换得到的三种角色:Query 是“我想找什么”,Key 是“我能被什么匹配”,Value 是“匹配后我传递什么”。
可以类比查字典:Query 是你要查的词,Key 是字典里的词条索引,Value 是词条解释。匹配用 Query 和 Key 算,取回的是 Value。在自注意力里,三者都来自同一段序列,所以每个位置既在提问,也在被别的位置检索。
二、缩放点积注意力
设输入矩阵 X 的形状是 n × d(n 个 token,每个 d 维),三个可学习矩阵 W_Q、W_K、W_V 的形状都是 d × d_k:
text
Q = X · W_Q # n × d_k
K = X · W_K # n × d_k
V = X · W_V # n × d_v
S = Q · Kᵀ / √d_k # n × n,每个位置对每个位置的分数
A = softmax(S) # 按行归一化,每行和为 1
O = A · V # n × d_v,每个位置的新表示S[i][j] 表示第 i 个位置对第 j 个位置的关注分数。softmax 按行做,所以每个位置分给所有位置的权重之和是 1;O 的第 i 行就是所有位置 Value 的加权平均,权重由第 i 行的 Query 决定。
三、为什么要除以 √d_k
假设 Q 和 K 的每个分量独立、均值 0、方差 1,那么点积 q · k 是 d_k 个独立项之和,方差是 d_k,标准差是 √d_k;d_k 越大,分数越容易落到绝对值很大的区间。softmax 对大的输入会饱和:最大项接近 1,其余接近 0,梯度也接近 0。模型会很难通过梯度调整“应该关注谁”。除以 √d_k 把分数的方差拉回 1,让 softmax 工作在梯度更健康的区间。这不是经验常数,而是从方差推导出来的。
四、掩码:谁能看见谁
原始分数矩阵 S 是 n × n,但并非所有位置都该互相看见:因果掩码把未来位置设成 -∞,解码器只能看当前及之前;填充掩码屏蔽 padding,避免短序列被无意义 token 稀释;编码器自注意力双向可见,不加上三角掩码,这是 BERT 类模型能利用右侧上下文的原因。
掩码必须在 softmax 之前加 -∞,而不是之后把权重乘 0。先乘 0 再归一化,其他位置的权重之和不再是 1,输出尺度会变。
五、多头:一次看多种关系
一组 Q/K/V 只能产生一种加权方式。把 d_model 切成 h 份(每个头 d_k = d_model / h),每份独立做一次注意力,再拼接成 n × d_model 并乘 W_O,就是多头注意力。总计算量和单头全维度注意力接近,但模型可以同时学习多种关系:一个头关注相邻词,另一个头关注指代,另一个头关注句法结构。这种分工是训练中涌现的,不保证每个头都有干净的人类可解释含义。
六、交叉注意力
自注意力的 Q、K、V 来自同一段序列;交叉注意力让 Query 来自解码器(n_dec × d_k),Key 和 Value 来自编码器(n_enc × d_k / d_v),分数矩阵是 n_dec × n_enc。这就是机器翻译里“解码器每一步去看编码器的哪些位置”的机制。只有 Encoder-Decoder 架构需要它;Decoder-only 模型只有自注意力。
图中最关键的是 n × n 的分数矩阵:它的行数和列数都随序列长度增长,所以显存和计算都是平方级。O = AV 则说明输出只是 Value 的加权平均,注意力本身不产生新内容,它决定“取哪些位置的信息、各取多少”。
七、复杂度:长上下文为什么贵
| 操作 | 时间 | 显存 |
|---|---|---|
QKᵀ | O(n² · d_k) | O(n²) 分数矩阵 |
softmax | O(n²) | O(n²) |
AV | O(n² · d_v) | O(n²) 或按块 |
序列长度翻倍,注意力计算量约翻四倍。FlashAttention 这类算法通过分块和重计算,把显存降到线性、减少 HBM 读写,但数学上仍然是平方级计算;它优化的是 IO,不是复杂度。真正要处理超长序列,还得靠稀疏注意力、线性注意力或检索式外挂,而这些都会在效果上做取舍。
最小上手
用 numpy 实现带因果掩码的缩放点积注意力,验证三件事:每行权重和为 1、第一个位置只能看见自己、输出形状不变。
python
import numpy as np
rng = np.random.default_rng(0)
n, d = 4, 8
X = rng.normal(size=(n, d))
Wq, Wk, Wv = (rng.normal(size=(d, d)) / np.sqrt(d) for _ in range(3))
Q, K, V = X @ Wq, X @ Wk, X @ Wv
scores = Q @ K.T / np.sqrt(d)
mask = np.triu(np.ones((n, n), dtype=bool), k=1)
scores = np.where(mask, -np.inf, scores)
A = np.exp(scores - scores.max(axis=1, keepdims=True))
A = A / A.sum(axis=1, keepdims=True)
out = A @ V
print(A.round(2))
print(A.sum(axis=1))
print(out.shape)输出中上三角全为 0,每行和为 1,out.shape 仍是 (4, 8):
text
[[1. 0. 0. 0. ]
[1. 0. 0. 0. ]
[0.31 0.34 0.35 0. ]
[0.22 0.28 0.21 0.29]]
[1. 1. 1. 1.]
(4, 8)A 的第 0 行只有第一个位置有权重,因为因果掩码不允许它看见未来;第 2 行只能看见前三个位置,权重在它们之间分配。把掩码去掉,上三角就会变成非零值,这就是双向编码器和因果解码器的区别。
生产动作
- 按任务定上下文,不按模型上限定。注意力是平方级,长度翻倍的代价接近四倍;能在检索层解决的,不要全塞进窗口。
- softmax 要数值稳定。实现时先减去每行最大值,再取指数;半精度下直接
exp容易溢出成NaN。 - 区分“模型看到了”和“模型用上了”。注意力权重高只说明加权多,不证明这个位置是因果原因;不要把注意力图当解释。
- 批处理时检查 padding 掩码。padding 位置必须被 mask 掉,否则短序列的表示会被无意义的填充 token 稀释。
- 长上下文优先用现成优化。FlashAttention、分块 prefill、前缀缓存先上,再考虑改模型结构;优化显存不等于降低计算量。
边界
- 注意力本身没有顺序感。把 token 打乱,注意力集合不变;顺序来自位置编码,见 04-位置编码与长上下文。
- 平方级成本不会消失。FlashAttention 省的是显存读写,不是
n²次比较;窗口开到多长,算力就要付多少。 - 注意力是稠密的。每个位置对所有位置都分配一点权重,长上下文里远处信息容易被稀释,出现“中间丢失”。
- 注意力不等于推理。它负责信息路由,事实知识主要存在 FFN 的参数里,见 03-Transformer Block。
- 交叉注意力只属于 Encoder-Decoder。Decoder-only 模型没有独立的编码器,所谓“看上下文”仍然是自注意力。
讲清楚
注意力把每个 token 映射成 Query、Key、Value 三个角色:Query 和所有 Key 做点积得到分数,除以
√d_k防止 softmax 饱和,加掩码决定谁能看见谁,softmax 归一化后对 Value 加权求和。它让任意两个位置一步直达,训练可以并行,代价是n²的分数矩阵,长上下文的成本就来自这里。多头是在多个子空间里并行做同一件事,交叉注意力则让解码器去查编码器。实际使用时要记住:注意力权重不是解释,FlashAttention 不改变平方级复杂度,顺序信息完全依赖位置编码。
适用:所有 Transformer 变体;掩码与头数细节以具体模型为准 | 初稿 2026-10