企业官网建设流程全解析

transformer你不学？05｜自注意力机制：QKV 的魔法

作 者：吴佳浩（Alben）

微某信公某众号：全栈架构师笔记

导读
论文标题里那句 “Attention Is All You Need”，说的就是本篇要手搓的东西。Q 查询、K 键、V 值——三个矩阵、一次矩阵乘法、一次 softmax，模型就学会了"读到一个词时，该看句子里哪些其他词"。
这是整个 Transformer 的心脏，本篇用 20 行代码把它跑透。

本篇你会学到什么

• 🔸 Q/K/V 各自扮演什么角色（用检索系统的类比）；
• 🔸 缩放因子√d_k为什么必须存在；
• 🔸 注意力矩阵的 shape 变换与行归一化性质。

一、为什么要学这个

前三篇解决的是"把词变成带位置的向量"，但向量之间还是孤立的。"我喜欢这个产品的设计"里，"这个"指什么？人类靠上下文指代消解，模型靠 Attention：每个位置的向量去"询问"所有其他位置，按相关性加权汇总信息。

一次 Attention 之后，每个位置的输出都融合了全句上下文——这就是"这个词在句子里是什么意思"的数学表达。

二、核心理论：用数据库查询理解 QKV

把每个 token 向量经过三个线性变换，得到三种角色：

检索过程：Q @ K^T算"每个查询和每个标签的匹配度" → softmax 变成概率 → 按概率加权取 Value。

缩放因子 √d_k：向量维度越大，点积数值越大，softmax 会饱和到 one-hot（梯度消失）。除以√d_k把方差拉回 1，保持 softmax "软"度适中——论文 3.2.1 节的原话是 “for larger values of d_k”。

三、流程图

四、真实代码

# 手搓 Scaled Dot-Product Attention（论文公式逐行实现）importtorchimporttorch.nnasnnimporttorch.nn.functionalasFimportmath torch.manual_seed(42)B,T,C=2,4,16# batch=2, 序列长=4, 特征维=16x=torch.randn(B,T,C)# 输入：第 3 篇+第 4 篇 产出的带位置向量# 三个线性变换：同一个 x 变出三种角色Wq=nn.Linear(C,C)Wk=nn.Linear(C,C)Wv=nn.Linear(C,C)Q,K,V=Wq(x),Wk(x),Wv(x)# 各 (B, T, C)# 第一步：相似度矩阵——每个位置和所有位置的两两点积att=Q @ K.transpose(-2,-1)# (B, T, T)print(f"原始相似度矩阵 shape:{att.shape}")print(f"缩放前标准差:{att.std():.2f}")# 随 C 增大会变大# 第二步：除以 √d_k 缩放att_scaled=att/math.sqrt(C)print(f"缩放后标准差:{att_scaled.std():.2f}")# 第三步：softmax 行归一化——相似度变权重att_weights=F.softmax(att_scaled,dim=-1)# 每行和为 1print(f"权重行和:{att_weights[0,0].sum().item():.4f}")# 第四步：按权重聚合 Valueout=att_weights @ V# (B, T, C)print(f"输出 shape:{out.shape}")

输出结果怎么看（真实运行输出）：

原始相似度矩阵 shape: torch.Size([2, 4, 4]) 缩放前标准差: 1.19 缩放后标准差: 0.30 权重行和: 1.0000 输出 shape: torch.Size([2, 4, 16])

• 🔸 缩放前标准差 1.19 → 缩放后 0.30：恰好缩小到 1/√d_k（d_k=16，√16=4，1.19/4≈0.30）——缩放项给点积方差装了"闸门"，d_k 越大不缩放越容易让 softmax 饱和成 one-hot，这就是"为什么开根号"的数值证据；
• 🔸 每行权重和精确为 1：softmax 保证每个位置的"注意力总量"守恒，只是分配比例不同；
• 🔸 输出 shape 与输入相同(B, T, C)：Attention 不改变形状，这正是它能被堆叠成深层网络的前提。

五、运行结果应该怎么看

把att_weights打印出来（print(att_weights[0].round(decimals=2))），你会看到每行一个概率分布：

• 🔸 对角线附近权重高：token 天然和自己最相关；
• 🔸 训练后的模型里会浮现模式：代词行在先行词列有高权重（指代消解）、动词行关注宾语列（句法结构）；
• 🔸 这张 T×T 图就是所谓"注意力热力图"，可视化 GPT 在"看"哪。

六、常见错误与排查

• 🔸忘除 √d_k：小模型看不出差别，d_model 上到 512+ 后 loss 降不动——缩放缺失导致 softmax 饱和；
• 🔸softmax 维度搞错：必须dim=-1（对每行归一化）。写成dim=0会对列归一化，物理意义全错；
• 🔸transpose 用错轴：K.transpose(-2, -1)只换最后两维，别用view硬捏——会打乱 token 对应关系。

七、练习题

1. 把C从 16 改到 512，观察缩放前后标准差的变化，验证"维度越大缩放越必要"；
2. 故意把softmax(dim=-1)改成dim=0，打印行和与列和，说明为什么错了；
3. 构造一个"指代消解"实验：手工设计 V 使"它"的位置输出偏向"猫"的向量，验证注意力权重可控。

一句话总结：自注意力 = 查询匹配键 → softmax 变权重 → 加权聚合值，√d_k 缩放保 softmax 不饱和——20 行代码，模型的"上下文理解"能力全部在这。

下一篇：transformer你不学？06｜多头注意力：八个视角看世界。