PaperlayerView exact revision on arXiv

Source-linked AI summary

DeepNet: Scaling Transformers to 1,000 Layers

Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, Furu Wei

arXiv:2203.00555v1cs.CLcs.LG

TL;DR

Training instability limits Transformer depth despite the success of large-scale models. DEEPNORM modifies residual connections with theoretically derived initialization, enabling stable optimization and Transformers exceeding 1,000 layers. Experiments show improved translation performance across benchmarks, including a 4.4 BLEU improvement over the baseline at 1,000 layers.

  • Problem

    Despite large parameter counts and strong task performance, Transformer depth remains limited by training instability.

  • Method

    DEEPNORM modifies residual connections, scales residual-branch weights with architecture-dependent constants, and theoretically bounds model updates by a constant.

  • Results

    DEEPNET improves translation performance across benchmarks and scales to 1,000 layers, outperforming the baseline by 4.4 BLEU.

  • Takeaways & Limitations

    DEEPNET combines the reported performance of Post-LN with the training stability of Pre-LN, supporting deeper Transformer scaling for machine translation.

  • Takeaways & Limitations

    ReZero is unstable with half precision in the reported experiments, even when the model is shallow.

Abstract

from arXiv · show

In this paper, we propose a simple yet effective method to stabilize extremely deep Transformers. Specifically, we introduce a new normalization function (DeepNorm) to modify the residual connection in Transformer, accompanying with theoretically derived initialization. In-depth theoretical analysis shows that model updates can be bounded in a stable way. The proposed method combines the best of two worlds, i.e., good performance of Post-LN and stable training of Pre-LN, making DeepNorm a preferred alternative. We successfully scale Transformers up to 1,000 layers (i.e., 2,500 attention and feed-forward network sublayers) without difficulty, which is one order of magnitude deeper than previous deep Transformers. Remarkably, on a multilingual benchmark with 7,482 translation directions, our 200-layer model with 3.2B parameters significantly outperforms the 48-layer state-of-the-art model with 12B parameters by 5 BLEU points, which indicates a promising scaling direction.

1 Introduction

Transformer models have grown to billions or trillions of parameters, but their depth remains limited by training instability. DeepNorm addresses this limitation by stabilizing optimization and enabling much deeper Transformers while retaining favorable performance.

  • Large-scale Transformers achieve strong performance across tasks, but their depth is constrained by training instability.Model capacity has expanded from millions to billions and even trillions of parameters.
  • Pre-LN improves Transformer stability, but its larger bottom-layer gradients can reduce performance relative to Post-LN.
  • DeepNorm modifies residual connections and initialization to bound model updates and stabilize extremely deep Transformers.The method requires only small code changes and is architecture-dependent.
  • 1,000 layers are reached with DeepNorm, while a 200-layer, 3.2B-parameter model improves multilingual translation by 5 BLEU over a 48-layer, 12B-parameter state-of-the-art model.

2 TL;DR for Practitioners

DeepNorm is designed as a simple modification to Post-LN Transformers. It scales residual connections and selected initialization weights using architecture-dependent factors.

  • DeepNorm is implemented by up-scaling the residual connection before layer normalization.
  • Initialization down-scales selected weights, including feed-forward weights and attention value and output projections.
  • The residual and initialization scales depend on the encoder-decoder architecture.Figure 2 provides parameters for N-layer encoders and M-layer decoders.

3 Instability of Deep Transformer

The analysis attributes deep Post-LN instability primarily to exploding model updates at the beginning of training. These updates enlarge layer-normalization inputs, causing progressively smaller gradients and trapping optimization.

  • Post-LN-init narrows lower-layer weight scales after Xavier initialization while preserving the Post-LN architecture.
  • Post-LN-init converges on IWSLT-14 De-En, whereas Post-LN does not, despite Post-LN-init retaining larger last-layer gradient norms.
  • Post-LN-init has larger last-layer gradient norms than Post-LN across depths from 6L-6L to 24L-24L, indicating that exploding deep-layer gradients are not the root cause.
  • Post-LN exhibits an exploding model update at the beginning of training, followed by nearly no update.Warm-up and better initialization help the model update more smoothly.
  • Large updates increase layer-normalization inputs, whose larger magnitude makes gradients through layer normalization smaller.
  • The resulting gradient vanishing can trap the model in a spurious local optimum and further destabilize optimization.

4 DEEPNET: Extremely Deep Transformers

DEEPNET stabilizes extremely deep Transformers by replacing Post-LN with DEEPNORM and architecture-dependent initialization, theoretically bounding model updates. Its updates remain nearly constant during early training, while vanilla Post-LN updates explode.

  • Architecture: DEEPNET replaces Post-LN with DEEPNORM in every Transformer sub-layer and introduces architecture-dependent constants for residual scaling and initialization.The residual branch is scaled by α, while its parameters are scaled by β; both depend only on the architecture.
  • Architecture: DEEPNET’s attention analysis identifies value and output projections as determining attention-output magnitude, while query and key projections do not change its bound.The analysis studies attention through its value and output projection matrices.
  • Expected Magnitude of Model Update: The model update is defined as ||∆F|| = ||F(x, θ*) − F(x, θ)|| and is analyzed for N-layer DEEPNET with attention and feed-forward sublayers.The theorem uses DEEPNORM-normalized sublayers of the form x_l+1 = LN(αx_l + G_l(x_l, θ_l)).
  • Expected Magnitude of Model Update: Vanilla Post-LN accumulates updates across sublayers, producing an exploding update magnitude that destabilizes early optimization; warm-up and smaller initialization alleviate this effect.The theoretical result explains why vanilla Post-LN has update growth with depth and why these interventions improve update smoothness.
  • Expected Magnitude of Model Update: With proper α and β, DEEPNET’s expected model updates are bounded by a constant, under an analysis based on SGD that is empirically verified with Adam.The analysis is presented for encoder-decoder models and can be extended to encoder-only and decoder-only architectures.
  • Empirical Evidence: On IWSLT-14 De-En, DEEPNET’s early-training update is nearly constant, whereas vanilla Post-LN’s update is exploding.Figure 5 visualizes this behavior for tiny Transformers ranging from 6L-6L to 100L-100L.

5 Neural Machine Translation

DEEPNET improves Transformer training stability across depths and settings, while maintaining or improving translation quality relative to deep-model baselines.

  • DEEPNET successfully scales to 100L-100L, reaching 28.9 BLEU, whereas Post-LN baselines encounter unstable optimization at 50L-50L.
  • DEEPNET achieves comparable performance with baselines on shallow models and avoids the performance harm observed when scaling deeper.
  • Pre-LN models are more stable than Post-LN baselines but incur a 0.5-1.0 BLEU drop, while DEEPNET outperforms all Pre-LN baselines.
  • Convergence with varying depth: DEEPNET remains stable from 10L-10L through 100L-100L, exceeds 30 BLEU within 8,000 steps, and improves as depth increases.
  • Large learning rate, batch size, and hidden dimension: DEEPNET trains without difficulty under the largest tested learning rate, batch size, and hidden dimension settings, with larger settings producing faster convergence and lower validation loss.
  • ReZero is unstable with half precision, even when the model is shallow.

6 Massively Multilingual Neural Machine Translation

DEEPNET extends deep Transformer scaling to massively multilingual translation, where increasing depth improves translation quality and outperforms a larger-parameter baseline across evaluation datasets.

  • DEEPNET scales to 1,000 layers using a 500-layer encoder and 500-layer decoder on the OPUS-100 multilingual corpus.
  • 4.4 BLEU is the improvement of 1,000-layer DEEPNET over the baseline, after the 48-layer baseline gains 3.2 points over the 12-layer model.
  • Scaling law in terms of depth: DEEPNET exhibits logarithmic BLEU growth with depth on multilingual NMT, trained at 12, 20, 100, 200, and 1,000 layers.
  • More data and language directions: DEEPNET uses 3.2B parameters compared with M2M-100's up to 12B parameters in the expanded multilingual comparison.
  • DEEPNET significantly outperforms M2M-100 on all reported WMT, OPUS, TED, and Flores evaluation datasets.

7 Conclusion and Future Work

DEEPNET improves Transformer stability and scales successfully to 1,000 layers. The experiments focus on machine translation, while future work will extend the approach to other tasks.

  • Conclusion: DEEPNET successfully scales Transformers to 1,000 layers using the DEEPNORM normalization function.DEEPNORM is theoretically justified to stabilize optimization with a constant upper bound for model updates.
  • Conclusion: Experimental results verify the effectiveness of the proposed methods across various machine-translation benchmarks.Machine translation is the current experimental test bed.
  • Future Work: Future work will extend DEEPNET to language-model pre-training, protein-structure prediction, and BEiT vision pre-training.

A.2 Proof of Theorem 4.2

The proof analyzes model-update magnitudes for DEEPNET under explicit simplifying assumptions and establishes bounds for both homogeneous and encoder-decoder architectures. It contrasts these bounds with vanilla Post-LN, whose updates grow with depth.

  • Theorem A.2: Theorem A.2 characterizes the model-update norm for an N-layer DEEPNET whose attention and feed-forward sublayers use DEEPNORM.The theorem defines each layer through xl+1 = LN(αxl + Gl(xl, θl)).
  • Assumptions: The derivation assumes hidden dimension d = 1 and positive weights with magnitude below 1, while DEEPNORM parameters α and β exceed 1.
  • Proof: The attention and feed-forward derivations estimate sublayer outputs and model updates using assumptions on projection weights and Taylor expansion.The analysis notes that query and key projections do not change the bound on attention-output magnitude.
  • Comparison: For vanilla Post-LN with standard initialization, the model-update norm grows as O(P2N).
  • Encoder-Decoder Extension: Theorem A.3 extends the model-update bound to an encoder-decoder DEEPNET with N encoder layers and M decoder layers.The encoder and decoder sublayers use separate residual scaling constants, αe and αd.
  • Comparison: For vanilla Post-LN in the encoder-decoder setting, the update norm is O(M P2N).

B Derivation for Encoder-Decoder Architecture

The encoder-decoder derivation selects architecture-dependent initialization and residual-scaling parameters to control update terms for encoder and decoder components.

  • Decoder: For an encoder-decoder DEEPNET, decoder parameters are set using M, with vd = wd = (12M)−1 and αd = (3M) 1 4.These choices bound the second term of Equation (13) to Θ(η).
  • Encoder: Encoder parameters are tied across layers by setting vei = ve and wei = we.
  • Encoder: The derivation sets ve = we = βe and gives αe = 0.81(N 4M) and βe = 0.87(N 4M)−1 16 to satisfy the stated condition.

C Derivation for Encoder-only (Decoder-only) Architecture

The encoder-only or decoder-only derivation starts from the update bound and imposes symmetric parameter choices. The accompanying tables list hyperparameters for the machine-translation experiments.

  • Derivation: The derivation begins from Theorem 4.2 for an N-layer DEEPNET.
  • Parameter Choices: By symmetry, the projection scales are set uniformly as vi = v and wj = w.
  • Parameter Choices: The construction uses v = w = (8N)−1 4 and α = (2N) 1 4 to satisfy the stated condition.
  • Experimental Settings: Tables 4–8 provide hyperparameters for IWSLT-14 De-En, WMT-17 En-De, OPUS-100, and 102-language machine-translation experiments.

D.5 Evaluation Details

The evaluation uses specified BLEU procedures and compares language-pair results across 7,482 directions for DEEPNET and M2M-100.

  • Evaluation metrics: BLEU scores for IWSLT-14 and WMT-17 use Fairseq scripts, while OPUS-100 uses case-sensitive detokenized sacreBLEU.The evaluation protocols differ by benchmark.
  • Evaluation metrics: WMT, OPUS, and TED use the same test sets and evaluation scripts as M2M.M2M results are taken directly from Fan et al. (2021).
  • Multilingual evaluation: Figures 9 and 10 report results for 12B M2M-100 and 3.2B DEEPNET across 7,482 directions.Each figure uses source languages as rows and target languages as columns, covering 87 languages.
Loading 2203.00555v1…