在随机优化的理论中，学习率往往设置为常数或者逐渐衰减 (decay)，从而保证算法的收敛，这种学习率的设置方法也与机器学习里很多任务上的实际经验类似，例如图像分类、语音识别等。然而，不管是设置学习率为常数还是使学习率逐渐衰减都不能让 Transformer 很好地收敛。
在优化 Transformer 结构时，除了设置初始学习率与它的衰减策略，往往还需要在训练的初始阶段设置一个非常小（接近0）的学习率，让它经过一定的迭代轮数后逐渐增长到初始的学习率，这个过程也被称作 warm-up 阶段。
Warm-up 是原始 Transformer 结构优化时的一个必备学习率调整策略。Transformer 结构对于 warm-up 的超参数（持续轮数、增长方式、初始学习率等）非常敏感，若调整不慎，往往会使得模型无法正常收敛。
图1：常见的 warm-up 学习率策略近日，来自伊利诺伊大学香槟分校的研究员们对 warm-up 问题进行了研究[2], 他们认为 Transformer 结构的优化中常用的 Adam 优化器[3]（自适应学习率优化器的一种）在训练前期由于数据样本不足，导致自适应学习率在调整学习率时方差过大，从而引起优化稳定性不足，所以提出了 Rectified Adam（RAdam）算法来控制 Adam 优化器前期的方差。我们在 IWSLT14 De-En 翻译数据集上使用 Transformer 模型对这一观点进行验证，采用 Adam 和 SGD 两种随机优化器，分别测试了两种学习率（1e^-3, 5e^-4）在没有 warm-up（0轮）, warm-up 迭代轮数不足（500轮）和迭代轮数充足（4000轮）情况下模型的验证集 Loss 与 BLEU 分数。
图2：IWSLT14 De-En 对比实验结果从实验中可以看出：（1）模型最终性能对无论是 warm-up 所需迭代轮数还是学习率的大小都非常敏感；（2）不仅仅是自适应学习率方法 Adam, warm-up 方法对非自适应性学习率方法 SGD 也同样十分重要。由此看来，“自适应学习率在训练前期方差过大”这个解释并不能完全解答 Transformer 依赖 warm-up 的问题。由于 Transformer 优化困难的阶段是在训练的初始阶段，warm-up 也只是在迭代的前若干轮起作用，因此我们从模型的初始化阶段开始探究原因。如图3（a）所示，原始 Transformer 结构的每一层中分别经过了带残差连接的 Multi-Head Attention 和 FFN 两个子层（sub-layer），在两子层之后分别放置了层归一化（Layer Normalization）层，即 Post-LN Transformer。
图3：(a) Post-LN Transformer；(b) Pre-LN Transformer当采用 Xavier[4] 方法对 Post-LN Transformer 进行初始化后，通过对各隐层梯度值进行分析可以证明，在初始化点附近的 Post-LN Transformer 结构最后一层的梯度值非常大，同时随着反向传播的前传会导致梯度值迅速衰减。这种在各层之间不稳定的梯度分布必然会影响优化器的收敛效果，导致训练过程初始阶段的不稳定。造成 Post-LN Transformer 梯度分布出现问题的核心原因在于各子层之后的 Layer Normalization 层会使得各层的输入尺度与层数 L 无关，因此当 Layer Normalization 对梯度进行归一化时，也与层数 L 无关。将 Layer Normalization 放到残差连接中的两个子层之前，并且在整个网络最后输出之前也增加一个 Layer Normalization 层来对梯度进行归一化，我们称这样的结构为 Pre-LN Transformer[5][6]，如图3(b)所示。
使用相同的方法对 Pre-LN Transformer 结构进行分析后，发现最后一层 Layer Normalization 层的输入尺寸的量级只有 Post-LN 的√(1/L)倍，并且每个 LN 层都会对梯度以 √L 的比例归一化。所以对于 Pre-LN 结构来说，其每层梯度范数都近似不变。
我们在 IWSLT 14 De-En 数据集对两种结构在初始化附近以及经过了 warm-up 阶段的梯度范数进行了验证：
图4：梯度范数验证实验可以看出，相比于 Post-LN 结构梯度分布的不稳定，Pre-LN 在各层之间梯度范数几乎保持不变，这种结构明显更利于优化器进行优化。而在进行一定轮数的 warm-up 后，Post-LN 的梯度范数也基本保持不变，并且其量级非常小（图4中绿色），这也验证了我们的结论。在 IWSLT、WMT 和 BERT 三个任务上验证了是否可以在 Pre-LN 结构中去掉 warm-up：
图5：IWSLT 与 WMT 实验结果
图6：BERT 验证集 Loss 与下游任务表现我们发现，当使用 Pre-LN 结构时，warm-up 阶段已经不再是必需，并且 Pre-LN 结构可以大幅提升 Transformer 的收敛速度。对于机器翻译任务（IWSLT/WMT)，不需要 warm-up 的 Pre-LN 结构可以比 Post-LN 收敛快1倍左右，而在 BERT 上，Pre-LN 在下游任务上达到和 Post-LN 相同的性能也只需要后者迭代轮数的1/3左右，并且最终的效果也更好。至此，我们通过对 warm-up 问题的本质进行分析，不仅彻底摆脱了恼人的 warm-up 以及敏感的调参过程，还令 Transformer 结构的训练过程更稳定，收敛速度也更快。
参考文献
[1] Ruibin X, Yunchang Y, Di H, et al. On Layer Normalization in the Transformer Architecture. ICML 2020
[2] Liu L, Jiang H, He P, et al. On the variance of the adaptive learning rate and beyond. ICLR 2019.
[3] Kingma D P, Ba J. Adam: A method for stochastic optimization[J]. arXiv preprint arXiv:1412.6980, 2014.
[4] Glorot X, Bengio Y. Understanding the difficulty of training deep feedforward neural networks[C]//Proceedings of the thirteenth international conference on artificial intelligence and statistics. 2010: 249-256.
[5] Baevski, A. and Auli, M. Adaptive input representations for neural language modeling. arXiv preprint arXiv:1809.10853, 2018.
[6] Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
