[论文解读] On Layer Normalization in the Transformer Architecture

Ruibin Xiong, Yunchang Yang+8 more|arXiv (Cornell University)|Feb 12, 2020

Power Transformer Diagnostics and Insulation被引用 109

一句话总结

该论文分析层归一化放置位置如何影响 Transformer 的优化，并显示 Pre-LN 在训练中无需预热阶段即可收敛并更快收敛，而 Post-LN 依赖预热以保持稳定。

ABSTRACT

The Transformer is widely used in natural language processing tasks. To train a Transformer however, one usually needs a carefully designed learning rate warm-up stage, which is shown to be crucial to the final performance but will slow down the optimization and bring more hyper-parameter tunings. In this paper, we first study theoretically why the learning rate warm-up stage is essential and show that the location of layer normalization matters. Specifically, we prove with mean field theory that at initialization, for the original-designed Post-LN Transformer, which places the layer normalization between the residual blocks, the expected gradients of the parameters near the output layer are large. Therefore, using a large learning rate on those gradients makes the training unstable. The warm-up stage is practically helpful for avoiding this problem. On the other hand, our theory also shows that if the layer normalization is put inside the residual blocks (recently proposed as Pre-LN Transformer), the gradients are well-behaved at initialization. This motivates us to remove the warm-up stage for the training of Pre-LN Transformers. We show in our experiments that Pre-LN Transformers without the warm-up stage can reach comparable results with baselines while requiring significantly less training time and hyper-parameter tuning on a wide range of applications.

研究动机与目标

• 激励为何对 Post-LN Transformer 的学习率预热至关重要，以及层归一化放置如何影响梯度行为。
• 理论性地利用均场理论分析 Post-LN 和 Pre-LN 变体在初始化时的梯度尺度。
• 实证验证在 Pre-LN 下是否可以去除预热，并在多项 NLP 任务中衡量训练速度和性能。

提出的方法

• 使用均场理论研究 Post-LN 和 Pre-LN Transformer 初始化时的梯度尺度。
• 对最后一个前馈网络层的梯度范数及其对深度 L 的依赖进行理论分析。
• 在 IWSLT14 De-En、WMT14 En-De 和 BERT 预训练上进行实证实验，以比较有预热与无预热的设置。
• 控制化初始化：单头注意力、Xavier 初始化、注意力中的零 Q/K、以及高斯输入。
• 使用 Adam 和 SGD/RAdam 变体，在有无预热的情况下比较 Post-LN 与 Pre-LN 架构。

实验结果

研究问题

• RQ1 初始化时 Pre-LN Transformer 的学习率预热阶段是否变得不必要？
• RQ2 层归一化的位置如何影响 Transformer 架构中的梯度尺度与训练稳定性？
• RQ3 无预热的 Pre-LN Transformer 是否能在翻译和预训练任务中实现与 Post-LN 基线相当或更快的收敛和最终性能？

主要发现

• Post-LN Transformer 在初始化时输出层附近存在较大的梯度，在没有预热时较大学习率不稳定。
• Pre-LN Transformer 在初始化时梯度行为良好，允许移除预热阶段。
• 在 IWSLT14 De-En、WMT14 En-De 与 BERT 预训练中，Pre-LN 无预热在速度和最终性能方面达到或超过带预热的 Post-LN。
• 在相同 lr_max 设置下，Pre-LN 的收敛速度快于 Post-LN，降低了超参数敏感性和训练时间。
• 移除预热带来显著的加速（例如更快的收敛和更少的超参数调优），同时保持竞争力的结果。