PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding the Difficulty of Training Transformers

Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, Jiawei Han

arXiv:2004.08249v3cs.LGcs.CLstat.ML

TL;DR

The paper asks what makes Transformer training unstable despite their effectiveness and analyzes this question empirically and theoretically. It identifies residual-branch dependency as an amplification mechanism, proposes Admin’s adaptive initialization, and reports improved stability, convergence, and performance. The study also outlines scope for generalization and larger-scale application.

  • Problem

    Transformer training requires specialized optimizers and learning-rate schedules, and the paper investigates the factors underlying its instability beyond unbalanced gradients.

  • Method

    The paper combines theoretical and empirical analysis of gradient behavior, residual-branch dependency, layer-normalization effects, and an adaptive initialization method called Admin.

  • Results

    Admin stabilizes training, converges faster, and achieves better performance; it stabilizes 72-layer WMT’14 En-Fr training without additional hyper-parameters and reaches 43.80 BLEU.

  • Takeaways & Limitations

    Strong residual-branch dependency amplifies parameter-change fluctuations, while Admin controls this dependency early and preserves flexibility after training stabilizes.

  • Takeaways & Limitations

    The analysis uses ReLU, and the authors identify generalization to other models and larger-scale training as future work.

Abstract

from arXiv · show

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 $\textit{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 ($\textbf{Ad}$aptive $\textbf{m}$odel $\textbf{in}$itialization) to stabilize 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 performance. Implementations are released at: https://github.com/LiyuanLucasLiu/Transforemr-Clinic.

1 Introduction

Transformers are effective but difficult to train, motivating empirical and theoretical analysis of the instability. The paper attributes instability to residual-branch dependency and proposes Admin to combine early stability with late-stage model potential.

  • Motivation: Transformers improve effectiveness, efficiency, and scalability through parallel computation, but conventional SGD can converge to bad or suspicious local optima.Removing warmup can also cause severe consequences, including model divergence.
  • Training difficulty: Post-LN Transformers are less robust than Pre-LN, and fixing gradient vanishing alone does not stabilize Post-LN training.This indicates that factors beyond unbalanced gradients substantially influence training.
  • Amplification effect: Post-LN layers depend more heavily on residual branches than Pre-LN layers, amplifying parameter-change fluctuations and destabilizing model training.Pre-LN’s lighter residual-branch reliance generally limits algorithmic potential and can produce inferior models.
  • Admin: Admin is an adaptive initialization method that restricts residual-branch dependency early and unleashes model potential later.It is designed to retain Pre-LN stability without hurting performance.
  • Empirical outcomes: Admin is more stable, converges faster, and achieves better performance across IWSLT’14 De-En, WMT’14 En-De, and WMT’14 En-Fr experiments.Without additional hyper-parameters, it stabilizes 72-layer Transformer training on WMT’14 En-Fr and achieves a 43.80 BLEU score.

2 Preliminaries

The preliminaries define Transformer modules, notation, and the distinct organization of Pre-LN and Post-LN architectures. They also introduce relative gradient norms for comparing sub-layer output gradients.

  • Architecture: Transformer layers contain Attention and Feedforward sublayers built mainly from Layer Norm, Multi-head Attention, and Feedforward Network modules.The Feedforward Network is a two-layer perceptron.
  • Layer normalization: Pre-LN and Post-LN organize Transformer modules differently, with Layer Norm placed inside Pre-LN residual blocks and outside Post-LN residual blocks.The architecture notation distinguishes network type, layer indexes, and sequence-by-hidden-dimensional inputs and outputs.
  • Residual notation: Residual branches are sublayer transformations, while layer or sublayer outputs integrate residual outputs with shortcut outputs.For a residual block x + f(x), the shortcut output is x and the residual-branch output is f(x).
  • Gradient notation: Relative gradient norms compare sub-layer output gradient magnitudes after scaling by the largest gradient norm in the same 18-layer WMT’14 En-De network.The histogram uses a logarithmic scale.
  • Attention: Multi-head Attention uses H heads and supports Encoder-Attention and Self-Attention settings.Self-Attention uses the same input for queries, keys, and values, while Encoder-Attention uses encoder outputs for keys and values.

3 Unbalanced Gradients

The analysis finds that unbalanced or vanishing gradients do not directly explain Transformer training instability. Instead, it examines how gradient behavior, parameter changes, and layer-normalization placement relate to stability across Pre-LN and Post-LN models.

  • Pre-LN converges in all 15 settings, whereas Post-LN diverges in 7 of 15; when Post-LN converges, it outperforms Pre-LN in 7 of 8 settings.
  • Only Post-LN decoders suffer from gradient vanishing; Post-LN encoders, Pre-LN encoders, and Pre-LN decoders do not.The analysis establishes this theoretically and empirically at initialization.
  • The analysis uses theoretical results and empirical measurements of gradient norms, parameter updates, and encoder output changes to study training stability.The output-shift analysis treats very large changes between F(x0, W) and F(x0, W*) as evidence of ill-conditioned training.
  • Fixing gradient vanishing alone fails to stabilize training, showing that gradient vanishing is not the direct cause of Post-LN instability.A hybrid model with a Post-LN encoder and Pre-LN decoder still does not achieve stable training after the decoder change.
  • Attention gradients remain unbalanced even for Pre-LN, while adaptive optimizers assign different learning rates and produce more consistent update magnitudes.This explains why standard SGD, which lacks this handling of unbalanced gradients, performs poorly for Transformers.
  • Pre-LN and Post-LN differ primarily in layer-normalization placement, which changes how residual outputs are aggregated.Pre-LN normalizes residual outputs once before they enter other layers, whereas Post-LN applies normalization differently.

4 Instability from Amplification Effect

The paper attributes Post-LN instability to an amplification effect caused by strong dependency on residual branches, rather than unbalanced gradients. Pre-LN reduces this dependency but can limit model potential, motivating Admin’s adaptive control of residual dependencies.

  • Amplification effect: Unbalanced gradients are not the root cause of Post-LN instability; layer dependency’s amplification effect intensifies updates and destabilizes training.The analysis identifies another factor beyond gradient imbalance: strong residual-branch dependency amplifies fluctuations from parameter changes.
  • Layer-normalization positions: Layer-normalization placement changes how residual outputs are aggregated: Pre-LN normalizes them once, whereas Post-LN normalizes most residual outputs multiple times.These differing normalization patterns assign different effective weights to residual outputs.
  • Layer dependency: The coefficient βi,j represents the proportion of the j-th residual branch in the i-th layer output and therefore reflects dependency among layers.The coefficients integrate the scaling operations of layer normalizations.
  • Layer dependency: Post-LN outputs retain stronger reliance on their own residual branches, while Pre-LN begins with more evenly distributed dependencies and remains less residual-dependent after training.This difference is visible from initialization through the final trained model.
  • Stability–capacity trade-off: Pre-LN’s lower residual-branch dependency improves stability but restricts higher-layer outputs from relying heavily on their own residual branches, limiting full model potential.Post-LN gives the network greater choice between wider and deeper representations.
  • Amplification effect: For parameter perturbations, output-change variance scales as O(log N) for Pre-LN but O(N) for Post-LN.Theoretical results attribute the larger Post-LN sensitivity to its stronger residual-branch dependency.
  • Empirical verification: Empirically, the predicted relationships achieve R2 = 0.99 for both architectures, while optimization updates produce weaker but still substantial correlations.With Adam updates, the correlation is R2 = 0.75, and Post-LN output shifts remain larger than Pre-LN by multiple magnitudes.
  • Admin: Admin adds parameters ω to control Post-LN residual dependencies and adaptively initializes them to target an O(log N) output change.Its construction uses profiling of residual-branch output variances before initialization.

5 Experiments

Experiments across IWSLT’14 and WMT’14 evaluate Admin against Post-LN and Pre-LN. Admin stabilizes deeper training, converges competitively, and achieves stronger performance while preserving late-stage model capacity.

  • Performance comparison: Admin achieves better performance than Post-LN and Pre-LN for 6-, 12-, and 18-layer WMT’14 Transformer-base models.The experiments use BLEU on WMT’14 En-De and En-Fr, with configurations varying in encoder and decoder depth.
  • Training stability: 12-layer and 18-layer Post-LN models diverge without adaptive initialization, whereas Pre-LN converges in all tested settings.Pre-LN nevertheless produces sub-optimal performance relative to Admin.
  • Performance comparison: Admin stabilizes deeper models and benefits more from increased capacity than Pre-LN.This supports the paper’s claim that Pre-LN stability can limit model potential.
  • Convergence behavior: Admin matches Pre-LN’s early-stage convergence speed while reaching stronger late-stage performance.This combines early stability with the capacity needed for better final performance.
  • Performance comparison: On IWSLT’14 De-En, all 6-layer Transformer-small methods perform similarly, with Admin outperforming the other two by a small margin.The authors associate the smaller stability difference with the shallower network.
  • Connection to warmup: Extending Post-LN warmup from 8,000 to 16,000, 24,000, or 32,000 updates still fails for the 18-layer model.The result indicates that learning-rate warmup alone does not neutralize Post-LN instability.
  • Comparison with other initializations: ReZero, FixUp, and LookLinear all produce divergent training for the tested 18-layer WMT’14 De-En Transformers.The authors link FixUp and ReZero instability in these experiments to half-precision training, dropout, and absent layer normalization.

6 Related Work

Prior work established Transformer effectiveness and identified training difficulties, including sensitivity to warmup and layer normalization. This paper adds a systematic analysis connecting those difficulties to residual-branch dependency and proposes adaptive initialization.

  • Transformer background: Transformers improve effectiveness, efficiency, and scalability through parallel computation without recurrent connections.Prior studies report breakthroughs across multiple deep learning tasks and domains.
  • Training difficulties: Removing warmup is more damaging for Transformers, especially Post-LN, than for other architectures.Warmup has also been reported as useful for vanilla SGD.
  • Deep network initialization: Deep-network initialization research addresses obstacles to training deeper models, whose depth can improve performance but hinder precise optimization.The paper situates Admin among methods designed to ease deep-network training.
  • Research gap: Although scaling residual outputs can stabilize training, prior work lacked a systematic analysis of Transformer-specific difficulty and its connection to residual-branch dependency.The paper identifies amplification effects and introduces an adaptive initialization method to address this gap.

7 Conclusion

The paper identifies strong residual-branch dependency as the root cause of Transformer instability and proposes Admin to control this dependency during training. Experiments show that Admin provides more stable training, faster convergence, and better performance without additional hyper-parameters.

  • Strong dependency on residual branches amplifies parameter-change fluctuations and destabilizes Transformer training.The analysis finds gradient vanishing is not the root cause, while residual-branch dependency is.
  • Admin controls residual-branch dependency early and preserves flexibility to capture these dependencies after training stabilizes.Admin is designed to retain Pre-LN stability without sacrificing model potential.
  • Without additional hyper-parameters, Admin achieves more stable training, faster convergence, and better performance.
  • The work motivates future research on generalizing Theorem 2, adapting deep networks to training configurations, upgrading Transformer architectures, and scaling Admin training.

Appendices

The appendix analyzes gradients at initialization and finds that Pre-LN gradients do not vanish through backpropagation. Under the stated independence assumptions, lower layers have larger gradients than higher layers.

  • Pre-LN does not suffer from gradient vanishing at initialization.The analysis defines gradients as ∆x = ∂L/∂x, where L is the training objective.
  • The derivation assumes randomly initialized parameters are independent of partial derivatives with respect to module inputs.
  • At initialization, the two terms in the Pre-LN gradient recurrence are approximately independent.
  • For Pre-LN, gradient variance does not decrease through backpropagation, so lower layers have larger gradients than higher layers.

A.2 Post-LN Encoder Analysis

The Post-LN encoder analysis shows that most encoder sublayers do not exhibit gradient vanishing, while encoder-attention sublayers do. It also examines how layer-normalization placement shapes gradient distribution and residual-output weighting.

  • Under Theorem 1’s initialization and independence assumptions, Post-LN encoder gradients satisfy Var[∆x_i−1] ≥ Var[∆x_i].The result implies lower encoder layers have at least as large gradient variance as higher layers.
  • Backpropagation through Post-LN encoder FFN sublayers does not suffer from gradient vanishing.
  • Backpropagation through Post-LN encoder self-attention sublayers also does not suffer from gradient vanishing.
  • The Post-LN encoder-attention sublayer suffers from gradient vanishing, with encoder-attention output gradients shorter than input gradients in Figure 3.
  • Attention-module gradients remain unbalanced in Pre-LN because softmax-related parameters, including W (K) and W (V1), have smaller gradients.
  • Adaptive optimizers convert unbalanced parameter gradients into more consistent update magnitudes, unlike vanilla SGD.

B Proof of Theorem 2

The proof of Theorem 2 relates layer dependency to output fluctuations caused by parameter changes. Under its assumptions, the output-change variance is estimated by a sum of layer-dependent terms proportional to a constant C.

  • Theorem 2 estimates output-change variance after parameter perturbation when layer dependencies remain unchanged.The perturbation is δ = W*−W, assumed independent of the randomly initialized parameters W.
  • The proof represents each sublayer as a module output combined through normalized residual coefficients β_i,j.
  • At initialization, independently initialized module outputs are independent across distinct layers.
  • Residual sharing between adjacent layers is used in the proof to analyze how layer contributions accumulate.
  • The derivation treats the variance of the relevant per-layer term as a common constant C.

C Admin Implementation Details

Admin introduces parameters ω to rescale module outputs and limit residual-branch dependency during training. After training, the model can be reparameterized as a conventional Post-LN structure by removing ω.

  • Admin introduces parameters ω to rescale module outputs when constructing Post-LN sub-layers.
  • The element-wise product is used in Admin’s module construction.
  • After training, Admin can be reparameterized into the conventional Post-LN structure by removing ω_i and changing associated parameters.
  • The implementation provides corresponding formulations for feedforward, self-attention, and encoder-attention sub-layers.
  • Adding ω_i is equivalent to rescaling some model parameters, but direct initialization rescaling is less stable in half-precision training.

D Experimental Setup

Experiments use fairseq-based implementations across IWSLT’14 De-En and WMT’14 En-De and En-Fr, with configurations ranging from standard Transformer models to very deep networks. Comparisons include ReZero, Pre-LN, and Admin, including a 72-layer WMT’14 En-Fr model.

  • The experiments use fairseq implementations and standard preprocessing, with WMT evaluations conducted on newstest14.
  • IWSLT’14 De-En uses Transformer-small settings, while WMT’14 En-De and En-Fr use Transformer-base configurations with additional layers.
  • Training uses RAdam, inverse square-root learning-rate scheduling, warmup, dropout, and dataset-specific regularization settings.
  • Comparison to ReZero: The study evaluates ReZero under original and Post-LN-normalized configurations, finding that ReZero initialization reduces performance while layer normalization improves it.
  • Performance on the WMT’14 En-Fr: 43.80 BLEU: Admin (60L–12L) on WMT’14 En-Fr achieves the reported new state of the art with a 60-layer encoder and 12-layer decoder.
Loading 2004.08249v3…