查看全课大纲（第 19 / 40 节）

Self-Attention：QKV 与缩放点积

注意力为什么除以 $\sqrt{d_k}$ ？

小象实战讲义 · 深度学习基础：从神经网络到 Transformer

上一章的 RNN 有个硬约束：第 $t$ 步必须等第 $t-1$ 步算完，序列一长就慢，远距离信息还会在层层传递中衰减。Transformer 换了个思路——让序列里任意两个位置直接对话，一步到位，这就是自注意力（Self-Attention）。本节回答三个问题：Q、K、V 到底在算什么？为什么点积之后要除以 $\sqrt{d_k}$ ，不除会怎样？以及怎么用几十行 PyTorch 把它写出来并与官方实现对上。

先修：矩阵乘法与形状规则、softmax 的定义。后续：下一节把单头扩成多头，再往后讲位置编码与残差/LayerNorm，最后从零拼出一个完整的 Transformer block。

💡 核心导读

注意力 = 加权平均：每个位置的输出，是所有位置「值」的加权和，权重由「查询」与「键」的相似度决定。
Q、K、V 来自同一份输入：三个线性投影把同一个 $X$ 变成三种角色，所以叫「自」注意力。
四步走：点积 → 缩放 → softmax → 加权求和，形状从 (B, L, d_k) 一路走到 (B, L, d_v)。
缩放不是玄学： $d_k$ 越大，点积的方差越大，softmax 越容易饱和；除以 $\sqrt{d_k}$ 把方差拉回 1 附近。

🧭 前置知识

矩阵乘法形状规则：(B, L, d) @ (B, d, S) -> (B, L, S)。
softmax： $\mathrm{softmax}(z)_i = \dfrac{e^{z_i}}{\sum_j e^{z_j}}$ ，沿指定维度归一化，输出非负且和为 1。
方差的可加性：独立同分布变量求和，方差相加（下文解释缩放时用到）。

一句话说明本节的位置：注意力的公式只有一行，但涉及「谁和谁相乘、在哪一维归一化、为什么缩放」三个容易含糊的地方。把这三个地方钉死，后面多头、位置编码、残差堆叠都只是在这行公式外面加壳。

注意力的直觉：一次带权检索

先把公式放一边，用检索的场景理解它。你在搜索框里输入一句话（查询 Query），系统拿它去和一堆文档的标题（键 Key）比相似度，得到每个文档的相关度分数；再按这些分数把文档正文（值 Value）加权混合，得到一段「为你定制」的摘要（输出）。

三个角色因此有了分工：

查询（Query）：当前位置「我想找什么」。
键（Key）：每个位置「我能被什么匹配到」。
值（Value）：每个位置「我实际提供的信息」。

关键点：权重不是学出来的固定参数，而是每次前向现算的，随输入而变。这正是注意力区别于全连接层的地方——它按内容动态地决定「看谁」。

一个补充说明，避免后面混淆：本节只讲自注意力，即 Q、K、V 三者的来源完全相同（都是同一个序列 $X$ ）。Transformer 里还有一种交叉注意力，它的查询来自解码器、键和值来自编码器输出，用来把两段不同的序列对齐。两者的计算公式一字不差，区别只在于「Q 从哪来、K 和 V 从哪来」。

Q、K、V：从同一份输入变出三种角色

Transformer 论文把注意力写成一个映射：把一个查询和一组键值对映射成输出（Vaswani et al., 2017，arXiv:1706.03762）。在自注意力里，查询、键、值都来自同一份输入 $X$ ，只是各自过了一个不同的线性投影：

                 ┌──► X @ W_q  ──►  Q  (B, L, d_k)   我要找什么
X (B, L, d_model)┼──► X @ W_k  ──►  K  (B, S, d_k)   我能被谁匹配
                 └──► X @ W_v  ──►  V  (B, S, d_v)   我提供什么

三个投影矩阵 $W_q, W_k, W_v$ 是可学习参数，形状都是 (d_model, d_k) 或 (d_model, d_v)。三个角色共用输入、不共用参数，所以同一个位置在不同角色下可以表达不同的东西。

缩放点积：四步与形状

论文给的公式是（arXiv:1706.03762，式 (1)）：

$\mathrm{Attention}(Q,K,V) = \mathrm{softmax}\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V$

拆成四步，每步都盯住形状：

步骤	运算	形状	说明
1 打分	$QK^{\top}$	`(B, L, S)`	每个查询对每个键的相似度
2 缩放	除以 $\sqrt{d_k}$	`(B, L, S)`	控制数值范围，理由见下节
3 归一	沿最后一维 softmax	`(B, L, S)`	每一行是一个概率分布，和为 1
4 聚合	权重乘 $V$	`(B, L, d_v)`	每个输出是值向量的加权平均

注意第 3 步的方向：softmax 必须沿「键」这一维（dim=-1）做，因为我们要的是「这个查询在各键上的分配比例」。做在 dim=-2 上就变成了「各查询在同一键上的比例」，语义完全错了（练习里会看到差别）。

下面一个完整的实现，同时打印关键中间量，便于你对照形状：

import torch
import torch.nn.functional as F

torch.manual_seed(0)

batch, seq_len, d_model, d_k = 2, 5, 8, 4
X = torch.randn(batch, seq_len, d_model)   # 一份输入，既当 Q 又当 K、V 的来源
W_q = torch.randn(d_model, d_k)
W_k = torch.randn(d_model, d_k)
W_v = torch.randn(d_model, d_k)

Q, K, V = X @ W_q, X @ W_k, X @ W_v        # (2, 5, 4)
scores = Q @ K.transpose(-2, -1)           # (2, 5, 5)：5 个查询 × 5 个键
weights = F.softmax(scores / d_k ** 0.5, dim=-1)   # 沿键维归一化
out = weights @ V                          # (2, 5, 4)

print("X.shape =", tuple(X.shape))
print("Q.shape =", tuple(Q.shape))
print("scores.shape =", tuple(scores.shape))
print("weights[0] 各行和 =", [round(v, 4) for v in weights.sum(dim=-1)[0].tolist()])
print("weights[0][0] =", [round(v, 4) for v in weights[0, 0].tolist()])
print("out.shape =", tuple(out.shape))
print("out[0][0] =", [round(v, 4) for v in out[0, 0].tolist()])

X.shape = (2, 5, 8)
Q.shape = (2, 5, 4)
scores.shape = (2, 5, 5)
weights[0] 各行和 = [1.0, 1.0, 1.0, 1.0, 1.0]
weights[0][0] = [0.0001, 0.0, 0.9977, 0.0002, 0.0019]
out.shape = (2, 5, 4)
out[0][0] = [-2.2439, -2.8877, -1.4509, -4.722]

读输出时留意两件事：weights[0] 的每一行加起来都是 1（这是 softmax 沿 dim=-1 的直接结果）；out 的形状回到了 (B, L, d_v)，与 V 的最后一维一致——所以注意力层可以原样接在下一层之前，不需要额外的形状修补。

顺便说清 d_k 与 d_v 为什么要分开写。打分只用到 d_k（查询与键的维度必须一致才能点积），而聚合用到的是 d_v（值的维度）。两者在论文里取相同值，但在工程实现中允许不同——输出维度由 d_v 决定，这也是为什么注意力的输出可以接到任意宽度的下一层。

还有一层容易被忽略但很有用的性质：因为权重非负且每行和为 1，每个输出都是所有值向量的凸组合，落在这组值向量张成的凸包内。这解释了两件事——注意力不会把输出「放大」到值域之外；当某个位置几乎独占权重时，输出就退化成一个具体的值向量，信息被「锁定」在那一处。

为什么除以 $\sqrt{d_k}$

论文给出的解释（arXiv:1706.03762 §3.2.1 脚注 4）：假设查询和键的每个分量都是均值 0、方差 1 的独立随机变量，那么它们的点积

$q\cdot k = \sum_{i=1}^{d_k} q_i k_i$

的均值是 0，方差是 $d_k$ 。也就是说， $d_k$ 越大，点积的典型绝对值越大（量级约 $\sqrt{d_k}$ ）。

这会带来一个具体的麻烦：softmax 的输入一旦很大，输出就会变成「一个位置接近 1、其余接近 0」的尖峰分布。此时 softmax 的雅可比矩阵几乎处处接近 0——梯度被压平，参数几乎收不到有效信号。这正是论文所说的「把 softmax 推进梯度极小的区域」。

除以 $\sqrt{d_k}$ 后，点积的方差被拉回 1 附近，softmax 的输入落在正常范围，梯度不再被压平。做一个 64 维的对照实验就能直观看到差别（这里用分布的熵衡量「尖锐程度」，熵越小分布越尖）：

import torch
import torch.nn.functional as F

torch.manual_seed(0)

d_k, n_keys = 64, 6
q = torch.randn(d_k)             # 一个查询
K = torch.randn(n_keys, d_k)     # 6 个键
scores = q @ K.t()               # 未缩放的原始分数

raw = F.softmax(scores, dim=-1)
scaled = F.softmax(scores / d_k ** 0.5, dim=-1)

def entropy(w):
    return round(float(-(w * w.log()).sum()), 4)

print("原始分数 =", [round(v, 4) for v in scores.tolist()])
print("未缩放：最大权重 =", round(float(raw.max()), 4), "熵 =", entropy(raw))
print("缩放后：最大权重 =", round(float(scaled.max()), 4), "熵 =", entropy(scaled))

原始分数 = [-11.4345, 2.2093, 0.3528, -14.6962, -4.5081, -6.3864]
未缩放：最大权重 = 0.8638 熵 = 0.4053
缩放后：最大权重 = 0.3486 熵 = 1.5694

$n=6$ 个键时，最大熵是 $\ln 6 \approx 1.7918$ 。缩放后的分布接近均匀（熵接近上限），未缩放的分布则明显更尖——这就是「除以 $\sqrt{d_k}$ 让注意力不至于过早锁死在某一个位置」的数值含义。

论文的实践取值可以直接对上这个结论： $h=8$ 个头， $d_k = d_v = d_{\text{model}}/h = 64$ （ $d_{\text{model}}=512$ ），且作者指出由于每头维度降低，多头总计算量与全维单头相当（arXiv:1706.03762 §3.2.2）。

为什么是点积，而不是加性注意力

注意力并不只有点积一种算法。Bahdanau 等人提出的加性注意力用一个单隐层前馈网络算兼容性分数：

$\mathrm{score}(q, k) = v^{\top}\tanh(W_1 q + W_2 k)$

论文 §3.2.1 对两者做了对比：理论复杂度相近，但点积注意力在实践中更快、更省内存，因为它可以直接写成矩阵乘法，吃到高度优化的 BLAS 实现；而加性注意力要过一遍非线性层，无法这样摊平。论文同时指出， $d_k$ 较小时两者表现接近， $d_k$ 较大且不做缩放时，点积反而不如加性注意力——正是这一现象促使作者引入 $1/\sqrt{d_k}$ 。

于是 Transformer 的选择可以概括成一句话：用点积换取矩阵乘法的速度，用缩放补回数值稳定性。理解了这一层，你在读实现时看到的那行 scores / d_k ** 0.5 就不再是「照抄的常数」，而是一个有明确动机的设计。

怎么读权重矩阵

weights 的形状是 (B, L, S)，把它想成一张表：行是查询位置，列是键位置，第 $i$ 行第 $j$ 列就是「位置 $i$ 把多少注意力放在位置 $j$ 上」。行内和为 1，所以每行是一个分配方案。

某一行很「平」：该位置的信息来自所有位置的均匀混合。
某一行很「尖」：该位置几乎只看某一个位置，信息高度集中。
整列都很大：该位置被很多查询同时关注，说明它在当前任务里是「信息枢纽」。

需要提醒的是，这张表描述的只是计算路径，不等于模型「解释」。同一组权重在不同层、不同输入下都会变，不能据此断言「模型因为这个原因给出了这个答案」。把它当作观察工具，而不是证据。

与 PyTorch 官方实现对账

手写一遍是为了理解，工程上直接用 F.scaled_dot_product_attention：它把四步封装成一次调用，并在可用时走更高效的内核实现。官方文档（PyTorch 2.14.1，torch.nn.functional.scaled_dot_product_attention）说明它接受 attn_mask、dropout_p、is_causal、scale 等参数，其中 scale 缺省时按 $1/\sqrt{d}$ 计算（ $d$ 为查询最后一维）。

验证两者等价只要一个块：

import torch
import torch.nn.functional as F

torch.manual_seed(1)

Q = torch.randn(2, 5, 4)
K = torch.randn(2, 5, 4)
V = torch.randn(2, 5, 4)

manual = F.softmax(Q @ K.transpose(-2, -1) / 4 ** 0.5, dim=-1) @ V
official = F.scaled_dot_product_attention(Q, K, V)

print("manual.shape =", tuple(manual.shape))
print("official.shape =", tuple(official.shape))
print("最大绝对差 =", round(float((manual - official).abs().max()), 8))
print("allclose =", torch.allclose(manual, official, atol=1e-6))

manual.shape = (2, 5, 4)
official.shape = (2, 5, 4)
最大绝对差 = 1.2e-07
allclose = True

输出里最大绝对差在 $10^{-7}$ 量级（float32 的舍入误差），allclose 为 True——在这个小规模、非因果、无 dropout 的设定下，两者算的是同一个函数。换到大规模张量时数值可能因内核实现略有差异，但这不影响「公式一致」这一结论。

官方文档还写清了一个容易踩坑的约定：F.scaled_dot_product_attention 的布尔掩码中 True 表示该位置参与注意力，这与 nn.MultiheadAttention 的 key_padding_mask（True 表示被屏蔽的填充位）语义正好相反。迁移代码时这一处必须显式取反，否则「屏蔽」会变成「只看填充」。

📝 动手练一练

把上面第一个代码块里的 F.softmax(..., dim=-1) 改成 dim=-2，观察输出的变化，并解释为什么错。

👉 查看参考答案

思路：dim=-1 是键维，归一化后「每个查询在各键上的权重和为 1」，这是加权的正确含义；dim=-2 是查询维，归一化后变成「各查询在同一键上的权重和为 1」，每一行不再和为 1，输出也不再是值的凸组合。

import torch
import torch.nn.functional as F

torch.manual_seed(0)

Q = torch.randn(1, 3, 4)      # 3 个查询
K = torch.randn(1, 5, 4)      # 5 个键
scores = Q @ K.transpose(-2, -1)        # (1, 3, 5)

right = F.softmax(scores, dim=-1)       # 沿键维
wrong = F.softmax(scores, dim=-2)       # 沿查询维（错误）

print("scores.shape =", tuple(scores.shape))
print("dim=-1 时各行和 =", [round(v, 4) for v in right[0].sum(dim=-1).tolist()])
print("dim=-2 时各列和 =", [round(v, 4) for v in wrong[0].sum(dim=-2).tolist()])
print("dim=-2 时第 0 行和 =", round(float(wrong[0, 0].sum()), 4))

scores.shape = (1, 3, 5)
dim=-1 时各行和 = [1.0, 1.0, 1.0]
dim=-2 时各列和 = [1.0, 1.0, 1.0, 1.0, 1.0]
dim=-2 时第 0 行和 = 0.8702

常见错误

错误理解或写法	现象或报错	原因	正确做法
忘记除以 $\sqrt{d_k}$	$d_k$ 较大时训练不收敛，损失长时间不动	点积方差随 $d_k$ 增大，softmax 饱和、梯度被压平	用 `F.scaled_dot_product_attention`，或手动除以 `d_k ** 0.5`
softmax 的 `dim` 写错	行和不为 1，输出不是值的凸组合，训练照样能跑但学不到东西	归一化维度必须是键维（最后一维），不是查询维	统一写 `dim=-1`，并打印行和做一次自查
把 `nn.MultiheadAttention` 的 `key_padding_mask` 直接喂给 `F.scaled_dot_product_attention`	该屏蔽的位置反而被关注	两者布尔语义相反：前者 `True` 表示屏蔽，后者 `True` 表示参与	传参前取反：`~key_padding_mask`
转置转错了轴（如写 `K.transpose(0, 1)`，或误写成 `Q.transpose(-2, -1) @ K`）	形状报错，或形状偶然对得上但结果错	只有对 `K` 的最后两维转置才得到 `(B, d_k, S)`；高维下转错轴不一定报错	固定写 `K.transpose(-2, -1)`，并核对 `scores.shape == (B, L, S)`

本章小结

注意力是一次「按内容动态加权」的检索：查询决定看谁，键提供匹配依据，值提供信息。
自注意力里 Q、K、V 都来自同一份输入 $X$ ，分别过三个不同的线性投影。
计算四步： $QK^{\top}$ → 除以 $\sqrt{d_k}$ → 沿键维 softmax → 乘 $V$ ，输出形状回到 (B, L, d_v)。
缩放的理由：点积方差是 $d_k$ ，不缩放会让 softmax 进入梯度极小的饱和区。
F.scaled_dot_product_attention 与手写公式等价，但布尔掩码语义与 nn.MultiheadAttention 相反，迁移时须取反。
权重矩阵的行是查询、列是键，行内和为 1；输出是值向量的凸组合，形状由 d_v 决定。

📋 行动清单

把第一个代码块的 d_k 从 4 改成 64，先预测 scores.shape 与 weights[0][0] 的长度，再跑一遍核对。
把 dim=-1 改成 dim=-2 跑一次，用打印出的行和确认「错在哪」。
用 F.scaled_dot_product_attention 重写第一个代码块的最后两行，确认输出与手写版本一致。

注意力为什么除以 dk\sqrt{d_k}dk​​？