论文阅读:Attention Is All You Need
《Attention Is All You Need》发表于 2017 年的 NIPS,提出了 Transformer 架构。它在机器翻译任务上取得了当时最好的效果,但更深远的影响是:此后的 BERT、GPT 系列、ViT 等几乎所有大规模预训练模型都建立在这一架构之上。
和多数论文不同,这篇工作不是为了解决某个具体痛点,而是重新思考了序列建模的方式——把”沿位置一步步递推”换成”让所有位置一次性互相关注”。
一句话概括:Transformer 抛弃了循环(recurrence)和卷积(convolution),完全依靠注意力机制来建模序列内部和序列之间的依赖关系。
1. 背景:从循环到注意力
要理解 Transformer 换掉了什么,先看它出现之前的主流做法。序列转换(sequence transduction)任务,例如机器翻译,主流模型是基于循环神经网络(RNN、LSTM、GRU)的编码器—解码器结构,通常再配合注意力机制。这类模型有一个根本性的限制:计算沿序列位置顺序进行。
RNN 在处理长度为 的序列时,会依次生成隐藏状态 ,其中 是前一个状态 和当前输入的函数:
这种顺序依赖意味着,第 步必须等第 步算完才能开始,无法在单个样本内部并行。序列越长,问题越严重;同时较长的序列会因为内存限制压缩批大小(batch size),进一步降低训练效率。
为绕过顺序计算,此前出现了基于卷积的方案,例如 Extended Neural GPU、ByteNet 和 ConvS2S。卷积可以对所有位置并行计算,但关联两个相距较远的位置所需的操作数会随距离增长:ConvS2S 是线性增长,ByteNet 是对数增长。距离越远,学习它们之间的依赖关系就越困难。
注意力机制此时已被广泛使用,它允许模型在不考虑距离的前提下建模依赖关系。但在当时,注意力几乎总是和 RNN 配合使用,而非独立承担建模任务。
于是一个自然的想法是:能否完全去掉循环和卷积,只用注意力机制来构建序列转换模型?Transformer 给出的回答是肯定的。
2. 核心思想
用自注意力(self-attention)完全替代 RNN 和卷积来建模序列依赖。
自注意力让序列中任意两个位置直接交互,二者之间的路径长度是常数 ,与它们的距离无关;同时一层内所有位置可以并行计算。这同时缓解了 RNN 的两个痛点——无法并行、长程依赖路径过长。
下图对比了循环网络与自注意力在连接序列中两个位置时的差异。
3. 方法细节
3.1 整体架构
Transformer 沿用编码器—解码器结构。编码器把输入符号序列 映射为连续表示 ;解码器以 为条件,逐个生成输出序列 。解码器是自回归的:生成每个符号时,会把之前已生成的符号作为额外输入。
编码器由 个相同的层堆叠而成。每层包含两个子层:一个多头自注意力子层,一个逐位置的前馈网络子层。每个子层都使用残差连接,再接层归一化(layer normalization),即每个子层的输出为:
为了让残差连接可以相加,模型中所有子层和嵌入层的输出维度统一为 。
解码器同样由 个相同的层堆叠而成。除了编码器中的两个子层,每层还插入第三个子层:对编码器输出做多头注意力(即”编码器—解码器注意力”)。解码器的自注意力子层经过修改,通过掩码(masking)阻止某个位置关注其后的位置;这一掩码加上输出嵌入右移一位,保证位置 的预测只能依赖位置小于 的已知输出,从而维持自回归性质。
3.2 注意力
注意力函数把一个查询(query)和一组键值对(key-value pairs)映射为输出。输出是各 value 的加权和,每个 value 的权重由 query 与对应 key 的相容性(compatibility)函数计算得到。
缩放点积注意力(Scaled Dot-Product Attention)。 这是论文采用的具体注意力形式。query 和 key 的维度为 ,value 的维度为 。计算 query 与所有 key 的点积,每个除以 ,再经 softmax 得到 value 上的权重。实际计算中把一批 query、key、value 分别打包成矩阵 :
两种常见的注意力是加性注意力(additive attention)和点积注意力。二者理论复杂度相近,但点积注意力可以用高度优化的矩阵乘法实现,更快也更省空间。
这里的关键设计是缩放因子 。当 较小时,加性注意力和未缩放的点积注意力表现接近;但当 较大时,点积的数值会变大。论文给出的解释是:若 query 和 key 各分量是均值 0、方差 1 的独立随机变量,则它们的点积 均值为 0、方差为 。点积过大会把 softmax 推入梯度极小的区域,因此用 把方差缩放回 1 附近。
多头注意力(Multi-Head Attention)。 与其用 维直接做一次注意力,论文把 query、key、value 用不同的、可学习的线性投影分别投影到 、、 维,并行做 次注意力,再把 个 维输出拼接起来,最后经一次线性投影:
多头机制让模型能在不同的表示子空间、不同的位置上同时关注信息;单头注意力会因为平均而削弱这种能力。论文取 ,并令 。由于每个头的维度降低,多头的总计算量与单头全维度注意力相近。
注意力在模型中的三种用法。 Transformer 在三处使用多头注意力,区别在于 query、key、value 的来源:
- 编码器—解码器注意力:query 来自解码器上一层,key 和 value 来自编码器输出。这让解码器每个位置都能关注输入序列的所有位置。
- 编码器自注意力:query、key、value 都来自编码器上一层的输出,每个位置可关注上一层的所有位置。
- 解码器自注意力:每个位置可关注解码器中包括自身在内、直到当前位置的所有位置;通过把 softmax 输入中对应非法连接的值置为 来屏蔽后续位置,保持自回归性质。
3.3 逐位置前馈网络
除了注意力子层,编码器和解码器的每一层还包含一个前馈网络,它对每个位置独立且相同地作用,由两个线性变换和中间的 ReLU 激活组成:
输入输出维度为 ,中间层维度为 。
3.4 嵌入与 Softmax
模型用可学习的嵌入把输入和输出 token 转成 维向量,并用线性变换加 softmax 把解码器输出转成下一个 token 的概率。两个嵌入层和 softmax 前的线性变换共享同一权重矩阵;嵌入层中的权重还会乘以 。
3.5 位置编码
模型不含循环和卷积,本身无法感知序列顺序,因此需要显式注入位置信息。论文在编码器和解码器底部的输入嵌入上加入”位置编码”(positional encoding),其维度与嵌入相同(均为 ),二者可以相加。位置编码采用不同频率的正弦和余弦函数:
其中 是位置, 是维度索引。选择正弦形式的原因是:对任意固定偏移 , 都可以表示为 的线性函数,模型因此更容易学习按相对位置关注;同时正弦形式可能让模型外推到比训练时更长的序列。论文也试验了可学习的位置嵌入,效果几乎相同。
4. 为什么用自注意力
论文从三个角度比较了自注意力层与循环层、卷积层:每层的计算复杂度、可并行的计算量(用所需的最少顺序操作数衡量)、网络中长程依赖的路径长度。
| 层类型 | 每层复杂度 | 顺序操作数 | 最大路径长度 |
|---|---|---|---|
| 自注意力 | |||
| 循环 | |||
| 卷积 | |||
| 自注意力(受限) |
其中 是序列长度, 是表示维度, 是卷积核大小, 是受限自注意力的邻域大小。
要点如下:
- 自注意力用常数个顺序操作连接所有位置,循环层则需要 个顺序操作。在路径长度上,自注意力是 ,循环层是 ,更短的路径让长程依赖更容易学习。
- 当序列长度 小于表示维度 时(这在机器翻译常用的句子表示中很常见),自注意力每层的计算量比循环层更小。
- 对极长序列,可以把自注意力限制在以当前位置为中心、大小为 的邻域内,把最大路径长度增加到 以换取更低的计算量;论文将此列为未来工作。
此外,自注意力可能带来更可解释的模型:论文在附录中展示,不同注意力头会学到不同的任务,部分头表现出与句法、语义结构相关的行为。
5. 训练设置
- 数据:WMT 2014 英德数据集约 450 万句对,用 byte-pair encoding,共享词表约 37000 个 token;英法数据集约 3600 万句对,词表约 32000 个 word-piece。按近似序列长度分批,每批约含 25000 个源 token 和 25000 个目标 token。
- 硬件与时长:在 8 张 NVIDIA P100 GPU 上训练。基础模型每步约 0.4 秒,共训练 10 万步(约 12 小时);大模型每步约 1.0 秒,共训练 30 万步(约 3.5 天)。
- 优化器:Adam,,,。学习率先在前 步线性升高,之后按步数平方根的倒数衰减,:
- 正则化:对每个子层输出在残差相加前做 dropout,并对嵌入与位置编码之和做 dropout,基础模型 ;标签平滑(label smoothing),它会略微提高困惑度,但提升准确率和 BLEU。
6. 效果
6.1 机器翻译
在 WMT 2014 英德翻译任务上,大模型达到 28.4 BLEU,比此前最好结果(包括集成模型)高出 2 BLEU 以上。在英法翻译任务上,大模型达到 41.8 BLEU,超过此前所有已发表的单模型,而训练成本不到此前最优模型的四分之一。即使是基础模型,也以远低于竞争对手的训练成本超过了此前所有已发表的模型和集成模型。
| 模型 | 英德 BLEU | 英法 BLEU | 训练成本(FLOPs) |
|---|---|---|---|
| GNMT + RL | 24.6 | 39.92 | |
| ConvS2S | 25.16 | 40.46 | |
| MoE | 26.03 | 40.56 | |
| ConvS2S(集成) | 26.36 | 41.29 | |
| Transformer(基础) | 27.3 | 38.1 | |
| Transformer(大) | 28.4 | 41.8 |
6.2 模型变体
论文在英德开发集上改动基础模型的各个部分,观察性能变化,主要结论有:
- 改变注意力头数 (保持计算量不变)时,单头比最佳设置低 0.9 BLEU,头数过多质量也会下降。
- 减小注意力 key 的维度 会损害质量,说明相容性的计算并不简单,可能需要比点积更复杂的函数。
- 模型越大效果越好;dropout 对抑制过拟合很有帮助。
- 用可学习的位置嵌入替换正弦位置编码,结果与基础模型几乎相同。
6.3 泛化到句法分析
为检验泛化能力,论文把 Transformer 用于英语成分句法分析(English constituency parsing)。该任务输出受强结构约束、且通常远长于输入,此前 RNN 序列到序列模型在小数据上难以达到最优。一个 4 层 Transformer 在仅用 WSJ 4 万句训练时达到 91.3 F1,半监督设置下达到 92.7 F1,即使在小数据情形下也超过了 BerkeleyParser。这说明 Transformer 不局限于翻译,能较好地迁移到其他任务。
7. 局限与后续发展
- 复杂度对长序列不友好:自注意力每层复杂度为 ,随序列长度平方增长。处理很长的序列时代价很高。论文提出的受限自注意力只是设想,被列为未来工作;后续的 FlashAttention、稀疏/线性注意力等方向才进一步缓解了这一问题。
- 评测任务有限:论文的实验集中在机器翻译和句法分析,主要是 NLP 序列转换任务,并未涉及更广的模态或任务类型。
- 规模有限:论文验证的是中等规模模型在翻译任务上的表现。Transformer 真正的威力——作为大规模预训练的骨干——是由后来的 BERT、GPT 等工作揭示的,并不在论文范围内。
- 位置编码的外推能力:正弦位置编码被设计为可能外推到更长序列,但实践中对极长序列的外推效果有限,这也催生了后续 RoPE、ALiBi 等相对位置方案。
其中关于 FlashAttention、BERT/GPT、RoPE 的内容属于这篇论文之后的发展,并非论文本身的论述。
8. 总结
Transformer 的核心贡献是证明:仅依靠注意力机制,不用循环和卷积,就能构建出在质量、并行性和训练效率上都更优的序列转换模型。它把”建模任意两位置之间的依赖”从随距离增长的操作,压缩为常数路径长度的一次注意力,从而既容易并行,又便于学习长程依赖。这一架构后来成为绝大多数大规模语言模型与多模态模型的基础。
参考
- Vaswani et al. Attention Is All You Need. NIPS 2017. arXiv:1706.03762