1. 引言
前文词嵌入和位置嵌入已经做了模型推理前对token序列的预处理工作，本文来重点讨论语言模型的核心组件——Attention。
语言模型最早面临的问题是翻译问题，如何从一种语言翻译成另一种语言，例如：德文翻译成英文，由于两种语言的语法结构不同，如果逐字翻译会造成语法错误。
为了解决无法逐字翻译文本的问题，出现了基于编码器-解码器架构的递归神经网络（RNN）。
- 编码器按顺序对输入文本进行处理，在每个步骤中更新其隐藏状态，并试图在最终的隐藏状态中捕获输入句子的全部含义。
- 解码器则以这个最终的隐藏状态作为输入，开始一个单词一个单词的生成翻译后的句子，并在每一步更新其隐藏状态，以确保下一个单词预测时携带了已经生成的词作为上下文。
但RNN只适用于翻译短句，而不适用于较长的文本，它最大的局限性在于：解码阶段完全依赖于编码器输出的最终隐藏状态，而无法访问编码器中更早期的隐藏状态，这可能会导致上下文丢失，尤其是在依赖关系可能跨越很长距离的复杂句子中。
为此，研究人员在 2014 年开发了 RNN 的 Bahdanau 注意力机制，该机制所作的重要修改是允许在解码步骤中选择性地访问输入序列的不同部分。
仅仅三年后，谷歌发表了论文《Attention is all you need》,指出RNN对于构建自然语言的神经网络并不是必需的，并提出了完全基于self-attention的transformer架构。
自注意力是一种更有效的进行输入表示的机制。它允许序列中的每个位置在计算它在序列中的表示时，都能关注序列中的所有位置。通俗来讲，就是每个单词在句子中的含义是由句子中其它单词组成的上下文决定的。
下面我们将从头开始编写这种自注意力机制。
2. 简单自注意力
2.1 单个token的注意力权重
自注意力机制中的自（self）是指注意力权重所要关注的是单一输入序列内部不同位置的联系。作为对比，传统注意力关注的是两个不同序列之间的元素关系。
自注意力的目标是为每个输入token计算一个上下文向量，该向量结合了序列中所有其他输入token的信息。
以一个输入文本 “Your journey starts with one step.” 为例，假设一个单词对应一个token，如果我们要计算第2个tokenjourney的上下文向量z(2)，那最终每个token x(1)……x(T)对z(2)的重要性将由注意力权重 α \alpha α 21到 α \alpha α 2T决定。
接下来，我们将用代码来演示这个计算的过程。
为了可显示的需要，我们降低了嵌入的维度，我们对每个token采用三维嵌入。
import torch
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)
)
计算注意力得分
实现自注意力机制的第一步是计算中间变量 ω，这些变量被称为注意力得分。方法是计算 journey 的嵌入向量x(2)与其它token的嵌入向量之间的点积。
query = inputs[1]
atten_scores_2 = torch.empty(inputs.shape[0])
for i, x_i in enumerate(inputs):
atten_scores_2[i] = torch.dot(query, x_i)
atten_scores_2
tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])
那通过点积计算的这些数值有什么含义呢？
在自注意力机制中，点积用于衡量序列中各token之间的关注程度，点积值越高，两个token之间的相似性和注意力得分就越高。
分数归一化
归一化就是将注意力得分中的每个数值进行变换，目标是各项数值的总和为1。归一化的最简单做法是每个元素除以所有元素之和，如下所示：
normalized_scores_2 = atten_scores_2/atten_scores_2.sum()
normalized_scores_2
tensor([0.1455, 0.2278, 0.2249, 0.1285, 0.1077, 0.1656])
查验这些注意力分数相加之和是否为1：
normalized_scores_2.sum()
tensor(1.0000)
但实际场景中，更推荐使用softmax进行归一化，与上面我们直接采用元素本身进行运算不同，softmax先对所有元素进行指数运算。
def softmax_normalize(x):
return torch.exp(x) / torch.exp(x).sum(dim=0)
atten_weights_2 = softmax_normalize(atten_scores_2)
atten_weights_2, atten_weights_2.sum()
(tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581]), tensor(1.))
使用softmax进行归一化的好处在于：指数运算确保了所有注意力得分为正，这意味着输出可以被解释为概率，高数值代表更大的重要性。
现在我们已经得到了注意力权重w, 下一步是将每个嵌入的向量 x(i) 与相应的注意力权重相乘，然后将结果向量求和，计算出上下文向量 z(2)。
context_vec_2 = torch.empty(query.shape)
for i, x_i in enumerate(inputs):
context_vec_2 += atten_weights_2[i] * x_i
context_vec_2
tensor([0.4419, 0.6515, 0.5684])
2.2 所有token的注意力权重
上面对第2个tokenjourney计算了上下文向量，下面会对代码作一些修改，以将计算过程扩展至所有输入token的注意力权重和上下文向量。
atten_scores = torch.empty(inputs.shape[0], inputs.shape[0])
for i, query in enumerate(inputs):
for j, x_i in enumerate(inputs):
atten_scores[i, j] = torch.dot(query, x_i)
atten_scores
tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],
[0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],
[0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],
[0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],
[0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],
[0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])
上面计算过程本质上是inputs矩阵和inputs矩阵的转置相乘，在pytorch中矩阵相乘有更简单的写法：
atten_scores = inputs @ inputs.T
atten_scores
tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],
[0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],
[0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],
[0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],
[0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],
[0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])
pytorch中的这种矩阵乘法运算不仅简洁，而且执行效率也比python中的for循环更高效。
下面对每行的注意力得分进行归一化，使其总和为1。同样，对于归一化也使用pytorch中封装的softmax来代替我们手工编写的softmax。
atten_weights = torch.softmax(atten_scores, dim=1)
atten_weights
tensor([[0.2098, 0.2006, 0.1981, 0.1242, 0.1220, 0.1452],
[0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581],
[0.1390, 0.2369, 0.2326, 0.1242, 0.1108, 0.1565],
[0.1435, 0.2074, 0.2046, 0.1462, 0.1263, 0.1720],
[0.1526, 0.1958, 0.1975, 0.1367, 0.1879, 0.1295],
[0.1385, 0.2184, 0.2128, 0.1420, 0.0988, 0.1896]])
简单确认下，各行相加之和是否为1。
atten_weights.sum(dim=1)
tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000])
最后通过矩阵乘法来生成所有的上下文向量，在得到的输出结果中，每一行都包含一个三维的上下文向量。
context_vecs = atten_weights @ inputs
context_vecs
tensor([[0.4421, 0.5931, 0.5790],
[0.4419, 0.6515, 0.5683],
[0.4431, 0.6496, 0.5671],
[0.4304, 0.6298, 0.5510],
[0.4671, 0.5910, 0.5266],
[0.4177, 0.6503, 0.5645]])
其中，第二行的输出张量[0.4419, 0.6515, 0.5683]与我们上面计算的context_vec_2完全相同，也说明代码计算无误。
到这里，我们就完成了一个简单自注意力的代码演示。接下来，我们将添加可训练的权重，使自己编写的注意力具备可学习性。
3.可训练权重的自注意力
3.1 初始化权重矩阵
首先，我们定义输入和输出的维度dim_in和dim_out。
dim_in = inputs.shape[1]
dim_out = inputs.shape[1]
在前面简单注意力的基础上，引入三个可训练的权重矩阵 Wq、Wk 和 Wv，这三个矩阵会用于将嵌入的输入向量 x(i) 投影为查询向量Q、键向量K和值向量V。
torch.manual_seed(123)
Wq = torch.nn.Parameter(torch.rand(dim_in, dim_out), requires_grad=False)
Wk = torch.nn.Parameter(torch.rand(dim_in, dim_out), requires_grad=False)
Wv = torch.nn.Parameter(torch.rand(dim_in, dim_out), requires_grad=False)
3.2 计算Q、K、V向量
同样，我们先以第二个tokenjourney的输入向量x(2)作查询向量，来计算与其对应的查询（q）、键（k）和值（v）向量。
x_2 = inputs[1]
q_2 = x_2 @ Wq
k_2 = x_2 @ Wk
v_2 = x_2 @ Wv
q_2, k_2, v_2
(tensor([0.8520, 0.4161, 1.0138]),
tensor([0.7305, 0.4227, 1.1993]),
tensor([0.9074, 1.3518, 1.5075]))
尽管只是计算一个token的上下文向量z(2)，但我们仍然需要所有token的键向量k和值向量v，因为查询向量q_2需要与序列中所有token的k向量和v向量运算，才能得到x_2的注意力权重和上下文向量。
keys = inputs @ Wk
values = inputs @ Wv
keys, values
(tensor([[0.6813, 0.2706, 1.0793],
[0.7305, 0.4227, 1.1993],
[0.7355, 0.4227, 1.1901],
[0.3363, 0.2225, 0.6077],
[0.6184, 0.3038, 0.6909],
[0.3178, 0.2383, 0.7426]]),
tensor([[0.4976, 0.9655, 0.7614],
[0.9074, 1.3518, 1.5075],
[0.8976, 1.3391, 1.4994],
[0.5187, 0.7319, 0.8493],
[0.4699, 0.7336, 0.9307],
[0.6446, 0.9045, 0.9814]]))
3.3 计算注意力得分
如同前面的简单自注意力实现一样，注意力得分是一个点积运算，有所不同的是，我们不再直接计算输入元素，而是经过权重矩阵变换过的查询向量q、键向量k和值向量v。
atten_scores_2 = q_2 @ keys.T
atten_scores_2
tensor([1.7872, 2.0141, 2.0091, 0.9952, 1.3538, 1.1227])
3.4 计算注意力权重
从上面可以知道，注意力得分到注意力权重的变换是通过softmax进行归一化运算。与之前不同的是，这里我们会引入缩放点积来计算注意力，具体操作就是在softmax之前，先将注意力得分除以键的嵌入维度的平方根来缩放注意力得分。
dim_k = keys.shape[-1]
atten_weights_2 = torch.softmax(atten_scores_2/dim_k ** 0.5, dim = -1)
atten_weights_2
tensor([0.1862, 0.2123, 0.2117, 0.1179, 0.1450, 0.1269])
通过嵌入维度的平方根进行缩放的这种注意力机制也被称为缩放点积注意力。
3.5 计算上下文向量
得到了注意力权重后，我们通过对所有值向量加权求和就能计算出上下文向量，而注意力权重就是衡量每个值向量重要性的权重因子。
context_vec_2 = atten_weights_2 @ values
context_vec_2
tensor([0.6864, 1.0577, 1.1389])
上面的整个计算过程可以用下面这个图来表示：
4. 实现selfAttention类
上面逐步计算自注意力的过程是为了理解和展示，实际应用中，会把上面的代码整合到一个python类中以便高效使用。
from torch import nn
class SelfAttention(nn.Module):
def __init__(self, dim_in, dim_out):
super().__init__()
self.dim_out = dim_out
self.Wq = nn.Parameter(torch.rand(dim_in, dim_out), requires_grad=True)
self.Wk = nn.Parameter(torch.rand(dim_in, dim_out), requires_grad=True)
self.Wv = nn.Parameter(torch.rand(dim_in, dim_out), requires_grad=True)
def forward(self, x):
q = x @ self.Wq
k = x @ self.Wk
v = x @ self.Wv
atten_scores = q @ k.T
atten_weights = torch.softmax(atten_scores/self.dim_out ** 0.5, dim=-1)
context_vecs = atten_weights @ v
return context_vecs
直接使用封装的自注意力类来计算上下文向量。
torch.manual_seed(123)
atten = SelfAttention(inputs.shape[-1], inputs.shape[-1])
atten(inputs)
tensor([[0.6692, 1.0276, 1.1106],
[0.6864, 1.0577, 1.1389],
[0.6860, 1.0570, 1.1383],
[0.6738, 1.0361, 1.1180],
[0.6711, 1.0307, 1.1139],
[0.6783, 1.0441, 1.1252]], grad_fn=<MmBackward0>)
小结：本节从一个简单文本的三维嵌入开始，一步一步复现了自注意力的详细计算过程，而后从简单自注意力逐步过滤到可训练权重的自注意力，从基础的矩阵计算逐步过渡到pytorch的高级API，最后将这个计算过程封装为一个组件以备复用。下一节，我们将结合因果关系和多头元素对自注意机制进行改进。