Attention is All You Need

摘要

现阶段的序列变换模型主要采用了具有编码器和解码器(encoder-decoder)的复杂循环或卷积神经网络,其中表现最好的模型通过注意力机制来连接的编码器和解码器。而我们提出了一种新的简单网络架构,即 Transformer 模型。它完全基于注意力机制,避免了循环和卷积的过程。通过两个机器翻译任务的实验表明,采用该模型可以在质量上优于平均水平,并且所需训练时间显著减少。我们的模型在 WMT 2014 英语至德语的翻译任务上 BLEU 分数达到了 28.4,超过现有的最佳结果(包含使用了模型融合的模型),且超过的数值达到了 2 分。在 WMT 2014 英语至法语的翻译任务中,我们的训练模型只采用单个模型,在 8 个 GPU 上训练 3.5 天后,其达到的最好的 BLEU 分数为 41.0,且训练成本只用了现阶段该任务的最佳模型的一小部分。

1. 介绍

循环神经网络,尤其是长短期记忆(LSTM)和门控循环(GRU)神经网络,已经成为了序列建模和变换问题的主要方法,这些网络模型常被用于解决语言建模和机器翻译问题。自此之后,业界一直努力尝试推动循环语言模型和编码器-解码器架构之间的界限。

循环模型通常沿着输入和输出序列的符号位置算出特征值。根据先前隐藏状态 $h_{t-1}$ 和当前时刻 $t$ 的输入,对齐需要计算的位置和时序,生成隐藏状态序列 $h_t$。因为内存约束限制了样本能够处理的长度,所以这种固有的顺序特性导致了循环神经网络无法并行化地训练样本。当输入序列的长度变得很长的时候,这一特性问题就会暴露出来。近期相关领域的研究通过分解序列和添加约束条件的方式在提高计算的效率方面取得了重大进展,尤其当在约束条件下进行训练的同时还能提高模型性能。但只要还是进行序列化的计算,那么这一基本特性约束(译者注:因为前后依赖导致的无法并行化训练)就依然存在。

注意机制已经成为序列建模和变换模型在不同任务中不可或缺的一部分,它允许基于依赖关系进行建模而不考虑输入或输出序列中的位置关系。当然也有少数情况除外,即注意机制与递归网络同时出现的情况。

在本文中,我们提出了 Transformer ,一种规避循环的模型架构,它完全依赖于注意机制来绘制输入和输出之间的全局依赖关系。Transformer 允许模型进行更多的并行化计算,并且在实验中,如果在八个 P100 GPU 上接受不少于十二小时的训练,该模型就可以具有当前最好的翻译质量。

2. 背景

减少序列计算这一目标构成了 Extended Neural GPU,ByteNet 和 ConvS2S 模型的基础,它们都使用卷积神经网络作为基本构筑块,在所有输入和输出位置的计算中都采用并行的方式计算隐藏层。在这些模型中,关联来自两个任意输入或输出位置的信号所需的操作次数会随着位置之间的距离的增加而增加,对于 ConvS2S 而言这一变化是线性的,而对于 ByteNet 则是呈对数形式的。这使得学习远距离位置之间的依赖关系变得困难。在 Transformer 中,尽管由于注意力机制对位置进行加权平均而降低了有效依赖长度,但这一结构可以减少所需的操作次数,这一效果源于我们将在 3.2 节中描述的多头注意力( Multi-Head Attention)。

自我注意力(Self Attention),有时又被称为内部注意力,是一种注意力机制,它将单个序列中不同的位置进行关联来用于计算序列的表示。自我注意力已经成功地用于各种任务,包括阅读理解,抽象概括,文本蕴涵和学习任务独立句子表示等。

端到端内存网络(End-to-end memory networks)基于一种循环的注意力机制,而不是对齐序列后采用循环的方式,并且已被证明在简单的语言问答任务和语言建模任务上表现良好。

然而,据我们所知,Transformer 是第一个使用自我注意力的变换模型来计算其输入和输出的表示,而不使用将序列对齐后采用 RNN 或卷积来进行。在接下来的部分中,我们将描述 Transformer,使用自我注意力的动机并讨论其优于[17,18]和[9]等模型的优势

3. 模型架构

大多数极具竞争性的神经序变换模型都具有编码器—解码器结构 。在本文中,编码器将符号表示 $(x_1,\ldots,x_n)$ 的输入序列映射到连续表示序列 $z=z_1,\ldots,z_n$。当给定 $z$,解码器按照元素一个接一个地生成输出符号序列 $(y_1,\ldots,y_m)$。在每一步中,模型都是自洽的,每当生成下一个时,就会使用先前生成的符号作为附加的输入。

Transformer 采用这种整体架构,分别使用图 1 左半部分和右半部分所示的编码器和解码器。他们采用了堆叠 self-attention 模型和 point-wise 模型,并通过全连接层将他们连接了起来。

3.1 堆叠的编码器和解码器

编码器:编码器由 $N = 6$ 个相同的层组成。每层都有两个子层。第一层是多头自我注意力机制,第二层是一个简单的、按位置进行全连接的前向反馈网络。我们在两个子层中的每一个周围都采用一个残差连接[10],然后对该层进行归一化[1]。也就是说,每个子层的输出是 ${\rm LayerNorm}(x+{\rm Sublayer}(x))$,其中 ${\rm Sublayer}(x)$ 是由子层本身实现的函数。为了更好地使用残差连接,模型中的所有子层以及嵌入层都会生成维度为 $d_{model}=512$ 的输出。

解码器:解码器也由 $N = 6$ 个的相同层组成。除了也有编码器层中的两个子层之外,解码器还插入第三个子层,该层对编码器堆叠的输出进行多头注意力操作。与编码器类似,我们在每个子层之间采用了残差连接,然后对层进行归一化。我们还修改解码器堆中的自我注意力子层,以防止该位置参与后续位置的计算。该掩码的产生基于输出嵌入具有一个位置偏移的事实,这一操作是为了确保 $i$ 位置的预测只能依赖于小于 $i$ 的位置的已知输出。

3.2 注意力

注意力函数可以描述为将 query 和由 key-value 对组成的集合映射到输出,其中 query、keys、values 和输出都是向量。输出的值为这些 values 的加权和,而分配给每个 value 的权重由 query 与相应 key 通过一致性函数来计算。

3.2.1 缩放点积注意力

我们特别关注“缩放点积注意力”(图2)。输入由 query、维度为 $d_k$ 的 key 和维度为 $d_v$ 的 value 所组成。我们计算 query 和所有的 key 的点积,然后除以 $\sqrt{d_k}$,最后使用 softmax 函数来获得 value 的权重。

在实践中,我们同时对 query 的集合计算注意力函数,并将它们组合在一起形成一个矩阵 $Q$。这些 key 和 value 也分别组合成矩阵 $K$ 和 $V$。我们通过下式计算输出矩阵:

$${\rm Attention}(Q,K,V)={\rm softmax}(\frac{QK^T}{\sqrt{d_k}})V$$

两个最常用的注意力函数是加法注意力[2]和点积(乘法)注意力。点积注意力的计算与我们的算法相同,只是增加了 $1/\sqrt{d_k}$ 缩放了因式。加法注意力使用了具有单个隐藏层的前馈网络计算一致性函数。虽然两者在理论上的复杂性相似,但在实践中点积注意力计算的速度更快,更节省空间,因为它可以使用具有高度优化的矩阵乘法代码来实现。

虽然对于较小的 $d_k$ 值,两种注意力机制的表现相似,但是在不对较大的 $d_k$ 值进行缩放[3]的情况下,加法注意力性能优于点积注意力。我们怀疑对于较大的 $d_k$ 值,点积的计算量将大幅度增大,一旦计算量增大就会将 softmax 函数推入梯度极小的区域。为了抵消这种影响,我们才在点积计算中乘以 $1/\sqrt{d_k}$。

3.2.2 多头注意力

我们发现相比于单单使用 $d_{model}$ 维度的 key,value 和 query 来执行单个注意力函数,利用不同的、训练过的线性模型多次映射 queries,keys 和 values 可以获得更好的效果。在每一版本的经过映射后的 query、key 和 value 上,我们并行地执行注意力函数,产生输出值,其维度为 $d_v$。这些输出值进行连接然后再次映射,产生最终的值,如图 2 所示。

多头注意力允许模型中的具有不同表示的子空间去共同关注来自不同位置的信息。如果只有单个注意力头部,平均值会抑制这一信息的表示。

$$ \begin{align} {\rm MultiHead}(Q,K,V) & ={\rm Concat}({\rm head_1,\ldots,head_n})W^O \\ {\rm where\ head_i} & = {\rm Attention}(QW_i^Q,KW_i^K,VW_i^V) \end{align} $$

其中 $W_i^Q\in R^{d_{model}\times d_k}$, $W_i^K\in R^{d_{model}\times d_k}$, $W_i^V\in R^{d_{model}\times d_V}$, $W^o\in R^{hd_v\times d_{model}}$ 。

在这项工作中,我们使用 $h=8$ 个并行的注意力层(头部)。对于其中的每一个头部,我们采用 $d_k=d_v=d_{model}/h=64$。由于每个头部的尺寸都进行了减少,所以总的计算成本与具有全维度的单个头部注意力模型相似。

3.2.3 注意力在 Transformer 中的应用

Transformer 以三种不同方式使用头部注意力:

  • 在“编码器—解码器的注意力”层,queries 来自前一个解码器层,key 和 value 来自编码器的输出。 这允许解码器中的每个位置都能注意到输入序列中的所有位置。这模仿了序列到序列模型中典型的编码器—解码器的注意力机制,例如[38, 2, 9]。
  • 编码器包含自我注意力层。 在自我注意力层中,所有的 key、value 和 query 来自同一地点,在这种情况下,它是编码器中前一层的输出。编码器中的每个位置都可以处理编码器前一层中的所有位置。
  • 类似地,解码器中的自我注意力层允许解码器中的每个位置都关注该解码器和之前的所有位置。我们需要防止解码器中的信息向左流动来保持自回归属性。通过我们在缩放的点积注意力中屏蔽来自 softmax 中的所有不合法连接的输入值(设置为 $-\infty$)来实现。见图 2。

3.3 Position-wise Feed-Forward Networks

除了注意力子层之外,我们的编码器和解码器中的每一层间都采用一个全连接前馈网络,该前馈网络将单独且相同地应用于每个位置。它由两个线性变换组成,其间具有 ReLU 激活函数。

$${\rm FFN}(x)=\max(0,xW_1+b_1)W_2+b_2$$

尽管线性变换在不同位置上是相同的,但它们在层与层之间有着不同的参数。描述这种情况的另一种方式是使用了两个核心大小为1的卷积层。输入和输出的维数为 $d_{model}=512$,内层的维数为 $d_{ff}=2048$。

3.4 Embedding 与 Softmax

与其他序列变换模型类似,我们使用经过训练的 Embedding 模型将输入和输出转换为维度 $d_{model}$ 的向量。此外为了完成预测目标标签的概率,我们还使用经过训练的线性变换与 softmax 函数的组合,这一操作较为常见。在我们的模型中,我们共享两个 Embedding 层(译者注:输入与输出所需的两个)以及线性变换的权重矩阵,该处线性变换是指进行 softmax 之前所用的,类似于[24]。在 Embedding 层中,我们用 $\sqrt{d_{model}}$ 来乘以这些权重。

3.5 位置编码

由于我们的模型不包含循环和卷积层,为了使模型能够利用序列的顺序关系,我们必须加入一些关于序列中每个元素的相对或绝对位置的信息。为此,我们在编码器堆和解码器堆的底部给输入 Embedding 添加“位置编码”。位置编码具有与 Embedding 相同的维数 $d_{model}$,因此这两者可以相加。此外,有若干种位置编码的方式可供我们选择、学习和修补[8]。

在这项工作中,我们使用不同频率的正弦和余弦函数作为位置编码:

$$ \begin{align} PE_{(pos,2i)} & =sin(pos/{10000}^{2i/d_{model}}) \\ PE_{(pos,2i+1)} & =cos(pos/{10000}^{2i/d_{model}}) \end{align} $$

其中 pos 是位置,i 是维度。也就是说,位置编码的每个维度对应于正弦曲线。波长形成从 2π 到10000·2π 的几何级数。我们之所以选择这个函数是因为它可以允许模型比较容易地得到相对位置,对于任何固定的偏移量 k,$PE_{pos+k}$ 都可以表示为 $PE_{pos}$ 的线性函数。

我们还尝试使用 Embedding 技术[8]来学习位置信息,并发现这两个版本产生了几乎相同的结果(参见表 3-行(E))。最终我们选择了正弦曲线版本,因为它可以允许模型使用比训练中的序列更长的序列。

3.6 为什么采用注意力

在本节中,我们从多个方面将 self-attention 与循环和卷积进行了比较,主要针对于将一个可变长度的符号表示序列 $(x_1,\cdots ,x_n)$ 映射到相同长度的另一个序列 $(z_1,\cdots ,z_n )$ 的问题,例如传统序列转导中编码器和解码器的隐藏层。激发我们采用 self-attention 的主要原因来自以下三个需求:

一个是每一层总的计算复杂度;另一个是可以并行化的计算量,通常以所需的最小连续操作数量来衡量。

第三个是网络中长依赖间的路径长度。学习长依赖是许多序列转换任务中的关键性挑战。影响学习这种依赖关系性能的一个关键因素是前向和后向网络中信号传播的路径长度。输入和输出序列中任意位置组合间的这些路径越短,学习长依赖就越容易。因此,我们还比较了由不同类型的层组成的网络中任意两个输入和输出位置之间的最大路径长度。

如表 1 所示,self-attention 层需要将所有位置连接到具有固定数量的顺序操作上,而循环层只需要 $O(n)$ 个连续操作。在计算复杂度方面,当序列长度 n 小于表示维数 d 时,self-attention 层比循环层要快,这在当前机器翻译最好模型所使用的句子表示中非常常见,例如在考虑单词和字节对的表示方法时。为了提高长序列任务的计算性能,可以将 self-attention 限制为只考虑以相应输出位置为中心在输入序列中的大小为 r 的邻域,这会将最大路径长度限制到 $O(n/r)$。我们计划在未来的工作中进一步研究这种方法。

内核宽度为 $k<n$ 的单个卷积层不能连接所有输入和输出的位置对。这样做需要内核连接一个复杂度为 $O(n/k)$ 的卷积层,或者扩大卷积,增加在任何两个位置之间最长路径的长度在网络中,其复杂度为 $O(logk(n))$。卷积层通常比循环层消耗大 k 倍,采用可分割的卷积显著降低了复杂度,使得最终复杂度为 $O(k·n·d+n·d2)$。然而,在 k = n 的情况下,可分割卷积的复杂度是等于我们在 Transformer 模型中采用 self-attention 层和 point-wise 前馈层组合所用复杂度。

此外作为附加的效果,self-attention 可以产生更多可解释的模型。我们检查了我们模型中的注意力分布,并在附录中展示和讨论了示例,发现这不仅独立的注意力头部可以清晰地训练用于不同任务,而且还似乎表现出与句子的句法和语义结构有关的信息。

5 训练

本节介绍我们的模型的训练方法。

5.1 训练集和训练批次

我们对标准的 WMT 2014 英语—德语数据集进行了训练,其中包含约 450 万个句子对。句子使用字节对编码,该编码具有大约 37000 个源 token 的共享目标词汇表。对于英法翻译问题,我们使用了由 36M 个句子组成的更大的 WMT 2014 英法数据集,并将 token 拆分为 32000 个单词词汇表。句子对按照近似的序列长度一起进行批处理。每个批次包含一组句子对,其中包含大约 25000 个源 token 和 25000 个目标 token。

5.2 硬件和流程

我们使用了 8 块 NVIDIA P100 GPU 在一台机器上训练我们的模型。对于依据我们在本文描述的超参数而建立的基本模型,每个训练步骤花费大约 0.4 秒,共训练了 10 万次约 12 小时。而对于大型模型(表 3 的底部所述),单步时间为 1.0 秒,大型模型总共训练了 30 万步(3.5 天)。

5.3 优化器

我们使用 Adam 优化器,$\beta_1 = 0.9,\beta_2= 0.98,\epsilon=10^{-9}$。我们在训练过程中采用变化学习率,公式为:

$$lrate=d_{model}^{-0.5}\cdot \min(step\_num^{-0.5},step\_num\cdot warmup\_step^{-1.5})$$

这对应于为在第一个 $warmup_steps$ 步数内线性地增加学习速率,之后学习速率与步数的反平方根成比例地减小。本文采用了 $warmup_steps = 4000$。

5.4 正则化

训练期间我们采用三种正则化方法:

残差 Dropout:在将子层的输出添加到下一子层输入并进行归一化之前,我们将 Dropout 应用于每个输出。此外,我们将 Dropout 也添加到编码器和解码器堆中对 Embedding 和位置编码求和阶段。对于基线模型,我们使用 $P_{drop}=0.1$。

标签平滑:在训练过程中,我们使用标签平滑的值为 $\epsilon_{ls} = 0.1$。模型学习变得更不确定,变得更加泛化,提高了准确性和 BLEU 分数。

6 实验结果

6.1 机器翻译

在 WMT 2014 英语至德语翻译任务中,大型 transformer 模型(表 2 中的 Transformer(Big))比以前报告的最优模型(包括模型融合)的 BLEU 分数高出 2.0 以上,成为了一个全新的 SOTA 模型, BLEU 分数优雅地达到了 28.4。该模型的配置列于表 3 的底部。训练在 8 台 P100 GPU 上花费 3.5 天。不仅我们的基础模型超过了所有以前发布的单个模型和融合模型,而且我们的训练成本只占所有参与竞争的模型的训练成本中的一小部分。

在 WMT 2014 英语到法语翻译任务中,我们的大型模型的 BLEU 成绩为 41.0,超过了之前发布的所有单一模型,比以前的最新技术训练成本低了 1/4。英语到法语训练的变种(大)模型使用的 Dropout 为 $P_{drop}=0.1$,而不是 0.3。

对于基本模型,我们使用经过平均后的最后 5 个检查点获得的单个模型,这些检查点以 10 分钟的间隔写入。对于大型模型,我们对最后 20 个检查点进行了平均。我们使用波束大小为 4,长度惩罚 $\alpha = 0.6$ 的波束搜索。这些超参数是在训练集上进行实验后选定的。我们推断输入长度 $+50$ 时设置为最大输出长度,但在必要时尽早终止(early stop)。

表 2 总结了我们的结果,并将我们的翻译质量和训练成本与文献中的其他模型架构进行了比较。我们通过将训练时间,使用的GPU数量以及每个GPU 的持续单精度浮点容量的估计相乘来估计用于训练模型的浮点运算的数量。

6.2 模型的变种

为了评估 Transformer 不同组件的重要性,我们以不同的方式改变了我们的基础模型,测量开发集 newstest2013 上英文到德文翻译的性能变化。我们使用前一节所述的波束搜索,但没有对检查点平均。我们在表 3 中列出这些结果。

在表 3 行(A)中,我们改变了注意力头部的数量和 key 和 value 的维度,保持计算量不变,如3.2.2 节所述。虽然单个头部注意力比最佳设置差 0.9 个 BLEU 分数,但翻译质量也会随着头部太多而下降。

在表 3 行(B)中,我们观察到减小注意力中 key 的大小 $d_k$ 会损害模型的质量。这表明确定一致性并不容易,并且比点积更复杂的一致性函数可能是有益的。我们在行(C)和(D)中进一步观察到,如预期的那样,更大(更复杂)的模型更好,并且提早结束对避免过度拟合非常有帮助。在行(E)中,我们用学习到的位置嵌入代替了我们的正弦位置编码,并观察到与基本模型几乎相同的结果。

6.3 英语成分句法分析

为了评估 Transformer 是否可以推广到其他任务,我们进行了英语成分句法分析任务的实验。此任务提出了更具体的挑战:输出受到强大的结构约束,并且明显长于输入。此外,RNN 的序列到序列模型还无法在小数据体系中获得最先进的结果。

我们在 Penn Treebank 的华尔街日报(WSJ)部分训练了一个 $d_{model} = 1024$ 的 4 层 Transformer,大约 40K 训练语句。我们还使用更大的高可信度的 BerkleyParser 语料库包括大约17M 的句子在半监督环境中训练它。我们使用 16K 的 token 的词汇表用于 WSJ 任务的设置,并使用 32K 的 token 词汇表用于半监督任务的设置。

我们仅进行了少量实验,来选择 dropout 大小(包括注意力和残差中的 dropout),学习率和波束大小,其他参数与英语-德语基础翻译模型保持不变。在训练过程中,我们将最大输出长度增加到输入长度的 $+300$。我们使用的波束尺寸为 21, $\alpha= 0.3$ 用于 WSJ 和半监督任务的设置。

我们在表 4 中的结果表明,尽管缺乏任务特定的调整,我们的模型表现出色得多,产生的结果比之前报道的除了 Recurrent Neural Network Grammar 外的所有模型都要好。

与 RNN 序列到序列模型相比,Transformer 的表现优于 BerkeleyParser,即使只训练 40K 句子的 WSJ 训练集也是如此。

7 结论

在这项工作中,我们提出了 Transformer,第一个完全基于注意力的序列转换模型,用多头自注意力取代了编码器—解码器架构中最常用的循环层。

对于翻译任务,Transformer 可以比基于循环或卷积层的架构更快地进行训练。在 WMT 2014 英语对德语和 WMT 2014 英语至法语翻译任务中,我们达到了了最顶尖的水平。在之前的实验中,我们最好的模型甚至胜过以前报道过的所有模型融合。

我们对基于注意力模型的未来感到兴奋,并计划将它们应用于其他任务。我们计划将 Transformer 扩展到涉及文本以外的输入与输出架构的问题,并调查局部的有限注意力机制以有效处理大型输入和输出,如图像,音频和视频。产生更少顺序信息是我们的另一个研究目标。我们用来训练和评估我们的模型的代码可以在 http://github.com/tensorflow/tensor2tensor 中找到。