Attention Is All You Need（论文）

按原文结构转述 Vaswani 等 2017 年论文：架构、注意力公式、训练与实验结果。

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin. Attention Is All You Need. NIPS 2017. arXiv:1706.03762（HTML、PDF）。八位作者均标注为同等贡献。文首注明：在注明出处的前提下，Google 允许在学术或新闻作品中复制该文的表与图。

中文译本：Arthur Chiao、GitHub LaTeX 中译、博客园。民间译本质量不一，英德/英法 BLEU 有的把摘要 41.8 和正文 41.0 抄混。

摘要

当时占主导的序列转换（sequence transduction）模型，建立在编码器—解码器结构上的复杂循环网络或卷积网络。表现最好的模型还会用注意力把编码器与解码器连起来。

本文提出一种简单网络架构 Transformer：只基于注意力机制，完全不用循环和卷积。在两项机器翻译任务上，质量更高、更易并行、训练时间显著更短。

  • WMT 2014 英→德：28.4 BLEU，比此前最佳结果（含模型集成）高 2 个 BLEU 以上。
  • WMT 2014 英→法：单模型 41.8 BLEU（正文 6.1 节叙述为 41.0；表 2 为 41.8），在 8 块 GPU 上训练 3.5 天，成本只占文献中最佳模型的一小部分。
  • 在英语成分句法分析上，无论训练数据多还是少，都表明该架构可以推广到翻译以外的任务。

1 引言

循环神经网络，尤其 LSTM 与门控循环网络，已被确立为语言建模、机器翻译等序列建模与转换任务的主流。循环模型把计算沿符号位置展开：位置 (t) 的隐状态 (h_t) 依赖 (h_{t-1}) 与当前位置输入。这种顺序性阻碍单个训练样本内部的并行；序列一长，内存又限制跨样本的 batch，问题更严重。因子分解技巧、条件计算可以提高效率，但顺序计算这一条约束仍然在。

注意力已经是序列模型里的常规组件：它可以建模依赖，而不必管依赖在输入或输出里隔了多远。但除少数工作外，注意力仍与循环网络一起用。

本文提出 Transformer：去掉循环，完全靠注意力建立输入与输出之间的全局依赖。并行度显著提高；在 8 块 P100 上训练约 12 小时，即可达到新的翻译质量水平（指 base 模型）。

2 背景

减少顺序计算，也是 Extended Neural GPU、ByteNet、ConvS2S 的出发点。它们用卷积对所有位置并行算隐表示。任意两个位置之间要建立联系，所需运算次数随距离增长：ConvS2S 线性，ByteNet 对数。远距离依赖因此更难学。

在 Transformer 里，这一次数变成常数。代价是对注意力加权位置做平均会降低有效分辨率；第 3.2 节用多头注意力抵消。

自注意力（self-attention，也称 intra-attention）把同一序列里不同位置联系起来，以计算该序列的表示。它已用于阅读理解、摘要、文本蕴含、与任务无关的句子表示等。端到端记忆网络用循环式注意力代替按序列对齐的循环，在简单问答与语言建模上有效。

据作者所知，Transformer 是第一个完全靠自注意力计算输入与输出表示、而不使用按序列对齐的 RNN 或卷积的转换模型。

3 模型架构

多数有竞争力的神经序列转换模型是编码器—解码器。编码器把输入符号序列 ((x_1,\ldots,x_n)) 映射为连续表示 (\mathbf{z}=(z_1,\ldots,z_n))。解码器据此一次一个符号地生成 ((y_1,\ldots,y_m))。每一步是自回归的：生成下一个符号时，把已经生成的符号当作额外输入。

Transformer 沿用这一总体结构。编码器、解码器都用堆叠的自注意力和逐位置全连接层（原文 Figure 1 左右两半）。

3.1 编码器栈与解码器栈

编码器。 (N=6) 层相同的层。每层两个子层：多头自注意力；逐位置全连接前馈网络。每个子层外包残差连接，再接层归一化：

[ \mathrm{LayerNorm}(x + \mathrm{Sublayer}(x)) ]

所有子层以及嵌入层的输出维度均为 (d_{\mathrm{model}}=512)，以便做残差。

解码器。 同样 (N=6) 层。在编码器那两个子层之外，插入第三个子层：对编码器栈的输出做多头注意力。同样使用残差与层归一化。解码器里的自注意力被改成不能看后续位置。这种掩码，再加上输出嵌入右移一个位置，保证对位置 (i) 的预测只依赖小于 (i) 的已知输出。

3.2 注意力

注意力：把一个 query 和一组 key–value 对映射到一个输出；四者都是向量。输出是 value 的加权和；每个 value 的权重由 query 与对应 key 的兼容性函数算出。

3.2.1 缩放点积注意力

输入：维度 (d_k) 的 query 与 key，维度 (d_v) 的 value。用 query 与所有 key 做点积，再除以 (\sqrt{d_k})，经 softmax 得到 value 上的权重。一组 query 打成矩阵 (Q)，key、value 打成 (K)、(V)：

[ \mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V ]

两种常用注意力：加性注意力（Bahdanau 等）与点积（乘性）注意力。点积注意力与上式相同，只是没有 (1/\sqrt{d_k}) 这个缩放。加性注意力用单隐层前馈网算兼容性。理论复杂度相近，实践上点积可用高度优化的矩阵乘，更快、更省空间。

(d_k) 较小时两者表现接近；(d_k) 较大且不缩放时，加性更好。作者认为 (d_k) 大时点积幅值变大，把 softmax 推到梯度极小的区域（若 (q,k) 各维独立、均值 0 方差 1，则 (q\cdot k) 的方差为 (d_k)）。因此除以 (\sqrt{d_k})。

3.2.2 多头注意力

不是用 (d_{\mathrm{model}}) 维做一次注意力，而是把 (Q,K,V) 用不同的、可学习的线性投影做 (h) 次，投影到 (d_k,d_k,d_v)，再并行做注意力，把 (d_v) 维输出拼接起来，再投影一次：

[ \mathrm{MultiHead}(Q,K,V)=\mathrm{Concat}(\mathrm{head}_1,\ldots,\mathrm{head}_h)W^{O} ]

[ \mathrm{head}_i=\mathrm{Attention}(QW_i^{Q},,KW_i^{K},,VW_i^{V}) ]

其中 (W_i^{Q}\in\mathbb{R}^{d_{\mathrm{model}}\times d_k}) 等。多头让模型能同时从不同表示子空间、不同位置取信息；单头会被平均抑制。

本文 (h=8)。每个头 (d_k=d_v=d_{\mathrm{model}}/h=64)。头变窄之后，总计算量与全维度单头相近。

3.2.3 在本模型中的三种用法

  1. 编码器—解码器注意力。 query 来自上一解码器层；key 与 value 来自编码器输出。解码器每个位置可以关注输入序列所有位置。对应以往 seq2seq 里的编码器—解码器注意力。
  2. 编码器自注意力。 (Q,K,V) 都来自编码器上一层输出。每个位置可以关注上一层所有位置。
  3. 解码器自注意力。 每个位置可以关注解码器中不超过该位置的位置。为保持自回归、禁止信息左流，在缩放点积注意力里把 softmax 输入中非法连接置为 (-\infty)。

3.3 逐位置前馈网络

每个编码器、解码器层除注意力子层外，还有对每个位置单独、相同地施加的全连接前馈网：两层线性，中间 ReLU。

[ \mathrm{FFN}(x)=\max(0, xW_1+b_1)W_2+b_2 ]

不同位置共享同一套线性变换；层与层之间参数不同。也可看成核大小为 1 的两次卷积。输入输出维 (d_{\mathrm{model}}=512)，内层 (d_{ff}=2048)。

3.4 嵌入与 Softmax

与其他序列转换模型一样，用可学习嵌入把输入、输出 token 变成 (d_{\mathrm{model}}) 维向量。解码器输出经线性变换与 softmax，得到下一个 token 的概率。两个嵌入层与 softmax 前的线性变换共享同一权重矩阵。嵌入层中把这些权重乘以 (\sqrt{d_{\mathrm{model}}})。

3.5 位置编码

模型没有循环也没有卷积，必须向序列注入相对或绝对位置信息。在编码器、解码器栈底部，把位置编码加到输入嵌入上；维度同为 (d_{\mathrm{model}})，以便相加。

本文用不同频率的正弦与余弦：

[ PE_{(pos,2i)}=\sin(pos/10000^{2i/d_{\mathrm{model}}}) ]

[ PE_{(pos,2i+1)}=\cos(pos/10000^{2i/d_{\mathrm{model}}}) ]

(pos) 是位置，(i) 是维度。波长从 (2\pi) 到 (10000\cdot 2\pi) 成几何级数。选择该函数是因为：对任意固定偏移 (k)，(PE_{pos+k}) 可写成 (PE_{pos}) 的线性函数，便于按相对位置做注意力。

也试过可学习的位置嵌入，与正弦版本结果几乎相同（表 3 行 E）。选用正弦，是因为它可能外推到训练时未见过的序列长度。

4 为什么用自注意力

比较自注意力、循环、卷积三类层：把变长符号序列 ((x_1,\ldots,x_n)) 映到等长 ((z_1,\ldots,z_n))，(x_i,z_i\in\mathbb{R}^{d})。三条标准：每层计算复杂度；可并行程度（最少顺序操作数）；网络中长程依赖的路径长度。路径越短，长程依赖越好学。

表 1. (n) 为序列长度，(d) 为表示维度，(k) 为卷积核宽，(r) 为受限自注意力的邻域大小。

层类型 每层复杂度 顺序操作 最大路径长度
自注意力 (O(n^{2}\cdot d)) (O(1)) (O(1))
循环 (O(n\cdot d^{2})) (O(n)) (O(n))
卷积 (O(k\cdot n\cdot d^{2})) (O(1)) (O(\log_{k} n))
受限自注意力 (O(r\cdot n\cdot d)) (O(1)) (O(n/r))

自注意力用常数次顺序操作连接所有位置；循环需要 (O(n))。当 (n<d) 时（机器翻译里 word-piece / BPE 句子表示常见如此），自注意力比循环更快。极长序列可以把注意力限制在半径 (r) 的邻域，最大路径变为 (O(n/r))；文中列为未来工作。

单层卷积在 (k<n) 时连不满所有位置对；要连满需要堆 (O(n/k)) 层（连续核）或 (O(\log_{k} n)) 层（膨胀卷积），从而加长最长路径。

附带好处：自注意力可能更可解释。附录可视化显示，不同头学会不同工作，不少头与句法、语义结构有关。

5 训练

5.1 数据与 batch

英德：WMT 2014，约 450 万句对；BPE，源/目标共享约 37,000 token 词表。英法：WMT 2014，3600 万句，32,000 word-piece 词表。按近似长度组 batch；每个训练 batch 约 25,000 个源 token 与 25,000 个目标 token。

5.2 硬件与日程

一台机器、8 块 NVIDIA P100。base：每步约 0.4 秒，共 100,000 步（12 小时）。big：每步 1.0 秒，300,000 步（3.5 天）。

5.3 优化器

Adam，(\beta_1=0.9)，(\beta_2=0.98)，(\epsilon=10^{-9})。学习率：

[ lrate = d_{\mathrm{model}}^{-0.5}\cdot \min(\mathrm{step_num}^{-0.5},; \mathrm{step_num}\cdot \mathrm{warmup_steps}^{-1.5}) ]

前 (\mathrm{warmup_steps}) 步线性升高，之后按步数平方根的倒数下降。(\mathrm{warmup_steps}=4000)。

5.4 正则

残差 Dropout。 每个子层输出在与输入相加、归一化之前做 dropout；嵌入与位置编码之和也做 dropout。base 模型 (P_{drop}=0.1)。

标签平滑。 (\epsilon_{ls}=0.1)。会损害困惑度（模型学得更不确定），但提高准确率与 BLEU。

6 结果

6.1 机器翻译

表 2（newstest2014；训练成本为估计 FLOPs）。Transformer (big) 英德 28.4，英法 41.8；base 英德 27.3，英法 38.1。英德 big 比此前最佳（含集成）高 2 BLEU 以上。base 也超过此前所有已发表模型与集成，训练成本只占竞争模型的一小部分。英法 big 的训练成本不到此前 SOTA 的 1/4。英法 big 使用 (P_{drop}=0.1) 而非 0.3。

推理：base 对最后 5 个 checkpoint 做平均（每 10 分钟一个）；big 平均最后 20 个。束搜索束宽 4，长度惩罚 (\alpha=0.6)。最大输出长度为输入长度 + 50，可提前结束。

6.2 模型变体

在英德开发集 newstest2013 上改 base（束搜索，不做 checkpoint 平均）。要点（表 3）：

  • (A) 头数与 (d_k,d_v)：单头比最佳设置差 0.9 BLEU；头太多质量也下降。
  • (B) 减小 (d_k) 损害质量，说明兼容性不容易，可能需要比点积更复杂的函数。
  • (C)(D) 更大的模型更好；dropout 对过拟合很有帮助。
  • (E) 正弦位置编码换成可学习位置嵌入，结果几乎相同。
  • big：(N=6)，(d_{\mathrm{model}}=1024)，(d_{ff}=4096)，(h=16)，(P_{drop}=0.3)，300K 步，约 2.13 亿参数。base 约 6500 万参数。

6.3 英语成分句法分析

输出有强结构约束，且比输入长得多。RNN seq2seq 在小数据上未达到当时 SOTA。4 层 Transformer，(d_{\mathrm{model}}=1024)。仅 WSJ（约 4 万句）F1 91.3；半监督（约 1700 万句）92.7。除 RNN Grammar 外，好于此前报告的模型。仅用 WSJ 4 万句时，也超过 BerkeleyParser。任务相关调参很少。

7 结论

Transformer 是第一个完全基于注意力的序列转换模型，用多头自注意力替换编码器—解码器里最常用的循环层。翻译任务上，训练显著快于基于循环或卷积的架构；英德、英法均达到新的 SOTA，英德上最佳模型超过此前所有集成。

作者计划把 Transformer 扩展到文本以外的输入输出模态，研究局部、受限注意力以处理图像、音频、视频等大输入输出，并使生成不那么顺序化。

训练与评测代码：https://github.com/tensorflow/tensor2tensor。

附录中的注意力可视化（论文所述）

编码器第 5 层（共 6 层）的自注意力：多个头关注动词 making 的远距离依赖，补全 making…more difficult。另有头表现为指代消解（如 its）。许多头表现出与句子结构相关的行为，且不同头做不同的事。