从零开始实现大语言模型（四）：简单自注意力机制

1. 前言

理解大语言模型结构的关键在于理解自注意力机制(self-attention)。自注意力机制可以判断输入文本序列中各个token与序列中所有token之间的相关性，并生成包含这种相关性信息的context向量。

本文介绍一种不包含训练参数的简化版自注意力机制——简单自注意力机制(simplified self-attention)，后续三篇文章将分别介绍缩放点积注意力机制(scaled dot-product attention)，因果注意力机制(causal attention)，多头注意力机制(multi-head attention)，并最终实现OpenAI的GPT系列大语言模型中MultiHeadAttention类。

2. 从循环神经网络到自注意力机制

解决机器翻译等多对多(many-to-many)自然语言处理任务最常用的模型是sequence-to-sequence模型。Sequence-to-sequence模型包含一个编码器(encoder)和一个解码器(decoder)，编码器将输入序列信息编码成信息向量，解码器用于解码信息向量，生成输出序列。在Transformer模型出现之前，编码器和解码器一般都是一个循环神经网络(RNN, recurrent neural network)。

RNN是一种非常适合处理文本等序列数据的神经网络架构。Encoder RNN对输入序列进行处理，将输入序列信息压缩到一个向量中。状态向量$h_0$包含第一个token $x_0$的信息，$h_1$包含前两个tokens $x_0$和$x_1$的信息。以此类推， Encoder RNN最后一个状态$h_m$是整个输入序列的概要，包含了整个输入序列的信息。Decoder RNN的初始状态等于Encoder RNN最后一个状态$h_m$。$h_m$包含了输入序列的信息，Decoder RNN可以通过$h_m$知道输入序列的信息。Decoder RNN可以将$h_m$中包含的信息解码，逐个元素地生成输出序列。

RNN的神经网络结构及计算方法使Encoder RNN必须用一个隐藏状态向量$h_m$记住整个输入序列的全部信息。当输入序列很长时，隐藏状态向量$h_m$对输入序列中前面部分的tokens的偏导数（如对$x_0$的偏导数$\frac{\partial h_m}{x_0}$）会接近0。输入不同的$x_0$，隐藏状态向量$h_m$几乎不会发生变化，即RNN会遗忘输入序列前面部分的信息。

本文不会详细介绍RNN的原理，大语言模型的神经网络中没有循环结构，RNN的原理及结构与大语言模型没有关系。对RNN的原理感兴趣读者可以参见本人的博客专栏：自然语言处理。

2014年，文章Neural Machine Translation by Jointly Learning to Align and Translate提出了一种改进sequence-to-sequence模型的方法，使Decoder每次更新状态时会查看Encoder所有状态，从而避免RNN遗忘的问题，而且可以让Decoder关注Encoder中最相关的信息，这也是attention名字的由来。

2017年，文章Attention Is All You Need指出可以剥离RNN，仅保留attention，且attention并不局限于sequence-to-sequence模型，可以直接用在输入序列数据上，构建self-attention，并提出了基于attention的sequence-to-sequence架构模型Transformer。

3. 简单自注意力机制

自注意力机制的目标是计算输入文本序列中各个token与序列中所有tokens之间的相关性，并生成包含这种相关性信息的context向量。如下图所示，简单自注意力机制生成context向量的计算步骤如下：

1. 计算注意力分数(attention score)：简单注意力机制使用向量的点积(dot product)作为注意力分数，注意力分数可以衡量两个向量的相关性；
2. 计算注意力权重(attention weight)：将注意力分数归一化得到注意力权重，序列中每个token与序列中所有tokens之间的注意力权重之和等于1；
3. 计算context向量：简单注意力机制将所有tokens对应Embedding向量的加权和作为context向量，每个token对应Embedding向量的权重等于其相应的注意力权重。

3.1 计算注意力分数

对输入文本序列 “Your journey starts with one step.” 做tokenization，将文本中每个单词分割成一个token，并转换成Embedding向量，得到$x_1,x_2,\cdots ,x_6$。自注意力机制分别计算$x_i$与$x_1,x_2,\cdots ,x_6$的注意力权重，进而计算$x_1,x_2,\cdots ,x_6$与其相应注意力权重的加权和，得到context向量$z_i$。

如下图所示，将context向量$z_i$对应的向量$x_i$称为query向量，计算query向量$x_2$对应的context向量$z_2$的第一步是计算注意力分数。将query向量$x_2$分别点乘向量$x_1,x_2,\cdots ,x_6$，得到实数${\omega }_{21},{\omega }_{22},\cdots ,{\omega }_{26}$，其中${\omega }_{2i}$是query向量$x_2$与向量$x_i$的注意力分数，可以衡量$x_2$对应token与$x_i$对应token之间的相关性。

两个向量的点积等于这两个向量相同位置元素的乘积之和。假如向量$x_1=(x_{11},x_{12},x_{13})$，向量$x_2=(x_{21},x_{22},x_{23})$，则向量$x_1$与$x_2$的点积等于$x_{11}\times x_{21}+x_{12}\times x_{22}+x_{13}\times x_{23}$。

可以使用如下代码计算query向量$x_2$与$x_1,x_2,\cdots ,x_6$的注意力分数：

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)
)query = inputs[1]
attn_scores_2 = torch.empty(inputs.shape[0])
for i, x_i in enumerate(inputs):attn_scores_2[i] = torch.dot(x_i, query)
print(attn_scores_2)

执行上面代码，打印结果如下：

tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])

3.2 计算注意力权重

如下图所示，将注意力分数${\omega }_{21},{\omega }_{22},\cdots ,{\omega }_{26}$归一化可得到注意力权重${\alpha }_{21},{\alpha }_{22},\cdots ,{\alpha }_{26}$。每个注意力权重${\alpha }_{2i}$的值均介于0到1之间，所有注意力权重的和${\sum }_i{\alpha }_{2i}=1$。可以用注意力权重${\alpha }_{2i}$表示$x_i$对当前context向量$z_2$的重要性占比，注意力权重${\alpha }_{2i}$越大，表示$x_i$与$x_2$的相关性越强，context向量$z_2$中$x_i$的信息量比例应该越高。使用注意力权重对$x_1,x_2,\cdots ,x_6$加权求和计算context向量，可以使context向量的数值分布范围始终与$x_1,x_2,\cdots ,x_6$一致。这种数值分布范围的一致性可以使大语言模型训练过程更稳定，模型更容易收敛。

可以使用softmax函数将注意力分数归一化得到注意力权重：

attn_weights_2 = torch.softmax(attn_scores_2, dim=0)
print("Attention weights:", attn_weights_2)
print("Sum:", attn_weights_2.sum())

执行上面代码，打印结果如下：

Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
Sum: tensor(1.)

3.3 计算context向量

简单注意力机制使用所有tokens对应Embedding向量的加权和作为context向量，context向量$z_2={\sum }_i{\alpha }_{2i}{x}_i$。

可以使用如下代码计算context向量$z_2$：

query = inputs[1] # 2nd input token is the query
context_vec_2 = torch.zeros(query.shape)
for i,x_i in enumerate(inputs):context_vec_2 += attn_weights_2[i] * x_i
print(context_vec_2)

执行上面代码，打印结果如下：

tensor([0.4419, 0.6515, 0.5683])

3.4 计算所有tokens对应的context向量

将向量$x_2$作为query向量，按照3.1所述方法，可以计算出注意力分数${\omega }_{21},{\omega }_{22},\cdots ,{\omega }_{26}$。使用softmax函数将注意力分数${\omega }_{21},{\omega }_{22},\cdots ,{\omega }_{26}$归一化，可以得到注意力权重${\alpha }_{21},{\alpha }_{22},\cdots ,{\alpha }_{26}$。Context向量$z_2$是使用注意力权重对$x_1,x_2,\cdots ,x_6$的加权和。

计算所有tokens对应的context向量，可以使用矩阵乘法运算，分别将各个$x_i$作为query向量，一次性批量计算注意力分数及注意力权重，并最终得到context向量$z_i$。

如下面代码所示，可以使用矩阵乘法，一次性计算出所有注意力分数：

attn_scores = inputs @ inputs.T
print(attn_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中的矩阵乘法运算符号，与函数torch.matmul运算逻辑相同。

将一个$n$行$m$列的矩阵$A$与另一个$m$行$n$列$B$的矩阵相乘，结果$C$是一个$n$行$n$列的矩阵。其中矩阵$C$的$i$行$j$列元素等于矩阵$A$的第$i$行与矩阵$B$的第$j$列两个向量的内积。

如下面代码所示，使用softmax函数注意力分数归一化，可以一次批量计算出所有注意力权重：

attn_weights = torch.softmax(attn_scores, dim=1)
print(attn_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]])

可以同样使用矩阵乘法运算，一次性批量计算出所有context向量：

all_context_vecs = attn_weights @ inputs
print(all_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]])

4. 结束语

自注意力机制是大语言模型神经网络结构中最复杂的部分。为降低自注意力机制原理的理解门槛，本文介绍了一种不带任何训练参数的简化版自注意力机制。

自注意力机制的目标是计算输入文本序列中各个token与序列中所有tokens之间的相关性，并生成包含这种相关性信息的context向量。简单自注意力机制生成context向量共3个步骤，首先计算注意力分数，然后使用softmax函数将注意力分数归一化得到注意力权重，最后使用注意力权重对所有tokens对应的Embedding向量加权求和得到context向量。

接下来，该去看看大语言模型中真正使用到的注意力机制了！