← 全部论文 / 2017

Attention Is All You Need

2017 TPU、加速器与 Transformer 时代 被引用 331 次 ★ 获奖
下载 PDF

Attention Is All You Need - 中文验证版

英文原文卡片：attention_is_all_you_need_2017.md

状态：已翻译。

元数据

  • Slug： attention_is_all_you_need_2017
  • 年份： 2017
  • 会议： NeurIPS
  • 作者： Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin
  • 阅读状态： read complete
  • 计算范式： TPU、加速器与 Transformer 时代
  • 主要来源： PDF、抽取文本

计算设置

论文明确列出主要训练设置：一台带 8 块 NVIDIA P100 GPU 的机器。这不是推断。对 base 模型，每训练步约 0.4 秒，训练 100,000 步，约 12 小时。对 big 模型，每步约 1.0 秒，训练 300,000 步，3.5 天。

数据管线是 token 批处理而非句子数批处理。WMT 2014 English-German 约有 450 万句对，使用约 37,000 token 的共享 BPE 词表。WMT 2014 English-French 有 3600 万句，使用 32,000 word-piece 词表。句对按近似序列长度分组；每个训练 batch 包含约 25,000 源 token 和 25,000 目标 token。这减少了填充浪费，并使密集矩阵乘保持更充分填充。

论文还通过将训练时间、GPU 数量和每 GPU 的假定持续单精度能力相乘，以 FLOPs 估算训练成本。对 P100 使用 9.5 TFLOPS，使加速器成本成为报告结果的一部分而非背景基础设施。

瓶颈

论文指出循环模型的序列瓶颈：循环阻止单个样本内部并行化。该瓶颈对长序列更为严重，因为内存限制迫使更小 batch，而顺序依赖增加关键路径上的操作数。

对 2017 年加速器而言，有利原语是大批量矩阵乘。RNN 使用 GPU 算术，但无法在序列维度暴露足够多的并行工作。

论文用三个准则形式化这一点：每层复杂度、可并行化计算和长程依赖的最大路径长度。循环层在长度为 n 的序列上需要 O(n) 次顺序操作。自注意力层以常数次顺序操作连接各位置，将序列建模从时间展开的循环转换为批量注意力和前馈矩阵操作。

代价是二次注意力。作者论证当序列长度 n 小于表示维度 d 时，自注意力比循环更快，这在其机器翻译句表示中通常成立。他们已指出极长序列可能需要受限注意力，预示了后来的内存高效注意力工作。

方法适配

Transformer 是对加速器结构的直接适配：

  • 自注意力取代循环，将顺序操作减少为常数个层步骤。
  • 缩放点积注意力用优化的矩阵乘实现，在可用硬件上比加法注意力更快且更省内存。
  • 多头注意力保持投影为密集批处理，同时允许不同表示子空间。
  • 位置编码加入顺序信息，而不重新引入循环计算。
  • 逐位置前馈块是密集层，也是加速器友好的原语。

论文论证了缩放点积注意力比加法注意力更快、更省空间，因为它使用了优化的矩阵乘。h = 8 头的多头注意力保持每个头的维度较小（d_k = d_v = d_model / h = 64），因此总计算量与单头全维注意力相似，同时暴露多个密集投影。逐位置前馈网络也易于跨位置和样本批处理。

优化器和调度也针对这一加速器范式调整。论文使用 Adam，beta 值为 0.9 和 0.98，4000 步线性预热，然后逆平方根衰减。预热很重要，因为模型在新架构上使用大 token batch 训练；它让运行在不稳定性上浪费昂贵的 P100 早期步之前先逐步加速。

证据

  • 在 WMT 2014 English-German 上，big Transformer 达到 28.4 BLEU。
  • 在 WMT 2014 English-French 上，摘要和表 2 报告 Transformer big 为 41.8 BLEU；结果文字中也陈述 big English-French 模型为 41.0 BLEU。两者均作为单模型结果报告，与先前系统对比。
  • 训练 big English-German 模型在 8 块 P100 GPU 上耗时 3.5 天。
  • 表 2 报告 base 模型估计训练成本为 3.3 × 10^18 FLOPs，Transformer big 为 2.3 × 10^19 FLOPs，低于若干先前循环、卷积和集成系统的较大成本。
  • base 模型以报告训练成本的很小一部分超越了所有先前发表的 English-German 模型和集成。
  • 推理/评估时，论文使用束搜索，束宽为 4，长度惩罚为 0.6，并对 base 模型取最后 5 个 checkpoint 平均、对 big 模型取最后 20 个平均。

历史影响

本文将主流序列建模从循环形状的计算转移到矩阵乘形状的计算。这是使后来加速器规模语言模型变得可行的核心架构步骤。

在计算结构主线中，Transformer 不仅是方法创新。它是序列建模与 GPU/TPU 吞吐对齐的节点，从而在后来 BERT、GPT 和大语言模型系统中启用密集缩放。

一旦移除了循环，缩放模型宽度、批量 token 和训练 FLOPs 就成为语言和序列建模的更直接工程路径。

局限

  • 完整自注意力具有二次序列长度成本。
  • 论文最强证据是机器翻译，尚非互联网规模通用语言建模。
  • 它使用单台多 GPU 机器，而非后来的数据中心规模训练系统。
  • 报告的成本模型使用估计的持续单精度 TFLOPS，而非实测端到端利用率。

链接