《从零构建大模型》
[美]塞巴斯蒂安·拉施卡
书中资料 https://github.com/rasbt/LLMs-from-scratch
第三章 注意力机制
3.1 长序列建模中的问题
Transformer出现之前，循环神经网络(recurrent neural network, RNN)是语言翻译中最流行的编码器-解码器架构。RNN是一种将前一步骤的输出作为当前步骤的输入的神经网络，它非常适合处理像文本这样的序列数据
编码器-解码器RNN中，输入文本被传递给编码器以逐步处理。编码器在每一步都会更新其隐藏状态（隐藏层的内部值），试图在最终的隐藏状态中捕捉输入句子的全部含义。然后，解码器使用这个最终的隐藏状态开始逐字生成翻译后的句子。解码器同样在每一步更新其隐藏状态，该状态应包含为下一单词预测所需的上下文信息
编码器部分会将整个输入文本处理成一个隐藏状态（记忆单元）。然后解码器会使用这个隐藏状态来生成输出。你可以将这个隐藏状态视为一种嵌入向量
问题：在解码阶段，RNN无法直接访问编码器中早期隐藏状态，它只能依赖当前的隐藏状态，这会导致上下文丢失，特别是复杂的句子，依赖关系跨越很长的距离。对于较长的文本，它无法直接访问输入中靠前的单词。
研究人员在2014年为RNN开发了Bahdanau注意力机制（以该研究论文的第一作者命名，更多信息请参见附录B），该机制对编码器-解码器RNN进行了修改，使得解码器在每个解码步骤中可以选择性地访问输入序列的不同部分
3.2 使用注意力机制捕捉数据依赖关系
自注意力是Transformer模型中的一种机制，它通过允许一个序列中的每个位置与同一序列中的其他所有位置进行交互并权衡其重要性，来计算出更高效的输入表示。
3.3 通过自注意力机制关注输入的不同部分
自注意力机制中，“自”指的是该机制通过关联单个输入序列中的不同位置来计算注意力权重的能力。它可以评估并学习输入本身各个部分之间的关系和依赖，比如句子中的单词或图像中的像素。
传统的注意力机制关注的是两个不同序列元素之间的关系，比如在序列到序列模型中，注意力可能在输入序列和输出序列之间
自注意力机制的目标是为每个输入元素计算一个上下文向量(context vector)，该向量结合了其他所有输入元素信息的嵌入向量
上下文向量在自注意力机制中起着关键作用。它们的目的是通过结合序列中其他所有元素的信息，为输入序列（如一个句子）中的每个元素创建丰富表示，因为这些模型需要理解句子中单词之间的关系和相关性。
类似我们做阅读理解，要理解一个单词在一句话中的含义，需要看这个单词和句子中其他单词的关系，例如Apple is a good food. 通过food，我们可以知道这里的Apple是苹果水果，而不是苹果公司。
简单的自注意力机制(没有可训练权重)
对于一句文本输入序列”Your journey starts with one step”，它有6个词元，且按照前一章节的方法计算出来了它的嵌入向量$x^{(1)}$ to $x^{(T)}$ ，它的嵌入向量维度为3。现在以第二个词元“journey”为例计算它的上下文向量。
计算注意力分数 $\omega$，把第二个输入作为查询$q^{(2)} = x^{(2)}$，让它依次与输入中所有词元向量进行点积计算得到对应的注意力分数。点积本质上是将两个向量逐个元素相乘然后对乘积求和的简洁方法
点积不仅被视为一种将两个向量转化为标量值的数学工具，而且也是度量相似度的一种方式，因为它可以量化两个向量之间的对齐程度：点积越大，向量之间的对齐程度或相似度就越高，角度也越接近。在自注意机制中，点积决定了序列中每个元素对其他元素的关注程度：点积越大，两个元素之间的相似度和注意力分数就越高。
- $\omega_{21} = x^{(1)} q^{(2)\top}$ 表示第二个输入与第一个元素的点积计算得到注意力分数
- $\omega_{22} = x^{(2)} q^{(2)\top}$
- …
- $\omega_{2T} = x^{(T)} q^{(2)\top}$
计算注意力权重，将得到的注意力分数进行归一化得到注意力权重，归一化的主要目的是获得总和为1的注意力权重。这种归一化是一个惯例，有助于解释结果，并能维持大语言模型的训练稳定性
在实际应用中，使用softmax函数进行归一化更为常见，而且是一种更可取的做法。这种方法更好地处理了极值，并在训练期间提供了更有利的梯度特性
计算上下文向量$z^{(2)}$，通过将嵌入的输入词元与相应的注意力权重相乘，再将得到的向量求和来计算上下文向量
1 | def simple_attention(): |
计算所有输入词元的上下文向量
最终计算出来上下文向量的维数和输入是完全相同的
在计算前面的注意力分数张量时，使用for循环通常较慢，因此可以使用矩阵乘法来得到相同的结果
torch.softmax这样的函数中的dim参数用于指定输入张量的计算维度。将dim设置为-1表示让softmax函数在attn_scores张量的最后一个维度上进行归一化。如果attn_scores是一个二维张量（比如形状为［行, 列］），那么它将对列进行归一化，使得每行的值（在列维度上的总和）为1。
1 | def simple_attention(): |
3.4 实现带可训练权重的自注意力机制
和之前简单自注意力机制差别在于，这里引入了在模型训练过程中会更新的权重矩阵，这些可训练的权重矩阵可以让模型学习生成很好的上下文向量。
三个权重矩阵$W_q$, $W_k$, and $W_v$用于将嵌入的输入词元$x^{(i)}$分别映射为$Q$查询向量、$K$键向量和$V$值向量
- Query vector: $q^{(i)} = x^{(i)},W_q $
- Key vector: $k^{(i)} = x^{(i)},W_k $
- Value vector: $v^{(i)} = x^{(i)},W_v $
在权重矩阵$W$中，“权重”是“权重参数”的简称，表示在训练过程中优化的神经网络参数，随着模型在训练中接触更多数据，它会调整这些可训练的权重。这与前面的注意力权重是不同的。正如我们已经看到的，注意力权重决定了上下文向量对输入的不同部分的依赖程度（网络对输入的不同部分的关注程度）。权重参数是定义网络连接的基本学习系数，而注意力权重是动态且特定于上下文的值。
缩放点积注意力(scaled dot-product attention) 是实际在GPT-2模型中使用的自注意力机制。核心公式如下：
$$
\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V
$$
基本流程
生成3个权重矩阵
输入的嵌入向量维度和查询向量的嵌入维度可以相同也可以不同。在类GPT模型中，输入和输出的维度通常是相同的，但为了便于理解计算过程，这里使用不同的输入维度
(d_in=3)和输出维度(d_out=2)计算每一个输入元素的权重向量，将输入与权重进行矩阵乘法，这里将词元从3维空间映射到了2维空间
计算注意力分数，使用输入元素的查询向量Q和每一个元素的键向量K点积计算
计算注意力权重（归一化），通过缩放注意力分数并应用
softmax函数来计算注意力权重。不过，此时是通过将注意力分数除以键向量的嵌入维度的平方根来进行缩放（取平方根在数学上等同于以0.5为指数进行幂运算）对嵌入维度进行归一化是为了避免梯度过小，从而提升训练性能。例如，在类GPT大语言模型中，嵌入维度通常大于1000，这可能导致点积非常大，从而在反向传播时由于
softmax函数的作用导致梯度非常小。当点积增大时，softmax函数会表现得更像阶跃函数，导致梯度接近零。这些小梯度可能会显著减慢学习速度或使训练停滞。 因此，通过嵌入维度的平方根进行缩放解释了为什么这种自注意力机制也被称为缩放点积注意力机制。计算上下文向量，通过对值向量进行加权求和。注意力权重作为加权因子，用于权衡每个值向量的重要性。和之前一样，可以使用矩阵乘法一步获得输出结果
1 | def weight_attention(): |
为什么要用查询、键和值
查询类似于数据库中的搜索查询。它代表了模型当前关注或试图理解的项（比如句子中的一个单词或词元）。查询用于探测输入序列中的其他部分，以确定对它们的关注程度。
键类似于用于数据库索引和搜索的键。在注意力机制中，输入序列中的每个项（比如句子中的每个单词）都有一个对应的键。这些键用于与查询进行匹配
值类似于数据库中键-值对中的值。它表示输入项的实际内容或表示。一旦模型确定哪些键以及哪些输入部分与查询（当前关注的项）最相关，它就会检索相应的值。
自注意类实现
在自注意力机制中，我们用3个权重矩阵$W_q$, $W_k$, and $W_v$来变换输入矩阵$X$中的输入向量。根据所得查询矩阵($Q$)和键矩阵($K$)计算注意力权重矩阵。然后，使用注意力权重矩阵和值矩阵($V$)计算上下文向量($Z$)。为了视觉清晰，我们关注具有$n$个词元的单个输入文本，而不是一批多个输入。因此，在这种情况下，三维输入张量被简化为二维矩阵，方便更直观地可视化和理解所涉及的过程。
- 输入6个词元，每个词元嵌入向量维度为3，对应矩阵为[6, 3]，假设输出嵌入维度为2，权重矩阵就是[3, 2]，因为要把输入映射到权重矩阵上，左矩阵的列数就是右矩阵的行数，二者相乘得到权重向量的维度为[6,2]
- 以输入的第二个词元为例，它的查询向量Q为[6,2]依次与第一个词元的键K向量[6, 2]点积后，得到标量值如图中的0.2，由于查询要和每一个词元的键都进行点积，所以对第二个词元最终会得到一个[1, 6]的向量，即下图6*6矩阵的第二行。所有的词元都作为查询计算权重矩阵的结果就是[6, 6]即[n,n]的矩阵
- 还以第二个词元为例，它对每一个其他词元(包括它自己)用上一步算出来的权重标量和对应词元的值向量V矩阵乘法计算得到中间向量[1,2]，再把6(n)个中间向量相加得到[1,2]的第二个词元最终的上下文向量。
无论输入词元的嵌入向量维度是多少，最终每个词元的上下文向量的维度都是输出的维度，一般这个维度和字典的个数相同，表示每个词出现的可能性。
1 | # 从nn.Module派生出来的类。nn.Module是PyTorch模型的一个基本构建块，它为模型层的创建和管理提供了必要的功能 |
3.5 利用因果注意力隐藏未来词汇
对于许多大语言模型任务，你希望自注意力机制在预测序列中的下一个词元时仅考虑当前位置之前的词元
因果注意力（也称为掩码注意力）是一种特殊的自注意力形式。它限制模型在处理任何给定词元时，只能基于序列中的先前和当前输入来计算注意力分数，而标准的自注意力机制可以一次性访问整个输入序列。
在因果注意力机制中，我们掩码了对角线以上的注意力权重，并归一化未掩码的注意力权重，使得每一行的权重之和为1，以确保在计算上下文向量时，大语言模型无法访问未来的词元。例如，对于第2行的单词“journey”，仅保留当前词(“journey”)和之前词(“Your”)的注意力权
在因果注意力中，获得掩码后的注意力权重矩阵的一种方法是对注意力分数应用
softmax函数，将对角线以上的元素清零，并对所得矩阵进行归一化
简单掩码处理流程
按照之前的方法，通过
softmax函数计算出注意力权重1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17inputs = torch.tensor(
[[0.43, 0.15, 0.89], # Your (x^1)
[0.55, 0.87, 0.66], # journey (x^2)
[0.57, 0.85, 0.64], # starts (x^3)
[0.22, 0.58, 0.33], # with (x^4)
[0.77, 0.25, 0.10], # one (x^5)
[0.05, 0.80, 0.55]] # step (x^6)
)
d_in = inputs.shape[1] # 输入嵌入维度, d=3
d_out = 2 # 查询嵌入维度, d=2
torch.manual_seed(789)
sa_v2 = SelfAttention_v2(d_in, d_out)
queries = sa_v2.W_query(inputs)
keys = sa_v2.W_key(inputs)
attn_scores = queries @ keys.T
attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)创建一个对角线以上元素为0的掩码矩阵，矩阵维数为词元个数
1
2
3
4
5
6
7
8
9
10
11
12
13# 输入的词元个数
context_length = attn_scores.shape[0]
# 生成一个下三角矩阵
mask_simple = torch.tril(torch.ones(context_length, context_length))
print(mask_simple)
'''
tensor([[1., 0., 0., 0., 0., 0.],
[1., 1., 0., 0., 0., 0.],
[1., 1., 1., 0., 0., 0.],
[1., 1., 1., 1., 0., 0.],
[1., 1., 1., 1., 1., 0.],
[1., 1., 1., 1., 1., 1.]])
'''把这个掩码矩阵和注意力权重矩阵相乘，使权重矩阵对角线上方的值变为0
1
2
3
4
5
6
7
8
9
10
11
12# 只保留下三角矩阵部分的权重
masked_simple = attn_weights*mask_simple
print(masked_simple)
'''
tensor([[0.1921, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.2041, 0.1659, 0.0000, 0.0000, 0.0000, 0.0000],
[0.2036, 0.1659, 0.1662, 0.0000, 0.0000, 0.0000],
[0.1869, 0.1667, 0.1668, 0.1571, 0.0000, 0.0000],
[0.1830, 0.1669, 0.1670, 0.1588, 0.1658, 0.0000],
[0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
grad_fn=<MulBackward0>)
'''重新归一化注意力权重，使每一行的总和再次为1。可以通过将每行中的每个元素除以每行中的和来实现这一点
1
2
3
4
5
6
7
8
9
10
11
12
13# 对每一行重新归一化
row_sums = masked_simple.sum(dim=-1, keepdim=True)
masked_simple_norm = masked_simple / row_sums
print(masked_simple_norm)
'''
tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.5517, 0.4483, 0.0000, 0.0000, 0.0000, 0.0000],
[0.3800, 0.3097, 0.3103, 0.0000, 0.0000, 0.0000],
[0.2758, 0.2460, 0.2462, 0.2319, 0.0000, 0.0000],
[0.2175, 0.1983, 0.1984, 0.1888, 0.1971, 0.0000],
[0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],
grad_fn=<DivBackward0>)
'''
信息泄露
当我们应用掩码并重新归一化注意力权重时，初看起来，未来的词元（打算掩码的）可能仍然会影响当前的词元，因为它们的值会参与
softmax计算。然而，关键的见解是，在掩码后重新归一化时，我们实际上是在对一个较小的子集重新计算softmax（因为被掩码的位置不参与softmax计算）softmax函数在数学上的优雅之处在于，尽管最初所有位置都在分母中，但掩码和重新归一化之后，被掩码的位置的效果被消除——它们不会以任何实际的方式影响softmax分数。注意力权重的分布就像最初仅在未掩码的位置计算一样，这保证了不会有来自未来或其他被掩码的词元的信息泄露
改进掩码方法
softmax函数会将其输入转换为一个概率分布。当输入中出现负无穷大$-\infty $值时，softmax函数会将这些值视为零概率。（从数学角度来看，这是因为 $ e^{-\infty} $无限接近于0），所以通过优化以下步骤，相对之前的方法减少一次归一化。
- 对未归一化的注意力分数对角线以上部分用负无穷进行掩码
- 再用
softmax函数进行归一化
1 | def causal_attention(): |
利用dropout掩码额外的注意力权重
dropout是深度学习中的一种技术，通过在训练过程中随机忽略一些隐藏层单元来有效地“丢弃”它们。这种方法有助于减少模型对特定隐藏层单元的依赖，从而避免过拟合。需要强调的是，dropout仅在训练期间使用，训练结束后会被取消。
- GPT在内的模型通常会在两个特定时间点使用注意力机制中的dropout：
- 计算注意力权重之后，一般都在这时使用dropout - 注意力权重与值向量相乘之后
代码示例中使用了50%的dropout率，这意味着掩码一半的注意力权重。（当我们在接下来的章节中训练GPT模型时，将使用较低的dropout率，比如10%或20%。）
1 | torch.manual_seed(123) |
- 对注意力权重矩阵应用50%的dropout率时，矩阵中有一半的元素会随机被置为0。为了补偿减少的活跃元素，矩阵中剩余元素的值会按1/0.5=2的比例进行放大。放大比例系数计算规则为
1 / (1 - dropout_rate)这种放大对于维持注意力权重的整体平衡非常重要，可以确保在训练和推理过程中，注意力机制的平均影响保持一致。
因果注意力类实现
相对之前增加了多个批次处理，因果掩码和dropout掩码
1 | class CausalAttention(nn.Module): |
类的使用
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39def test_CausalAttention():
inputs = torch.tensor(
[[0.43, 0.15, 0.89], # Your (x^1)
[0.55, 0.87, 0.66], # journey (x^2)
[0.57, 0.85, 0.64], # starts (x^3)
[0.22, 0.58, 0.33], # with (x^4)
[0.77, 0.25, 0.10], # one (x^5)
[0.05, 0.80, 0.55]] # step (x^6)
)
# 把输入重复两遍，模拟两个批次
batch = torch.stack((inputs, inputs), dim=0)
# 2行输入，每个输入6个词元，每个词元的嵌入维度为3
print(batch.shape) # torch.Size([2, 6, 3])
d_in = inputs.shape[1] # 输入嵌入维度, d=3
d_out = 2 # 查询嵌入维度, d=2
torch.manual_seed(123)
context_length = batch.shape[1] # 上下文长度为6，每一个输入6个词元
ca = CausalAttention(d_in, d_out, context_length, 0.0)
context_vecs = ca(batch)
# 输出为2个批次，每个批次6个词元，每个词元2维向量表示
print("context_vecs.shape:", context_vecs.shape) # torch.Size([2, 6, 2])
print(context_vecs)
'''
tensor([[[-0.4519, 0.2216],
[-0.5874, 0.0058],
[-0.6300, -0.0632],
[-0.5675, -0.0843],
[-0.5526, -0.0981],
[-0.5299, -0.1081]],
[[-0.4519, 0.2216],
[-0.5874, 0.0058],
[-0.6300, -0.0632],
[-0.5675, -0.0843],
[-0.5526, -0.0981],
[-0.5299, -0.1081]]], grad_fn=<UnsafeViewBackward0>)
'''
3.6 将单头注意力扩展到多头注意力
“多头”这一术语指的是将注意力机制分成多个“头”，每个“头”独立工作。在这种情况下，单个因果注意力模块可以被看作单头注意力，因为它只有一组注意力权重按顺序处理输入。
实现多头注意力需要构建多个自注意力机制的实例（参见3.4 自注意类实现中的图），每个实例都有其独立的权重，然后将这些输出进行合成。虽然这种方法的计算量可能会非常大，但它对诸如基于Transformer的大语言模型之类的模型的复杂模式识别是非常重要的。
多头注意力的主要思想是多次（并行）运行注意力机制，每次使用学到的不同的线性投影——这些投影是通过将输入数据（比如注意力机制中的查询向量、键向量和值向量）乘以权重矩阵得到的。通过多个不同的、经过学习得到的线性投影，多次（并行地）运行注意力机制，这样可以使模型能够共同关注来自不同位置、不同表示子空间的信息。
简单的叠加多个单头注意力层
1 | class MultiHeadAttentionWrapper(nn.Module): |
测试
结果中的context_vecs张量的第一维是2，因为我们有两个批次的输入文本（输入文本是重复的，所以这些上下文向量完全相同）。第二维表示每个输入中的6个词元。第三维表示每个词元的四维嵌入。
因为通过d_out=2指定了Q，K，V和上下文向量的嵌入维度为2，我们沿着列维度连接这些上下文向量向量得到最终的矩阵。由于我们有2个注意力头并且嵌入维度为2，因此最终的嵌入维度是2×2=4。
1 | def test_MultiHeadAttentionWrapper(): |
改进的多头注意力类
主要思想是把多个头的向量矩阵放在一个大的矩阵向量中计算，从而减少计算过程中矩阵乘法的次数。
不同于之前的方法创建每个头创建一个权重矩阵$W_{q1}$和$W_{q2}$，新方法初始化了一个更大的权重矩阵$W_q$，并只与输入矩阵进行一次矩阵乘法操作，得到一个查询矩阵$Q$。
根据输出维度d_out按头数num_heads除后得到每个头的输出维度head_dim，这里测试代码例子中head_dim就是2/2 = 1，公式为head_dim = d_out / num_heads
通过增加一个head_dim维度隐式的将一个形状为(b, num_tokens, d_out)的张量通过view函数重塑形状为(b, num_tokens, num_heads, head_dim)，这里num_heads为2，所以隐含的就有两个查询矩阵$Q_1$和$Q_2$。其他矩阵处理类似。
然后转置张量，使num_heads维度置于num_tokens维度之前，从而形成一个(b, num_heads, num_tokens, head_dim)的形状。这种转置对于正确对齐不同头的查询矩阵、键矩阵和值矩阵，以及有效地执行批处理矩阵乘法至关重要。接着就可以使用批处理矩阵乘法，queries @ keys.transpose(2, 3)来计算注意力分数。
最后对计算得到的上下文向量(b, num_tokens, num_heads, head_dim)接着重塑（展平）为(b, num_tokens, d_out)的形状，从而有效地整合所有头的输出。
使用批量矩阵乘法的效率更高。原因是我们只需进行一次矩阵乘法来计算键矩阵，例如，keys = self.W_key(x)（查询矩阵和值矩阵也是如此）。在MultiHeadAttentionWrapper中，我们需要对每个注意力头重复进行这种矩阵乘法，而矩阵乘法是计算资源消耗较大的操作之一。
1 | class MultiHeadAttention(nn.Module): |
- 测试函数，这里总输出维数为2，即每个头的输出维数为1
1 | def test_MultiHeadAttention(): |
pytorch中也有多头注意力的实现 torch.nn.MultiheadAttention
最小的GPT-2模型（参数量为1.17亿）有12个注意力头，上下文向量嵌入维度为768，而最大的GPT-2模型（参数量为15亿）有25个注意力头，上下文向量嵌入维度为1600。请注意，在GPT模型中，词元输入和上下文嵌入的嵌入维度是相同的(d_in = d_out)
批处理矩阵乘法
PyTorch的矩阵乘法实现处理了四维输入张量，使得矩阵乘法在最后两个维度（num_tokens和head_dim）之间进行，并对每个头重复这一操作
1 | def batch_matrix_mul(): |