注意力机制

“小林把红杯子放到桌上，过了一会儿，又拿起了它。”

读到“它”时，我们会回头找前面的对象。对于模型，问题也很具体：当前这个位置，怎样从已有位置取出有用信息？只把前文压进一个固定大小的循环状态，有时不容易保留需要的细节。

注意力提供了另一种做法：让当前位置给其他位置打分，再按这些分数混合信息。这不是给计算机装上人的注意力，而是一组可以训练、可以逐项算清的向量运算。

1. 先把“怎样匹配”与“取回什么”分开

设我们有三个位置，每个位置提供两类向量：

Key，键：用来与查询比较，决定这个位置得到多少权重。
Value，值：真正被取出并混合的内容。

当前位置另有一个 Query，查询，表达当前计算要匹配什么。用 $\boldsymbol q$ 与各个 $\boldsymbol k_j$ 计算匹配分数，得到权重 $a_j$ 后，输出就是：

$\boldsymbol o=\sum_j a_j\boldsymbol v_j$

在这里， $a_j\ge0$ 且 $\sum_j a_j=1$ ，所以单个注意力头的输出是这些值向量的加权平均。它没有一定“选中一个词”，也不会直接输出一个词元编号。

查询、键和值通常由输入表示经过不同的可学习线性变换得到。一个位置会同时有自己的查询、键和值；这三种角色不是三组预先写好的词义标签。

例如，在自注意力中，把各位置的行向量堆成 $X$ ，可以写成：

$\begin{aligned} Q&=XW_Q\\ K&=XW_K\\ V&=XW_V \end{aligned}$

它们都来自同一序列，却可以经过不同参数变成不同表示。这里暂时省略偏置，后面的代码也采用这个约定。

2. 用两维向量算一遍

先不训练参数，直接指定一组小向量，观察计算：

$\begin{aligned} \boldsymbol q&=(1,0)\\[4pt] K&=\begin{bmatrix}1&0\\0&1\\1&1\end{bmatrix},\quad V=\begin{bmatrix}10&0\\0&10\\10&10\end{bmatrix} \end{aligned}$

查询与三把键的点积分别为 1、0、1。若键有 $d_k$ 维，缩放点积注意力会先除以 $\sqrt{d_k}$ 。这里 $d_k=2$ ，分数为：

$\boldsymbol s=\frac{(1,0,1)}{\sqrt2} \approx(0.7071,0,0.7071)$

接着用 Softmax 把分数变为权重：

$a_j=\frac{\exp(s_j)}{\sum_r\exp(s_r)}$

算出来约为 $(0.4011,0.1978,0.4011)$ 。第一个和第三个位置得到相同权重，中间位置较少，但并没有被完全忽略。

最后按这些权重混合值向量：

$\begin{aligned} \boldsymbol o &\approx0.4011(10,0)+0.1978(0,10)\\ &\quad+0.4011(10,10)\\ &\approx(8.0222,5.9889) \end{aligned}$

最后一行用未四舍五入的权重计算。这个结果是一个新向量，后面的网络还会继续处理它。这里所有向量都是教学设定，不能把权重大小解释成模型已经理解了某个句子。

3. 为什么除以维数的平方根

点积把许多乘积相加。维数变大时，分数的波动也可能变大；一旦 Softmax 输入差得很远，概率就会非常尖锐，部分梯度也会变小。

可以用一个简化假设理解缩放：若查询和键的各分量相互独立、均值为零、方差为 1，那么每个乘积的方差为 1， $d_k$ 个乘积之和的方差就是 $d_k$ 。除以 $\sqrt{d_k}$ ，方差便回到 1。

这不是在声称训练中的向量始终满足这些假设，而是解释一个有用的尺度选择。它也不是把点积变成余弦相似度；余弦相似度需要除以两个向量各自的长度。

4. 把所有查询一起算

一个序列里的每个位置都可以发出查询。把它们堆在一起，就得到矩阵形式：

$O=\operatorname{softmax} \left(\frac{QK^\mathsf T}{\sqrt{d_k}}\right)V$

Softmax 沿每一行计算：固定一个查询，在它能访问的键之间分配权重。

对象	形状	每个维度在表示什么
$Q$	$T_q\times d_k$	查询位置数、查询特征数
$K$	$T_k\times d_k$	键位置数、键特征数
$V$	$T_k\times d_v$	同一批位置提供的值特征
$QK^\mathsf T$	$T_q\times T_k$	每个查询对每把键的分数
$O$	$T_q\times d_v$	每个查询得到的新表示

查询和键的特征维度要一致，键和值的位置数要一致；查询位置数却不一定等于键位置数。

如果查询、键和值来自同一序列，叫自注意力。如果查询来自一条序列，键和值来自另一条序列，叫交叉注意力。例如翻译时，目标句正在生成的表示可以查询源句的表示；源句和目标句不必一样长。

5. 不能让模型偷看答案

训练下一词元预测时，我们常把整条真实序列一起送入网络。如果每个位置随意访问所有输入，就可能直接看见自己要预测的那个词元。

假设输入是 BOS 我爱茶，对应目标为 我爱茶 EOS。输入位置“我”要预测“爱”，所以它可以访问 BOS 和“我”，不能访问后面的“爱”和“茶”。

按输入位置从 0 编号，允许访问的条件就是 $j\le i$ 。这叫因果掩码。可以把它写成附加到分数上的矩阵：

$M_{ij}= \begin{cases} 0,&j\le i\\ -\infty,&j>i \end{cases}$

于是：

$A=\operatorname{softmax} \left(\frac{QK^\mathsf T}{\sqrt{d_k}}+M\right)$

因为 $\exp(-\infty)=0$ ，未来位置得到零权重，剩余位置的权重重新归一化。

相同查询与键在无掩码和因果掩码下的注意力权重，因果矩阵右上角为零。

图 1：行表示查询位置，列表示被访问的位置。数字来自本章三位置实验；每一行加起来为 1，因四舍五入可能有微小差异。

把未来分数乘零不等于屏蔽。 分数为 0 时，指数仍是 1，照样会分到概率。另一种错误是 Softmax 后才把未来权重清零，却忘了重新归一化。

补齐批次长度时，还可能需要屏蔽 PAD 对应的键。填充掩码和因果掩码解决不同问题，可以组合使用。若某个查询整行都被屏蔽，朴素 Softmax 会遇到未定义的归一化；本章实现选择直接报错。实际项目要先约定如何处理这种查询，而不是等 NaN 出现后再猜。

不同库的布尔掩码含义也可能相反。本章 allowed=True 表示允许访问；PyTorch 的缩放点积注意力接口也采用这一含义，但不要据此推断所有接口都相同。

6. 多个头，分别学习不同的混合方式

只用一套查询、键和值，会把信息混合成一个结果。多头注意力用多组投影分别计算，再把结果拼起来：

$\begin{aligned} O_r&=\operatorname{Attention}(XW_{Q,r},XW_{K,r},XW_{V,r})\\ O&=\operatorname{Concat}(O_1,\ldots,O_h)W_O \end{aligned}$

$h$ 是头数， $r$ 是其中一个头。每个头都可以访问允许范围内的全部位置；不是把前几个词交给第一个头、后几个词交给第二个头。

如果模型宽度 $D=8$ ，使用两个头，常见做法是每头使用 4 维查询、键和值。对 $B=2,T=3$ 的批次，形状可以这样追踪：

输入表示                    (2, 3, 8)
分别投影得到 Q、K、V         各为 (2, 3, 8)
拆成两个头并调整维度顺序      各为 (2, 2, 3, 4)
每个头的分数和权重           (2, 2, 3, 3)
混合值向量                  (2, 2, 3, 4)
把头拼回去                  (2, 3, 8)
输出投影                    (2, 3, 8)

两个头拥有不同的投影参数，因此可以学出不同关系，但不能预先保证“某个头专门负责语法，另一个专门负责事实”。多头输出还会经过投影、残差和后续层，单看一张权重图也不足以完整解释最终预测。

朴素实现显式存下 $B\times h\times T\times T$ 的权重。序列长度翻倍，这部分元素数会变成四倍；后面的系统篇会讨论怎样改进内存和计算，而不在这里把所有实现都等同于这张完整矩阵。

7. 用代码检查信息究竟流向哪里

环境沿用第十四章的 Python 3.12 与固定版 PyTorch。这里没有训练，只做确定的矩阵运算。可以下载脚本并运行：

python src/examples/attention.py

"""第 19 章：亲手计算缩放点积注意力与因果掩码；PyTorch CPU。"""

from math import sqrt
import torch


def scaled_attention(query, key, value, allowed=None):
    scores = query @ key.transpose(-2, -1) / sqrt(query.shape[-1])
    if allowed is not None:
        if not allowed.any(dim=-1).all():
            raise ValueError("每个查询至少需要一个允许访问的位置")
        scores = scores.masked_fill(~allowed, float("-inf"))
    weights = scores.softmax(dim=-1)
    return weights @ value, weights


def example():
    query = torch.tensor([[1, 0], [0, 1], [1, 1]], dtype=torch.float64)[None, None]
    key = query.clone()
    value = torch.tensor([[10, 0], [0, 10], [10, 10]], dtype=torch.float64)[None, None]
    allowed = torch.ones(3, 3, dtype=torch.bool).tril()
    return query, key, value, allowed


def rows(tensor):
    return [[round(float(x), 4) for x in row] for row in tensor[0, 0]]


def main():
    torch.set_num_threads(1)
    query, key, value, allowed = example()
    full, full_weights = scaled_attention(query, key, value)
    causal, causal_weights = scaled_attention(query, key, value, allowed)
    print(f"Q/K/V 形状：{tuple(query.shape)}")
    print(f"未加掩码的权重：{rows(full_weights)}")
    print(f"第一行加权结果：{rows(full)[0]}")
    print(f"因果权重：{rows(causal_weights)}")
    print(f"因果输出：{rows(causal)}")
    changed_key, changed_value = key.clone(), value.clone()
    changed_key[:, :, 2] += 100
    changed_value[:, :, 2] -= 100
    changed, _ = scaled_attention(query, changed_key, changed_value, allowed)
    print(f"改变第三个位置后，前两个输出不变："
          f"{torch.equal(causal[:, :, :2], changed[:, :, :2])}")
    # 乘零不是屏蔽：分数 0 经指数函数后仍然贡献 1。
    scores = query @ key.transpose(-2, -1) / sqrt(2)
    wrong = (scores*allowed).softmax(-1)
    print(f"错误地将未来分数乘零，第一行未来权重合计：{wrong[0, 0, 0, 1:].sum():.4f}")


if __name__ == "__main__":
    main()

输出为：

Q/K/V 形状：(1, 1, 3, 2)
未加掩码的权重：[[0.4011, 0.1978, 0.4011], [0.1978, 0.4011, 0.4011], [0.2483, 0.2483, 0.5035]]
第一行加权结果：[8.0222, 5.9889]
因果权重：[[1.0, 0.0, 0.0], [0.3302, 0.6698, 0.0], [0.2483, 0.2483, 0.5035]]
因果输出：[[10.0, 0.0], [3.3024, 6.6976], [7.5174, 7.5174]]
改变第三个位置后，前两个输出不变：True
错误地将未来分数乘零，第一行未来权重合计：0.4965

最前面的两个维度分别是批次和头数，这里都取 1。代码随后故意改动第三个位置的键和值，前两个输出仍然完全相同：它们没有访问未来的通路。

这个检查比只看一个“下三角矩阵”更进一步。真实模型里，还要确认其他计算没有跨位置泄露信息；下一章会对堆叠后的完整网络做同样的检查。

8. 自测与小实验

注意力权重是和键比较得到的，最终混合的又是什么？
查询有 5 个、键有 7 个，每个值向量 3 维，输出是什么形状？
为什么分数为零的位置仍可能得到非零权重？
因果注意力能否访问输入的当前位置？为什么这不等于偷看下一词元？
两个头是否分别处理序列的前半段和后半段？
不改变查询和键，只把所有值向量乘 2，权重与输出怎样变化？
改变第三个查询，但保持前两个查询以及键、值不变，前两个输出会变化吗？
假设忽略所有其他内存，批次和头数不变，序列从 128 增至 512，朴素权重矩阵的元素数增加几倍？
将一行 allowed 全部设成 False，本章代码为什么拒绝计算？
去掉因果掩码，再运行“修改未来位置”的实验，解释结果。

参考答案与提示

混合值向量。键决定权重，值提供被混合的信息。
$5\times3$ ；权重矩阵为 $5\times7$ 。
因为 $\exp(0)=1$ ，它仍参与分母与归一化。
可以。输入与目标已错开一位，当前位置的输入属于已知前文。
不是。每个头都在允许的位置集合上分配权重，只是投影不同。
权重不变，输出乘 2；若后面还有非线性层，其最终变化需另行分析。
单层这一步计算中不会；各查询行独立计算自己的权重。
长度变为 4 倍，元素数变为 16 倍。
没有可用位置，朴素 Softmax 无法产生有效的归一化权重。
前两个查询也能访问第三个位置，输出一般会改变。这正是下一词元训练中需要禁止的通路。

下一章将这些运算组装成完整的 Transformer。注意力负责位置之间的信息混合，网络还需要位置、非线性变换和稳定的信息传递路径。

继续阅读：Attention Is All You Need介绍了 Transformer 中的缩放点积与多头注意力。本章图表由教学向量重新计算，没有使用论文中的实验权重。

Keyboard shortcuts

从机器学习到大模型：原理、实践与前沿