Understanding the Difficulty of Training Transformers

By Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, Jiawei Han. 2020

315 citations recorded by the source corpus.

Transformers have been proved effective for many deep learning tasks. Training transformers, however, requires non-trivial efforts regarding carefully designing learning rate schedulers and cutting-edge optimizers (the standard SGD fails to train Transformers effectively). In this paper, we study Transformer training from both theoretical and empirical perspectives. Our analysis reveals that unbalanced gradients are not the root cause of the instability of training. Instead, we identify an amplification effect that substantially influences training. Specifically, we observe that 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 result in significant disturbances in the model output, yet a light dependency limits the potential of model training and can lead to an inferior trained model. Inspired by our analysis, we propose Admin ($\mathbf{Ad}$aptive $\mathbf{m}$odel $\mathbf{in}$itialization) to stabilize the training in the early stage and unleash its full potential in the late stage. Extensive experiments show that Admin is more stable, converges faster, and leads to better performance.

Subfields: Artificial Intelligence, Model Architectures, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Model Architectures

In Optimization & Training

Citations in the corpus

Zero-Shot Text-to-Image Generation (2021, methodology): “Similar to prior work (Liu et al., 2020), we found that the norms of the activation gradients from the resblocks decrease monotonically as we move from the earlier resblocks Effective Parameter Count Compression Rank Compression Rate”
Group Equivariant Stand-Alone Self-Attention For Vision (2020, background): “These are in fact open questions for self-attention in general (Xiong et al., 2020; Liu et al., 2020; Zhao et al., 2020) and developments in this direction are of utmost importance.”
ImageNet-21K Pretraining for the Masses (2021, background): “Previous works have observed that MLP-based models can be harder and less stable to use in transfer learning since they don’t have inherent translation inductive bias [6, 39, 35].”
MetaFormer Baselines for Vision (2022, methodology): “Similar to LayerScale, [89, 43, 60] attempt to stabilize architectures by scaling the residual branch (ResScale [60]): X ′ = λr ⊙X + F(Norm(X)), (12)”
Foundation Transformers (2022, unknown): “It is also veriﬁed by Liu et al. (2020).”
ActionFormer: Localizing Moments of Actions with Transformers (2022, background): “The warm-up stage turned out to be critical for model convergence and to obtain good performance, as also pointed out by [42].”
Adaptive Gradient Methods at the Edge of Stability (2022, background): “The Post-LN Transformer is known to be less stable than the Pre-LN variant [46, 27]; this instability can be seen in the spike in λ1(PH) (and training loss) towards the end of the warmup period.”
End-to-End Instance Edge Detection (2022, background): “However, despite its high accuracy, the transformer architecture suffers from slow convergence [42] and quadratic computation and memory consumption [32,3] necessitating a high number of GPUs and up to weeks for training.”
A Survey of Transformers (2021, background): “Although Post-LN often results in unstable training and divergence, it usually outperforms preLN variants after convergence [83].”
Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer (2022, background): “Many previous works proposed different initialization or parametrizations with favorable properties, such as better stability for training deep neural networks [5, 10, 13, 21, 30, 47, 48, 51].”
Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers (2021, methodology): “Whether to use pre-norm or post-norm was chosen on a per-dataset basis, and depended on whether the model overfit; recent results have shown that pre-norm architectures are more stable [15, 35], so we used it on harder datasets with less overfitting.”
Are Transformers More Robust Than CNNs? (2021, background): “As explained in [25], compared to SGD, adaptive optimizers (like AdamW) are capable of assigning different learning rates to different parameters, resulting in consistent update magnitudes even with unbalanced gradients.”
Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training (2023, unknown): “First, the benefit of the first-order (gradient-based) pre-conditioner in Adam is not yet well understood (Liu et al., 2020; Zhang et al., 2020; Kunstner et al., 2023).”
Fine-Tuning Language Models with Just Forward Passes (2023, methodology): “As such, ZO methods have been applied in deep learning settings to find adversarial examples or tune input embeddings [91, 90] but not to directly optimize large-scale models (see Liu et al. [61] for a survey).”
Stabilizing Transformer Training by Preventing Attention Entropy Collapse (2023, background): “(2022); Liu et al. (2020a) are recent works that proposed to rescale residual connections.”
Solving ImageNet: a Unified Scheme for Training any Backbone to Top Results (2022, methodology): “Due to lack of inductive bias, Transformerbased [8, 24, 11] and MLP-only [36, 37] models for vision are harder to train, and less stable [29, 23].”
On the Transformer Growth for Progressive BERT Training (2020, methodology): “The original BERT models use the AdamW (Loshchilov and Hutter, 2019) optimizer with learning rate decay from 0.0001 to 0 and 10K steps of warmup (Liu et al., 2020a).”