1.2.2 模型架构
在本小节中，我们将介绍Transformer模型的详细架构，为了更容易理解，我们首先将模型简化为一个黑盒子，然后一层层深入，并一一介绍相关概念。我们以Transformer最开始的应用——机器翻译任务为例进行讲解。在机器翻译任务中，它的输入是一种语言（比如汉语）的一个句子，输出是表示相同意思的另一种语言（比如英语）的一个句子，如图1-11所示。
图1-11 Transformer模型在机器翻译中的输入与输出示意
打开Transformer模型，可以发现它是由一个编码器组件、一个解码器组件以及它们之间的连接组成的。其中，编码器组件由一系列编码器组成，解码器组件由一系列相同数量的解码器组成。
所有的编码器在结构上都是相同的，如图1-12左图所示，它分为两个部分：自注意力层和前馈神经网络。编码器的输入首先经过自注意力层，对输入句子中的特定单词向量计算注意力分数，结合所有单词向量的注意力分数，编码为一个新的隐藏状态向量。自注意力层的输出被发送到前馈神经网络，进行线性映射。自注意力层根据输入句子中不同的单词得出不同的注意力分数，因此自注意力层的权重参数不同，而前馈神经网络对输入句子中不同的单词应用完全相同的权重参数。
解码器的结构与编码器类似，如图1-12右图所示，不过分为了3个部分：自注意力层、注意力层和前馈神经网络。其中自注意力层和前馈神经网络的结构与编码器中相同，在这两层之间是一个“编码器-解码器”注意力层，使用来自编码器压缩的上下文向量，让解码器不仅可以关注输入句子中的相关单词，还可以关注输出句子中的相关单词。
图1-12 Transformer模型中编码器和解码器的结构细节示意
1.自注意力层
正如前面提到的，编码器接收向量列表作为输入，然后将向量列表传递给自注意力层进行处理，之后传递给前馈神经网络，最后将输出发送给下一个编码器。在编码器计算的过程中，每个位置的单词都会经过自注意力层的计算，然后通过一个完全相同的前馈神经网络。
实现编码器中的自注意力层需要四个步骤，如图1-13所示。
图1-13 自注意力层的计算流程示意
1）为每个单词创建查询（query）键（key）和值（value）三个向量。对于输入序列中的每个单词，将其表征向量分别和对应的权重矩阵（WQ，WK，WV）相乘，映射成查询向量、键向量和值向量。
2）使用查询向量对其他单词的键向量进行评分。对于每个输入单词，使用其对应的查询向量与其他所有单词的键向量进行内积运算，获得注意力分数，这个注意力分数就表示当前单词与其他单词之间的相关性。
3）对注意力分数进行标准化。每个单词对于其他单词的注意力分数除以键向量维度的平方根
，然后通过softmax函数进行标准化，这样有助于提高梯度的稳定性。
4）将值向量乘以标准化后的注意力分数然后加权求和。对于每个输入单词，将其对应的值向量与上一步得到的标准化后的注意力分数相乘，然后将所有乘积结果相加，得到经过自注意力机制修正后的表征。
为什么需要三个向量来计算注意力分数呢？query、key和value的概念来自信息检索系统。例如，当我们想要搜索某一个问题时，会打开浏览器在搜索栏输入关键字，这个关键字就是查询问题的query，搜索结果可能包含很多文章，这个文章的标题就是key，文章的内容就是value，对搜索结果进行排序时，其实就是用问题query与标题key进行匹配，计算相关度，即注意力分数。
为什么要除以
进行缩放呢？假设query向量和key向量中的元素都是相互独立、均值为0、方差为1的随机变量，那么这两个向量的内积
的期望为0，方差为向量的维度dk。因此，当dk较大时，点积的方差也较大，这样不同的查询向量与不同的键向量计算出来的分数会相差较大。softmax函数的公式为
对xi求偏导得
对xj求偏导得
这里的
。当方差较大时，可能会出现某个xi远大于或者远小于其他的情况，进而可能导致梯度消失和梯度爆炸。
以如下代码为例，我们先创建一个从-10到9的数组x1，然后在其尾部添加上一个值100，表示可能出现的远大于其他数字的值，之后x1将除以100进行缩放，定义为x2，最后分别绘制两个数组的softmax值，结果如图1-14所示。可以发现，不进行缩放的蓝色的曲线比较陡峭，在两端可能会出现梯度接近于0或梯度过大的值；而进行缩放的红色的曲线相对平缓，梯度的变化不会太剧烈，这样有助于保持模型训练的稳定性。
python
import numpy as np
import matplotlib.pyplot as plt
#定义softmax函数
def softmax(x):
e_x=np.exp(x-np.max(x)) # for numerical stability
return e_x/e_x.sum()
#创建一个从-10到9的数组，并在末尾添加一个值100
x1=np.arange(-10, 10)
x1=np.append(x1, 100)
#计算softmax值
y1=softmax(x1)
#创建另一个从-10到9的数组，在末尾添加一个值100，然后除以100进行缩放
x2=np.arange(-10, 10)
x2=np.append(x2, 100)
x2=x2/100
#计算softmax值
y2=softmax(x2)
plt.figure(figsize=(10, 6))
plt.plot(range(1, len(x1)+1), y1, marker="o", color="blue", label="Original")
plt.plot(range(1, len(x2)+1), y2, marker="o", color="red", label="Scaled")
plt.xlabel("Index")
plt.ylabel("Softmax Value")
plt.title("Softmax Values of Two Vectors")
plt.grid(True)
plt.legend()
plt.show()
图1-14 关于softmax的输入是否进行缩放的函数值对比
为了更直观地理解注意力分数，我们来看一个例子。假如我们要翻译一句话“The animal didn’t cross the street because it was too tired.”，其中文意思是“动物没有过马路，因为它太累了。”，那么在翻译的时候，“it”指的到底是什么呢，是animal还是street？对于人类来说，我们可以很轻易地意识到“it”指的是动物（animal）。然而，对于算法来说，确定“it”指的是什么可能是一个复杂的问题。引入了注意力机制之后，在处理这个单词时，模型就可以使“it”与“animal”建立更强的联系。事实上，如图1-15所示，我们可以通过谷歌提供的Tensor2Tensor Notebook来交互查看注意力分数，训练完成后，结果确实如我们前面所述。
图1-15 注意力机制中“it”的注意力分数可视化示意
最后需要注意的是，每个单词都会对应查询、键和值三个向量，但是在实际的代码中使用的是整个输入序列的矩阵。输入X矩阵，它的每一行就对应输入句子中的每一个单词，然后分别乘以WQ、WK、WV这3个矩阵，得到Q、K、V三个矩阵，接着计算注意力分数矩阵，之后经过softmax函数，再乘以V矩阵，就得到了Z矩阵，这个Z矩阵就是要发送给前馈神经网络的矩阵。整个流程如图1-16所示。
图1-16 自注意力机制的计算流程
2.多头注意力机制
在Transformer模型的论文中还使用了一种加强注意力的设计——多头注意力（MHA）机制，进一步完善了注意力层。多头注意力机制可以让模型同时关注不同表示子空间的信息。例如，当我们看一篇文章的时候，会更注意标题和粗体的文字，而不是正文那种又小又密集的文字，也会更注意颜色鲜艳的文字，比如红色的标题。这里的字体和颜色就是两个表示子空间，我们如果能同时关注字体和颜色，那么就可以更有效地定位文章中强调的内容。同理，多头注意力机制就是综合利用各个方面的信息，从多个表示子空间中汇总重要特征。
要使用多头注意力机制，就要有多组query、key、value权重矩阵。在标准的Transformer模型中使用了8个注意力头，因此每个编码器和解码器的注意力层都有8个权重集合，这些集合中的每一个参数矩阵在开始训练时都是随机初始化的。模型训练完成后，每个权重集合可以将输入特征向量投影到不同的表示子空间中，模型可以通过Q、K、V在不同的空间去学习特征，从而避免使用同一种注意力机制时可能产生的偏执，让语义拥有更多元的表达。
多头注意力要为每个头维护单独的WQ、WK、WV权重矩阵，首先用X乘以WQ、WK、WV矩阵以产生Q、K、V矩阵，然后进行与前面提到的注意力机制相同步骤的计算，只是使用的权重矩阵不同，最终会得到8个不同的Z矩阵。但是前馈神经网络不能处理8个矩阵，因此我们需要再做一个矩阵运算以将这8个矩阵变换成1个矩阵。整个流程如图1-17所示。
图1-17 多头注意力机制流程
在结构上，多头注意力就是由多个点积注意力模块组合而成的，可以表示为
，多头中的h个点积注意力模块是可以并行的，它们之间没有依赖关系，因此可以进一步提升效率。多头注意力机制还可以提升模型效果，不过论文中并没有给出更清晰的理论支持，只是在实际训练时发现效果更好。这也是AI科研论文的一个特点，研究员经常凭借强烈的科研意识和敏锐性，发现一些新的研究方向，并通过实验证实其有效性。而这些方向并不一定能够得到完美的理论支持，这也为后续研究者提供了改进的空间。MQA和GQA就是其中的两种优化方案。
如图1-18a所示，多头注意力机制的每个头都有自己独立的查询（Q）、键（K）和值（V），在进行计算时，每个头都要依据自己的Q、K、V进行计算，这会占用大量的存储空间，并且其占用空间规模是随着模型隐藏层维度的增加而成倍增加的。
为了解决这个问题，有一种新提出的改进模型，称为多查询注意力（MQA）机制。MQA的设计思想是，让查询保持原来的多头设计，但键和值则只有一个头，如图1-18c所示。在这种设计中，所有的Q头共享一组KV头，因此得名“多查询”注意力。尽管这种设计在某些情况下可能会对模型性能产生一定的影响，但基于其各种优势，这种微小的性能降低是可以接受的。实验发现，MQA模型通常可以提高30%到40%的处理效率。
MQA模型为何能提高处理效率呢？主要原因在于它降低了KV缓存的大小。虽然从运算量上看，MQA和MHA的计算复杂度是差不多的，但由于MQA模型只需要读取一组KV头然后供所有Q头使用，因此，这种设计在内存和计算之间存在不对称性的情况下具有明显的优势。具体来说，MQA模型的这种设计可以减少需要从内存中读取的数据量，从而缩短了计算单元的等待时间，提高了计算效率。同时，由于KV缓存的大小减小了，显存中需要保存的张量的大小也相应减小，这为增大批处理量留出了空间，进一步提高了显存利用率。
除了MQA模型外，还有一种折中的解决方案，即分组查询注意力（GQA）机制，如图1-18b所示。GQA是MHA和MQA的一种混合模型，其目的是在不过分损失性能的前提下，尽可能地获取MQA模型的推理加速优势。在GQA模型中，不是所有的Q头共享一组KV，而是分组的一定数量的Q头共享一组KV。这种设计既兼顾了性能，也考虑了推理速度。
图1-18 MHA、GQA和MQA的对比
在后面我们将介绍的Llama 2的论文中，给出了三种多头注意力机制的效果对比，如图1-19所示，可以发现，GQA的效果相对来说是比较好的，MHA的效果次之，MQA的效果则要差一些。
图1-19 Llama 2论文中MHA、GQA和MQA在不同任务上的效果
3.注意力层
Transformer模型的解码器使用的是注意力层，基本的计算步骤与自注意力层相似，只不过其中的键矩阵和值矩阵是由编码器提供的上下文向量。编码器首先处理输入序列，然后将顶部编码器的输出转换为一组注意力矩阵——K和V，这两个矩阵将在每个解码器的注意力层中使用。
在解码阶段，我们需要给定一个初始的随机向量，用于表示开始生成序列。类似于我们在编码器阶段对输入序列进行处理的方式，解码器在生成序列的过程中也会添加位置编码向量，用于指示每个单词在序列中的位置。解码器每一步生成序列时只输出一个元素，而该元素会作为下一步的输入再次输入到解码器中。这样，解码器会逐步生成完整的序列。如图1-20所示，在注意力层解码器每次都会根据输入元素构建Q矩阵，与编码器发送过来的K矩阵和V矩阵共同进行注意力计算，这样有助于解码器在生成时注意到输入序列中的元素位置。解码器不断重复该生成过程，直到出现一个表示生成结束的特殊符号，表明解码器已完成输出。
图1-20 Transformer模型解码器输出步骤示意
解码器中的自注意力层的操作方式与编码器中的自注意力层的计算方式略有不同：在解码器中，自注意力层只允许关注输出序列中出现较早的位置。这个限制是通过在计算注意力分数时，在应用softmax函数之前，对未来位置的分数设置为负无穷（-inf）来实现的。这样一来，在softmax步骤中，未来位置的注意力分数会趋于0，从而使得解码器只关注过去和当前的位置。
输入矩阵和掩码矩阵具有相同的维度，在掩码矩阵中，遮挡位置的值被设置为0。如图1-21所示，对于单词0，它只能使用自身的信息；而对于单词1，它可以使用单词0和自身的信息。换句话说，掩码矩阵的作用是限制在注意力计算中使用的信息范围。当某个位置的掩码值为0时，该位置的相关信息应被排除在注意力计算之外。因此，掩码矩阵确保了每个位置只能关注它之前的位置，不会使用未来位置的信息。
图1-21 输入矩阵与掩码矩阵示意
通过输入矩阵X，我们计算得到了Q、K和V三个矩阵。然后我们将Q矩阵与K的转置矩阵KT进行矩阵乘法，得到了Q·KT矩阵。接下来，我们对Q·KT矩阵与掩码矩阵按位相乘，得到了Mask Q·KT矩阵。然后，我们对Mask Q·KT矩阵进行softmax操作，使得Mask Q·KT矩阵的每一行的值相加为1。最后，我们将完成softmax操作的Mask Q·KT矩阵与V矩阵进行矩阵乘法运算，得到最终的输出矩阵Z。整个流程如图1-22所示。
图1-22 注意力机制中的Q、K、V矩阵计算与掩码操作流程