ICML 2020 | 摆脱warm-up!巧置LayerNorm使Transformer加速收敛

2020-07-24 | 作者:贺笛、郑书新

编者按:Transformer 网络结构存在 warm-up 阶段超参数敏感、优化过程收敛速度慢等问题。为此,中科院、北京大学和微软亚洲研究院机器学习组的研究员们在 ICML 2020 的最新论文[1]中,从理论上分析了 Transformer 优化困难的原因,并提出了让 Transformer 摆脱 warm-up 阶段并加快收敛速度的解决方法。


近年来,Transformer 网络结构已经在自然语言处理的各项任务中都取得了“屠榜”的成绩。然而 Transformer 结构的优化非常困难,其具体表现有 warm-up 阶段超参数敏感、优化过程收敛速度慢等问题。近日,中科院、北京大学和微软亚洲研究院的研究员们在国际机器学习大会 ICML 2020 上发表了题为“On the Layer Normalization in the Transformer Architecture”的论文(点击阅读原文查看),从理论上详细分析了 Transformer 结构优化困难的原因,并给出了解决方法,可以让 Transformer 彻底摆脱 warm-up 阶段,并且大幅加快训练的收敛速度。

在随机优化的理论中,学习率往往设置为常数或者逐渐衰减 (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.