从 Bit 到 AGI第五章:Transformer Q、K、V 是什么
从 Bit 到 AGI · 第五章:Transformer · Query、Key、Value 三投影的几何含义
作者 David·更新于 ·第 69 篇(共 186 篇) Self-Attention 公式里出现 Q(Query)、K(Key)、V(Value)。它们不是三套独立数据,而是同一输入 X 经三个线性层得到的三种「视角」:问什么、匹配什么、传递什么。理解这三者,就读懂了 attention 权重的来源。
- Query:你手里的检索词——「我想找关于 attention 的书」
- Key:每本书 spine 上的标签——用于与 Query 比对相似度
- Value:书里的实际内容——相似度高就多读几本的内容进摘要
每个 token 同时扮演「读者」和「书」:自己的 Query 去匹配所有人的 Key,再按权重取所有人的 Value。
Q=XWQ,K=XWK,V=XWV
WQ,WK,WV∈Rdmodel×dmodel(单头全维写法;多头时投影到 dk)。
Attention 权重:
Aij=∑j′exp(qi⋅kj′/dk)exp(qi⋅kj/dk)
输出第 i 行:zi=∑jAijvj。
例 1:BERT-base 单头维度
- 输入 X:(n,768)
- WQ:(768,64) → Q:(n,64)(单头)
- 12 头时各有独立的 WhQ,最后 concat 成 (n,768)
三个投影层参数量(单头):3×768×64=147,456;12 头合计 3×768×768=1,769,472≈1.77M(含输出投影 WO 则再 7682)。
例 2:数值小例
n=3,dk=2,某头:
q1⋅k1=2.0,q1⋅k2=5.0,q1⋅k3=1.0
缩放(2≈1.41)后 softmax → 权重约 [0.10,0.85,0.05]。token 1 的输出主要来自 token 2 的 Value。
| token | A1j | 含义 |
|---|
| 1 | 0.06 | 自环,弱 |
| 2 | 0.87 | 强相关 |
| 3 | 0.07 | 弱 |
| 组件 | 参数量(BERT-base MHA) |
|---|
| WQ,WK,WV | 3×7682=1.77M |
| 输出 WO | 7682=0.59M |
| 合计 | ≈2.36M / 层 |
QKV 投影 FLOPs:6nd2(乘加各算一次);n=512 时约 1.5×109 / 层。
Cross-Attention(Encoder-Decoder):Q 来自 decoder 隐藏状态,K,V 来自 encoder 输出。Query 问「译文中这个词该看源句哪一段?」
- K 和 V 必须相同:否。K 只参与打分,V 参与加权求和;分工不同。
- Q·K 越大越好:未经 softmax 的 raw score 可任意大;dk 缩放是为避免 softmax 饱和、梯度消失。
- 三个 W 可以合并成一个:推理优化里常把 QKV 合成一次 GEMM(768→2304),数学等价,省 kernel launch。
import torch
import torch.nn as nn
n, d_model, n_heads = 128, 768, 12
d_k = d_model // n_heads # 64
X = torch.randn(n, d_model)
W_q = nn.Linear(d_model, d_model, bias=False)
W_k = nn.Linear(d_model, d_model, bias=False)
W_v = nn.Linear(d_model, d_model, bias=False)
Q = W_q(X).view(n, n_heads, d_k).transpose(0, 1) # (12, 128, 64)
K = W_k(X).view(n, n_heads, d_k).transpose(0, 1)
V = W_v(X).view(n, n_heads, d_k).transpose(0, 1)
scores = (Q @ K.transpose(-2, -1)) / (d_k ** 0.5) # (12, 128, 128)
Q、K、V 是同一隐藏状态的三次线性投影:Query 发起匹配,Key 被匹配,Value 被汇聚。BERT-base 每层约 2.36M 参数花在 Multi-Head Attention 的 QKV+输出上。下一篇看训练稳定性的关键——残差连接。
下一篇:Residual 连接