图解 Transformer

注:原文作者 Jay Alammar,地址 https://jalammar.github.io/illustrated-transformer/

先前文章中,我们介绍了能大大提高神经机器翻译的效果 Attention 结构,它目前也普遍存在于现代深度学习模型中。本文将要介绍的 Transformer,则是利用 Attention 来提高这类模型训练速度的方法。 在特定任务中,Transformers 表现已经优于谷歌的神经机器翻译。而它最大的优势则是来自于本身可以并行化。谷歌也推荐使用它作为 Cloud TPU 的教学范例。因此本文将结构这一模型的各个部分,并逐一解释其工作原理。

Transformer 起初由 Attention is All You Need 这篇文章提出。Tensor2Tensor 库中则提供了基于 TensorFlow 的实现,而哈佛的 NLP 组给出了 PyTorch 版本的指南,以注释形式解读原文。本文我们将高度简化地逐一介绍模型原理,使对这一主题涉猎不多的读者也能理解。

观其大略

我们不妨想把模型看作一个黑盒子,在机器翻译场景下,我们输入一个句子,另一端则会输出翻译的结果。

img

查看模型内部结构,我们发现有一个编码组件,一个解码组件,二者之间有连接线。

img

编码组件是一批堆叠在一起的编码器(文章中使用了 6 个编码器,选择 6 个没有什么玄学原因,换其他数量未尝不可)。而解码组件则是由相同数量的解码器堆叠而成。

img

所有编码器都有着相同的结构(然而它们并不共享权重),每个都可以拆分为两个子结构:

img

编码器的输入首先会进入 self-attention 层,在编码输入句子的某一个词语时,该层会纳入句中其他单词的信息。下文会详述这一过程。

self-attention 层的输出会被喂给一个前馈神经网络,句中不同位置的输出会各自独立地被送往相同的前馈网络中。

解码器也包含了这两层结构,但在二者之间多了一个 attention 层来帮助其聚焦于输入句中的相关部分(attention 的作用同其在seq2seq 模型中的类似)。

img

是时候来点张量了

既然我们已经了解模型的主要结构,下面我们来关注下其中涉及的各种向量,张量,看看它们在组件间如何传递使得我们能得到想要的模型输出。

同大多数 NLP 模型一样,我们先要利用embedding algorithm将输入的文本转化为数值向量。

img

每个词语都被嵌入成了一个长度为 512 的向量,图中将用小方块来表示它们。

词嵌入过程只在最底层的编码器中执行一次,而每个编码器都接受一个由长度为 512 的向量组成的列表作为输入。对于最底层的编码器而言,输入就是句中每个词嵌入后的向量列表。其他上层编码器的输入则是其正下方编码器的输出,输出列表的长度是我们可以设置的一个超参数,一般而言我们会将其设置成训练集中最长句子的长度。

经过词嵌入过程后的每一个输入句,都会流经编码器的这两层结构。

img

现在我们发现了 Transformer 的一个核心特性,输入句中不同位置的词语在编码器中的计算路径是不同的。在 self-attention 层中这些路径之间相互关联,而在前馈网络层中则是独立的,因此前馈层中的计算过程是可以并行的。

接下来我们将用一个更短的句子为例,看看每个编码器的子层都做了什么。

来编码吧!

正如上文提到的,编码器接收由向量组成的列表,然后将向量依次输送至 self-attention 层和前馈神经网络,再将输出传至下一个编码器。

img

每个位置上的词语都要经过这样的 self-attention 处理,然后各自传向一个同一个前馈神经网络。

Self-Attention 概述

假定我们想翻译如下句子:

The animal didn't cross the street because it was too tired

句中的 “it” 指代什么呢,是街道还是动物? 对于人类而言这个问题轻而易举但对算法来说并不容易。

当模型处理 “it” 这个词时, self-attention 会让其和 “animal” 产生联系。模型逐个处理句中的单词,self attention 将被处理词和句中其他词语关联起来已获得更好的编码效果。

如果你对 RNN 比较熟悉,可以类比其中的 hidden state,它也让 RNN 当前处理的内容可以用上先前处理过的词语或向量的表示结果。Transformer 则是通过 Self-attention 机制来讲关联词的语义信息纳入当前处理的词中。

img

当编号为 #5 的编码器(栈顶编码器)开始编码 “it” 一词时,注意力机制的一部分落在了 “The Animal” 上,并把其表示结果纳入 “it” 的编码中。

你可以参考 Tensor2Tensor notebook 中的内容,以交互式可视化的方式理解这部分内容。

Self-Attention 计算过程详解

让我们先学习如何向量来计算 self-attetion 结果,之后再看看如何利用矩阵应用它。

计算的第一步是为编码器中的每个输入向量(此处为每个词语的嵌入结果)创建三个向量,分别为 Query 向量,Key 向量和 Value 向量。这些向量是通过将输入向量和三个矩阵相乘得到的,而这三个矩阵也会在模型训练过程中不断更新。

注意这些向量的维度都比输入的词向量的维度要低,它们的长度为64,而输入输出向量维度为 512。这些向量也不是非得降维,这种选择是为了让计算更稳定。

img

将向量 x1 同 WQ 矩阵相乘会生成和输入词语关联的 query 向量,同理 key 和 value 向量也是这么生成的。它们是 attention 机制中所需的三类抽象,阅读后续内容后你就会理解它们的作用机制。

self-attention 的第二步是计算 score 值,假定women在计算例子中的 “Thinking” 一词,我们需要计算句中每个词相对该词的 score。score 的值代表了在编码 “Thinking” 时该给其他词分配多少注意力。

score 是通过将 query 向量 和 key 向量点乘得到的,当我们编码句中第一个单词时,第一个 score 由 q1 和 k1 点乘得到,第二个由 q1 和 k2 点乘得到。

img

第三步是将 score 除以 8(8 是 key 向量维数的平方根,这会使得梯度更加稳定,除以别的数字也可以,平方根是默认值),之后将结果进行 softmax,让向量的元素和为1。

img

softmax 之后的 score 决定了每个词对当前位置编码的贡献,该位置词语本身显然有最高的 score,但纳入和当前词相关的词的信息也是大有裨益的。

第五步是把每个 value 向量同 softmax socre 相乘作为赋权过程(为后续求和做准备),把值得注意的词赋予高权重,无关紧要的词赋予一个极小的权重。

第六步是将赋权后的 value 向量累加,这就是 self-attention 层在该位置的编码输出。

img

以上就是 self-attention 的计算全过程,输出的结果矩阵将被输入到前馈神经网络中。实际情况下整个计算过程是以矩阵计算的方式来提速的,下一节将阐述这个过程。

Self-Attention 的矩阵计算

现在的第一步是计算 Query, Key 和 Value 矩阵了,先把词嵌入向量拼成矩阵,然后分别和 WQ, WK, WV 三个权重矩阵相乘。

img

矩阵 X 中的每一行为输入的句子中的一个词向量,我们可以清晰地看到输入维度(512 维,图中用 4 个方块表示)和 q/k/v 向量(64 维,图中为 3 个方块)的差异。

最后,由于是矩阵计算,我们把 6 个计算步骤压缩成一个来获得 self-attention 层的结果。

img

多头怪兽

论文后续有又将 self-attention 层融入到名为多头注意力的机制中,从两个层面提升了注意力机制的效果:

  1. 它拓展了模型的注意力分配能力。上面的例子里,z1 包含了所有其他词编码的信息,但依旧是被自己位置的词语的结果主导。当我们翻译 “The animal didn’t cross the street because it was too tired” 时,我们希望知道 “it” 指代哪个词。(让 “it” 的编码以其指代对象词的信息为主,译者注)
  2. 多头扔注意力层有了更多的表示子空间。下文我们会看到在多头注意力机制下,会有多组 Query/Key/Value 权重矩阵(Transformer 使用了 8 个注意力头,所以对应也有 8 个编码器和解码器)。每组矩阵都是随机初始化的,在模型训练后,每组都用来将输入的向量投影至表示子空间。

img

多头注意力机制下我们将各个 WQ/WK/WV 矩阵分别保存,依次让输入矩阵 X 与其计算生成每个头的 Q/K/V 矩阵。

如果继续采用之前的方式,就是重复 8 次计算流程,得到 8 个矩阵 Z。

img

这给我们带来了一些些挑战,后续的前馈层并不接受 8 个矩阵作为输入,所以我们要把 8 个压缩成一个。具体做法是把 Z 矩阵们横向拼接,然后和一个权重矩阵 WO 相乘。

img

多头注意力机制就这些内容,里面涉及到了大量的矩阵,下图我试着用一个图来阐明他们的关系和作用。

img

了解多头的机理后,我们回顾之前的例子,看看对 “it” 的编码有什么不同:

img

编码过程中,一个注意力头权重主要落在 “the animal” 上,另一个落在 “tired” 上,某种意义上, “it” 的表示结果包含了这两个词的信息。

如果我们把所有注意力头都放在一个图里,事情就变得更难解释了:

img

用位置编码表示剧中词语的顺序

目前我们一直没有提到句中词语的顺序问题,为了表示它,transformer 在每个输入的嵌入层增加了一个向量。这些向量用特定的模式表示词语的顺序或是词语之间的距离。使得计算 Q/K/V 向量或是点乘生成注意力层时可以涵盖距离信息。

img

假定输入的词嵌入层维度为 4,实际的位置形式编码如下(来自实际案例):

img

那么高维下的位置编码什么样呢?

下图中,每一行代表一个词向量的位置编码,第一行就是句子里第一个词的位置信息,向量维度为 512,每个元素的取值咋吃 -1 到 1 之间,图中用颜色表示数值大小。

img

一个由 20 个单词构成的句子的真实案例,位置编码维度为 512,可以看到图从中间被一分为二,这是因为左右的值由不同的函数生成( 左侧为 sine,右侧围 cosine )。双边的结果直接被拼接到了一起共同编码位置信息。

论文的 3.5 部分也详细阐述了位置编码的计算公式,代码中的 get_timing_signal_1d() 函数实现了这一过程。当然,位置编码的生成方法远不止这一种,但它的优势在于可以编码任意长度的句子。

残差部分

在继续讲解之前还有一个细节需要阐明,即每个编码器中的每一个子层(self-attention, 前馈神经网络)都有一个残差连接结构,之后还有 layer-normalization](https://arxiv.org/abs/1607.06450) 步骤。

img

将向量和 self-attention 的 layer-norm 操作一起可视化如下:

img

解码器部分的 sub-layers 也有这样的结构,下图展示了一个包含两个编码器和解码器的 Transformer 的结构:

img

解码器

关于编码器,我们已经说得够多了,解码器部分的组件和它也基本一致。但我们还需要了解一下它们是如何协同工作的。

编码器的从输入序列开始处理,最后输出的是 attention 转化得到的一系列向量 K 和 V。它们会在每个解码器的 “encoder-decoder attention” 层被使用以辅助解码器关注输入序列的正确位置:

img

解码步骤的每一步都生成输出序列的一部分(本例中为英语句子)。

接下来的工作就是重复这一过程直到生成了告知解码器工作结束的特定符号。而在下一个时间步,这些输出都会喂给最底层的解码器。就像编码器一样,解码器的输出也向上层层传递。位置编码也一同输入来表名每个词语的位置。

img

解码器中的 self-attention 层和编码器中的稍稍有些不同,只纳入当前输出序列前的位置编码信息。这个限制通过在 softmax 处理前把后续位置编码设为 -inf 实现。

“Encoder-Decoder Attention” 层和多头 self-attention 层的工作方式大体一致,但它只从位于下方的层来生成 Query 矩阵,Key 和 Value 矩阵则用编码器的输出计算得到。

最后的 Linear 和 Softmax 层

编码器栈会产出一个元素为浮点数的向量,怎么把它转换为一个词呢?这就是 后续的 Linear 层和 Softmax 层的工作。

Linear 层是一个简单的全连接网络,它将解码器栈产出的向量投影成一个名为 logits 的大向量。

假设我们的模型从训练集中学习了 10000 个不同的英文单词(输出词典),那么 logits 向量的长度就是 10000,每个元素的值是一个词的概率得分。

而紧随其后的 Softmax 层则把这些得分归一化为概率值(全部为正数且和为1),概率最大的那个词会被用作当前时间步的输出。

img

上图演示了解码其栈输出转化为词语的过程。

训练过程回顾

我们已经了解了训练好的 Transformer 是如何工作的了,现在就来看看它的训练过程。

训练过程中模型的处理过程和上文提到的一样,但训练集是有标注的,我们可以比较输出和正确结果。为了方便图示,我们假设输出词典只有 6 个词(“a”, “am”, “i”, “thanks”, “student”, 和代表句子结尾的 “” ),对应索引如下:

img

之后英语句子就可以 one-hot 编码了, 例如单词 “am” 就可以表示为如下向量:

img

简要回顾后, 就该看看损失函数了。

损失函数

假设我们现在正在训练模型,样本非常简单 - 把 “merci” 翻译为 “thanks”。我们希望最终输出 “thanks” 的概率最大,但在模型没有训练好的情况下,这个情况一半不会发生。

img

由于模型的初始参数(权重)是随机生成的,未训练的模型输出也是很随机的。我们可以将其与实际结果比较,之后利用反向传播来调节参数使得输出和标注靠拢。那么问题来了,怎么比较两个离散概率分布呢,我们可以简单地采用相减的方式,更多细节可以参考 cross-entropyKullback–Leibler divergence

我们现在使用的是高度简化的例子,实际情况中一个句子往往不止一个单词。同时我们希望训练好的模型满足以下性质:

  • 使用和词典词语数等长的向量表示每一个位置的概率分布(例子里为6,实际上会是 30000 或 50000 这么大的数字)
  • 第一个位置的概率分布中单词 “i” 对应的概率值最大
  • 第一个位置的概率分布中单词 “am” 对应的概率值最大
  • And so on, until the fifth output distribution indicates ‘<end of sentence>’ symbol, which also has a cell associated with it from the 10,000 element vocabulary.以此类推,直到第五个位置输出 “<end of sentence>”

上图为训练集标注转换为向量的结果。

当我们在一个足够大的数据集上训练足够长的时间之后,我们期望输出的概率分布应当接近下图所示:

img

我们希望模型在经过训练之后能输出我们期望的结果,当然并不是指只在训练集上能很好的拟合( 详见: cross validation )。请注意,不论是哪个位置,每个单词相应的概率都不为 0,哪怕这本身是个彻底的不可能事件。这是 softmax 函数的性质决定的,它对于模型训练很有帮助。

目前模型的是把每个位置概率最高的词语输出,这种方式称为贪心解码( greedy decoding )。另一种方式则是保留每个位置概率最大的两个词(比如 “I” 和 “a”),下一个位置则假定前一个词语分别为二者跑两次模型,同样保留概率最大的两个结果,在对后续位置采用相同策略。这个方法叫做集束搜索 ( beam search )。在我们的例子里搜索的范围为 2 (即每个时间步会保留概率最大的前两个路径),这是我们可以通过实验效果来确定超参数。

继续深入

我希望本文能让你建立起对 transformer 的初步了解,如果你想更深入的研究,我建议如下方式:

之后的工作:

感谢

Thanks to Illia Polosukhin, Jakob Uszkoreit, Llion Jones , Lukasz Kaiser, Niki Parmar, and Noam Shazeer for providing feedback on earlier versions of this post.