注意力机制：让模型看见上下文

Bigram 模型的局限：只有一个字的记忆

第三篇那个 Bigram 模型，局限集中在一点：预测下一个字时，它只查看前一个字。对于"今天天气真"这五个字，它只取最后的"真"去查表，前面四个字全部被舍弃。

这就是它的能力上限。语言中的信息是前后连贯的——要恰当地接续"今天天气真"之后的字，需要知道前文正在谈论天气;要妥善地写好一个长句的结尾，需要记得开头说了什么。一个只有一个字记忆的模型，注定无法写出通顺的文字。

因此，接下来几篇的任务十分明确：让模型能够看见整段上下文。本篇要讲解的注意力机制AttentionAttentionAttention lets every position in the sequence look back over all earlier positions and decide, on its own, how much to draw from each one. Rather than treating prior tokens equally, it computes a data-dependent weight for each pair of positions. This is exactly what frees the model from the single-token memory of a bigram. ，正是解决这一问题的核心部件，也是 GPT 之中 T（Transformer）的灵魂所在。我们不会直接给出公式，而是从一个最初步的设想出发，逐步把它推导出来。

先升级表示：用向量表示每个字

在动手之前，先做一项关于"表示方式"的升级。

在 Bigram 模型中，一个字就是一个编号。但一个孤立的编号所能携带的信息十分有限。从本篇起，我们改用一串数字来表示每个字，也就是一个向量向量向量在这里指用一串固定长度的数字来表示一个字。单个编号只能区分这是哪个字，而一串数字构成的向量还能在不同维度上分别编码字的特征，例如偏书面还是偏口语、词性、语义倾向。训练的过程也包含不断调整这些数字，使含义相近的字在向量空间中彼此靠近。 。例如用 32 个数字来代表一个字，这一串数字便有足够的空间来编码诸如"这个字偏书面还是偏口语、属于名词还是动词"之类的特征。这一串数字的长度，我们称之为 n_embd（embedding 维度）。

具体做法仍然是查表，只是表的形状有所变化：建立一张 vocab_size × n_embd 的表，每个字的编号对应表中的一行，那一行的 n_embd 个数值就是这个字的向量。这张表本身也是参数，会在训练过程中一同被调整。

import torch
import torch.nn as nn

n_embd = 32
token_embedding_table = nn.Embedding(vocab_size, n_embd)

# 假设 idx 是一批文字的编号，形状为 (B, T)
# 查表之后，每个字便从一个编号变成了一个 n_embd 维的向量
x = token_embedding_table(idx)      # 形状为 (B, T, n_embd)

从现在起，模型内部流动的就是这种 (B, T, n_embd) 形式的数据：B 个样本，每个样本 T 个字，每个字是一个 n_embd 维的向量。注意力机制所要处理的，正是它。

一个初步的做法：对前文取平均

现在来让每个字"看见前文"。一个最初步、但方向正确的做法是：让每个位置的向量，等于它自身与它之前所有字向量的平均。

例如，第 5 个字的新向量，等于第 1 个到第 5 个字向量的平均。这样一来，第 5 个位置的向量中就掺入了前文的信息，不再只代表它自身。

直接用循环来求平均效率较低。这里可以借助一个简洁的技巧：求平均的操作，能够通过一次矩阵乘法来完成。构造一个下三角矩阵，对它的每一行做归一化，再让它与字向量相乘，得到的就是逐位置的累计平均。

import torch
from torch.nn import functional as F

T = 8
# 下三角的全 1 矩阵：第 i 行只有前 i+1 个位置为 1
tril = torch.tril(torch.ones(T, T))
# 对每一行做归一化，使一行的数值之和等于 1
weights = tril / tril.sum(dim=1, keepdim=True)
print(weights)

打印出来可以看到，weights 的第 1 行是 [1, 0, 0, ...]，第 2 行是 [0.5, 0.5, 0, ...]，第 3 行是 [0.33, 0.33, 0.33, 0, ...]。用它去乘字向量 x，第 i 个位置便自动得到了前 i 个字的平均。

# x 的形状为 (B, T, n_embd)，weights 的形状为 (T, T)
x_averaged = weights @ x      # 每个位置变为"自身与前文的平均"

请记住这里的加权求和结构：每个位置的新向量，是把前文所有位置的向量，按一组权重相加而成。这个结构，就是注意力机制的骨架。目前的权重是"平均"——前文各字的权重均等。注意力机制所要做的，只是把这一组权重换得更为合理。

取平均过于粗略：前文各字的重要性并不相同

"取平均"为何不够好。因为它把前文中的每个字都视为同等重要。

但真实情况并非如此。在预测"今天天气真"之后的字时，"天气"这个词显然比"今"这个字更为关键。一个良好的机制，应当让模型自行判断：对于当前这个位置，前文中哪些字重要、哪些字不重要，进而做到重要的字多参考、不重要的字少参考。

也就是说，加权求和所用的那一组权重，不应当是固定的平均，而应当是模型依据内容动态计算出来的。这正是注意力机制要完成的工作。

自注意力：查询、键、值

注意力机制如何计算这一组权重。它为每个字都设定了三个角色，各用一个向量来表示，分别称为查询（query）、键（key）、值（value）。

可以借助一个配对的设想来理解。每个字一方面在"寻求":我这个位置，希望从前文中获取什么样的信息——这个需求记录在它的查询向量中。另一方面，每个字也在"提供":我这个字，能够提供什么样的信息——这记录在它的键向量中。

当某个位置需要判断"前文中某个字对我有多重要"时，它就用自己的查询，去与那个字的键做匹配。匹配度高，说明那个字恰好具备我所需要的信息，权重就大;匹配度低，权重就小。匹配度通过两个向量的点积点积点积是把两个向量对应位置的数字相乘再求和，得到的一个数。它的几何含义是：两个向量方向越接近，点积越大;方向越无关，点积越接近零。注意力正是借此把查询与键是否吻合，量化为一个可比较的分数——无需额外规则，方向的接近程度本身就成了重要性的度量。 来计算——点积大代表两个向量的方向接近，也就意味着"需求与供给相吻合"。

权重确定之后，真正参与加权求和的，并不是字向量本身，而是每个字的第三个向量——值。键负责"匹配",值负责"匹配成功之后实际交付的内容"，两者各司其职。

下面用代码把这套机制实现出来。三个角色各用一个线性层从字向量变换得到：

import torch
import torch.nn as nn
from torch.nn import functional as F

torch.manual_seed(1337)
B, T, n_embd = 4, 8, 32
head_size = 16          # 查询、键、值向量的长度

x = torch.randn(B, T, n_embd)     # 此处用随机数模拟一批字向量

# 三个线性层，分别产出查询、键、值
query = nn.Linear(n_embd, head_size, bias=False)
key = nn.Linear(n_embd, head_size, bias=False)
value = nn.Linear(n_embd, head_size, bias=False)

q = query(x)        # (B, T, head_size)，每个字"希望获取什么"
k = key(x)          # (B, T, head_size)，每个字"能够提供什么"
v = value(x)        # (B, T, head_size)，每个字"实际交付的内容"

# 每个字的查询，与所有字的键做点积，得到一张"重要性"评分表
weights = q @ k.transpose(-2, -1)        # (B, T, T)
weights = weights * head_size ** -0.5    # 缩放，原因见下文

weights 的形状是 (B, T, T)：对每个样本而言，它是一张 T 行 T 列的表，第 i 行第 j 列代表"第 i 个字对第 j 个字的关注程度"。这张表就是我们所需要的那一组动态权重——由内容计算得出，不再是固定的平均。

代码中那一行乘以 head_size ** -0.5 是缩放操作。点积的数值会随着向量长度的增大而增大，数值过大的话，下一步的 softmaxsoftmaxsoftmax 把任意一组实数换算成一组概率：先对每个数取指数，再除以它们的总和，使结果都落在 0 到 1 之间且相加为 1。在注意力中，它把一行原始的匹配分数变为一组规范的权重。由于事先取了指数，原本的分数差距会被进一步放大，这正是此处所说会变得过于极端的来由。 会变得过于极端——几乎把全部权重集中到一个字上，其余字的权重归零。除以 head_size 的平方根，能把数值拉回到温和的范围，使权重的分布较为柔和。

不能参考未来：因果掩码

还差一步，而且这一步至关重要。

我们的模型要做的是预测下一个字。在训练时，第 3 个位置的任务是"依据前 3 个字，预测第 4 个字"。那么它在计算注意力时，绝对不能看到第 4 个、第 5 个字——那是答案，看到了就等同于作弊，这样训练出的模型一旦进入真实的生成场景（此时并没有答案可看），便会失效。

因此需要为注意力施加一项限制：每个位置只允许参考它自身以及它之前的字，不允许参考它之后的字。这项限制称为因果掩码causal maskcausal maskA causal mask forces the model, when it sits at position i, to attend only to positions up to and including i. In code it sets every future entry of the score matrix to negative infinity before softmax, so their weights collapse to zero. Without it the model would peek at the very token it must predict, and would break the moment it is asked to generate text with no answer at hand. 。

实现方式同样借助下三角矩阵：把权重表中"属于未来"的那些位置（上三角部分），强制设为负无穷。负无穷经过 softmax 之后会变成 0，这就等于把这些未来位置的权重彻底清零。

tril = torch.tril(torch.ones(T, T))
# 把上三角（未来位置）填充为负无穷
weights = weights.masked_fill(tril == 0, float("-inf"))
# softmax：把每一行的评分换算成相加为 1 的权重
weights = F.softmax(weights, dim=-1)

经过 masked_fill 与 softmax，weights 的每一行都只有"当前位置及之前"为非零值，且一行之和等于 1。这正是一组合法的、只能看见过去的加权系数。

一个完整的注意力头

把上述各个步骤连接起来——计算查询键值、做点积、缩放、施加因果掩码、做 softmax，最后用权重对值做加权求和——就构成了一个完整的注意力头（attention head）：

# 沿用上文的 q、k、v、B、T、head_size
weights = q @ k.transpose(-2, -1) * head_size ** -0.5   # 计算重要性
tril = torch.tril(torch.ones(T, T))
weights = weights.masked_fill(tril == 0, float("-inf")) # 屏蔽未来
weights = F.softmax(weights, dim=-1)                    # 换算成权重

out = weights @ v       # 用权重对"值"做加权求和

print(f"输出的形状：{out.shape}")     # (B, T, head_size)

最后这一行 weights @ v，其结构与本篇开头那个"取平均"完全相同——依然是加权求和。唯一的、也是全部的进步在于：那一组权重不再是固定的平均，而是模型依据每个字的查询与键、自行计算得出的。重要的字权重大，不相关的字权重小。

这就是注意力机制。每个字不再只关注前一个字，而是回顾整段前文，自行判断应当重点参考其中的哪几个字。Bigram 模型那个"只有一个字记忆"的局限，由此得到了彻底的突破。

本篇要点

• Bigram 只查看前一个字，能力上限很低;模型需要能够看见整段上下文。
• 从本篇起，每个字用一个 n_embd 维向量表示，比单个编号能携带更多信息。
• 让每个位置"看见前文",本质上是对前文向量做加权求和;最初步的权重是取平均。
• 注意力机制把平均换成动态权重：每个字有查询、键、值三个向量，查询与键的点积给出关注程度。
• 因果掩码用下三角把"未来位置"的权重清零，保证模型预测时不会参考到答案。
• 一个注意力头的完整流程是：计算查询键值、做点积、缩放、施加掩码、做 softmax、对值加权求和。

下一篇

一个注意力头能让模型看见上下文，但还不足以构成强大的模型。下一篇会把多个注意力头、前馈网络、残差连接、层归一化、位置编码这些部件组合起来，构成一个完整的 GPT 结构。

参考资料

• Attention Is All You Need（Transformer 原始论文）
• 图解 Transformer（The Illustrated Transformer）
• 本博客 ml-basics 系列《注意力机制》《Transformer》