PaperlayerView exact revision on arXiv

Source-linked AI summary

On Layer Normalization in the Transformer Architecture

Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, Tie-Yan Liu

arXiv:2002.04745v2cs.LGcs.CLstat.ML

TL;DR

The paper asks why Transformer training commonly requires learning-rate warm-up and whether its necessity depends on layer-normalization placement. Using mean field theory and experiments, it finds that Post-LN has large output-adjacent gradients at initialization, whereas Pre-LN has well-behaved gradients and can train without warm-up while retaining comparable performance with less training time and tuning.

  • Problem

    Post-LN Transformer training usually requires carefully configured learning-rate warm-up, which slows optimization and increases hyperparameter tuning.

  • Method

    The paper combines mean field theory of initialization gradients with experiments comparing Post-LN and Pre-LN Transformers.

  • Results

    Pre-LN Transformers can omit warm-up across translation and BERT pretraining tasks, achieving comparable final performance with much less training time.

  • Takeaways & Limitations

    Layer-normalization placement controls initialization gradient behavior, making warm-up removable for the studied Pre-LN training settings.

  • Takeaways & Limitations

    The warm-up configuration affects final performance, requiring careful hyperparameter tuning that is computationally expensive for large-scale NLP tasks.

Abstract

from arXiv · show

The Transformer is widely used in natural language processing tasks. To train a Transformer however, one usually needs a carefully designed learning rate warm-up stage, which is shown to be crucial to the final performance but will slow down the optimization and bring more hyper-parameter tunings. In this paper, we first study theoretically why the learning rate warm-up stage is essential and show that the location of layer normalization matters. Specifically, we prove with mean field theory that at initialization, for the original-designed Post-LN Transformer, which places the layer normalization between the residual blocks, the expected gradients of the parameters near the output layer are large. Therefore, using a large learning rate on those gradients makes the training unstable. The warm-up stage is practically helpful for avoiding this problem. On the other hand, our theory also shows that if the layer normalization is put inside the residual blocks (recently proposed as Pre-LN Transformer), the gradients are well-behaved at initialization. This motivates us to remove the warm-up stage for the training of Pre-LN Transformers. We show in our experiments that Pre-LN Transformers without the warm-up stage can reach comparable results with baselines while requiring significantly less training time and hyper-parameter tuning on a wide range of applications.

1. Introduction

The paper explains why Post-LN Transformers require learning-rate warm-up and shows that Pre-LN Transformers can remove it while preserving comparable performance and reducing training time and tuning.

  • Motivation: Learning-rate warm-up is typically required for Post-LN Transformers, but it slows optimization and increases hyperparameter tuning.The warm-up begins with a very small learning rate and gradually increases it to a predefined maximum.
  • Approach: Mean field theory analyzes optimization behavior at initialization to explain why warm-up is needed.The analysis focuses on gradient scales under different layer-normalization placements.
  • Post-LN analysis: Post-LN Transformers have large expected gradients for parameters near the output layer, making large learning rates potentially unstable.Warm-up and initially small learning rates practically avoid this instability.
  • Pre-LN analysis: Pre-LN Transformers place layer normalization inside residual connections and have well-behaved initialization gradients without exploding or vanishing.This design also adds a final layer normalization before prediction.
  • Experimental findings: Across translation and BERT pretraining tasks, Pre-LN training can omit warm-up, achieve comparable final performance, and use much less training time.The experiments also report faster loss decay and fewer hyperparameters.

2. Related work

Prior work established that warm-up is unusually important for Transformers, while explanations based specifically on Adam do not fully account for its benefits across optimizers.

  • Warm-up in prior work: Unlike many convolutional and recurrent architectures, Transformers commonly require learning-rate warm-up during optimization.Prior studies found that removing or shortening warm-up can cause Post-LN optimization to diverge.
  • Warm-up in prior work: Popel and Bojar reported that Post-LN optimization diverges without warm-up or with relatively few warm-up iterations.Their work investigated the influence of different warm-up strategies.
  • Alternative architecture: Pre-LN was proposed in recent work as an alternative architecture intended to alleviate optimization difficulties associated with Post-LN.The supplied related-work passage introduces Pre-LN as a newer Transformer variant.
  • Explanations for warm-up: The claim that warm-up benefits arise from Adam’s adaptive-learning-rate variance does not explain why other optimizers also benefit from warm-up.This suggests Adam is not necessary for warm-up to matter.

3. Optimization for the Transformer

The paper explains why Post-LN Transformers require learning-rate warm-up and connects this need to layer-normalization placement and initialization-time gradient scales. Experiments and theory show that Post-LN gradients near the output can be large, whereas Pre-LN gradients are better behaved.

  • Transformer architectures: Post-LN places layer normalization after the residual connection, while self-attention and feed-forward sub-layers use scaled dot-product attention and position-wise transformations.The Transformer layer applies residual connection and layer normalization to both sub-layers individually.
  • Learning-rate warm-up: For both Adam and SGD, removing warm-up harms Post-LN optimization; with Adam, BLEU reaches only 8.45 without warm-up versus around 34 with warm-up.The same qualitative pattern appears in validation-loss curves, although SGD performs worse than Adam overall.
  • Learning-rate warm-up: Warm-up length is an important Post-LN hyper-parameter: with Twarmup = 500, Adam obtains BLEU scores of 31.16 at lrmax = 5e−4 and 2.77 at lrmax = 1e−3.These results show sensitivity to the interaction between warm-up duration and maximum learning rate.
  • Initialization and gradient scales: Mean field analysis attributes the instability to large initialization-time gradients near the Post-LN output layer, while Pre-LN gradients are smaller and better behaved.For Post-LN, the last FFN gradient scale is O(d ln d) and independent of depth; experiments examine gradient norms across layers and model sizes.
  • Initialization and gradient scales: A fixed learning rate of 1e−4 can mitigate Post-LN instability, but convergence is significantly slower than with larger learning rates after warm-up.After warm-up, gradient scales are very small, allowing training with large learning rates.

4. Experiments

Experiments show that Pre-LN Transformers can be trained without learning-rate warm-up across translation, BERT pre-training, and downstream tasks, while retaining competitive performance and converging faster than Post-LN models.

  • Machine Translation: On IWSLT14 De-En, Pre-LN reaches approximately 34 BLEU and 4 validation loss, comparable to Post-LN.
  • Machine Translation: The Pre-LN model converges faster: its ninth IWSLT14 checkpoint nearly matches the Post-LN model’s fifteenth checkpoint, with similar observations on WMT14 En-De.
  • Machine Translation: Changing layer-normalization placement has more effect than changing Adam to RAdam: RAdam helps Post-LN without warm-up but differs little from Adam for Pre-LN.
  • Unsupervised Pre-training (BERT): In BERT pre-training, Pre-LN reaches a validation loss similar to Post-LN’s 1.69 at 700k rather than 500k updates, corresponding to a reported 40% speed-up.
  • Unsupervised Pre-training (BERT): Pre-LN also converges faster on the MRPC and RTE downstream tasks.

5. Conclusion and Future Work

The paper argues that layer-normalization placement determines Transformer optimization stability and that Pre-LN can avoid warm-up while converging faster. Experiments span translation and pretraining settings, with specific implementation configurations reported for these evaluations.

  • Conclusion and Future Work: Layer-normalization placement matters: Post-LN has large output-layer gradients at initialization, whereas Pre-LN can train without warm-up and converges faster.The paper’s conclusion attributes the Post-LN instability to large gradients under large learning rates and reports faster convergence for the inside-residual-block design.
  • Conclusion and Future Work: The IWSLT14 German-English setup uses about 153K/7K/7K training/validation/test sentence pairs, a 10K joint BPE vocabulary, and 6-layer encoder and decoder.The architecture uses 512-dimensional embeddings and hidden nodes in the attention and feed-forward sublayers.
  • Conclusion and Future Work: The WMT14 English-German setup uses about 4.5M training sentence pairs, a 37K joint BPE vocabulary, and the base Transformer configuration.The base configuration has 6-layer encoder and decoder, 512 hidden nodes and embeddings, 8 heads, label smoothing ε = 0.1, and 8192 tokens per GPU on 16 P40 GPUs.
  • Conclusion and Future Work: BERT pretraining uses English Wikipedia and BookCorpus2, whose concatenation contains roughly 3.4B words.Documents are sentence-segmented, normalized, lower-cased, tokenized with Moses, and BPE-encoded; experiments use 32 NVIDIA Tesla P40 GPUs.
  • Conclusion and Future Work: 35.5 test BLEU is reported for Pre-LN on IWSLT14 De-En with a 7.5e−4 initial learning rate and inverse square root decay after 8000 updates.The reported configuration also uses dropout 0.3, attention dropout 0.1, and batch size 8192.

A.3. GLUE Dataset

The GLUE appendix describes MRPC and RTE as sentence-pair classification tasks and specifies validation-based fine-tuning with repeated random seeds. Evaluation uses accuracy for both datasets.

  • A.3. GLUE Dataset: MRPC evaluates whether sentence pairs from online news sources are semantically equivalent.The corpus includes human annotations, and performance is evaluated by accuracy.
  • A.3. GLUE Dataset: RTE evaluates whether sentences in a pair stand in an entailment relation.The datasets come from annual textual-entailment challenges, with accuracy as the evaluation metric.
  • A.3. GLUE Dataset: Fine-tuning uses validation evaluation, searches batch sizes 16/32, learning rates 1e−5 - 1e−4, and 3-8 epochs, then repeats each task 6 times.The repeated runs use different random seeds because validation accuracy is sensitive to seed choice.

B. Proof of Lemma 1

The lemma proof computes the expected squared norm of a ReLU-transformed Gaussian vector by exploiting independence and Gaussian symmetry. The resulting expression reduces the calculation to a one-dimensional conditional expectation.

  • B. Proof of Lemma 1: For i.i.d. Gaussian coordinates, the squared norm of ReLU(X) is decomposed into a sum of coordinate-wise expectations.Independence and identical distributions allow the d-dimensional expectation to be reduced to one representative coordinate.
  • B. Proof of Lemma 1: The proof therefore expresses the expected ReLU squared norm through d, the nonnegative-half probability, and a conditional Gaussian second moment.This is the lemma’s core reduction from a vector norm to a scalar Gaussian calculation.
  • B. Proof of Lemma 1: The probability that a centered Gaussian coordinate is nonnegative contributes a factor of 1/2.The proof rewrites the coordinate expectation using the conditional second moment given Xi ≥ 0.

C. Proof of Lemma 2

The lemma proofs track intermediate-vector norms under layer normalization and ReLU-based transformations. They use the layer-normalization norm constraint, Gaussian Xavier weights, and induction to bound Pre-LN quantities.

  • C. Proof of Lemma 2: At initialization, layer normalization projects vectors onto a sphere with controlled norm.The proof states that the layer-normalized vector has squared norm d.
  • C. Proof of Lemma 2: The proof estimates intermediate-vector norms by combining Xavier Gaussian weights with the ReLU expectation from Lemma 1.The weight entries are sampled from N(0, 1/d), and the resulting norm calculations are applied to the relevant transformations.
  • C. Proof of Lemma 2: Induction yields bounds on Pre-LN intermediate norms, including (1 + l/2)d ≤ E(∥xpre_l,i∥2).The proof proceeds through successive norm expressions and bounds involving d and layer index l.

D. Proof of Lemma 3

The proof establishes the eigenstructure needed to analyze layer normalization and explicitly derives its Jacobian. The key auxiliary result is that I − α^Tα has eigenvalues only 0 and 1.

  • Lemma 4 states that I − α^Tα has eigenvalues 0 or 1 when α is a unit vector.The eigenvalue 0 corresponds to α, while the orthogonal directions have eigenvalue 1.
  • The proof constructs an orthonormal basis containing α and shows directly how I − α^Tα acts on each basis vector.
  • The analysis rewrites layer normalization in matrix form before calculating its Jacobian explicitly.
  • The Jacobian’s eigenvalue behavior follows from the auxiliary eigenvalue lemma.

E. Proof of Theorem 1

The theorem proof estimates gradients with respect to the final feed-forward parameters using backpropagation, bounded derivatives, and concentration arguments. It contrasts large output-layer gradients in Post-LN with better-behaved gradients in Pre-LN.

  • The proof uses Lemma 5 to control the relevant random-variable tail in the Pre-LN analysis.
  • The proof estimates each element of the gradient matrix with respect to the final feed-forward parameter W 2,L.The Post-LN loss and backpropagated gradient are written explicitly for this parameter.
  • Chernoff and union bounds control ReLU activations with probability at least 0.99 across coordinates.The bound applies to the coordinatewise squared ReLU outputs in the analyzed layer.
  • The Pre-LN loss and last-layer gradient are analyzed analogously, with each gradient component bounded using the same concentration strategy.
  • Post-LN output-layer gradients remain large at initialization, whereas the corresponding Pre-LN bound is smaller and depends on model width.The supplied theorem statements include the Post-LN and Pre-LN scaling expressions, though parts of the equations are truncated.

F. Extension to other layers

The extension analyzes gradients in layers beyond the output layer by composing layer Jacobians. It finds that Post-LN gradients decrease sharply toward earlier layers, while Pre-LN gradients remain comparatively stable across layer position.

  • The gradient for a layer parameter is expressed through products of Jacobians from that layer to the output.The analysis applies this formulation to W 2,l in both normalization placements.
  • The extension relies on norm inequalities and concentration properties of Gaussian-derived random variables.
  • Post-LN gradient norms include a factor O(3^(L−l)/2), which decreases exponentially as l moves toward earlier layers.
  • In Pre-LN, layer-normalization Jacobians are close to identity when l is sufficiently large because their deviations are O(1/√j).
  • The Pre-LN gradient with respect to W 2,l is independent of l in the analyzed regime.

H. Small learning rate experiment

The experiment tests whether a small fixed learning rate can mitigate the instability associated with training Post-LN without warm-up. It evaluates model performance on IWSLT14 German-English translation.

  • The experiment uses a fixed initial learning rate of 1e−4, compared with lrmax = 1e−3 in the paper.
  • Model performance is evaluated on the IWSLT14 De-En task.
Loading 2002.04745v2…