Transformer原理

Transformer由论文《Attention is All You Need》提出。学（烤）习（贝）一下。实现声明，本文大部分内容来源于Transformer模型详解（图解最完整版），对于不理解的地方，我会加上个人注解。

Attention简介

Attention注意力机制

从“Attention”这个名字可以读出，Attention机制主要是对注意力的捕捉。Attention的原理与大脑处理信息有一些相似。比如看到下面这张图，短时间内大脑可能只对图片中的“锦江饭店”有印象，即注意力集中在了“锦江饭店”处。短时间内，大脑可能并没有注意到锦江饭店上面有一串电话号码，下面有几个行人，后面还有“喜运来大酒家”等信息。

所以，大脑在短时间内处理信息时，主要将图片中最吸引人注意力的部分读出来了，类似下面。

Attention的输入由三部分构成：Query、Key和Value。其中，(Key, Value)是具有相互关联的KV对，Query是输入的“问题”，Attention可以将Query转化为与Query最相关的向量表示。

Attention的计算主要分3步，如下图所示。

第一步：Query和Key进行相似度计算，得到Attention Score；

第二步：对Attention Score进行Softmax归一化，得到权值矩阵；

第三步：权重矩阵与Value进行加权求和计算。

Query、Key和Value的含义是什么呢？我们以刚才大脑读图为例。Value可以理解为人眼视网膜对整张图片信息的原始捕捉，不受“注意力”所影响。我们可以将Value理解为像素级别的信息，那么假设只要一张图片呈现在人眼面前，图片中的像素都会被视网膜捕捉到。Key与Value相关联，Key是图片原始信息所对应的关键性提示信息，比如“锦江饭店”部分是将图片中的原始像素信息抽象为中文文字和牌匾的提示信息。一个中文读者看到这张图片时，读者大脑有意识地向图片获取信息，即发起了一次Query，Query中包含了读者的意图等信息。在一次读图过程中，Query与Key之间计算出Attention Score，得到最具有吸引力的部分，并只对具有吸引力的Value信息进行提取，反馈到大脑中。就像上面的例子中，经过大脑的注意力机制的筛选，一次Query后，大脑只关注“锦江饭店”的牌匾部分。

再以一个搜索引擎的检索为例。使用某个Query去搜索引擎里搜索，搜索引擎里面有好多文章，每个文章的全文可以被理解成Value；文章的关键性信息是标题，可以将标题认为是Key。搜索引擎用Query和那些文章们的标题（Key）进行匹配，看看相似度（计算Attention Score)。我们想得到跟Query相关的知识，于是用这些相似度将检索的文章Value做一个加权和，那么就得到了一个新的信息，新的信息融合了相关性强的文章们，而相关性弱的文章可能被过滤掉。

cross-attention与self-attention

看完下文之后，就知道了Encoder部分中只存在self-attention，而Decoder部分中存在self-attention和cross-attention。这里提前解释一下，看完下文之后可以再回味这一小节。

self-attention：encoder中的self-attention的query、key、value都对应了源端序列，decoder中的self-attention的query、key、value都对应了目标端序列。
cross-attention：decoder中的cross-attention的query对应了目标端序列，key、value对应了源端序列(每一层中的cross-attention用的都是encoder的最终输出)

也就是说，Attention connecting between the encoder and the decoder is called cross-attention since keys and values are generated by a different sequence than queries. If the keys, values, and queries are generated from the same sequence, then we call it self-attention. The cross-attention mechanism allows output to focus attention on input when producing output while the self-attention model allows inputs to interact with each other.

下图是cross-attention的示意图。

下图是self-attention的示意图。

核心思想

RNN缺陷

在没有Transformer以前，大家做神经机器翻译用的最多的是基于RNN的Encoder-Decoder模型：在没有Transformer以前，大家做神经机器翻译用的最多的是基于RNN的Encoder-Decoder模型：

上图中的 Encoder-Decoder 模型将德语短语转换为英语短语。让我们把它分解一下：

编码器和解码器都是递归神经网络。
在编码器中的每个时间步骤，递归神经网络从输入序列获取词向量（$x_i$），从前一个时间步骤中获取一个隐状态（$H_i$）。
隐状态在每个时间步骤中更新。
最后一个单元的隐状态称为语境矢量（context vector）。它包含有关输入序列的信息。
然后将该语境矢量传递给解码器，然后使用它生成目标序列（英文短语）。

Encoder-Decoder模型当然很成功，在2018年以前用它是用的很多的。而且也有很强的能力。但是RNN天生有缺陷，只要是RNN，就会有梯度消失问题，核心原因是有递归的方式，作用在同一个权值矩阵上，使得如果这个矩阵满足条件的话，其最大的特征值要是小于1的话，那就一定会出现梯度消失问题。后来的LSTM和GRU也仅仅能缓解这个问题Encoder-Decoder模型当然很成功，在2018年以前用它是用的很多的。而且也有很强的能力。但是RNN天生有缺陷，只要是RNN，就会有梯度消失问题，核心原因是有递归的方式，作用在同一个权值矩阵上，使得如果这个矩阵满足条件的话，其最大的特征值要是小于1的话，那就一定会出现梯度消失问题。后来的LSTM和GRU也仅仅能缓解这个问题

Transformer为何优于RNN

那如果要获得比RNN更进一步能力的提升，那怎么办呢？Transformer给了我们一种答案。

Transformer中抛弃了传统的CNN和RNN，整个网络结构完全是由Attention机制组成。作者采用Attention机制的原因是考虑到RNN（或者LSTM，GRU等）的计算限制为是顺序的，也就是说RNN相关算法只能从左向右依次计算或者从右向左依次计算，这种机制带来了两个问题：

顺序计算的过程中信息会丢失，尽管LSTM等门机制的结构一定程度上缓解了长期依赖的问题，但是对于特别长期的依赖现象，LSTM依旧无能为力。
时间片$t$的计算依赖$t-1$时刻的计算结果，这样限制了模型的并行能力。

Transformer的提出解决了上面两个问题：

首先它使用了Attention机制，将序列中的任意两个位置之间的距离缩小为一个常量。因为在集成信息的时候，当前单词和句子中任意单词都发生了联系，所以一步到位就把这个事情做掉了。不像RNN需要通过隐层节点序列往后传，也不像CNN需要通过增加网络深度来捕获远距离特征，Transformer在这点上明显方案是相对简单直观的。
其次它不是类似RNN的顺序结构，因此具有更好的并行性，符合现有的GPU框架。注意这里的并行性，Transformer encoder与decoder在训练过程中都可以并行化，主要借助了teacher force与masked self attention，具体可以参考下文。至于在推理过程中，encoder可以并行化，但是decoder是step by step的，无法并行化。

整体结构

首先介绍 Transformer 的整体结构，下图是 Transformer 用于中英文翻译的整体结构：

可以看到 Transformer 由 Encoder 和 Decoder 两个部分组成，各包含 6 个 block。它的工作流程大体如下：

第一步：获取输入句子的每一个单词的表示向量 x，x由单词的 Embedding（从原始数据提取出来的Feature）和单词位置的 Embedding 相加得到。

第二步：将得到的单词表示向量矩阵 (如上图所示，每一行是一个单词的表示 x) 传入 Encoder 中，经过 6 个 Encoder block 后可以得到句子所有单词的编码信息矩阵 C，如下图。单词向量矩阵用$X_{n,d}$表示， n 是句子中单词个数，d 是表示向量的维度 (论文中 d=512)。每一个 Encoder block 输出的矩阵维度与输入完全一致。

第三步：将 Encoder 输出的编码信息矩阵 C 传递到 Decoder 中，Decoder 依次会根据当前翻译过的单词 1~ i 翻译下一个单词 i+1，如下图所示。在使用的过程中，翻译到单词 i+1 的时候需要通过 Mask (掩盖) 操作遮盖住 i+1 之后的单词。

上图 Decoder 接收了 Encoder 的编码矩阵 C，然后首先输入一个翻译开始符 ““，预测第一个单词 “I”；然后输入翻译开始符 ““ 和单词 “I”，预测单词 “have”，以此类推。这是 Transformer 使用时候的大致流程，接下来是里面各个部分的细节。

输入

Transformer 中单词的输入表示 x由单词 Embedding 和位置 Embedding （Positional Encoding）相加得到。

单词 Embedding

单词的 Embedding 有很多种方式可以获取，例如可以采用 Word2Vec、Glove 等算法预训练得到，也可以在 Transformer 中训练得到（比如fairseq的翻译模型就利用了torch.nn.Embedding进行Embedding并训练）。

位置 Embedding

Transformer 中除了单词的 Embedding，还需要使用位置 Embedding 表示单词出现在句子中的位置。因为 Transformer 不采用 RNN 的结构，而是使用全局信息，不能利用单词的顺序信息，而这部分信息对于 NLP 来说非常重要。所以 Transformer 中使用位置 Embedding 保存单词在序列中的相对或绝对位置。

位置 Embedding 用 PE表示，PE 的维度与单词 Embedding 是一样的。PE 可以通过训练得到，也可以使用某种公式计算得到。在 Transformer 中采用了后者，计算公式如下：

其中，pos 表示单词在句子中的位置，d 表示 PE的维度 (与词 Embedding 一样)，2i 表示偶数的维度，2i+1 表示奇数维度 (即 2i≤d, 2i+1≤d)。使用这种公式计算 PE 有以下的好处：

使 PE 能够适应比训练集里面所有句子更长的句子，假设训练集里面最长的句子是有 20 个单词，突然来了一个长度为 21 的句子，则使用公式计算的方法可以计算出第 21 位的 Embedding。
可以让模型容易地计算出相对位置，对于固定长度的间距 k，PE(pos+k) 可以用 PE(pos) 计算得到。因为 Sin(A+B) = Sin(A)Cos(B) + Cos(A)Sin(B), Cos(A+B) = Cos(A)Cos(B) - Sin(A)Sin(B)。

将单词的词 Embedding 和位置 Embedding 相加，就可以得到单词的表示向量 x，x 就是 Transformer 的输入。

Self-Attention（自注意力机制）

上图是论文中 Transformer 的内部结构图，左侧为 Encoder block，右侧为 Decoder block。红色圈中的部分为 Multi-Head Attention，是由多个 Self-Attention组成的，可以看到 Encoder block 包含一个 Multi-Head Attention，而 Decoder block 包含两个 Multi-Head Attention (其中有一个用到 Masked)。Multi-Head Attention 上方还包括一个 Add & Norm 层，Add 表示残差连接 (Residual Connection) 用于防止网络退化，Norm 表示 Layer Normalization，用于对每一层的激活值进行归一化。

因为 Self-Attention是 Transformer 的重点，所以我们重点关注 Multi-Head Attention 以及 Self-Attention，首先详细了解一下 Self-Attention 的内部逻辑。

Self-Attention 结构

上图是 Self-Attention 的结构，在计算的时候需要用到矩阵Q(查询),K(键值),V(值)。在实际中，Self-Attention 接收的是输入(单词的表示向量x组成的矩阵X) 或者上一个 Encoder block 的输出。而Q,K,V正是通过 Self-Attention 的输入进行线性变换得到的。

Q, K, V 的计算

Self-Attention 的输入用矩阵X进行表示，则可以使用线性变阵矩阵WQ,WK,WV计算得到Q,K,V。计算如下图所示，注意 X, Q, K, V 的每一行都表示一个单词。

Self-Attention 的输出

得到矩阵 Q, K, V 之后就可以计算出 Self-Attention 的输出了，计算的公式如下（先计算括号内+softmax，然后乘上V）：

公式中计算矩阵Q和K每一行向量的内积，为了防止内积过大，因此除以$d_k$的平方根。Q乘以K的转置后，得到的矩阵行列数都为 n，n 为句子单词数，这个矩阵可以表示单词之间的 attention 强度。下图为Q乘以$K^T$，1234 表示的是句子中的单词。

得到 $QK^T$ 之后，使用 Softmax 计算每一个单词对于其他单词的 attention 系数，公式中的 Softmax 是对矩阵的每一行进行 Softmax，即每一行的和都变为 1。

得到 Softmax 矩阵之后可以和V相乘，得到最终的输出Z。

上图中 Softmax 矩阵的第 1 行表示单词 1 与其他所有单词的 attention 系数，最终单词 1 的输出$Z_1$等于所有单词 $i$ 的值 $V_i$ 根据 attention 系数的比例加在一起得到，如下图所示：

Multi-Head Attention

在上一步，我们已经知道怎么通过 Self-Attention 计算得到输出矩阵 Z，而 Multi-Head Attention 是由多个 Self-Attention 组合形成的，下图是论文中 Multi-Head Attention 的结构图。

从上图可以看到 Multi-Head Attention 包含多个 Self-Attention 层，首先将输入X分别传递到 h 个不同的 Self-Attention 中（也就是说会有8个不同的V、K、Q），计算得到 h 个输出矩阵Z。下图是 h=8 时候的情况，此时会得到 8 个输出矩阵Z。

得到 8 个输出矩阵 Z1 到 Z8 之后，Multi-Head Attention 将它们拼接在一起 (Concat)，然后传入一个Linear层，得到 Multi-Head Attention 最终的输出Z。

可以看到 Multi-Head Attention 输出的矩阵Z与其输入的矩阵X的维度是一样的。

Encoder 结构

上图红色部分是 Transformer 的 Encoder block 结构，可以看到是由 Multi-Head Attention, Add & Norm, Feed Forward, Add & Norm 组成的。刚刚已经了解了 Multi-Head Attention 的计算过程，现在了解一下 Add & Norm 和 Feed Forward 部分。

Add & Norm

Add & Norm 层由 Add 和 Norm 两部分组成，其计算公式如下（以下分别表示上图红色部分第一个和第二个Add & Norm 层）：

其中 X 表示 Multi-Head Attention 或者 Feed Forward 的输入，MultiHeadAttention(X) 和 FeedForward(X) 表示输出 (输出与输入 X 维度是一样的，所以可以相加)。

Add指 X+MultiHeadAttention(X)或者X+FeedForward(X)，是一种残差连接，通常用于解决多层网络训练的问题，可以让网络只关注当前差异的部分，在 ResNet 中经常用到：

Norm指 Layer Normalization，通常用于 RNN 结构，Layer Normalization 会将每一层神经元的输入都转成均值方差都一样的，这样可以加快收敛。

Feed Forward

Feed Forward 层比较简单，是一个两层的全连接层，第一层的激活函数为 Relu，第二层不使用激活函数，对应的公式如下。

X是输入，Feed Forward 最终得到的输出矩阵的维度与X一致。

组成 Encoder

通过上面描述的 Multi-Head Attention, Feed Forward, Add & Norm 就可以构造出一个 Encoder block，Encoder block 接收输入矩阵 $X_{(n\times d)}$，并输出一个矩阵 $O_{(n\times d)}$。通过多个 Encoder block 叠加就可以组成 Encoder。

第一个 Encoder block 的输入为句子单词的表示向量矩阵，后续 Encoder block 的输入是前一个 Encoder block 的输出，最后一个 Encoder block 输出的矩阵就是编码信息矩阵 C，这一矩阵后续会用到 Decoder 中。

Decoder 结构

上图红色部分为 Transformer 的 Decoder block 结构，与 Encoder block 相似，但是存在一些区别：

包含两个 Multi-Head Attention 层。
第一个 Multi-Head Attention 层采用了 Masked 操作。
第二个 Multi-Head Attention 层的K, V矩阵使用 Encoder 的编码信息矩阵C进行计算，而Q使用上一个 Decoder block 的输出计算。
最后有一个 Softmax 层计算下一个翻译单词的概率。

第一个 Multi-Head Attention

Decoder block 的第一个 Multi-Head Attention 采用了 Masked 操作，因为在翻译的过程中是顺序翻译的，即翻译完第 i 个单词，才可以翻译第 i+1 个单词。通过 Masked 操作可以防止第 i 个单词知道 i+1 个单词之后的信息。下面以 “我有一只猫” 翻译成 “I have a cat” 为例，了解一下 Masked 操作。

下面的描述中使用了类似 Teacher Forcing 的概念，在 Decoder 的时候，是需要根据之前的翻译，求解当前最有可能的翻译，如下图所示。首先根据输入 ““ 预测出第一个单词为 “I”，然后根据输入 “ I” 预测下一个单词 “have”。

Decoder 可以在训练的过程中使用 Teacher Forcing 并且并行化训练，即将正确的单词序列 ( I have a cat) 和对应输出 (I have a cat) 传递到 Decoder。那么在预测第 i 个输出时，就要将输出序列第 i+1 之后的单词掩盖住（输入序列不需要遮挡），注意 Mask 操作是在 Self-Attention 的 Softmax 之前使用的，下面用 0 1 2 3 4 5 分别表示 “ I have a cat“。

第一步：先介绍 Decoder 的输入矩阵和 Mask 矩阵，输入矩阵包含 “ I have a cat” (0, 1, 2, 3, 4) 五个单词的表示向量，Mask 是一个 5×5 的矩阵。在 Mask 可以发现单词 0 只能使用单词 0 的信息，而单词 1 可以使用单词 0, 1 的信息，即只能使用之前的信息（从图中可以看出，Mask操作必须放在Softmax之前，否则还是利用到了之后的信息）。

第二步：接下来的操作和之前的 Self-Attention 一样，通过输入矩阵X计算得到Q,K,V矩阵。然后计算Q和 $K^T$ 的乘积 $QK^T$ 。

第三步：在得到 $QK^T$ 之后需要进行 Softmax，计算 attention score，我们在 Softmax 之前需要使用Mask矩阵遮挡住每一个单词之后的信息，遮挡操作如下：

也有可能不是按位相乘，而是让 $QK^T$ 中mask==0的对应位置，是一个极小值，这样这些位置在经过softmax后，值仍然很小。

得到 Mask $QK^T$ 之后在 Mask $QK^T$上进行 Softmax，每一行的和都为 1。但是单词 0 在单词 1, 2, 3, 4 上的 attention score 都为 0。

第四步：使用 Mask $QK^T$与矩阵 V相乘，得到输出 Z，则单词 1 的输出向量 $Z1$ 是只包含单词 1 信息的。

为什么说这样可以保证单词 1 的输出向量 $Z1$ 是只包含单词 1 信息的，可以看如下这个图理解这个矩阵相乘。

第五步：通过上述步骤就可以得到一个 Mask Self-Attention 的输出矩阵 $Z_i$ ，然后和 Encoder 类似，通过 Multi-Head Attention 拼接多个输出$Z_i$ 然后计算得到第一个 Multi-Head Attention 的输出Z，Z与输入X维度一样。

第二个 Multi-Head Attention

Decoder block 第二个 Multi-Head Attention 变化不大，主要的区别在于其中 Self-Attention 的 K, V矩阵不是使用上一个 Decoder block 的输出计算的，而是使用 Encoder 的编码信息矩阵 C 计算的。

根据 Encoder 的输出 C计算得到 K, V，根据上一个 Mask Multi-Head Attention的输出 Z 计算 Q （因为此时得到的Z已经经过了mask，所以无需再次mask了），后续的计算方法与之前描述的一致。

这样做的好处是在 Decoder 的时候，每一位单词都可以利用到 Encoder 所有单词的信息 (这些信息无需 Mask)。

Softmax 预测输出单词

Decoder block 最后的部分是利用 Softmax 预测下一个单词，在之前的网络层我们可以得到一个最终的输出 Z，因为 Mask 的存在，使得单词 0 的输出 Z0 只包含单词 0 的信息，如下：

Softmax 根据输出矩阵的每一行预测下一个单词：

下动图能够更好的解释这个过程：

给Decoder输入Encoder对整个原始句子embedding的结果和一个特殊的开始符号</s>。Decoder 将产生预测，在我们的例子中应该是 ”I”。
给Decoder输入Encoder的embedding结果和</s> I，在这一步Decoder应该产生预测 am。
给Decoder输入Encoder的embedding结果和</s> I am，在这一步Decoder应该产生预测a。
给Decoder输入Encoder的embedding结果和</s> I am a，在这一步Decoder应该产生预测student。
给Decoder输入Encoder的embedding结果和</s> I am a student, Decoder应该生成句子结尾的标记，Decoder 应该输出</eos>。
然后Decoder生成了</eos>，翻译完成。

这就是 Decoder block 的定义，与 Encoder 一样，Decoder 是由多个 Decoder block 组合而成。

Transformer 总结

Transformer 与 RNN 不同，可以比较好地并行训练。
Transformer 本身是不能利用单词的顺序信息的，因此需要在输入中添加位置 Embedding，否则 Transformer 就是一个词袋模型了。
Transformer 的重点是 Self-Attention 结构，其中用到的 Q, K, V矩阵通过输出进行线性变换得到。
Transformer 中 Multi-Head Attention 中有多个 Self-Attention，可以捕获单词之间多种维度上的相关系数 attention score。

Attention在Transformer中的应用

文章一开始解释了Self-Attention和Multi-Head Attention。通过对Transformer模型的深入解读，可以看到，模型一共使用了三种Multi-Head Attention：

Encoder Block中使用的Attention。第一个Encoder Block的Query、Key和Value来自训练数据经过两层Embedding转化，之后的Encoder Block的Query、Key和Value来自上一个Encoder Block的输出。
Decoder Block中的第一个Attention。与Encoder Block中的Attention类似，只不过增加了Mask，在预测第 $i$个输出时，要将第$i+1$ 之后的单词掩盖住。第一个Decoder Block的Query、Key和Value来自训练数据经过两层Embedding转化，之后的Decoder Block的Query、Key和Value来自上一个Decoder Block的输出。
Decoder Block中的第二个Attention。这是一个 Encoder-Decoder cross Attention，它建立起了 Encoder 和 Decoder 之间的联系，Query来自第2种 Decoder Attention的输出，Key和Value 来自 Encoder 的输出。

Beam Search

前文讨论的主要是训练过程，在推理时，Decoder侧最后的 Softmax 将 logit 转化为概率，选择概率最大的词作为预测词。Encoder侧输入是源语言句子。第一个时间步，Decoder侧首先以开始符”<Begin>”作为输入，预测下一个概率最大的词；第二个时间步，Decoder侧以开始符和第一个预测词作为输入，来预测下一个概率最大词。每次基于上一次预测结果，选择最优的解，这是一种贪心搜索算法。贪心搜索从局部来说可能是最优的，但是从全局角度并不一定是最优的。

Beam Search对贪心算法进行了改进。在每一个时间步预测时，不再只选择最优解，而是保留num_beams个结果。当num_beams=1时 Beam Search 就退化成了贪心搜索。

下图中，每个时间步有ABCDE共5种可能的输出，设置num_beams=2，也就是说每个时间步都会保留到当前步为止条件概率最优的2个序列。

在第一个时间步，A和C是最优的两个，因此得到了两个结果[A],[C]，其他三个就被抛弃了；第二步会基于这两个结果继续进行生成，在A这个分支可以得到5个候选人，[AA],[AB],[AC],[AD],[AE]，C也同理得到5个，此时会对这10个进行统一排序，再保留最优的两个，即图中的[AB]和[CE]；第三步同理，也会从新的10个候选人里再保留最好的两个，最后得到了[ABD],[CED]两个结果。可以发现，Beam Search在每一步需要考察的候选人数量是贪心搜索的num_beams倍，因此是一种牺牲时间换取准确率的方法。

Decoder加速

文小节主要来自于faster-decoder之 decoder解码加速。

背景

Transformer 模型舍弃了 step by step 的时序循环，依靠 self-attention 的并行结构获得了远超同量级 LSTM 网络的训练速度。即使做auto-regresisve 任务时，通过attention-mask 机制依然可以像encoder 一样并行计算。然而在解码时，却任然需要step by step 的进行，即需要知道上一个time step 的结果后才能进行下一个time step 的解码。此外，通常我们的解码策略是在获得模型结果后在内存中计算的，需要不停的将结果从GPU load 进 CPU 然后计算，这就进一步的拖慢了解码速度。而通常我们在部署时，首选的tf-serving 需要将结果通过网络传输，这将进一步的拖慢解码速度。而针对解码慢的问题，主要的加速方案有：

将解码策略放在GPU 上计算，这样将避免结果在GPU/CPU 之间转换与等待；
attention cache，根据attention 层的特点，对attention 中的 $K$ / $V$ 进行cache，避免对之前的time step 进行重复计算，将attention 层的计算由 $O(n^2)$ 降低到 $O(n)$。
transformer 计算最耗时的是attention 层中的softmax，尝试使用一些线性函数进行近似替换。

三种方案中，GPU 上进行解码需要一些底层技术进行开发，暂时没能力，而替换softmax 方案则会或多或少的损失一些精度，本文都不做进一步讨论。本文聚焦在attention cache 方案上，加速的同时又“不会”损失精度。

原理

attention 的计算公式：

$A = Softmax(QK^{T})* V$

在解码时，我们是step by step 进行的，所以，我们将时刻 t 的attention 写出来：

$A = Softmax(Q_{t}K^{T})* V$

即：对于时刻t 来说，attention 只需要当前的 $Q_{t}$ 时刻信息，$K$ / $V$ 的所有时刻信息进行计算。而 $Q_{t}$ 的计算只需要 $Token_{t}$ 即可，如何加速计算的关键就剩下如何更加高效的计算 $K$ / $V$.

encoder-decoder cross-attention

对于encoder-decoder cross-attention 来说，对应的 $K$ / $V$ 都来自encoder 的outputs，所以直接将其整个进行cache 即可，而无需每步都重新计算。

self-attention

而当attention 是self-attention 时，对于时刻 $t$ 来说，此时的 $K$ / $V$ 为 $K$ / $V$ 的前 $t$ 时刻信息，即 $K_{\leq t}$ / $V_{\leq t}$ .此时的 attention 计算为：

$A = Softmax(Q_{t}K_{\leq t}^{T})* V_{\leq t}$

而 $K_{t}$/$V_{t}$ 的计算只与 $Token_t$ 有关，与其他时刻的 $Token$ 无关，且不论是时刻 $t$ 还是时刻 $t+1$,对应的 $K_{t-1}$ / $V_{t-1}$ 的计算结果都是一样的。因此，每个时刻都对 $K_{\leq t}$ / $V_{\leq t}$ 全部计算是低效且浪费的。

由于 $K_t$ / $V_t$ 有只需 $Token_t$ 计算且不同时刻结果“一致”的特点，对于$t$时刻，我们只需要根据当前$Token_t$ 计算 $K_t$ / $V_t$ ，然后拼接cache 中的 $K_{< t}$ / $V_{< t}$上，得到完整序列的$K_{\leq t}$ / $V_{\leq t}$ 。用公式表示如下：

$\begin{array}{c}q_{t}^{l}, k_{t}^{l}, v_{t}^{l}=Q x_{t}^{l}, K x_{t}^{l}, V x_{t}^{l} \\ h_{t}^{l}=\left[h_{t-1}^{l},\left(k_{t}^{l}, v_{t}^{l}\right)\right]\end{array}$

此外，由于使用了attention cache 后，每次解码输入只需要 $Token_t$ 而非 $Token_{\leq t}$ ，这样将其他层的计算量也会随之降低。

PS：由于decoder 中为了实现auto-regressive 而采用了下三角的attention mask，因此，不同时刻的attention mask 是不同的，这会导致不同时刻的 $K_t$ / $V_t$ 的结果略有不同（约e-10)，但是这并不影响最终端到端的结果。

不要小看这样的工程化技巧，在实际测试时候，解码速度可以提高到 4~8倍。

实现

attention 层在实现时，除了进行attention 计算的同时，还会包含attention mask 和 position bias 两种信息，其中，attention mask 来实现auto regressive，即当前位置的attention 只能包含当前位置及之前的信息；position bias 则包括各种position 信息的实现，所以在使用attention cache 后，还需要对这两种信息进行“纠正”。

attention 层修改

具体实现时，对于encoder-decoder cross-attention, 我们之间将encoder outputs 计算一次后进行cache，每次进行解码时作为inputs 送人decoder；

对于self-attention ，我们在得到 $Q_{t}$/$K_{t}$/$V_{t}$ 后，将 $K_t$/$V_t$ 与之前的 $K_{\leq t-1}$/ $V_{\leq t-1}$ cache 进行拼接，构造出完整的$K_{\leq t}$ / $V_{\leq t}$, 然后将$Q_t$ / $K_{\leq t}$ / $V_{\leq t}$ 进入self-attention 层进行计算。

attention mask 的“纠正”

由于attention mask 的作用是防止当前位置看到其后位置的信息，而在使用cache 后，当前位置即最后时刻的位置，所以此时的attention mask 已没有存在的必要，直接取消即可；PS: 由于这里直接取消了attention mask，而attention mask 的实现通常是通过加上一个负无穷(-e12) 来实现的，所以加了cache 后的outputs 与没加之前会有一定的差异，大概在e-10 量级。

position bias 的“纠正”

由于position bias 通常是通过inputs 的长度进行计算的，而加了attention cache 后，每次的inputs 的长度变为1 了（当前时刻的$Token_t$），所以此时的position bias 恒等于长度为1 的序列。为了还原他原始的position bias，我们使用拼接了cache 后的$K_{\leq t}$ 进行计算完整序列的position bias，然后取出当前query 在完整序列中位置对应的position bias 即可。

解码实现

此外，在解码函数上，也需要进行相应的修改，以获得当前时刻的$K_t$/$V_t$ , 然后与之前时刻的所有 $K_{\leq t-1}$ / $V_{\leq t-1}$ cache 进行拼接，为下一个时刻计算做准备。

onnx

由于tensorflow 会对当前显卡的显存全部占用，所以一个显卡只能启动一个tensorflow 进程，这样就导致当一个模型的显存不需要占用所有显存即可解码时，使用tensorflow 会浪费一部分显存，这里我们将其转为onnx ，这样只需要占用模型需要的显存即可，避免显存浪费。即一个显卡可以起多个解码进程。

demo

在bert4keras 的基础上，对 T5/Roformer 进行了实现，具体代码参考：faster-decoder。

关于TensorFlow的官方实现可以参考tensor2tensor。

疑问解决

Decoder中是否用了真实标签

训练时：第i个decoder的输入 = encoder输出 + ground truth embedding

预测时：第i个decoder的输入 = encoder输出 + 第(i-1)个decoder输出

训练时因为知道ground truth embeding，相当于知道正确答案，网络可以一次训练完成。

预测时，首先输入start，输出预测的第一个单词，然后start和新单词组成新的query，再输入decoder来预测下一个单词，循环往复直至end。

Softmax怎么将输出矩阵的行向量映射到相应的单词？

行向量代表着单词的类型，输出概率最大的那个位置就是预测的单词。行向量中单词的位置是固定的，只需要找位置信息就能找到相应的单词了。

Transformer中的mask

在《the annotated transformer》中有多个mask，这里总结一下。

整个模型中使用到的mask主要就是source mask和target mask，其各自的作用如下所示：

source mask：

source长短不一而无法形成batch，因此引入了pad。将source mask传入到encoder中，让attention在计算$\mathrm{softmax}(\frac{QK^T}{\sqrt{d_k}})$时，pad位置的值不起作用。
同时这个mask还需要传入每个decoderLayer第二个multi-head attention模块中，就是防止来自encoder的key和来自decoder的query在计算多头注意力的时候算了target中的词和source中pad的权重

target mask：需要分training和testing进行讨论

训练时，用于防止target的ground truth长短不一引入pad造成的误差，以及避免在自回归时看到正在预测的字和以后字的ground truth
测试时，逻辑上decoder不需要target mask，但出于编程方便的考虑引入mask，假装用于防止看到后面的ground truth，target mask的最后两维的shape和目前生成出来的序列长度相同，但实际上每次都会有一些重复运算在里面，比如目前在预测第10个词时，第1-9个词还需要重新算一遍。核心原因是：模型在写的时候主要考虑的是训练，执行一次attention函数翻译完一个batch的所有句子，而测试时必须是单个或多个句子word by word进行计算

Teacher Forcing是什么

我把这个概念放在了这里，前文没提到这个，现在知识不成体系，不知道放哪合适。

Teacher Forcing 是一种用于序列生成任务的训练技巧，与Autoregressive模式相对应，这里阐述下两者的区别：

Autoregressive 模式下，在 timesteps $t$ decoder模块的输入是 timesteps $t-1$ 的输出 $y_{t-1}$。这时候我们称 $y_{t-1}$为当前预测步的 context;
Teacher-Forcing 模式下，在 timestep $t$ decoder模块的输入是 Ground-truth 语句中位置的 $y_{t-1}^$ 单词。这时候我们称 $y_{t-1}^$为当前预测步的 context；

Teacher-Forcing 技术之所以作为一种有用的训练技巧，主要是因为：

Teacher-Forcing 能够在训练的时候矫正模型的预测，避免在序列生成的过程中误差进一步放大。
Teacher-Forcing 能够极大的加快模型的收敛速度，令模型训练过程更加快&平稳。
Teacher-Forcing 技术是保证 Transformer 模型能够在训练过程中完全并行计算所有token的关键技术。

如果要用比较不太严谨的比喻来说，Teacher-Forcing 技术相当于就是小明学习的时候旁边坐了一位学霸，当发现小明在做序列生成题目的时候，每一步都把上一步的正确答案给他偷看。那么小明当然只需要顺着上一步的答案的思路，计算出这一步的结果就行了。这种做法，比起自己每一步都瞎猜，当然能够有效的避免误差进一步放大，同时在学习前期还能通过学霸辅导的这种方式快速学到很多的知识。

Teacher Forcing 最常见的问题就是 Exposure Bias 了。

模型在训练时基于真实的描述句来生成下一个词，但在测试的时候，它只能基于模型自己的生成结果来生成下一个词，模型自己的生成结果有可能是错的，一旦出错，就会让模型处于一个在训练阶段没有见过的情况，从而导致在生成过程中的误差逐渐累积

上面的『比喻』，其实就是不太严谨的 Exposure Bias 现象了。更严谨的表述，由于训练和预测的时候decode行为的不一致，导致预测单词（predict words）在训练和预测的时候是从不同的分布中推断出来的。而这种不一致导致训练模型和预测模型直接的Gap，就叫做 Exposure Bias。