注意力机制与掩码注意力详解(附带例子说明)

注意力机制的公式只有一行，但「读懂公式」和「把每个数字算出来」是两回事。这篇文章拿一个四个 token 的句子，把 Q、K、V 的每一个分量、每一次内积、每一次 softmax 都算一遍——小到每个格子里的数都能手算核对——然后在同一组数上接着讲掩码注意力。

算例：一句话，四个 token

后面所有图、所有表，用的都是这一个例子：

真实模型里 token 是子词，一句话会被切成十几个甚至几十个 token。这里刻意用四个单字 token，是为了让每个中间结果都能在一屏里排下，而且能手算核对。结构上它和真实模型完全一样，只是小到可以逐个数字盯着看。

记号与形状

注意 $d_k$ 和 $d_v$ 是两个可以不相等的量：注意力矩阵 $A$ 是 $n \times n$ 的，而输出 $AV$ 的形状是 $n \times d_v$。输出的宽度由 V 决定，与 Q、K 无关。 本例里 $d_k = 4$、$d_v = 3$，就是为了让这件事在图上看得出来。

Q、K、V 从哪来

标准流程是先把每个 token 的 embedding 拼成 $X \in \mathbb{R}^{n \times d_{\text{model}}}$，再用三个可学的投影矩阵投影出三份表示：

$$Q = XW^Q,\qquad K = XW^K,\qquad V = XW^V$$

这三步都是普通的矩阵乘法，不涉及注意力机制本身。本文直接从投影之后开始，把 Q、K、V 当作给定输入——因为注意力真正做的事情，全在它们之后的四步里。

下面这三个矩阵是为了能手算而手工设定的（取值都是 $-1$ 到 $2$ 的小整数），不代表真实模型学到了什么。你要看的是计算过程，不是这几个数的语义。

$$Q=\begin{bmatrix} 1 & 1 & 1 & -1\\ 0 & 1 & 0 & 1\\ 0 & 0 & 1 & -1\\ 1 & 0 & 0 & 2 \end{bmatrix} \qquad K=\begin{bmatrix} 0 & -1 & -1 & -1\\ 0 & 1 & 0 & 2\\ 1 & 0 & 1 & 2\\ 1 & 1 & 1 & -1 \end{bmatrix} \qquad V=\begin{bmatrix} 2 & 0 & 1\\ 0 & 1 & 0\\ 1 & 0 & 2\\ 0 & 2 & 1 \end{bmatrix}$$

行序和 token 顺序一致：第 0 行是「我」，第 3 行是「苹果」。

全貌：四步

先看一眼整个流程，再逐步展开。每一步旁边的形状是这一步的产出：

图 1：注意力的四个步骤。Q 和 K 在前三步里用掉，V 一直等到第 4 步才登场——这是注意力机制最容易被忽略的结构特点：先决定「看谁」，再决定「拿什么」。

第 1 步：打分 $S = QK^\top$

注意力要回答的第一个问题是「第 $i$ 个 token 该看第 $j$ 个 token 多少」。问的方式是拿第 $i$ 行的查询和第 $j$ 行的键做内积：

$$S = QK^\top,\qquad S_{ij} = q_i \cdot k_j = \sum_{t=1}^{d_k} q_{i,t}\,k_{j,t}$$

内积大，表示「我在找的东西」和「你能提供的东西」对上了，这个位置就该被多看几眼。

先手算三个格子

拿「苹果」（$i = 3$）这一行做例子。它的查询向量是 $q_3 = [1, 0, 0, 2]$。

和「吃」（$j = 2$）比，键是 $k_2 = [1, 0, 1, 2]$：

$$S_{3,2} = 1 \times 1 + 0 \times 0 + 0 \times 1 + 2 \times 2 = 5$$

和「爱」（$j = 1$）比，键是 $k_1 = [0, 1, 0, 2]$：

$$S_{3,1} = 1 \times 0 + 0 \times 1 + 0 \times 0 + 2 \times 2 = 4$$

和「我」（$j = 0$）比，键是 $k_0 = [0, -1, -1, -1]$：

$$S_{3,0} = 1 \times 0 + 0 \times (-1) + 0 \times (-1) + 2 \times (-1) = -2$$

第 4 个分量在这三个结果里都是 $2 \times k_{j,3}$——苹果的查询在那一维上权重很高，于是「谁的键在第 4 维也大，谁得分就高」。内积就是这么逐个分量相乘再求和的。

完整的打分矩阵

四个 token 两两打分，得到 $4 \times 4 = 16$ 个数：

图 2：打分矩阵 $S$。颜色越深分数越高；框出的第 3 行是「苹果」的查询，它给出的最高分 5 落在「吃」那一列。

完整的表长这样：

每一行是「一个 token 对所有 token 的评价」。到这里为止，分数还是可正可负、大小没有上限的任意实数——不能直接当权重用。后面两步就是把它变成一组和为 1 的权重。

第 2 步：缩放，除以 $\sqrt{d_k}$

$$S' = \frac{S}{\sqrt{d_k}}$$

本例 $d_k = 4$，所以 $\sqrt{d_k} = 2$，把 $S$ 里每个数除以 2：

这一步只是缩放，不改变每一行里谁大谁小——排序不变，所以「看谁」的结论不变，变的是分布的陡峭程度。为什么非得除以 $\sqrt{d_k}$ 而不是别的数，后面单独一节用实测数据说。

第 3 步：按行 softmax，变成权重

$$\alpha_{ij} = \frac{\exp(S'_{ij})}{\sum_{t=1}^{n} \exp(S'_{it})}$$

注意分母对 $t$ 求和——归一化是沿着每一行做的，不是整个矩阵一起归一化。第 $i$ 行的四个权重加起来是 1，代表「第 $i$ 个 token 把它的注意力 100% 分配给了这 4 个位置」。

手算「苹果」这一行

这一行的缩放后分数是 $[-1,\ 2,\ 2.5,\ -0.5]$。

实现上不会直接对这四个数取指数，而是先减去行内最大值（这里最大值是 $2.5$）再取：

减最大值这一步不改变结果——分子分母同时乘上了 $e^{-2.5}$，会约掉——但它把最大的那个指数变成 $\exp(0) = 1$，其余都是比 1 小的数，从根本上避免了 $\exp$ 溢出。这是 softmax 的标准实现方式，不是这里为了好算才这么写的。

四个权重加起来：$0.0179 + 0.3596 + 0.5929 + 0.0295 = 1$。

完整的权重矩阵

四行都算完，得到注意力矩阵 $A$：

图 3：注意力权重矩阵 $A$。和图 2 是同一组数，只是每一行都过了 softmax——每行的四个数现在都在 $[0,1]$ 之间，且加起来恰好是 1。

几个可以直接读出来的现象：

• 「苹果」把 59.29% 的注意力给了「吃」，而给它自己的只有 2.95%——它更关心谁在吃它，而不是它自己。
• 「爱」把 52.31% 留给了自己，这是自注意力（self-attention）最常见的形态之一。
• 所有行和精确为 1，一行里的权重构成一个概率分布。这就是为什么它叫「软」注意力：不是硬选一个位置，而是按权重在所有位置上取一个加权平均。

第 4 步：加权求和，得到输出

$$o_i = \sum_{j=1}^{n} \alpha_{ij}\, v_j$$

到这里 V 才登场：权重已经定好了「看谁」，这一步决定「从每个位置拿走多少」。

手算「苹果」的输出

「苹果」那一行的权重是 $[0.0179,\ 0.3596,\ 0.5929,\ 0.0295]$。逐个乘上对应的值向量，再加起来：

第 3 个分量 $1.233$ 里，有 $1.186$ 来自「吃」——因为权重最高的就是「吃」（0.5929），而「吃」的值向量在那一维上恰好是 2。权重决定了谁说了算，V 决定了它说了什么。

图 4：把「苹果」的输出向量拆回每一维，看每个 token 贡献了多少。第 1 维和第 3 维几乎全由「吃」构成，第 2 维则主要来自「爱」——不同维度取决于不同的来源 token，这正是多头注意力要利用的结构。

四个 token 的输出拼起来，就得到整个输出矩阵 $O$：

形状是 $4 \times 3$。回到图 1：进来的 Q、K 是 $4 \times 4$，出去的 $O$ 是 $4 \times 3$——输出宽度跟着 V 走。这一步之后，$O$ 会送去后面的前馈网络，或者当作下一层的输入。

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

前面说排序不变、只改变陡峭程度。那「不缩放」到底会怎样？实测一下就清楚了。

$q \cdot k = \sum_{t=1}^{d_k} q_t k_t$ 是 $d_k$ 个乘积的和。如果每个分量都独立、均值 0、方差 1，那么每一项的方差是 1，$d_k$ 项加起来方差就是 $d_k$，标准差是 $\sqrt{d_k}$。抽样验证（每个维度取 20000 个随机向量对）：

对上了。于是维度一大，进 softmax 的分数就散得很开——而 softmax 对大的输入差异极其敏感。

在 $d_k = 512$ 下，取一个查询向量和 20 个键算分数，分别做「不缩放」和「除以 $\sqrt{512}$」的 softmax：

不缩放时最大权重是 $1.0000$、熵是 $0.000$——注意力彻底塌到一个位置上，其他 19 个位置的权重全是 0。熵为 0 意味着这个分布已经没有任何不确定性，softmax 的梯度在这个区域几乎消失，模型学不动。缩放之后最大权重降到 0.27 左右，熵接近均匀分布的 $\ln 20 \approx 2.996$，分布还有形状、还有梯度。

所以 $\sqrt{d_k}$ 不是一个拍脑袋的常数，它正好是让内积方差回到 1 的那个因子。

多头：把上面四步并行做几遍

单个注意力头只能输出一种加权方式。多头注意力把每个头的表示维度切小，并行跑几遍：

$$\text{head}_h = \mathrm{Attention}(QW_h^Q,\ KW_h^K,\ VW_h^V)$$ $$\mathrm{MultiHead}(Q, K, V) = \mathrm{Concat}(\text{head}_1, \dots, \text{head}_H)\,W^O$$

关键在「切」：$H$ 个头把 $d_k$ 维切成 $H$ 份，每个头只看自己那份。本例 $d_k = 4$、取 $H = 2$，每个头就是 2 维：

图 5：多头把 4 维切成两个 2 维的头，各算各的四步，结果拼接回 4 维再过一个输出投影 $W^O$。

每个头都完整地跑一遍前面那四步——打分、缩放、softmax、加权求和——但只在它那 2 维上跑。头与头之间在计算过程中完全不通气，直到最后拼接才合到一起。

多头为什么有用？直观说法是：一个头学会一种关注模式，比如某个头盯相邻位置、另一个头盯动词和宾语的关系。这是经验观察，不是设计时的硬性保证——真正有保证的是，它给了模型 $H$ 组互相独立的「查询 / 键 / 值」投影，而不是逼着一组投影同时干所有事。代价是计算量不变、但需要多存中间结果，且每个头的维度变小。

掩码注意力

前面四步算的是「全都能看」的注意力。但有一大类场景里，有些位置必须看不见。

为什么需要掩码

最典型的是自回归生成（GPT 那一类模型）：训练时把整句话一次性喂进去，让每个位置预测下一个 token。但第 0 个位置要预测的是第 1 个 token——如果它能在打分阶段看见位置 1 的内容，那就等于直接把答案抄给它。模型会学会抄答案，然后在推理时发现没有答案可抄。

所以规则是：第 $i$ 个位置只能看见 $j \le i$ 的位置（含自己）。位置 $i$ 看到自己不算作弊——它要预测的是 $i+1$。

因果掩码怎么加

做法极其简单：在 softmax 之前，给不该看见的位置加上 $-\infty$。

$$M_{ij} = \begin{cases} 0, & j \le i \\ -\infty, & j > i \end{cases} \qquad S'_{\text{masked}} = S' + M$$

图 6：因果掩码 $M$。虚线是分界——左下方（$j \le i$）加 0、分数不变；右上方（$j > i$）加 $-\infty$。注意每一行可见的格子数不同：第 0 行只有 1 个，第 3 行有 4 个。

把它加到缩放后的分数上，再照常 softmax。$\exp(-\infty) = 0$，所以被屏蔽的位置权重精确为 0，既不贡献输出，也不参与归一化——分母里那几项也是 0，行和依然是 1。

手算：加了掩码之后的「爱」

「爱」在位置 $i = 1$，只能看位置 0 和 1。它缩放后的分数本来是 $[-1,\ 1.5,\ 1,\ 0]$，加上掩码后变成 $[-1,\ 1.5,\ -\infty,\ -\infty]$：

对照没有掩码时「爱」那一行（$[0.0429,\ 0.5231,\ 0.3173,\ 0.1167]$）：本来分给「吃」的 31.73% 和分给「苹果」的 11.67% 全部被收回，按比例重新分配给了还看得见的两个位置。「爱」对「我」的注意力从 0.0429 涨到 0.0759，几乎翻倍——不是模型改了主意，是分母变小了。

完整的掩码注意力

四行都用各自的掩码（第 $i$ 行挡掉 $j > i$）算一遍：

对比几个可以直接观察到的点：

• 第 0 行（「我」）的输出是 $[2, 0, 1]$，和 $v_{\text{我}}$ 一模一样。 这不是巧合：位置 0 没有任何前文可看，权重必然是 $[1,0,0,0]$，输出就是它自己的值向量。这是加掩码后最该先做的 sanity check。
• 最后一行（「苹果」）和没加掩码时完全相同。 它本来就是最后一个位置，本来就能看见所有 token，掩码对它没有任何影响。
• 中间两行被改动，因为它们的可用上下文变少了。

另一个 mask：padding

除了因果掩码，还有一类完全不同来源的掩码。一个 batch 里的句子长短不一，短的要在右边补齐 <pad>。这些 pad token 是凑数的，不该被任何位置看见——被 pad 的列整列屏蔽。

它的形状和因果掩码不一样，值得分清：

两者常常同时使用，实现上就是把两个 mask 相加或取逻辑或，再一起加到分数上。

一个真实的坑：整行都被屏蔽会得到 NaN

如果某一行的全部位置都被屏蔽（比如一个序列从头到尾都是 padding），那么这一行加完掩码后全是 $-\infty$。softmax 要先减去行最大值，而这一行的最大值也是 $-\infty$：

$$-\infty - (-\infty) = \text{NaN}$$

结果是 NaN，不是 0。 这个 NaN 会顺着反向传播污染整个 batch 的梯度。这也是为什么很多实现用当前 dtype 能表示的最小有限值（比如 torch.finfo(dtype).min）去填充，而不是直接用 -inf——有限值在整行被屏蔽时至少能得到一个数值结果，不会炸出 NaN。

只要序列里每个样本都至少有一个非 pad token，这行就不会出现。但它确实是 padding 处理里最容易漏掉的一个边界。

复杂度

自注意力的计算量落在 $S = QK^\top$ 上：$n \times n$ 个格子，每个格子要做 $d_k$ 次乘法，所以是 $O(n^2 d)$；$AV$ 那一步是 $O(n^2 d_v)$。合起来 $O(n^2 d)$，$n$ 是序列长度。

平方项来自打分矩阵本身——每个 token 都要和每个 token 比一次，这是注意力机制的定义决定的，不是实现问题。序列翻倍，计算量翻四倍；这也是长上下文昂贵的根本原因。各种稀疏注意力、线性注意力的工作，改写的都是这个平方项。

顺便说一句：掩码是把 $S$ 的上三角置成 $-\infty$，并不节省计算——那些格子照样算出来了，只是算完被丢掉。真正省掉上三角半步计算的，是专门的 fused kernel 干的事。

小结

回到那句话：注意力就是把「我该看谁」和「我该拿什么」拆成了两步。

1. 打分 $S = QK^\top$：每个 token 拿自己的查询和所有 token 的键做内积，得到一张 $n \times n$ 的关系表。
2. 缩放 $S / \sqrt{d_k}$：把内积的方差拉回 1，让 softmax 停在有梯度的区间。
3. 归一化 $\mathrm{softmax}$：按行变成和为 1 的权重——这是「软」的地方，不是硬选一个位置。
4. 加权求和 $AV$：用权重把值向量加权平均，得到输出。输出宽度由 V 决定。

掩码只是在这四步的第 3 步之前插了一刀：把不许看的位置变成 $-\infty$，softmax 之后它们的权重精确为 0。