Understanding the Difficulty of Training Transformers
Liyuan Liu†‡†‡\dagger\ddagger Xiaodong Liu‡‡\ddagger Jianfeng Gao‡‡\ddagger Weizhu Chen§§\mathsection Jiawei Han††\dagger {ll2, , ††\daggerUniversity of Illinois at Urbana-Champaign ‡‡\ddaggerMicrosoft Research §§\mathsection Microsoft Dynamics 365 AI
Abstract
Transformers have proved effective in many NLP tasks. However, their training requires non-trivial efforts regarding designing cutting-edge optimizers and learning rate schedulers carefully (e.g., conventional SGD fails to train Transformers effectively). Our objective here is to understand what complicates Transformer training from both empirical and theoretical perspectives. Our analysis reveals that unbalanced gradients are not the root cause of the instability of training. Instead, we identify an amplification effect that influences training substantially–for each layer in a multi-layer Transformer model, heavy dependency on its residual branch makes training unstable, since it amplifies small parameter perturbations (e.g., parameter updates) and results in significant disturbances in the model output. Yet we observe that a light dependency limits the model potential and leads to inferior trained models. Inspired by our analysis, we propose Admin (Adaptive model initialization) to stabilize the early stage’s training and unleash its full potential in the late stage. Extensive experiments show that Admin is more stable, converges faster, and leads to better performance111Impleme
中文速览
训练Transformer模型时，普通SGD会失败、必须精心调学习率预热等问题困扰着研究者，但背后的根本原因一直不清楚。作者通过理论推导和实验发现，梯度消失并非训练不稳定的根本原因，真正的"幕后黑手"是一种放大效应——在Post-LN结构中，每一层对自身残差分支（residual branch）的依赖过重，导致参数的微小扰动被层层放大，最终使模型输出剧烈震荡、训练崩溃；而Pre-LN虽然稳定，但对残差分支依赖太轻又限制了模型的表达能力，最终性能不如Post-LN。基于这一洞察，作者提出了Admin（自适应模型初始化，Adaptive Model Initialization）方法，通过在初始化阶段调整各层对残差分支的依赖权重，在训练早期保持稳定、训练后期释放模型全部潜力。在机器翻译基准测试上，Admin不引入额外超参数，就能稳定训练72层的超深Transformer，并在WMT'14英法翻译任务上取得43.80的BLEU分数，收敛更快、性能更强，为Transformer的稳定训练提供了清晰的理论解释和实用解决方案。
原文 arXiv:2004.08249;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2004.08249v3