← Back to list
Transformer

Attention Is All You Need

other Google
Abstract 10 Reading 10 Rating —
2017-06-12
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin
Google Brain, Google Research, University of Toronto
提出完全基于注意力机制的Transformer架构,摒弃循环和卷积,在机器翻译任务上以更低训练成本达到SOTA
评分原因
精读评分:深度学习领域最具影响力的论文之一,提出了完全基于注意力机制的Transformer架构,彻底改变了NLP及几乎所有AI领域的研究范式,实验充分且开创性极强
transformer academic

Attention Is All You Need — Transformer 架构精读

研究动机与背景

在 Transformer 提出之前,主流的序列转导(sequence transduction)模型基于循环神经网络(RNN),特别是 LSTM 和 GRU 及其变体。这些模型沿着输入和输出序列的符号位置逐步计算,生成隐藏状态序列 $h_t = f(h_{t-1}, x_t)$。这种固有的顺序特性带来两个根本性问题:

  1. 训练无法并行化:由于每个时间步的计算依赖前一步的结果,RNN 无法在时间维度上并行,当序列较长时训练效率极低,且内存限制了跨样本的 batching。
  2. 长距离依赖建模困难:尽管 LSTM 在一定程度上缓解了梯度消失问题,但信号在网络中传播的路径长度仍与序列距离成正比,难以高效捕获远距离依赖关系。

注意力机制(attention mechanism)已经成为序列建模的重要组成部分,能够不受位置距离限制地建模依赖关系。但在本文之前,注意力机制几乎总是与循环网络结合使用。

本文的核心贡献是提出 Transformer 架构——一个完全基于注意力机制的序列转导模型,彻底抛弃了循环和卷积结构。Transformer 允许显著更高的并行化程度,在 8 块 P100 GPU 上训练仅 12 小时即可达到翻译任务的新 SOTA。

核心方法/模型架构

整体架构

Transformer 采用经典的 encoder-decoder 结构。Encoder 将输入符号序列 $(x_1, \ldots, x_n)$ 映射为连续表示序列 $\mathbf{z} = (z_1, \ldots, z_n)$;Decoder 在给定 $\mathbf{z}$ 的条件下,自回归地逐个生成输出序列 $(y_1, \ldots, y_m)$,每一步消费之前已生成的所有符号作为额外输入。

Figure 1: The Transformer - model architecture.

Encoder 由 $N = 6$ 个完全相同的层堆叠而成。每一层包含两个子层(sub-layer):

  1. Multi-Head Self-Attention 机制
  2. Position-wise 全连接前馈网络(Feed-Forward Network)

每个子层外围使用残差连接(residual connection)并紧跟 Layer Normalization,即每个子层的输出为:

$$\text{LayerNorm}(x + \text{Sublayer}(x)) \tag{1}$$

为了方便残差连接,所有子层及 embedding 层的输出维度统一为 $d_{\text{model}} = 512$。

Decoder 同样由 $N = 6$ 个相同的层堆叠而成。除了 Encoder 中的两个子层外,Decoder 额外插入第三个子层,对 Encoder 的输出执行 Multi-Head Attention(即 cross-attention)。同样在每个子层使用残差连接和 Layer Normalization。此外,Decoder 的 self-attention 子层经过修改(masked),防止当前位置关注到未来位置,确保位置 $i$ 的预测只依赖于位置 $\lt i$ 的已知输出。

Scaled Dot-Product Attention

注意力函数的本质是将一个查询(query)和一组键值对(key-value pairs)映射到一个输出。输出是 value 的加权和,权重由 query 与对应 key 的兼容性函数决定。

本文采用的注意力形式称为 "Scaled Dot-Product Attention"。输入由维度为 $d_k$ 的 query 和 key 以及维度为 $d_v$ 的 value 组成。计算 query 与所有 key 的点积,除以 $\sqrt{d_k}$,再通过 softmax 获得 value 的权重。实际计算中,query、key、value 分别打包为矩阵 $Q$、$K$、$V$:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V \tag{2}$$

为什么要除以 $\sqrt{d_k}$? 两种常见的注意力函数是 additive attention 和 dot-product attention。当 $d_k$ 较小时,两者表现相近;但当 $d_k$ 较大时,点积的量级增大(假设 $q$ 和 $k$ 的分量是均值为 0、方差为 1 的独立随机变量,则它们的点积均值为 0、方差为 $d_k$),将 softmax 推入梯度极小的区域。缩放因子 $\frac{1}{\sqrt{d_k}}$ 正是用来抵消这个效应。

Multi-Head Attention

与其使用 $d_{\text{model}}$ 维的单个注意力函数,不如将 query、key、value 分别用不同的学习到的线性投影投射 $h$ 次到 $d_k$、$d_k$、$d_v$ 维,然后在每组投影上并行执行注意力函数,得到 $d_v$ 维输出。这些输出拼接后再通过一次线性投影得到最终结果:

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O \tag{3}$$

$$\text{where } \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) \tag{4}$$

其中投影矩阵为 $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}}}$。

Figure 2: (left) Scaled Dot-Product Attention. (right) Multi-Head Attention

本文使用 $h = 8$ 个并行的注意力头,每个头的维度为 $d_k = d_v = d_{\text{model}} / h = 64$。由于每个头的维度降低,总计算成本与全维度的单头注意力相近。

Multi-Head Attention 的优势:允许模型在不同的表示子空间中同时关注来自不同位置的信息。单个注意力头的平均化操作会抑制这种能力。

Attention 在模型中的三种应用

  1. Encoder-Decoder Attention(cross-attention):Query 来自前一个 Decoder 层,Key 和 Value 来自 Encoder 的输出。这使得 Decoder 的每个位置可以关注输入序列的所有位置。

  2. Encoder Self-Attention:所有 Key、Value、Query 均来自 Encoder 前一层的输出。Encoder 中的每个位置可以关注前一层的所有位置。

  3. Decoder Masked Self-Attention:允许 Decoder 的每个位置关注包括自身在内的所有先前位置。为了保持自回归性质,通过将 softmax 输入中对应非法连接的位置设为 $-\infty$ 来实现 masking。

Position-wise Feed-Forward Networks

Encoder 和 Decoder 的每一层都包含一个全连接前馈网络,对每个位置独立且相同地应用。它由两个线性变换组成,中间用 ReLU 激活:

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

虽然不同位置共享同一组参数,但不同层之间使用不同的参数。这等价于两个 kernel size 为 1 的卷积。输入输出维度为 $d_{\text{model}} = 512$,内层维度为 $d_{ff} = 2048$。

Embeddings 与 Softmax

与其他序列转导模型类似,使用学习到的 embedding 将输入和输出 token 转换为 $d_{\text{model}}$ 维向量。同时使用线性变换和 softmax 将 Decoder 输出转换为预测的下一个 token 概率。两个 embedding 层和 pre-softmax 线性变换共享相同的权重矩阵。在 embedding 层中,权重乘以 $\sqrt{d_{\text{model}}}$。

Positional Encoding

由于 Transformer 不包含循环和卷积,为了让模型利用序列的顺序信息,需要注入位置信息。本文在 Encoder 和 Decoder 底部的 embedding 上加入 "positional encoding",其维度与 embedding 相同($d_{\text{model}}$),以便两者可以相加。

本文采用不同频率的正弦和余弦函数:

$$PE_{(pos, 2i)} = \sin(pos / 10000^{2i/d_{\text{model}}}) \tag{6}$$

$$PE_{(pos, 2i+1)} = \cos(pos / 10000^{2i/d_{\text{model}}}) \tag{7}$$

其中 $pos$ 是位置,$i$ 是维度索引。每个维度对应一个正弦曲线,波长从 $2\pi$ 到 $10000 \cdot 2\pi$ 构成几何级数。选择正弦函数的原因是假设它能让模型轻松学习关注相对位置——对于任何固定偏移 $k$,$PE_{pos+k}$ 可以表示为 $PE_{pos}$ 的线性函数。

实验发现使用可学习的 positional embedding 与正弦版本效果几乎完全相同(见消融实验 Table 3 行 (E))。作者选择正弦版本是因为它可能允许模型外推到比训练时更长的序列。

为什么选择 Self-Attention

本节从三个维度比较 self-attention 层与循环层和卷积层的优劣:

Layer Type Complexity per Layer Sequential Operations Maximum Path Length
Self-Attention $O(n^2 \cdot d)$ $O(1)$ $O(1)$
Recurrent $O(n \cdot d^2)$ $O(n)$ $O(n)$
Convolutional $O(k \cdot n \cdot d^2)$ $O(1)$ $O(\log_k(n))$
Self-Attention (restricted) $O(r \cdot n \cdot d)$ $O(1)$ $O(n/r)$

关键分析:

  1. 计算复杂度:当序列长度 $n$ 小于表示维度 $d$ 时(机器翻译中的常见情况),self-attention 比循环层更快。对于极长序列,可以限制 self-attention 只关注大小为 $r$ 的邻域,将复杂度降至 $O(r \cdot n \cdot d)$。
  2. 并行化程度:Self-attention 的所有位置可以同时计算($O(1)$ 顺序操作),而循环层需要 $O(n)$ 步顺序操作。
  3. 最大路径长度:Self-attention 将任意两个位置之间的距离缩短到 $O(1)$,而循环层为 $O(n)$,卷积层为 $O(\log_k(n))$。更短的路径使得学习长距离依赖更加容易。

此外,self-attention 还带来更具可解释性的模型——注意力分布可以被可视化检查,不同的 attention head 学习了不同的语言结构模式。

实验设置

训练数据

  • WMT 2014 English-German:约 450 万句对,使用 byte-pair encoding (BPE),共享源-目标词表约 37,000 token
  • WMT 2014 English-French:约 3600 万句对,使用 32,000 word-piece 词表
  • 句对按近似序列长度分批,每个训练 batch 包含约 25,000 源 token 和 25,000 目标 token

硬件与训练时间

  • 8 块 NVIDIA P100 GPU
  • Base model:每步约 0.4 秒,共训练 100,000 步(约 12 小时)
  • Big model:每步约 1.0 秒,共训练 300,000 步(约 3.5 天)

优化器

使用 Adam 优化器,参数为 $\beta_1 = 0.9$,$\beta_2 = 0.98$,$\epsilon = 10^{-9}$。学习率采用 warmup + inverse square root 调度:

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

前 $warmup\_steps = 4000$ 步线性增加学习率,之后按步数的逆平方根衰减。

正则化

  1. Residual Dropout:对每个子层的输出(加到子层输入并归一化之前)施加 dropout,同时对 embedding 和 positional encoding 之和施加 dropout。Base model 使用 $P_{drop} = 0.1$。
  2. Label Smoothing:使用 $\epsilon_{ls} = 0.1$ 的标签平滑。这会损害困惑度(perplexity),但提升准确率和 BLEU 分数。

主要实验结果

机器翻译

Model BLEU (EN-DE) BLEU (EN-FR) Training Cost (FLOPs) EN-DE Training Cost (FLOPs) EN-FR
ByteNet [18] 23.75
Deep-Att + PosUnk [39] 39.2 $1.0 \cdot 10^{20}$
GNMT + RL [38] 24.6 39.92 $2.3 \cdot 10^{19}$ $1.4 \cdot 10^{20}$
ConvS2S [9] 25.16 40.46 $9.6 \cdot 10^{18}$ $1.5 \cdot 10^{20}$
MoE [32] 26.03 40.56 $2.0 \cdot 10^{19}$ $1.2 \cdot 10^{20}$
Deep-Att + PosUnk Ensemble [39] 40.4 $8.0 \cdot 10^{20}$
GNMT + RL Ensemble [38] 26.30 41.16 $1.8 \cdot 10^{20}$ $1.1 \cdot 10^{21}$
ConvS2S Ensemble [9] 26.36 41.29 $7.7 \cdot 10^{19}$ $1.2 \cdot 10^{21}$
Transformer (base model) 27.3 38.1 $3.3 \cdot 10^{18}$
Transformer (big) 28.4 41.8 $2.3 \cdot 10^{19}$

结果分析

  • EN-DE 任务:Transformer (big) 以 28.4 BLEU 超越所有此前公开的模型和集成模型,比最好的集成模型(ConvS2S Ensemble, 26.36)高出 2.0 BLEU。即使是 base model (27.3) 也已经超越所有此前的单模型。
  • EN-FR 任务:Transformer (big) 达到 41.8 BLEU(当时的单模型 SOTA),仅使用不到之前最好单模型 1/4 的训练成本。EN-FR 的 big model 使用 $P_{drop} = 0.1$ 而非 0.3。
  • 训练成本:Transformer base model 的训练 FLOPs 仅为 $3.3 \times 10^{18}$,是所有竞争模型中最低的。

推理设置:Base model 取最后 5 个 checkpoint(间隔 10 分钟)平均;Big model 取最后 20 个 checkpoint 平均。Beam search 使用 beam size = 4,length penalty $\alpha = 0.6$。最大输出长度设为输入长度 + 50。

消融与分析

模型变体消融 (Table 3)

在 EN-DE 开发集 newstest2013 上评估,使用 beam search 但不使用 checkpoint 平均。

Table 3: Variations on the Transformer architecture

变体 $N$ $d_{\text{model}}$ $d_{ff}$ $h$ $d_k$ $d_v$ $P_{drop}$ $\epsilon_{ls}$ train steps PPL (dev) BLEU (dev) params ($\times 10^6$)
base 6 512 2048 8 64 64 0.1 0.1 100K 4.92 25.8 65
(A) h=1 1 512 512 5.29 24.9
(A) h=4 4 128 128 5.00 25.5
(A) h=16 16 32 32 4.91 25.8
(A) h=32 32 16 16 5.01 25.4
(B) $d_k$=16 16 5.16 25.1 58
(B) $d_k$=32 32 5.01 25.4 60
(C) $N$=2 2 6.11 23.7 36
(C) $N$=4 4 5.19 25.3 50
(C) $N$=8 8 4.88 25.5 80
(C) $d_{\text{model}}$=256 256 32 32 5.75 24.5 28
(C) $d_{\text{model}}$=1024 1024 128 128 4.66 26.0 168
(C) $d_{ff}$=1024 1024 5.12 25.4 53
(C) $d_{ff}$=4096 4096 4.75 26.2 90
(D) $P_{drop}$=0.0 0.0 5.77 24.6
(D) $P_{drop}$=0.2 0.2 4.95 25.5
(D) $\epsilon_{ls}$=0.0 0.0 4.67 25.3
(D) $\epsilon_{ls}$=0.2 0.2 5.47 25.7
(E) learned PE 4.92 25.7
big 6 1024 4096 16 0.3 300K 4.33 26.4 213

消融分析

(A) 注意力头数量:保持计算量恒定,改变 $h$。单头注意力($h=1$)BLEU 下降 0.9 到 24.9;头数过多($h=32$)时 BLEU 也下降到 25.4。$h=8$ 和 $h=16$ 表现最佳。

(B) 注意力 key 维度:减小 $d_k$ 会损害模型质量。$d_k=16$ 时 BLEU 降至 25.1,说明确定注意力兼容性并非简单的操作,更复杂的兼容性函数可能有益。

(C) 模型规模:更大的模型更好。增加层数($N=8$,25.5)、模型维度($d_{\text{model}}=1024$,26.0)、FFN 维度($d_{ff}=4096$,26.2)均带来提升。减小这些维度则显著降低性能。

(D) 正则化:Dropout 对避免过拟合非常重要,无 dropout 时 BLEU 仅 24.6。Label smoothing 带来小幅提升。

(E) Positional Encoding:将正弦 positional encoding 替换为可学习的 positional embedding 后,结果几乎完全相同(25.7 vs 25.8)。

English Constituency Parsing

为了验证 Transformer 的泛化能力,在英语句法分析任务上进行了实验。使用 4 层 Transformer,$d_{\text{model}} = 1024$,在 WSJ Penn Treebank 的约 40K 训练句子上训练。

Parser Training WSJ 23 F1
Vinyals & Kaiser et al. (2014) WSJ only, discriminative 88.3
Petrov et al. (2006) WSJ only, discriminative 90.4
Zhu et al. (2013) WSJ only, discriminative 90.4
Dyer et al. (2016) WSJ only, discriminative 91.7
Transformer (4 layers) WSJ only, discriminative 91.3
Zhu et al. (2013) semi-supervised 91.3
Huang & Harper (2009) semi-supervised 91.3
McClosky et al. (2006) semi-supervised 92.1
Vinyals & Kaiser et al. (2014) semi-supervised 92.1
Transformer (4 layers) semi-supervised 92.7
Luong et al. (2015) multi-task 93.0
Dyer et al. (2016) generative 93.3

结果分析

  • 在仅 WSJ 数据的设置下,4 层 Transformer 达到 91.3 F1,超过除 Dyer et al. (2016) 外的所有判别式模型。
  • 在半监督设置下(使用约 17M 句子的 BerkeleyParser 语料),Transformer 达到 92.7 F1,超过所有此前的半监督方法。
  • 值得注意的是,这些结果是在几乎不做任务特定调优的情况下取得的——除 dropout、学习率、beam size 外,其余超参数直接沿用机器翻译 base model 的配置。

Attention 可视化

Figure 3: An example of the attention mechanism following long-distance dependencies

论文附录展示了多个 attention head 的可视化。不同的 head 学习了不同的语言结构:

  • 某些 head 学会了跟踪长距离依赖,例如动词 "making" 的注意力指向了远距离的 "more difficult",完成了 "making...more difficult" 这个搭配结构。
  • 某些 head 参与指代消解(anaphora resolution),例如 "its" 的注意力精确指向了 "Law" 和 "application"。
  • 不同 head 展现出关注句子不同层面结构的行为,有些关注局部语法,有些关注全局语义。

Figure 4: Two attention heads involved in anaphora resolution

讨论与局限性

核心贡献

Transformer 是第一个完全基于 attention 的序列转导模型,用 multi-headed self-attention 替代了 encoder-decoder 架构中最常用的循环层。其核心洞察是:

  1. Self-attention 可以同时解决并行化和长距离依赖两个问题,而不是像 RNN 那样在两者之间折中。
  2. Multi-head attention 通过在不同的子空间并行计算注意力,赋予模型同时捕获多种不同类型依赖关系的能力。
  3. 位置编码用简洁的正弦函数即可提供足够的位置信息,且可能具有外推到更长序列的潜力。

值得借鉴的设计

  • Scaled Dot-Product:缩放因子 $1/\sqrt{d_k}$ 的设计基于对点积方差的精确分析,是一个简洁而关键的工程决策。
  • 权重共享:输入 embedding、输出 embedding 和 pre-softmax 线性变换三者共享权重,大幅减少参数量。
  • Warmup 学习率调度:线性 warmup + inverse sqrt decay 的策略后来被广泛采用。
  • Label Smoothing:虽然损害困惑度,但提升了实际的 BLEU 分数,体现了评价指标与训练目标之间的微妙差异。

局限性

  1. Self-attention 的二次复杂度:$O(n^2)$ 的计算和内存复杂度限制了 Transformer 处理超长序列的能力。论文提到了 restricted self-attention 的可能性,这成为后续大量高效 Transformer 研究(Linformer、Performer、Flash Attention 等)的出发点。
  2. 位置编码的设计空间:虽然正弦位置编码与可学习位置编码效果相当,但论文并未深入探索更多选择。后续工作(RoPE、ALiBi 等)在这个方向上取得了显著进展。
  3. 仅在翻译和句法分析上验证:论文发表时仅在两个 NLP 任务上验证了有效性。Transformer 后来被证明在几乎所有模态和任务上都极其有效,这在当时尚未被揭示。

历史影响

这篇论文是现代深度学习最具影响力的工作之一。Transformer 架构直接催生了 BERT、GPT、T5 等一系列里程碑模型,彻底改变了 NLP 的研究范式,并逐步扩展到计算机视觉(ViT)、语音、蛋白质结构预测(AlphaFold 2)、推荐系统等几乎所有领域。论文标题 "Attention Is All You Need" 已成为深度学习领域最广为人知的短语之一。