[Paper Review] NormFormer: Improved Transformer Pretraining with Extra Normalization

Sam Shleifer, Jason Weston+1 more|arXiv (Cornell University)|Oct 18, 2021

Topic Modeling34 references28 citations

TL;DR

NormFormer adds three lightweight normalization-based operations per transformer layer to Pre-LN models, reducing gradient mismatch and accelerating pretraining while improving perplexity and downstream tasks across causal and masked language models.

ABSTRACT

During pretraining, the Pre-LayerNorm transformer suffers from a gradient magnitude mismatch: gradients at early layers are much larger than at later layers. These issues can be alleviated by our proposed NormFormer architecture, which adds three normalization operations to each layer: a Layer Norm after self attention, head-wise scaling of self-attention outputs, and a Layer Norm after the first fully connected layer. The extra operations incur negligible compute cost (+0.4% parameter increase), but improve pretraining perplexity and downstream task performance for both causal and masked language models ranging from 125 Million to 2.7 Billion parameters. For example, adding NormFormer on top of our strongest 1.3B parameter baseline can reach equal perplexity 24% faster, or converge 0.27 perplexity better in the same compute budget. This model reaches GPT3-Large (1.3B) zero shot performance 60% faster. For masked language modeling, NormFormer improves fine-tuned GLUE performance by 1.9% on average. Code to train NormFormer models is available in fairseq https://github.com/pytorch/fairseq/tree/main/examples/normformer .

Motivation & Objective

• Identify gradient magnitude mismatches in Pre-LN transformers during pretraining.
• Propose lightweight normalization-based additions to stabilize and accelerate training.
• Evaluate NormFormer on causal and masked language models across multiple scales.
• Demonstrate improvements in pretraining perplexity and downstream task performance.
• Provide ablations and analyses to understand the contribution of each addition.

Proposed method

• Introduce three additions per layer: head-wise scaling of MHA outputs (HeadScale), a LayerNorm after the attention module, and a LayerNorm after the first FFN layer.
• Apply an additional LayerNorm inside the MHA path and a second LN after FFN, with small learnable parameters γ per head and per residual path.
• Optionally include residual scaling (ResScale) on the FFN path, analyzed for its impact at different scales.
• Train causal and masked language models across sizes 125M, 355M, 1.3B, and 2.7B, comparing NormFormer to compute-matched baselines under equal compute budgets.
• Experiment with zero-shot evaluations on GPT-3-like tasks and GLUE benchmarks to assess generalization.

Experimental results

Research questions

• RQ1 Does adding NormFormer’s extra normalization operations stabilize Pre-LN transformers and close gradient gaps across layers?
• RQ2 Do NormFormer gains persist across model scales from 125M to 2.7B parameters?
• RQ3 How do the added operations impact pretraining perplexity and downstream task performance (GLUE) compared to tuned Pre-LN baselines?
• RQ4 What is the effect of residual scaling in NormFormer across different model scales?
• RQ5 Are the gains robust to ablations removing any of the added components?

Key findings

• NormFormer improves pretraining perplexity and downstream task performance for both causal and masked language models across sizes 125M–2.7B.
• For 1.3B models, NormFormer matches baseline perplexity faster and can reach equal perplexity 24% faster under compute parity; it can converge 0.27 perplexity better in the same compute budget.
• Zero-shot evaluation shows NormFormer outperforms GPT-3 at all sizes on the tested tasks.
• GLUE fine-tuning results show NormFormer MLMs outperform Pre-LN baselines across tasks, with average gains.
• Ablation studies show removing any added operation degrades performance; HeadScale and post-attn LN are particularly impactful.
• Learned scaling parameters (γ) reduce early-layer FG gradients and downscale early FFN inputs, while HeadScale can emphasize certain heads, aiding stability and performance.