在随机优化的理论中,学习率往往设置为常数或者逐渐衰减 (decay),从而保证算法的收敛,这种学习率的设置方法也与机器学习里很多任务上的实际经验类似,例如图像分类、语音识别等。然而,不管是设置学习率为常数还是使学习率逐渐衰减都不能让 Transformer 很好地收敛。
在优化 Transformer 结构时,除了设置初始学习率与它的衰减策略,往往还需要在训练的初始阶段设置一个非常小(接近0)的学习率,让它经过一定的迭代轮数后逐渐增长到初始的学习率,这个过程也被称作 warm-up 阶段。
Warm-up 是原始 Transformer 结构优化时的一个必备学习率调整策略。Transformer 结构对于 warm-up 的超参数(持续轮数、增长方式、初始学习率等)非常敏感,若调整不慎,往往会使得模型无法正常收敛。

我们在 IWSLT14 De-En 翻译数据集上使用 Transformer 模型对这一观点进行验证,采用 Adam 和 SGD 两种随机优化器,分别测试了两种学习率(1e^-3, 5e^-4)在没有 warm-up(0轮), warm-up 迭代轮数不足(500轮)和迭代轮数充足(4000轮)情况下模型的验证集 Loss 与 BLEU 分数。

由于 Transformer 优化困难的阶段是在训练的初始阶段,warm-up 也只是在迭代的前若干轮起作用,因此我们从模型的初始化阶段开始探究原因。如图3(a)所示,原始 Transformer 结构的每一层中分别经过了带残差连接的 Multi-Head Attention 和 FFN 两个子层(sub-layer),在两子层之后分别放置了层归一化(Layer Normalization)层,即 Post-LN Transformer。

将 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 阶段的梯度范数进行了验证:

在 IWSLT、WMT 和 BERT 三个任务上验证了是否可以在 Pre-LN 结构中去掉 warm-up:

图5:IWSLT 与 WMT 实验结果

至此,我们通过对 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.