Self-Attention Transformer完全指南：像Transformer的创作者一样思考

本文从RNN到self-attention，再到Transformer来讲清楚整个算法。

近半年来有大量同学来找我问Transformer的一些细节问题，例如Transformer与传统seq2seq RNN的区别、self-attention层的深入理解、masked self-attention的运作机制；以及各种Transformer中的思路如何运用到自己的算法中，例如Transformer是怎么实现并行化的，decoder是怎么用cross-attention把context vector整合进来的，等等。在这篇文章中，我们将从RNN到self-attention，再到Transformer来讲清楚整个算法。注意，这不是一篇纯粹的科普文，更像是一个hands-on的作业，文中会大量出现要求读者思考的问题，尽力让同学们自己去想想如何一步步做出Transformer的整个结构，以至于看完后你不仅能从头手撸一个Transformer，还能理解它的各个细节是如何想出来的。

1 RNN的长距离依赖问题

RNN的长距离依赖问题一直非常棘手。RNN的结构是这样的：

上图中，输入序列是x，RNN会先生成一个初始的隐藏状态 $a^{<0>}$ ，然后这个把这个隐藏状态和对应的输入序列中的token做几个矩阵乘法，得到一个个输出序列的token $y^{<i>}$ ，同时更新这个隐藏状态。注意, RNN本质是一个MLP，也就是说，上图中的蓝色方框, 其实不是多个模型，而是一个模型。输入的序列逐个token流经同一个MLP。

相信大家都敏锐地发现了，这个模型的最大问题，就是这个隐藏状态的更新。由于只存在一个隐藏状态张量，当序列很长的时候，模型很自然地就会忘记了输入序列里面之前的token，因此信息对于整个输入序列是不对等的。我从头手撸了一下RNN并进行了实验，发现RNN总会将注意力放到token的后面位置，结论与我们的推断是相符的。这个观察其实和online learning中涉及的灾难性遗忘现象是同源的。

2 LSTM的时序依赖问题

发现这个问题后，后人用LSTM和GRU这样的模型进行了优化。LSTM的大致示意图如下：

这个模型的宗旨很简单：留住前面看过token的信息。除了隐藏状态和输入，LSTM还在顶部开了一个“后门”，专门用于存储之前看到过的知识。这个记忆张量本质上是另一种形式的隐状态张量。这种双隐方法有时候会减慢模型的收敛速度，但是确实可以有效保证长距离情况下相对持久的记忆。GRU的思想和LSTM差不多，在此不再赘述。

尽管LSTM有效减轻了长距离依赖问题，它依然存在信息不对等的问题。如果你仔细看上面这幅图，会发现尽管保留了记忆，输入序列里面比较靠前的token蕴含的信息仍然不可避免地被压缩了。LSTM确实减轻了问题，但是没有从根本上解决问题。如果仔细思考一下，会发现要从根本上解决问题，需要平等地看待输入序列中的每一个token，却又知道它们的位置信息。如果让你去设计一个模型，你会怎么做？没错，需要完全抛弃这个与时序有关的隐藏状态。

LSTM的另一个问题在于，它的扩展性和并行性是很差的。自己训练过RNN一族模型的同学都知道，RNN等模型在GPU上的加速是非常有限的，因为需要一个token一个token往后迭代，这一个时刻的隐状态必须由上一个时刻的隐状态推出，这就是时序依赖问题。由于时序依赖问题，训练RNN一族模型没法大规模扩展，因为模型参数量一旦大起来，就需要大量的训练数据和epoch。RNN一族根本快不起来，而部署讲究一个高效，做研究也要占坑，动作慢是非常吃亏的。这个时候整个领域都在迫切期待可以并行化的RNN的出现。

3 Self-Attention：并行的RNN

注意，attention和self-attention不是一个东西。attention是一个广泛的思想，包含self-attention、cross-attention、bi-attention等。本文着眼于Transformer，主要讨论self-attention和cross-attention。

self-attention机制其实在Transformer提出的时候已经应用在各种模型上了，但是效果一直有待提升。self-attention也有很多分类，比如单层attention，多层attention，还有多头attention。它们本质上都是一样的，无非是加几层或者加几个分支，我们从单层attention入手讨论。

3.1 self-attention的单层算法

单层self-attention的示意图如下所示。假设输入序列是 $x_{1,2,3}$ , 先通过embedding层得到序列的 embedding $e_{1,2,3}$ , 然后注意了, 我们定义三个矩阵c来作为模型参数。对于token embedding $e_{1}$ , 要先将 $e_{1}$ 与 $Q,K$ 矩阵相乘得到两个向量 $q_{1},k_{1}$ , 这两个向量再互相乘得到一个初步的attention值 $at_{11}$ （attention score，常写作 $\alpha$ ）。注意, 这个 $at_{11}$ 是一个值, 不是一个向量。对所有的attention score做一个softmax, 可以得到归一化的 $st_{11}$ 。同时, $e_{1}$ 要与矩阵 $V$ 相乘得到值向量 $v_{1}$ 。将归一化后的qk token和 token相乘, 就能得到一个qkv token $wt_{11}$ 。对 $q_{1}$ 去乘输入序列中的第二个、第三个token重复用的 $K$ , 就能得到 $wt_{12},wt_{13}$ 。加和这三个 token, 就能得到 $b_{1}$ 。对 $q_{2},q_{3}$ 重复这个过程，就能得到 $b_{2},b_{3}$ 。这就是self-attention的核心算法。

上面这张图有几个值得注意的地方：

self-attention layer的所有参数只有 $W^{Q},W^{K},W^{V}$ 这三个矩阵。
对于每一个输入token, 它对应的输出token其实是所有token（包括自己）的key-value pair和自己的query的乘加。
$q,k,v$ 本质上都是相应token $x$ 在隐空间(latent space)的某种语义表示。token $x$ 通过 $W^{Q},W^{K},W^{V}$ 进入隐状态 $q,k,v$ 。
输出的token并不依赖任何上一时刻的隐层状态(hidden state)。

其中，第四点肯定会让很多同学有疑问。相对于RNN，self-attention不需要任何随token顺序传递的隐层状态。这是因为，它使用一种positional encoding方法直接修改token的embedding，从而让模型感知token在序列中的相对位置，我们在第4节会详细说明。

3.2 self-attention的矩阵化

我们惊喜地发现, 根据第四点, 这个算法可以用矩阵并行化。 把 $e_{1,2,3}$ 这三个向量拼到一起成为一个矩阵, 就能得到下面这张图。其中, 由于我们的输入序列有3个token, 左边的Inputs矩阵有 $n=3$ 。这个输入矩阵被三个参数矩阵乘后会得到 $Q,K,V$ , 其中 $Q,K$ 可以使用 $QK^{T}$ 相乘来得到attention score matrix。这里要注意, 在上图中, self-attention看上去是分别对 $x_{1},x_{2},x_{3},$ 这三个token做的, 但事实上把它们的embedding vector拼起来就可以并行了。同理, 上图中的 $at_{11,12,13}$ 拼起来就是这个attention score matrix的一行（思考一下, 为什么是一行不是一列? ）。接下来对这个attention score matrix做row-wise softmax, 就能得到每一行加和均为1的归一注意力分数矩阵 $A$ 。这个矩阵 $A$ 的形状是 $n\ast n$ , 非常有意思, 我们接下来会重点看看这个 $A$ 矩阵到底有什么数学含义。最后我们将矩阵 $A$ 和 $V$ 相乘得到最后的矩阵 $Z$ , 形状是 $N\ast d_{v}$ , 意义是 $n$ 个token, 每个token此时都获得了一个长度为 $d_{v}$ 的值。

现在如果再放出self-attention的官方公式，相信大家就很好理解了：

其中, $Q,K,V$ 是输入矩阵 $X$ 分别乘以三个参数矩阵后的隐状态矩阵, $d_{k}$ 是 $W_{k}$ 矩阵的列数。

这里有一个小细节: 为什么需要除以一个 $\sqrt{d_{k}}$ 呢? 其实经常调参的同学已经很熟悉了, 这就是为了去掉方差的影响, 其实stable diffusion也借鉴了这个思想（见如何通俗理解扩散模型一文）。我们拿 $Q$ 矩阵中的任意一列 $q_{i}$ ， $K$ 矩阵中的任意一行 $k_{j}$ 出来, 如果 $q_{i}$ 和 $k_{j}$ 中每一个元素均为均值为0, 方差为1的独立同分布的随机变量, 则 $q_{i}k\tfrac{T}{j}$ 这个随机变量（也就是token $i$ 看token $j$ 的注意力分数）中的每个元素也有均值为 0 , 方差为1, 这是根据公式 $Var(x_{1},x_{2})=Var(x_{1})\cdot Var(x_{2})+Var(x_{1})\cdot E(x_{2})^{2}+Var(x_{2})\cdot E(x_{1})^{2}$ 得到的。而由于 $q$ 和 $k$ 独立同分布, 我们有 $E\left [X +Y \right ] = E\left [X \right ]+E\left [Y \right ] =0$ 和 $Var\left [X +Y \right ] = Var\left [X \right ]+Var\left [Y \right ] =\sum_{m}^{d_{k}}q_{i,m}k_{d,m}=\sum_{m}^{d_{k}}1=d_{k}$ 。也就是说, 每一对这样的qk对都会产生一个期望为 0 , 方差为 $d_{k}$ 的新向量, 这个新向量表示token $i$ 看token $j$ 的注意力分数。当 $d_{k}$ 很大的时候, 这个向量的方差会很大, 而特别大的值会在softmax层被挤到边缘去, 这样回传的梯度就很小了。这个时候就可以根据 $Var(kx)=k^{2}Var(x)$ 来用1/ $\sqrt{d_{k}}$ 将方差从 $\sqrt{d_{k}}$ 缩小为 1 。

3.3 self-attention的本质

透过方法看本质, 我们来讨论一下上面这个self-attention官方公式到底怎么理解。首先要确认一个事实, $Q,K,V$ 这三个矩阵其实都是输入矩阵 $X$ 的某种线性变换, 也就是 $X$ 在隐空间的某种语义表示。换句话说, 丢掉 $W^{Q},W^{K},W^{V}$ 这三个矩阵也是可以训练的, 但是复杂度不够, 效果会受到影响。因此, 为了方便说明, 我们用 $X$ 来表示这三个矩阵。另外, 1/ $\sqrt{d_{k}}$ 是用于scale的, 所以也不涉及本质，扔掉。那么，官方公式又可以写为：

假设一个句子Welcome to Starbucks，如果embedding层采用简单的2-hot编码（例如编码Welcome为1010），可以把输入矩阵表示为下图左边所示的样子。这个矩阵和自身的转置相乘，就能得到一个本质上是attention matrix的矩阵，如下图右边那个矩阵所示。

这个attention matrix有什么含义呢？因为向量相乘的本质是计算相似度，如果我们看第一行，可以发现这个row本质上在计算Welcome这个token和句子中所有token的相似度。到这里其实接触过推荐的同学已经懂了，词向量之间相似度的本质是关注度。如果一个token A和另一个token B经常一起出现，那A和B之间的相似度往往很高。例如，上图中Welcome这个词和自身“Welcome”、第三个token “Starbucks”的相似度就高，说明在推测Welcome这个token时应当对这两个token给予更高的关注。

将这个结果用softmax归一化，就能得到下图右侧所示的归一化attention matrix矩阵。归一化后，这个attention matrix就变成了系数矩阵，可以直接和原矩阵相乘了。

最后一步是将归一化后的attention matrix $\alpha$ 右乘输入 matrix $X$ , 得到结果 $\hat{X}$ , 如下图所示。这一步本质上做了什么呢? 我高亮了左边矩阵的第一行, 这一行会和输入矩阵 $X$ 的每一列乘加, 然后算出输出矩阵第一行的每个值。由于 $\alpha$ 矩阵的第一行是Welcome这个token对于所有token的注意力值, 输出矩阵 $\hat{X}$ 的第一行也就变成了注意力加权后的Welcome这个token的embedding。

总结一下，输入一个矩阵 $X$ , self-attention会输出一个矩阵 $\hat{X}$ ,这个矩阵是输入矩阵的注意力加权后的语义表示矩阵。上面这个过程，超详细图解Self-Attention这篇文章也讲的很清楚，如果还有疑问可以去这篇文章看看。

3.4 self-attention的Q, K, V思想

虽然 $W^{Q},W^{K},W^{V}$ 这三个矩阵都不是必要的，我们知道可以全部用 $X$ 来实现注意力的加权，但是最后效果会差很多。做算法的同学肯定会好奇QKV这一套思想到底怎么去理解、可能是怎么提出的，这样也许可以找到一些灵感放到自己的工作里面去。

相信大家一直都有个疑问，为什么要取QKV这些名字。Q是Query（查询），K是Key（键/钥匙），V是Value（值）。在数据库中，我们希望通过一个查询Q去数据库中找出相应的值V（for 做推荐的同学：这个值V在推荐中叫Document，就是双塔里面的Query-Document架构那个Document）。问题是，直接通过Query去搜索数据库中的V往往效果不好，我们希望每个V都对应一把可以找到它的钥匙K，这个K提取了V的特征，使得它更容易被找到。注意，一把钥匙对应一个值，也就是说K和V的数量相同。

参考下图。K提取的特征是如何获得的呢？根据attention的思想，需要先看过所有项才能准确地定义某一个项。因此对于每一个查询语句Q，我们取出所有钥匙K（ $\alpha$ 的第一行）作为系数，将它乘上对应的V（ $V$ 的第一列）算出加权值，得到这个查询语句期望的结果 0.67 。这样反向传播后，K就能逐渐学习到V的特征。

因此，self-attention中的 $Q,K,V$ 思想，本质是一个具有全局语义整合功能的数据库。而最后得到的矩阵Y , 就是输入矩阵X注意了上下文信息得到的在隐空间的语义矩阵，每一行代表一个token。

4 Positional Encoding：位置信息的整合

仔细思考的同学可以发现，self-attention本身只是三个存储语义映射的矩阵，现在还不能捕捉token的位置信息，整个输入序列仍然处于天涯若比邻的状态。这个时候已经可以在输出矩阵 Y后面接一个MLP来做分类任务了，但是缺乏位置信息使得效果并不好。如果让你自己设计一个算法来解决这个问题，你会怎么做？

首先是如何将positional encoding融合到self-attention算法中。显然，我们可以直接改输入矩阵X也可以改self-attention的算法。其中，改输入矩阵更直观一些，只要让positional encoding生成的positional embedding和token embedding一样长，就能直接相加获得与原来相同大小的新输入矩阵，不需要任何self-attention算法里面的修改。我们采用这种方法继续思考：如何进行编码呢？

没错，一个很直观的解决方法就是将token在序列中的位置从0到1编码，然后以某种方式融入到self-attention层中。但是，从0到1编码的序列有一个重大问题：不同长度的句子，相邻的token之间的距离是不同的。例如Welcome to Starbucks的编码是[0, 0.5, 1]，间隔是0.5，而Today is a very clear day的编码是[0, 0.2, 0.4, 0.6, 0.8, 1]，间隔是0.2。这样非归一化的间隔很明显会降低模型的性能，因为self-attention层中的参数是用序列长度无关方法训练出来的，引入序列长度相关的任何参数都会导致参数目标的异化。

还有一个想法，是固定间隔为1，然后逐渐往上增长。例如Welcome to Starbucks的编码是[0, 1, 2]，Today is a very clear day的编码是[0, 1, 2, 3, 4, 5]。这样的方法虽然解决了序列长度相关参数的问题，但是当序列很长时参数会非常大，导致梯度消失，其实和上面说的加一个1/ $\sqrt{d_{k}}$ 项是一个问题。

这里就引出了Transformer使用的方法，其实也是一个在应用数学界非常常用的位置编码方法：余弦位置编码。对于一个给定的token在序列中的位置t，和embedding长度d，余弦位置编码的数学表达是：

其中, c是token在序列中的位置除以2，d是embedding的维数。注意, 这个 $p_{t}$ 的长度是和embedding的长度相同的, 所以 $p_{t}$ 又叫positional embedding, 往往可以直接和 semantic embedding相加。我们可以将t扩展到整个序列, 得到一个矩阵。例如, 如果输入句子有4个token "I am a Robot", 且我们取embedding size d=4, 则可以得到如下图所示的4x4的位置编码矩阵。其中, 每一行都是上面这个公式的一个 $p_{t}$ 。

第一行是"I"这个token的positional embedding。由于在序列中的位置为 0 , 有 t=0, 所以 $p_{0}$ 的第一个item就是 $sin(w_{1}t)$ 也就是 $sin(\frac{1}{10000^{2*0/4}}*0)=0$ , 第二个item就是 $cos(w_{1}t)$ 也就是 $cos(\frac{1}{10000^{2*0/4}}*0)=1$ 。以此类推, 就能得到每个token的positional embedding。

如何理解这个方法呢？其实这个方法就是利用了余弦函数的周期性。实际上，只要能找到多个周期不同但特征相同（例如余弦函数的波形都是相同的，这就是特征相同）的函数，理论上都可以用来做positional encoding。例如，我们可以用二进制位置编码法来代替余弦相似编码法。如下图，假设有一个序列有16个token，且embedding size为4，我们就能用二进制编码法得出每个token的positional embedding：

显然，最低位（红色）的01变化非常快（周期为2），而最高位（黄色）的01变化最慢（周期为8）。也就是说，每个positional embedding位的周期不一样，因此不会产生重复的embedding。实验后，这种方法也有效果，但没有余弦位置编码法的效果好。

5 Transformers：打通seq2seq任督二脉

到此为止，self-attention已经可以做分类任务了，但是还没法做seq2seq任务。要做分类任务，只需要在self-attention前把positional embedding加到 XX 上来获得新的embedding，再在self-attention后接一个MLP头就行。然而，seq2seq任务往往需要encoder-decoder架构。如何用self-attention layer来实现encoder-decoder架构呢？想一想，如果你是Transformer算法的设计者，你很可能首先会去借鉴RNN一族是怎么实现seq2seq架构的。

下图表示RNN的encoder-decoder架构。可以发现，在encoder部分，RNN做的事和分类任务是一模一样的。但是在最后一个token处，RNN并没有接上MLP，而是把context vector丢给decoder。decoder是另一个RNN（参数不同），输入是context vector、上一个时刻预测的输出token、一个隐层向量。注意，这个隐层向量是在decoder阶段重新初始化的，不是encoder阶段拿过来的。每次decode时，我们将context vector和input token embedding拼起来一起送到RNN中，同时送入上一个时刻的隐向量，这个RNN会输出一个隐向量和一个输出向量。输出向量会过softmax获得每个token的probability，然后取得probability最大的token的输出，而隐向量会重新被输入这个RNN里面，从此RNN进入下一个时刻。当RNN输出<eos> token的时候，decode过程结束，我们就能得到一个生成的序列，而这个序列与输入序列不一定是等长的。

总结一下，RNN的seq2seq solution其实就是在encoder阶段获得整个句子的表征context vector，然后把这个表征送到decoder里面，在生成每个token的时候都要看一遍这个表征。另外，由于RNN必须一个一个token去看，隐向量依旧是不可缺少的。

现在我们再来想想，如何借鉴RNN来设计出用self-attention搭的encoder-decoder架构，来完成seq2seq任务呢？首先，encoder肯定也要输出看过的句子的表征。其次，decoder的每个token位也要算一个probability，从而取出可能行最大的那个。但是有一个问题，encoder输出的context vector应该如何交给decoder呢？在RNN里，context vector直接和输入的embedding拼在一起放到RNN里；如果用self-attention，是不是也可以把context vector和embedding拼到一起送到self-attention中呢？还是说，要用额外的网络来融合信息？如果使用self-attention，因为它有三个输入，我们又应该将context vector送到哪些输入里面呢？

让我们看一下Transformer的官方实现（建议截一下这个图，在下文中会反复用到）：

这张图相信大家已经看烂了，左边是encoder，右边是decoder。Multi-head Attention层是attention的一个升级，左下角的是multi-head self-attention，右上角的是multi-head cross-attention，右下角的是masked multi-head attention。蓝色的层是FFN，Feed Forward Layer，其实就是个MLP，包括几个全连接层和激活函数（如ReLU）。encoder和decoder都可以往上叠层数，如图中 $N\times$ 的提示。接下来我们具体分析一下每个部分。

5.1 Cross-Attention：Self-Attention的双塔实践

很多同学看到cross-attention会感到害怕，觉得比较难理解。但是如果你已经看完了self-attention的部分，其实cross-attention是一点就通的。为什么self-attention要加一个self？看Transformer的结构图，你会发现输入左下角和右下角attention layer的三个箭头是从一个地方分出来的，说明 $Q,K,V$ 三个矩阵都是从输入矩阵 $X$ 变化来的。那为什么cross-attention要加一个cross？没错，因为这三个矩阵要从不同的输入矩阵变化来。

如下图所示, 我们现在获得了两个输入矩阵 $X_{1},X_{2}$ , 其中, $X_{1}$ 提供线性变换 $Q$ ， $X_{2}$ 提供线性变换 $K,V$ 。cross-attention和self-attention的不同在图中用new注明了。

仔细观察最后得到的矩阵 $Z$ , 会发现它的行数和矩阵 $X_{1}$ 是一致的。思考一下, 如果让你设计 Transformer结构图右上角的cross-attention模块, 你会把encoder的context vector作为 $X_{1}$ 还是 $X_{2}$ ? 如果你看了 3.4节关于self-attention的解释, 会很自然地得出答案： $X_{1}$ 是 decoder上一个 token的embedding, $X_{2}$ 是encoder最后产生的context vector。这是因为encoder的context vector本质是聚合了输入序列信息的一个数据库（V），而decoder的每一个输入token本质是一条查询语句 (Q), 负责查询数据库中与之最相似的（最需要注意的）token。这样的话, 上图中的 $QK^{T}$ 这个矩阵, 每一行都表示decoder的一个输入token对context vector中所有token的注意力。这个attention matrix叫做他注意力矩阵。

5.2 Transformer Decoder的训练和预测

在seq2seq任务中，Transformer encoder显然需要和decoder联合训练，因为这不是个分类问题，encoder没法单独接MLP进行训练。其中，token流经encoder的过程比较简单，用self-attention就能解决，难点是这个decoder。首先要思考一个问题，如果让你设计如何使用这个decoder，你可以做到让decoder一次性得到所有的输出token吗？答案是不可能。一个序列的生成必须有一个结束的token（如[EOS]）来提示，而这个结束的token的生成必须依赖前面已经生成的token，因为打算结束一个句子，必须是因为这个句子已经完成了任务；也就是说，最后一个token的生成必须由前面生成的token来condition。以此类推，前面的每一个token都必须由更前面的token来condition，直到推到起始token（如[SOS]）。这和人说话的逻辑是相似的，人也许可以在脑中构思整个句子，但是表述一定是一个词一个词说出来的，而且后面的词一定会被前面的词所影响，这就是说话的逻辑。因此，decoder在预测的时候依然是一个一个token去预测的，不可能一次性全部输出。这个逐个token输出的过程如下图所示。也就是说，decoder的预测算法无法并行。

但是decoder的训练采用了一个不同的、很精妙的方式，叫做teacher forcing，也就是有老师带着做训练。这其实不是Transformer的创新点，在那个时代已经是一个老生常谈的话题了。teacher forcing是什么意思呢？我们拿RNN decoder的teacher forcing举一个例子。如下图所示，我们假设一个seq2seq的输入是What do you see，正确的输出(label)是Two people running，那么下图左边的训练过程就叫做free-running方法，下图右边的就叫做teacher forcing方法。如果你仔细观察Transformer结构图的右下角，会发现有一个提示是把Outputs做一个Shifted Right，也就是将所有输入token都往右移一个位置。这就对应着下图中右边所示的teacher forcing方法，是把label的所有token向右移一个位置，然后在最左边放上一个表示开始的token（例子中是<Start>）。这种方法和free-running有什么区别呢？很显然，free-running的decoder RNN的输入是上一时刻自己输出的值，而这不一定是正确的；而teacher forcing的decoder RNN的输入是上一时刻的label，这一定是正确的。这种训练方法可以防止错误的累积，从而提高训练效果。

teacher forcing的另一个好处是可以并行化。上面这幅图是RNN的，由于有一个需要传递的隐向量，所以训练阶段的decoder其实无法并行计算。然而，self-attention和cross-attention并不依赖隐向量来传递状态，而是通过输入阶段的positional embedding一步到位，且训练阶段使用teacher forcing，不需要上一状态的输出，所以可以把全量的输入和label直接投到由masked self-attention和cross-attention组成的decoder里面，直接并行计算进行训练。因此，通过使用teacher forcing，decoder的训练算法可以并行。

在decoder的预测和训练阶段要注意一个细节上的差异。在预测阶段，假设我已经预测了Welcome to，那么接下来在预测Starbucks的时候，decoder必须看到之前的所有张量Welcome to，而不仅仅是上一个张量to。很多同学都忘了这个细节，没有去stack之前的张量，导致生成的句子效果很差，而且decoder一直不输出结束的token。至于训练阶段，由于采用teacher forcing且已经有了一个attention mask，只要shift right后把整个句子扔到decoder里面就已经实现了看到之前所有token的目的了。

5.3 Masked Self-Attention：防止偷看答案

如果你仔细思考，会发现仅仅使用Shifted Right，并不能实现teacher forcing，这是因为如果右下角的attention模块是没有mask的self-attention，会造成数据泄漏的问题。我们拿出3.4节讲过的例子，看看问题出在什么地方。

在decode过程中，每次decoder应当只能看到之前自己生成的token（预测阶段）或者label中已经给出的这个时刻以前的teacher token（训练阶段）。总之，decoder在注意某个token时，不可能注意到它之后出现的token。如果注意到后面的token，说明decoder偷看了答案，这样会影响训练效果。

由于self-attention是矩阵运算，我们需要用mask的方法去强制去除这种bad case。如下图所示，我们可以将attention matrix的上三角矩阵用-inf来mask掉。通过将不希望模型注意到的token设为负无穷，可以让梯度不再流经这个位置（梯度为0），从而消除偷看后面的token的问题。

现在思考一个问题：masked attention后面的cross-attention需要也加一个attention mask吗？答案是不需要，因为encoder可以看到整条输入序列，已经获得了全量信息，所以decoder这一条Q可以看到context vector全部的K和V。换句话说，在训练和预测的时候，我们是允许decoder看到输入的全部信息的。

5.4 Multi-head Attention：扩大参数量和语义分化

前面已经讨论过，单层attention要学的参数其实就是三个矩阵 $W^{Q},W^{K},W^{V}$ ，这个参数量往往是很小的。虽然已经能够满足表示语义的需求，但是当语义逐渐复杂后，容易因为参数量的问题达到容量上限。multi-head attention的提出解决了这个问题。接下来我们来看一下Transformer结构图里的multi-head attention和我们前面说的单层attention有什么区别。

multi-head attention的结构如下图所示。多头注意力机制有几个额外的参数矩阵, 将获得的QKV矩阵进一步线性变换成多个矩阵。例如下图的双头注意力, 就是将 $Q,K,V$ 变换成 $Q1,K1,V1$ 和 $Q2,K2,V2$ 。在获得多个头后，对应的 $Q,K,V$ 头分别做单层注意力, 得到多个输出, 如双头会得到 $b_{11},b_{12},$ 两个输出。不同的头的输出会进一步聚合成一个输出向量 $b_{1}$ 。注意图中的一个细节, 在计算 $b_{1}$ 的时候, 是不会用到 $q_{21}$ 的。思考一下这是为什么? 因为 $b_{1}$ 是 $a_{1}$ 的查询在隐空间的特征表示, 所以不可能用到了 $a_{2}$ 的查询。如果不理解, 需要再复习一下第3.4节。

multi-head attention的结构如下图所示。多头注意力机制将 $W^{Q},W^{K},W^{V}$ 三个矩阵等分为多个小矩阵, 例如如果是双头注意力, $W^{Q}$ 会被分割成两个小矩阵 $W^{q_{1}},W^{q_{2}}$ 。这样, $a_{1}$ 生成的 $q$ 矩阵就也可以被分成两个小矩阵 $q_{11},q_{12}$ , 我们称之为两个注意力头。在获得多个头后, 对应的 $q,k,v$ 头分别做单层注意力, 得到多个输出, 如双头会得到 $b_{11},b_{12}$ 两个输出。不同的头的输出会进一步聚合成一个输出向量 $b_{1}$ 。注意图中的一个细节, 在计算 $b_{1}$ 的时候, 是不会用到 $q_{21}$ 的。思考一下这是为什么? 因为 $b_{1}$ 是 $a_{1}$ 的查询在隐空间的特征表示, 所以不可能用到了 $a_{2}$ 的查询。如果不理解, 需要再复习一下第3.4节。在全部头的计算都结束后, 还需要加一个仿射矩阵 $W^{O}$ , 用于聚合各个头的信息。各个矩阵的形状都在图中标出来了, 我们取句子最长长度为 256 个token, 且 embedding size为1024, 就能得到 $a_{1}$ 的形状为(256,1024)。其他矩阵的形状也在图中标明了。

明白了大致的思想，我们需要深入探索一下多头注意力的作用。头的数量是越多越好吗？我们来看一下Transformer论文中做的ablation，其中表示头的数量：

显然, 当h为 8 时效果最好, 而h太大 (16/32) 并不会显著提升性能, 而 h太小 (1/4) 往往会降低性能, 有研究说明这是因为参数量下来了。但是仔细想想, 增加头的作用真的只是提升参数量吗? 如果只是为了提升参数量, 完全可以扩大 $W^{Q},W^{K},W^{V}$ 的hidden size吧, 为什么要通过增加头来实现呢?

下图是一份研究可视化的不同的头对于注意力矩阵的影响。这份研究是基于encoder的，一共有4层encoder(0-3)，每层encoder都包括6个头(0-5)。图中每一行表示第几层encoder，每一列表示第几个头。

从图中可以发现，同一层中往往有几个头负责选择相同的特征（如第2层第1-4个头，第3层1-4个头），有几个头比较特立独行（如第2层第1个头，第三层第0、第5个头）。不同的头为什么会选择不同的特征？

我们来思考一下多头注意力机制的训练过程。观察本节开始时出现的multi-head attention结构图。由于参数初始化方式不同，有 $q_{11}\neq q_{12}$ 。同理有 $st_{111}\neq s_{121},b_{11}\neq b_{12}$ 。但是由于 $b_{11}, b_{12}$ 是concat起来的，所以梯度回流的时候，这两条路是对称的。那是什么导致了不同的头最后分化出不同的特征选择方式？是不同的初始化方式，导致最后训练出来的头有不同的特征选择功能。 以这个启发为出发点，有研究对初始化方式做了详细的ablation study，最后证明了可以通过改变初始化方法来降低层方差，从而获得更好的训练效果。

也有研究分析了不同的头到底都在关注什么信息。Adaptively Sparse Transformers一文指出，头的功能主要是三个：关注语法、关注上下文、关注罕见词。关注语法的头能够有效抑制不合语法的词的输出；关注上下文的头负责理解句子，并且会更加关注自己附近的词；关注罕见词的头负责抓住句子重点，例如Starbucks就比to罕见的多，因此往往蕴含更多的信息。

我们再思考一个问题，用多层attention（stack几个单层attention）来代替多头attention有什么明显的缺陷吗？有很大的并行化上的缺陷。多头attention是可以轻松并行化的，因为不同的头拿到相同的输入，进行相同的计算；而多层attention由于是stack的结构，必须等下层attention计算完成后才能传导到上层，所以是无法并行的，有几层就有几倍的时间上限。因此，从并行化的角度看，我们往往使用多头attention。

总结一下，multi-head attention可以增大attention层的参数量，同时增加特征提取器的分化程度，从而有效提升attention的性能。头的个数并不是越多越好，但是往往多头会比单头的效果更好。对比多层attention，多头attention方便并行化。

5.5 Feed Forward Layer：非线性能力的引入

如果你第一次看Transformer的文章，可能会不知道Feed Forward Layer是什么，毕竟这个名词已经比较有年代感了。FFN其实就是MLP，也就是几个矩阵和激活函数拼起来的层。在Transformer的结构图里面，FNN貌似是微不足道的，因为只有那么扁扁的一条。但是，如果你仔细计算一下Transformer的参数，会发现FFN原来占了整个Transformer结构的一半还多的参数。

Transformer采用的FNN可以用下面公式来表达, 其中 $W_{1}, b_{1}$ 是第一个全连接层的参数, $W_{2}, b_{2}$ 是第二个全连接层的参数。顺便附上一张比较好理解的FFN构造图。

根据这个公式，很容易完成FFN的实现。下面是PyTorch的实现。其中，d_model是embedding size（Transformer中为512），d_ff是FFN的hidden size（Transformer中为2048）。

class FFN(nn.Module):
    def __init__(self, d_model, d_ff):
        super(FFN, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
    def forward(self, x):
        return self.w_2(F.relu(self.w_1(x)))

现在我们思考一下这个FFN层的作用。输入张量的embedding size这个维度 (512) 被映射到一个更大的维度hidden size（2048），然后在下一层又映射回原来的embedding size（512）。比较显然的是，FFN明显可以用于增加模型的非线性，因为中间夹了一个relu激活函数。另外，FFN明显增加了模型的容量, 因为参数量极大增长。我们计算一下这个FFN的参数量, 发现竟然有 $2*512*2048=2,097,152$ 这么大（忽略bias）。再回头看看multi-head attention的参数量, 文章(https://towardsdatascience.com/how-to-estimate-the-number-of-parameters-in-transformer-models-ca0f57d8dff0）指出, Transformer论文提出的 8 头注意力网络需要的参数量为 $1,050,624$ , 计算公式如下所示：

如果想要自己求证，也可以使用PyTorch快速获得答案：

d_model = 512
n_heads = 8
multi_head_attention = nn.MultiheadAttention(embed_dim=d_model, num_heads=n_heads)
print(count_parameters(multi_head_attention))  # 1050624
print(4 * (d_model * d_model + d_model))  # 1050624

也就是说，两个multi-head attention的大小才赶上一个FFN的大小。由于Transformer结构中有3个multi-head attention和2个FFN，所以FFN占了过半的参数。

5.6 Residual Network与Layer Norm

观察Transformer的结构图，会发现大量出现了Add & Norm的模块，如下图所示。事实上，每经过一个计算层，Transformer就会过一个Add & Norm模块。这里的Add指的是residual addition（残差相加），也就是经过一个模块前的输入加上经过一个计算层后的输入来得到一个新的向量，用公式可以表达成 $y=x+f(x)$ , 这个就是计算层的函数。

在Transformer提出的时候，残差连接已经成为了主流的防止梯度消失的方法，最早出现在ResNet中，相信不管是做CV还是NLP的同学都已经对ResNet耳熟能详了。例如，FFN中出现了ReLU函数，在梯度回传的时候，有差不多一半的信号是会变成0的，这样就造成了比较大的信息损失。残差连接保留了经过ReLU之前的向量信息，因此能够有效缓解这样的情况。

Add & Norm中，在残差相加后是一个layer normalization模块。layer normalization也不是什么新鲜事，在Transformer提出的时候已经被广泛采纳了，主要用于向量的归一化。这里我们主要将layer normalization与batch normalization做一个简单的对比。

如下图所示, 给定一个三维张量(embedding size, token number, batch size), batch norm是对一个batch做归一化, 而layer norm是对一条序列做归一化。例如在batch size为2的情况下，我们输入了两个序列: “你好"与“机器学习”。假设“你好"的embedding为 $\left [ 4,6,0,0 \right ]$ , “机器学习”的 embedding为 $\left [ 1,2,3,3 \right ]$ 。batch norm对这一个batch中的每一个对应维度做归一化, 也就是说, 会对两个序列中对应的每个token和embedding做归一化, 这显然是不合理的, 因为这两个序列的token并不是一一对应的。如果我们使用和为1 的归一化方法, “你好"会获得归一化后的 embedding $\left [ 0.8,0.75,0,0 \right ]$ , 而“机器学习”会获得归一化后的embedding $\left [ 0.2,0.25,1,1 \right ]$ 。这反而破坏了原本embedding内的语义信息, 例如本来“你好”的embedding是 $4<6$ , 被batch norm后却变成了 $0.8>0.75$ , 语义发生了改变。

layer norm对每一条序列内部做归一化, 而每条序列内部显然是对应的。在上面的例子中, 如果使用layer norm, “你好"会获得归一化后的embedding $\left [ 0.4,0.6,0,0 \right ]$ , 而“机器学习"会获得归一化后的embedding $\left [ \frac{1}{9},\frac{2}{9},\frac{1}{3},\frac{1}{3} \right ]$ 。这样, 归一化既保留了原来的语义, 又能够有效防止值过大导致的梯度问题。