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

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$