2.2 Self-Attention

从 Query-Key-Value 看 Attention

Self-Attention 的核心是把每个 token 同时扮演三个角色：

• Query（Q）——"我想找什么样的信息？"
• Key（K）——"我身上有什么样的信息？"
• Value（V）——"如果你找到我，我实际能贡献什么？"

这和数据库检索（或者字典查找）的类比非常贴切：Query 是查询词，Key 是索引，Value 是真实内容；只不过在 Attention 中，匹配不是硬匹配，而是连续的、可导的"软匹配"。

对一个输入序列 $X \in \mathbb{R}^{n \times d_{\text{model}}}$（$n$ 为序列长度，$d_{\text{model}}$ 为隐维度），用三组投影矩阵得到：

其中 $W^Q, W^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$，$W^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$。通常取 $d_k = d_v = d_{\text{model}} / h$（$h$ 为头数，见下节）。

Scaled Dot-Product Attention 公式

Vaswani 等人给出的核心公式是：

这个公式可以拆成四步：

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

这个缩放因子是 Transformer 原论文的重要细节，也是考试和面试中的高频考点。

直觉推导：假设 $Q$ 和 $K$ 的每一维都是均值为 0、方差为 1 的独立随机变量，那么它们的点积

的方差是 $d_k$（$d_k$ 个方差为 1 的独立项之和）。当 $d_k = 64$ 时，点积的标准差约为 8；当 $d_k = 128$ 时，标准差约为 11.3。

问题：softmax 对输入的数值量级极为敏感。如果输入的量级过大，softmax 的输出就会集中在一个峰值上，接近 one-hot 分布——梯度几乎全为零，无法训练。

当 $\max(z) - \text{second}(z) \gg 1$ 时，梯度对大多数输入都几乎为零。

解决：除以 $\sqrt{d_k}$ 把点积的方差拉回 1，让 softmax 的输入停留在梯度良好的区间。

矩阵视角：三步矩阵乘法

把所有公式连起来，Self-Attention 本质上就是 3 个矩阵乘法 + 1 个 softmax：

这张图非常关键——所有 token 在一次矩阵乘法里同时完成 Attention。对比 RNN 一步一步算隐状态，差距一目了然。

一个数值小例子：3 个 token，$d_k = 2$

为了把公式落地，我们算一个最小的例子。设 $n = 3$, $d_k = d_v = 2$，并给出 $Q, K, V$（假设投影已完成）：

第一步：$QK^\top$

读一下第 3 行：$q_3 = (1, 1)$ 与 $k_1, k_2, k_3$ 分别点积得 1, 1, 1——token 3 对三个 token 的原始相似度都相同。

第二步：缩放 $\sqrt{d_k} = \sqrt{2} \approx 1.414$

第三步：softmax 按行归一化

对第 1 行 $(0.707, 0, 0.707)$：指数为 $(2.028, 1, 2.028)$，和为 $5.056$，归一化后得 $(0.401, 0.198, 0.401)$。

第 2 行 $(0, 0.707, 0)$：指数为 $(1, 2.028, 1)$，和为 $4.028$，归一化得 $(0.248, 0.504, 0.248)$。

第 3 行 $(0.707, 0.707, 0.707)$：三项相等，归一化得 $(1/3, 1/3, 1/3)$。

第四步：$AV$

逐行计算：

• 第 1 行：$0.401 \cdot (10, 0) + 0.198 \cdot (0, 10) + 0.401 \cdot (5, 5) = (6.02, 3.99)$
• 第 2 行：$0.248 \cdot (10, 0) + 0.504 \cdot (0, 10) + 0.248 \cdot (5, 5) = (3.72, 6.28)$
• 第 3 行：$\frac{1}{3}(10 + 0 + 5, 0 + 10 + 5) = (5.0, 5.0)$

读图解义：token 3 的 Query 对所有 Key 同样相似，最终输出是三个 Value 的均值 $(5, 5)$；token 1 的 Query 和 Key 1、Key 3 最匹配，输出就偏向 $V_1$ 和 $V_3$。

极简 PyTorch 实现

把公式直接翻译成 PyTorch（不含 Multi-Head 和 Mask，后续章节再加）：

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: (batch, n, d_k)
    K: (batch, n, d_k)
    V: (batch, n, d_v)
    """
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)  # (b, n, n)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))

    attn = F.softmax(scores, dim=-1)
    out = torch.matmul(attn, V)  # (b, n, d_v)
    return out, attn

本节小结

下一节讨论为什么一个头不够——引入 Multi-Head Attention 和 Transformer Block 的完整构造。