arXiv:1706.03762v7 [cs.CL] 2 Aug 2023

Attention Is All You Need

基于注意力机制的全新序列转导模型架构 —— Transformer

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin

Google Brain, Google Research, University of Toronto

NIPS 2017 机器翻译 自注意力 Transformer

摘要

当前主流的序列转导模型基于复杂的循环或卷积神经网络，包含编码器与解码器，并通过注意力机制连接二者。

本文提出一种全新的简单网络架构 —— Transformer，它完全基于注意力机制，彻底摒弃了循环与卷积。

在两个机器翻译任务上的实验表明：Transformer 模型在质量上更优、并行度更高、训练时间显著更短。在 WMT 2014 英德翻译任务上取得 28.4 BLEU，英德翻译取得 41.8 BLEU，且训练成本远低于已有最优模型。

28.4
EN-DE BLEU
+2.0 超越最优
41.8
EN-FR BLEU
单模型 SOTA
12h
8x P100 训练
极低成本

研究背景：RNN 的固有局限

序列计算的瓶颈

  • RNN、LSTM、GRU 等模型按时间步顺序计算隐藏状态 ht = f(ht-1, xt)
  • 这种固有的顺序性阻碍了训练样本内的并行化
  • 长序列下内存限制严重影响批处理效率

长距离依赖的挑战

  • 长程信号需穿越大量时间步，梯度易消失或爆炸
  • 位置相距越远，模型建立依赖关系越困难
  • 尽管已有分解技巧与条件计算改进效率，顺序约束依然存在
现有注意力机制几乎都与循环网络结合使用。我们提出一个关键问题：能否完全用注意力替代循环？

核心思想：Attention Is All You Need

Transformer 核心主张

  • 完全基于自注意力机制计算输入与输出的全局依赖关系
  • 彻底移除循环与卷积，所有位置可并行计算
  • 任意两位置间的路径长度缩短至 O(1)
  • 训练 12 小时即可在翻译质量上达到新 SOTA

与已有工作的区别

1
Neural GPU / ByteNet / ConvS2S：基于卷积，长距离依赖仍需 O(log n) 或 O(n) 操作
2
Self-attention 已有应用：阅读理解、摘要、文本蕴含等，但均未完全取代 RNN
3
Transformer：首个完全依赖自注意力的序列转导模型

模型架构总览

Transformer Architecture

Figure 1: The Transformer —— 模型架构图（Encoder 左半部，Decoder 右半部）

Encoder

N = 6 层堆叠，每层含 Multi-Head Self-Attention + Position-wise FFN，残差连接与层归一化

Decoder

N = 6 层堆叠，额外增加 Encoder-Decoder Attention，Masked Self-Attention 防止看到未来信息

Scaled Dot-Product Attention

Scaled Dot-Product Attention

Figure 2 (left): Scaled Dot-Product Attention

计算过程

Attention(Q, K, V) = softmax(QKT / sqrt(dk)) V
  • 计算 Query 与所有 Key 的点积，除以 sqrt(dk) 进行缩放
  • 应用 softmax 获得权重，加权求和得到输出
  • 缩放防止 dk 较大时点积幅度过大，softmax 梯度消失
Dot-Product Attention 比 Additive Attention 更快、更省内存，因可用高度优化的矩阵乘法实现

Multi-Head Attention

多头并行的动机

  • 单头注意力会平均抑制不同子空间的信息
  • 多头允许模型在不同表示子空间、不同位置联合关注信息
  • 每个头独立学习不同的语义/句法关系
MultiHead(Q,K,V) = Concat(head1,...,headh)WO

其中 headi = Attention(QWiQ, KWiK, VWiV)

本文使用 h = 8 头，dk = dv = dmodel / h = 64，总计算量与单头全维度注意力相近
Multi-Head Attention

Figure 2 (right): Multi-Head Attention

Attention 的三种应用

01

Encoder-Decoder Attention

Query 来自 Decoder 上一层，Key/Value 来自 Encoder 输出。Decoder 每个位置可关注输入序列的所有位置，模拟经典 Seq2Seq 注意力。

02

Encoder Self-Attention

Query、Key、Value 均来自 Encoder 上一层同一位置。Encoder 每个位置可 attending 到前一层的所有位置，实现全局上下文编码。

03

Decoder Masked Self-Attention

Decoder 每个位置只能 attending 到该位置及之前的所有位置。通过将非法连接置为 -inf 的 Mask 实现，保持自回归特性。

其他关键组件

Position-wise Feed-Forward Networks

FFN(x) = max(0, xW1 + b1)W2 + b2

每层含两个线性变换与 ReLU 激活，dmodel = 512，内层维度 dff = 2048。对不同位置使用相同参数，层间参数不同。

Embeddings & Softmax

输入/输出 token 均通过可学习嵌入映射到 dmodel 维。解码器输出经线性变换 + softmax 预测下一 token。嵌入层与 pre-softmax 线性层共享权重，并乘以 sqrt(dmodel)。

Positional Encoding（位置编码）

PE(pos,2i) = sin(pos / 100002i/dmodel)
PE(pos,2i+1) = cos(pos / 100002i/dmodel)
  • 模型不含循环与卷积，必须显式注入位置信息
  • 正弦/余弦函数使模型能学习相对位置：PEpos+k 可表示为 PEpos 的线性函数
  • 可外推到训练时未见过的更长序列

为什么选择 Self-Attention？

Layer Type Complexity per Layer Sequential Operations Maximum Path Length
Self-Attention O(n2 d) O(1) O(1)
Recurrent O(n d2) O(n) O(n)
Convolutional O(k n d2) O(1) O(logk(n))
Self-Attention (restricted) O(r n d) O(1) O(n / r)

Table 1: 复杂度、顺序操作数与最大路径长度对比（n=序列长度, d=表示维度, k=卷积核大小, r=邻域大小）

1
计算复杂度：当 n < d 时（机器翻译常见情况），Self-Attention 比 RNN 更快
2
可并行性：Self-Attention 仅需 O(1) 顺序操作，RNN 需要 O(n)
3
路径长度：任意两位置间信号路径为 O(1)，显著利于学习长距离依赖

实验设置

数据集

  • WMT 2014 English-German：约 450 万句对，BPE 共享词表约 37K tokens
  • WMT 2014 English-French：3600 万句，32K word-piece 词表
  • 每批次约 25K source + 25K target tokens

硬件与训练时间

  • 单台 8 块 NVIDIA P100 GPU
  • Base 模型：每步约 0.4s，共 100K 步（12 小时）
  • Big 模型：每步约 1.0s，共 300K 步（3.5 天）

优化器

Adam (beta1=0.9, beta2=0.98, eps=1e-9)，学习率随步数变化：

lrate = dmodel-0.5 * min(step-0.5, step * warmup-1.5)

warmup_steps = 4000

正则化

  • Residual Dropout：Pdrop = 0.1，应用于子层输出、嵌入与位置编码之和
  • Label Smoothing：epsilonls = 0.1，降低困惑度但提升 BLEU

实验结果：机器翻译

Model EN-DE BLEU EN-FR BLEU EN-DE FLOPs EN-FR FLOPs
ByteNet23.751.0e20
GNMT + RL24.639.922.3e191.4e20
ConvS2S25.1640.469.6e181.5e20
GNMT + RL Ensemble26.3041.161.8e201.1e21
ConvS2S Ensemble26.3641.297.7e191.2e21
Transformer (base)27.338.13.3e18
Transformer (big)28.441.82.3e19

Table 2: Transformer 在 WMT 2014 newstest2014 上的 BLEU 与训练成本对比

28.4
EN-DE BLEU
超越所有已有模型与集成
41.8
EN-FR BLEU
单模型新 SOTA
<1/4
训练成本
仅为先前 SOTA 的零头

消融实验：模型变体

Variant N dmodel dff h dk Pdrop BLEU (dev)
base651220488640.125.8
(A) single head151224.9
(A) 16 heads163225.4
(B) dk=323225.1
(C) big model1024409616640.326.4
(D) dropout=0.00.024.6
(E) learned PE25.7

Table 3: Transformer 架构变体在 newstest2013 开发集上的表现

关键发现

  • 单头注意力比最优设置低 0.9 BLEU，头数过多也会下降
  • 减小 attention key 维度 dk 损害质量，说明兼容性判断并不简单
  • 更大的模型更好，dropout 对防止过拟合非常关键

位置编码

正弦位置编码与可学习位置嵌入效果几乎相同（25.8 vs 25.7）。选择正弦版本是因为它可能允许模型外推到训练时未见过的更长序列。

总结与展望

核心贡献

  • 提出 Transformer：首个完全基于注意力的序列转导模型
  • 用 多头自注意力 替代编码器-解码器架构中的循环层
  • 训练速度显著快于 RNN/CNN 架构，翻译质量达到新 SOTA
  • 在英语成分句法分析任务上也展现出良好的泛化能力

局限与未来方向

  • Self-Attention 的复杂度为 O(n2)，对极长序列仍有挑战
  • 计划探索 局部、受限注意力机制 以高效处理图像、音频、视频等大输入/输出
  • 使生成过程更少序列化是未来的研究目标之一
  • 将 Transformer 扩展到文本以外的模态
代码开源：https://github.com/tensorflow/tensor2tensor —— 这一工作奠定了 GPT、BERT 等后续模型的基础架构

感谢观看

Attention Is All You Need

NIPS 2017 arXiv:1706.03762
01 / 15