Transformer原理

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

Attention简介

Attention注意力机制

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

img

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

img

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

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

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。这里提前解释一下,看完下文之后可以再回味这一小节。

  1. self-attention:encoder中的self-attention的query、key、value都对应了源端序列,decoder中的self-attention的query、key、value都对应了目标端序列。
  2. 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的示意图。

img

下图是self-attention的示意图。

img

核心思想

RNN缺陷

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

img

上图中的 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相关算法只能从左向右依次计算或者从右向左依次计算,这种机制带来了两个问题:

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

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

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

整体结构

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

img

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

第一步:获取输入句子的每一个单词的表示向量 xx单词的 Embedding(从原始数据提取出来的Feature) 和单词位置的 Embedding 相加得到。

img

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

img

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

img

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

输入

Transformer 中单词的输入表示 x单词 Embedding位置 Embedding (Positional Encoding)相加得到。

img

单词 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 中采用了后者,计算公式如下:

img

其中,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 相加,就可以得到单词的表示向量 xx 就是 Transformer 的输入。

Self-Attention(自注意力机制)

img

上图是论文中 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 结构

img

上图是 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 的每一行都表示一个单词。

img

Self-Attention 的输出

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

img

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

img

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

img

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

img

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

img

Multi-Head Attention

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

img

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

img

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

img

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

Encoder 结构

img

上图红色部分是 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 层):

img

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

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

img

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

Feed Forward

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

img

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 中。

img

Decoder 结构

img

上图红色部分为 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”。

img

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之前,否则还是利用到了之后的信息)。

img

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

img

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

img

也有可能不是按位相乘,而是让 $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 信息的。

img

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

aaa

第五步:通过上述步骤就可以得到一个 Mask Self-Attention 的输出矩阵 $Z_i$ ,然后和 Encoder 类似,通过 Multi-Head Attention 拼接多个输出$Z_i$ 然后计算得到第一个 Multi-Head Attention 的输出ZZ与输入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 的信息,如下:

img

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

img

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

img

  1. 给Decoder输入Encoder对整个原始句子embedding的结果和一个特殊的开始符号</s>。Decoder 将产生预测,在我们的例子中应该是 ”I”。
  2. 给Decoder输入Encoder的embedding结果和</s> I,在这一步Decoder应该产生预测 am。
  3. 给Decoder输入Encoder的embedding结果和</s> I am,在这一步Decoder应该产生预测a。
  4. 给Decoder输入Encoder的embedding结果和</s> I am a,在这一步Decoder应该产生预测student。
  5. 给Decoder输入Encoder的embedding结果和</s> I am a student, Decoder应该生成句子结尾的标记,Decoder 应该输出</eos>。
  6. 然后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:

  1. Encoder Block中使用的Attention。第一个Encoder Block的Query、Key和Value来自训练数据经过两层Embedding转化,之后的Encoder Block的Query、Key和Value来自上一个Encoder Block的输出。

  2. Decoder Block中的第一个Attention。与Encoder Block中的Attention类似,只不过增加了Mask,在预测第 $i$个输出时,要将第$i+1$ 之后的单词掩盖住。第一个Decoder Block的Query、Key和Value来自训练数据经过两层Embedding转化,之后的Decoder Block的Query、Key和Value来自上一个Decoder Block的输出。

  3. Decoder Block中的第二个Attention。这是一个 Encoder-Decoder cross Attention,它建立起了 Encoder 和 Decoder 之间的联系,Query来自第2种 Decoder Attention的输出,Key和Value 来自 Encoder 的输出。

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

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

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

img

在第一个时间步,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 需要将结果通过网络传输,这将进一步的拖慢解码速度。而针对解码慢的问题,主要的加速方案有:

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

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

原理

attention 的计算公式:

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

即:对于时刻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 计算为:

而 $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}$ 。用公式表示如下:

此外,由于使用了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,其各自的作用如下所示:

  1. 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的权重

  1. 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。

参考

Transformer模型详解(图解最完整版)
the annotated transformer中的关于mask的问题 - lumino的文章 - 知乎
各种各样的语言生成模型训练算法
Seq2Seq中Exposure Bias现象的浅析与对策
关于Teacher Forcing 和Exposure Bias的碎碎念
Attention 注意力机制
Transformer
What’s the Difference Between Attention and Self-attention in Transformer Models?
浅析Transformer训练时并行问题
【深度学习】NLP之Transformer (1) Encoder
一文理解 Transformer 的工作原理
Self-Attention和Transformer
faster-decoder之 decoder解码加速
碎碎念:Transformer的解码加速

------ 本文结束------
坚持原创技术分享,您的支持将鼓励我继续创作!

欢迎关注我的其它发布渠道