论文笔记《Attention Is All You Need》

Attention Is All You Need

原文链接:论文笔记《Attention Is All You Need》 | Karl的博客

CSDN链接:论文笔记《Attention Is All You Need》-CSDN博客

论文链接:[1706.03762] Attention Is All You Need (arxiv.org)

代码链接:tensorflow/tensor2tensor: Library of deep learning models and datasets designed to make deep learning more accessible and accelerate ML research. (github.com)

Abstract

现阶段,占主导地位的序列转换模型基于复杂的循环或卷积神经网络,包括编码器和解码器。性能最好的模型还通过注意力机制连接编码器和解码器。我们提出了一种新的简单网络架构——Transformer,它完全基于注意力机制,完全摒弃了递归和卷积。在两个机器翻译任务上的实验表明,这些模型在质量上更优,同时具有更强的并行性,需要的训练时间显著减少。我们的模型在WMT 2014英德翻译任务中的BLEU值为28.4,比现有的最佳BLEU值(包括集合)提高了2。在WMT 2014英法翻译任务中,我们的模型在8个GPU上经过3.5天的训练后,建立了一个新的SOTA单模型,其BLEU得分为41.0,这只是文献中最好的模型的一小部分。

1 Introduction

循环神经网络,特别是长短期记忆[12]和门控循环[7]神经网络,已被牢固地确立为序列建模和转换问题(例如语言建模和机器翻译)中最先进的方法[29, 2, 5]。此后,人们做出了许多努力,不断突破循环语言模型和编码器-解码器架构的边界[31, 21, 13]。

循环模型通常根据输入和输出序列的符号位置进行计算。将位置与计算时间中的步骤对齐,它们生成一个隐藏状态序列$h_t$,作为前一个隐藏状态$h_{t − 1}$和位置$t$的输入的函数。这种固有的顺序性妨碍了训练样例中的并行化,这在较长的序列长度时变得非常关键,因为内存限制限制了样例之间的批处理。最近的工作通过分解技巧[18]和条件计算[26]在计算效率上取得了显著的提高,同时也提高了后者的模型性能。然而,顺序计算的基本约束仍然存在。

注意力机制已经成为各种任务中引人注目的序列建模和转换模型的组成部分,允许对依赖关系进行建模,而不考虑它们在输入或输出序列中的距离[2, 16]。然而,除了少数情况外[22],这种注意力机制都是与循环网络结合使用的。

在本工作中,我们提出了Transformer,它是一种模型架构,避免了递归,完全依赖于注意力机制来绘制输入和输出之间的全局依赖关系。Transformer支持更多的并行化,在8个P100 GPU上经过12个小时的训练后,可以在翻译质量上达到一个新的水平。

2 Background

减少顺序计算的目标也构成了Extended Neural GPU[20]、ByteNet[15]和ConvS2S[8]的基础,它们都使用卷积神经网络作为基本的构建块,并行计算所有输入和输出位置的隐藏表示。在这些模型中,将来自两个任意输入或输出位置的信号关联起来所需的操作数量,按照位置之间的距离增长,对ConvS2S来说是线性增长,对ByteNet来说是对数增长。这使得学习遥远位置[11]之间的依赖关系变得更加困难。在Transformer中,这被减少为一个固定数量的操作,尽管因为平均注意力加权位置的影响,有效分辨率会降低,所以我们会用3.2节中描述的多头注意力抵消这一影响。

自注意力,有时也被称为内部注意力,是一种将单个序列的不同位置联系起来以计算序列的表示形式的注意力机制。自注意力已成功地应用于各种任务,包括阅读理解、抽象摘要、文本蕴涵和学习任务独立的句子表征[4, 22, 23, 19]。

端到端记忆网络基于循环注意力机制而不是序列对齐循环,并且在简单语言问答和语言建模任务[28]中表现良好。

然而,据我们所知,Transformer是第一个完全依赖自注意力来计算其输入和输出表示而不使用序列对齐RNN或卷积的转换模型。在接下来的章节中,我们将描述Transformer,自注意力的动机并讨论其相对于[14, 15]和[8]等模型的优势。

3 Model Architecture

大多数具有竞争力的神经序列转换模型都有编码器-解码器结构[5, 2, 29]。这里,编码器将符号表示的输入序列$(x_1, \cdots, x_n)$映射到连续表示序列$\mathbf{z} = (z_1, \cdots, z_n)$,给定$\mathbf{z}$,解码器每次生成一个元素的符号输出序列$(y_1, \cdots, y_m)$。在每个步骤中,模型都是自回归的[9],在生成下一个步骤时,使用之前生成的符号作为附加输入。

Transformer遵循这种整体架构,为编码器和解码器使用堆叠的自注意力层和逐点完全连接的层,分别如图1的左、右两部分所示。

image-20240209112650024

3.1 Encoder and Decoder Stacks

Encoder:编码器由$N = 6$个相同的层堆叠而成。每个层有两个子层。第一个子层是多头自注意力机制,第二个子层是简单的、位置明智的全连接前馈网络。我们在每个子层都使用了一个残差连接[10],然后是层归一化[1]。也就是说,每个子层的输出都是$\text{LayerNorm}(x + \text{Sublayer}(x))$,其中$\text{Sublayer}(x)$是子层本身实现的函数。为了便于这些残差连接,模型中的所有子层以及嵌入层都产生了维度为$d_{\text{model}} = 512$的输出。

Decoder:解码器也由$N = 6$个相同的层堆叠而成。除了每个编码器层中的两个子层外,解码器还插入了第三个子层,该子层对编码器的输出执行多头注意力。与编码器类似,我们在每个子层周围使用残差连接,然后进行层归一化。我们还修改了解码器中的自注意力子层,以防止位置对后续位置的影响。这种掩蔽,结合输出嵌入偏移一个位置的事实,确保了位置$i$的预测只能依赖于位置小于$i$的已知输出。

3.2 Attention

注意力函数可以描述为将查询和一组键值对映射到输出,其中查询、键、值和输出都是向量。输出是作为值的加权和计算的,其中分配给每个值的权重是通过查询与相应键的兼容性函数计算的。

3.2.1 Scaled Dot-Product Attention

我们称我们的特别注意力为“缩放的点积注意力”(图2)。

image-20240209115434068

输入由维度$d_k$的查询、键和维度$d_v$的值组成。我们计算查询与所有键的点积,每个键除以$\sqrt{d_k}$,并应用softmax函数来获得值的权重。

在实践中,我们同时计算一组查询的注意力函数,这些查询被打包成一个矩阵$Q$。键和值也打包成矩阵$K$和$V$,我们计算输出矩阵如下:

最常用的两个注意力函数是加性注意力[2]和点积(乘法)注意力。点积注意力与我们的算法相同,除了比例因子为$\frac{1}{\sqrt{d_k}}$。加性注意力使用带有单个隐藏层的前馈网络计算兼容性函数。虽然两者在理论上的复杂性相似,但点积注意力在实践中要快得多,而且空间效率更高,因为它可以使用高度优化的矩阵乘法代码来实现。

当$d_k$值较小时,两种机制表现相似,当$d_k$值较大时,不需缩放的加性注意力优于点积注意力[3]。我们猜想,当$d_k$值较大时,点积的幅度会变大,从而使softmax函数进入其梯度极小的区域。为了抵消这种影响,我们将点积乘以$\frac{1}{\sqrt{d_k}}$。

3.2.2 Multi-Head Attention

我们发现,将查询、键和值分别用不同的、学习过的线性投影投影到$d_k$、$d_k$和$d_v$维$h$次后,比用单个维度为$d_{\text{model}}$的注意力函数更有效。然后,在查询、键和值的每一个投影版本上,我们并行执行注意力函数,产生$d_v$维的输出值。这些输出被连接起来并再次投影,最终得到图2所示的值。

多头注意力允许模型在不同位置共同注意来自不同表示子空间的信息。对于单一注意力头,后续的平均操作就会抑制这一点。

其中投影是参数矩阵$W_i^Q \in \mathbb{R}^{d_\text{model} \times d_k}, W_i^K \in \mathbb{R}^{d_\text{model} \times d_k}, W_i^V \in \mathbb{R}^{d_\text{model} \times d_v}$,$W^O \in \mathbb{R}^{hd_v \times d_\text{model}}$。

在这项工作中,我们采用了$h = 8$个并行注意力层或头。对于每一个注意力层或头,我们使用$d_k = d_v = \frac{d_\text{model}}{h} = 64$。由于每个头的维度减小了,总的计算成本与全维度的单头注意力相似。

3.2.3 Applications of Attention in our Model

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

  • 在“编码器-解码器注意力”层中,查询来自先前的解码器层,而记忆键和值来自编码器的输出。这使得解码器中的每个位置都能处理输入序列中的所有位置。这模仿了序列到序列模型中典型的编码器-解码器注意力机制,如[31, 2, 8]。
  • 编码器包含自注意力层。在自注意力层中,所有的键、值和查询都来自同一个地方,在本例中,即编码器中前一层的输出。编码器中的每个位置可以处理所述编码器的前一层中的所有位置。
  • 类似地,解码器中的自注意力层允许解码器中的每个位置注意到解码器中的所有位置直至并包括该位置。我们需要防止信息在解码器中向左流动,以保持自回归特性。我们通过屏蔽(设置为$-\infty$)softmax输入中与非法连接对应的所有值来实现缩放点乘注意力。参见图2。

3.3 Position-wise Feed-Forward Networks

除了注意力子层外,我们的编码器和解码器中的每一层都包含一个全连接的前馈网络,该网络分别和相同地应用于每个位置。这包括两个线性转换,中间有一个ReLU激活。

虽然在不同的位置上线性转换是相同的,但它们在不同的层中使用不同的参数。另一种描述它的方法是两个卷积,卷积核大小为1。输入输出维数为$d_{model} = 512$,内层维数为$d_{ff} = 2048$。

3.4 Embeddings and Softmax

与其他序列转换模型类似,我们使用学习过的嵌入将输入token和输出token转换为维数为$d_\text{model}$的向量。我们还使用通常学到的线性变换和softmax函数将解码器输出转换为预测的下一token概率。在我们的模型中,我们共享两个嵌入层之间的权值矩阵和pre-softmax线性变换,类似于[24]。在嵌入层中,我们将这些权重乘以$\sqrt{d_\text{model}}$。

3.5 Positional Encoding

由于我们的模型不包含递归和卷积,为了让模型利用序列的顺序,我们必须注入一些关于序列中token的相对或绝对位置的信息。为此,我们将“位置编码”添加到编码器和解码器底部的输入嵌入中。位置编码与嵌入具有相同的维度$d_\text{model}$,因此可以将两者相加。有许多位置编码的选择,比如可学习的位置编码和固定的位置编码[8]。

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

其中$pos$是位置,$i$是维数。也就是说,位置编码的每个维度都对应一个正弦信号。波长从$2\pi$到$10000 \cdot 2\pi$,呈几何级数。我们选择这个函数是因为我们假设它可以让模型很容易通过相对位置进行学习,因为对于任何固定偏移量$k$,$PE_{pos + k}$可以表示为$PE_{pos}$的线性函数。

我们还使用学习过的位置嵌入[8]进行了实验,发现两个版本产生了几乎相同的结果(见表3行(E))。我们选择正弦版本是因为它可以让模型外推到比训练中遇到的序列更长的序列长度。

4 Why Self-Attention

在本节中,我们将自注意力层的各个方面与循环层和卷积层进行比较,这些层通常用于将符号表示的一个可变长度序列$(x_1, \cdots, x_n)$映射到另一个等长度序列$(z_1, \cdots, z_n)$,其中$x_i, z_i \in \mathbb{R}^d$,例如典型序列转换编码器或解码器中的隐藏层。关于使用自注意力的动机,我们考虑三个需求。

一个是每一层的总计算复杂度。另一个是可以并行化的计算量,由所需的最小顺序操作数来衡量。

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

如表1所示,自注意力层用固定数量的顺序执行操作连接所有位置,而循环层需要$O(n)$个顺序操作。在计算复杂度方面,当序列长度$n$小于表示维数$d$时,自注意力层比循环层更快,这是最常用的情况,在最先进的机器翻译模型中使用的句子表示,如词块[31]和字节对[25]表示。

image-20240209132432948

为了提高涉及非常长的序列的任务的计算性能,可以将自注意力限制在只考虑输入序列中大小为$r$的以各自输出位置为中心的邻域。这将增加最大路径长度到$O(\frac{n}{r})$。我们计划在未来的工作中进一步研究这种方法。

核宽度为$k < n$的单个卷积层并不连接所有的输入和输出位置对。在相邻核的情况下,这样做需要$O(\frac{n}{k})$个卷积层的堆叠,在扩张卷积[15]的情况下,需要$O(\log_k(n))$个卷积层的堆叠,增加网络中任意两个位置之间的最长路径的长度。卷积层通常比递归层贵$k$倍。而可分离卷积[6]则大大降低了卷积的复杂度,达到$O(k \cdot n \cdot d + n \cdot d^2)$。然而,即使在$k = n$的情况下,可分离卷积的复杂性也等于自注意力层和点前馈层的结合,这是我们在模型中采用的方法。

作为附加的好处,自注意力可以产生更多可解释的模型。我们检查了我们模型中的注意力分布,并在附录中展示和讨论了一些例子。不仅个别注意力头清楚地学会了执行不同的任务,许多注意力头似乎表现出与句子的句法和语义结构有关的行为。

5 Training

本节描述了我们模型的训练方法。

5.1 Training Data and Batching

我们使用由450万对句子组成的WMT 2014英德标准数据集进行训练。使用[3]编码的字节对对句子进行编码,该编码具有大约37000个标记的共享源目标词汇表。对于英法两种语言,我们使用更大的WMT 2014英法数据集,其中包含3600万个句子,并将token分解为32000个单词块词汇[31]。句子对按照近似的序列长度进行分组。每个训练批包含一组句子对,其中包含大约25000个源标记和25000个目标标记。

5.2 Hardware and Schedule

我们在一台装有8个NVIDIA P100 GPU的机器上训练我们的模型。对于使用本文中描述的超参数的基本模型,每个训练步骤大约需要0.4秒。我们对基本模型进行了总共10万步或12小时的训练。对于我们的大型模型(如表3所示),步长为1.0秒。这些大型模型被训练了30万步(3.5天)。

5.3 Optimizer

我们使用了Adam优化器[17],其中$\beta_1 = 0.9, \beta_2 = 0.98, \epsilon = 10^{-9}$。根据公式,我们在训练过程中改变学习率:

这对应于第一个$warmup_steps$训练步骤的学习速率线性增加,然后与步骤数的平方根倒数成比例减少。我们使用$warmup_steps = 4000$。

5.4 Regularization

在训练过程中,我们采用了三种正则化方法:

Residual Dropout,我们对每个子层的输出应用dropout[27],然后将其添加到子层输入并进行规范化。此外,我们在编码器和解码器中对嵌入和位置编码的和应用了dropout。对于基础模型,我们使用$P_{drop} = 0.1$的速率。

Label Smoothing,在训练过程中,我们使用了$\epsilon_{ls} = 0.1$[30]的标签平滑。这会降低模型的困惑度,因为模型会变得更加不确定,但会提高准确性和BLEU分数。

6 Results

6.1 Machine Translation

在2014年WMT英德翻译任务中,big transformer模型(表2中的Transformer (big))的性能比之前报告的最佳模型(包括集成模型)在BLEU评分上高出了2.0分以上,从而达到了一个SOTA的BLEU评分28.4。

image-20240209170517873

该模型的配置列于表3的底部。在8个P100 GPU上训练了3.5天。

image-20240209170606266

甚至我们的基础模型超过了所有以前发表的模型和集成模型,而训练成本只是任何竞争模型的一小部分。

在WMT 2014英法翻译任务中,我们的大模型获得了41.0的BLEU分数,超过了之前发布的所有单个模型,而训练成本不到之前最先进的模型的$\frac{1}{4}$。Transformer (big)模型的英语-法语训练的dropout概率为$P_{drop} = 0.1$,而不是$0.3$。

对于基本模型,我们使用了通过平均最后5个检查点获得的单个模型,这些检查点每10分钟记录一次。对于大型模型,我们计算了最后20个检查点的平均值。我们采用了束大小为4,长度惩罚$\alpha = 0.6$[31]的束搜索。这些超参数是在开发集上经过实验后选定的。在推理期间,我们将最大输出长度设置为输入长度$+50$,但在可能的情况下提前终止[31]。

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

6.2 Model Variations

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

在表3行(A)中,我们改变注意力头的数量以及注意力键和值维度,保持计算量不变,如章节3.2.2所述。虽然单头注意力比最好的设置差0.9 BLEU,但过多的头也会降低质量。

在表3行(B)中,我们观察到降低注意力键大小$d_k$会影响模型质量。这表明确定兼容性并不容易,一个比点积更复杂的兼容性函数可能是有益的。我们在(C)和(D)行进一步观察到,正如预期的那样,更大的模型更好,并且dropout在避免过拟合方面非常有帮助。在行(E)中,我们将正弦位置编码替换为学习过的位置嵌入[8],并观察到与基础模型几乎相同的结果。

7 Conclusion

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

对于翻译任务,Transformer的训练速度比基于循环或卷积层的体系结构快得多。在WMT 2014英语到德语和WMT 2014英语到法语的翻译任务中,我们都达到了一个新的水平。在前一个任务中,我们的最佳模型甚至优于所有之前报告模型的集成。

我们对基于注意力的模型的未来感到兴奋,并计划将其应用于其他任务。我们计划将Transformer扩展到涉及文本以外的输入和输出模式的问题,并研究局部受限的注意力机制,以有效地处理图像、音频和视频等大型输入和输出。使生成过程更少地依赖于序列是我们的另一个研究目标。

我们用来训练和评估模型的代码可以在https://github.com/tensorflow/tensor2tensor中找到。