NormFormer: Improved Transformer Pretraining with Extra Normalization

1 INTRODUCTION

• The original Transformer applies LayerNorm after each sublayer’s residual connection, whereas Pre-LN moves it to the beginning of each sublayer.Post-LN reduces variance in the inputs to the following sublayer.Source: §1 · p. 1
• Pre-LN permits larger learning rates and shorter warmup than Post-LN, often improves performance, and is used by most recent large pretrained language models.Source: §1 · p. 1
• Pre-LN tends to give earlier layers larger gradients than later layers; NormFormer adds three normalization operations that bring their gradient magnitudes closer together.The operations reduce gradients to early layers and increase gradients to later layers.Source: §1 · p. 1
• Figure 1 shows the baseline and NormFormer layers and places trainable HeadScale parameters γi before the attention output projection.Source: Figure 1 · p. 2
• Compared with compute-matched, well-tuned Pre-LN baselines, NormFormer reaches target pretraining perplexities faster and achieves better pretraining perplexities and downstream task performance.Source: §1 · p. 2

2.1 NORMFORMER

• NormFormer adds head-wise attention scaling and LayerNorm after attention and the first fully connected layer, with learnable parameters that adjust feature and subsequent-gradient magnitudes.The modifications are applied within each layer; the overall method is expressed through NormScaledMHA and NormFFN.Source: §2.1 · pp. 2–3
• Each attention head’s output is multiplied by a learned scalar initialized to 1 before the heads are concatenated and projected.The standard multi-head attention operation concatenates head outputs and applies an output projection.Source: §2.1 · p. 2

2.2 EXPERIMENTS

• Causal language models span 125M to 2.7B parameters, while 2.7B experiments compare a replicated GPT-3 baseline with two ReLU2 variants.The 2.7B baseline uses the published 1.6e-4 learning rate and 286K steps; the NormFormer variant uses 6e-4, while the authors did not exhaustively tune rates at this size.Source: §2.2 · pp. 3–4
• Models are compared at equal GPU time, so NormFormer typically receives 2–6% fewer training steps and tokens than the baseline.NormFormer adds fewer than 0.07% parameters but increases memory use by 2–6%, depending on model size; the authors report that moving HeadScale outside attention reduces overhead without noticeable performance degradation.Source: §2.2 · pp. 3–4
• Learning rates selected on the authors’ dataset are 3–5 times larger than those in GPT-3, and the search supplies stronger baselines for comparison.Except for 2.7B models, baseline and NormFormer rates are tuned using 50,000-step runs; the authors note that architectural differences may partly explain the learning-rate gap.Source: §2.2 · p. 3, Table 1 · p. 3
• Masked-language-model experiments use RoBERTa-base settings and fine-tune baseline and NormFormer models across seven learning rates on GLUE.The baseline trains for 2 million batches of 1 million tokens, about one quarter of RoBERTa-base’s original budget; NormFormer completes 1.92 million batches in the same time. Pretraining uses about 450GB of English text and 110B BPE tokens, with 40M held out for validation.Source: §2.2 · p. 4

3 RESULTS

• 60% and 57% as much compute lets NormFormer match Pre-LN models’ lowest perplexity for CLM and MLM, respectively.NormFormer models take 2–6% longer per training step and see less data in this comparison.Source: §3 · pp. 4–5, Figure 2 · p. 5
• NormFormer outperforms GPT-3 at all model sizes on causal-LM zero-shot tasks, including HellaSwag, PIQA, WinoGrande, StoryCloze, and OpenBookQA.The table caption identifies these tasks and groups compute-matched runs with horizontal lines.Source: §3 · p. 5, Table 2 · p. 5
• 0.83 perplexity improvement comes from NormFormer-125M with 100,000 added parameters, versus 0.08 from widening the baseline to 127 million parameters.The widened model changes hidden dimension from 768 to 780.Source: §3 · p. 6
• NormFormer MLM models outperform their Pre-LN counterparts on every GLUE task in equal-compute comparisons.The table caption notes that these models receive less compute than publicly released RoBERTa-base models.Source: §3 · p. 6, Table 3 · p. 6

4.1 ANALYSIS OF GRADIENT NORMS BY LAYER

• At the beginning of training, Post-LN gradients to later layers were much larger than to earlier layers, Pre-LN showed the opposite pattern, and NormFormer brought average gradient norms closer together across layers.The comparison measured average L1 gradient norms to the second fully connected weight in a 12-layer, 125M-parameter CLM model.Source: §4.1 · p. 6
• Earlier-layer FFN-LN γ values are smaller, while post-attention LN γ values are below 1 and HeadScale γ varies without a depth relationship.The authors interpret HeadScale variation as dynamically increasing the importance of well-initialized attention heads.Source: §4.1 · pp. 6–7, Figure 4 · p. 7
• NormFormer survives more updates under an increasing-learning-rate stability test; attention-output downscaling avoids the baseline 125M model’s layer-0 FP16 query-key overflow.The authors report that λresid reduces the stability improvement at all tested sizes.Source: §4.1 · p. 7, Figure 5 · p. 7

4.2 RESIDUAL SCALING

• Residual scaling improves perplexity and zero-shot performance for small-scale CLMs, but at 125M MLM scale its marginal pretraining gain hurts fine-tuned performance.At 1.3B scale, adding ResScale to NormFormer does not improve performance.Source: §4.2 · p. 8
• At 125M and 355M, later-layer residual weights are lower, whereas at 1.3B they are larger deeper in the network.Figure 6 presents average λresid weights by layer for different-sized CLMs.Source: §4.2 · p. 8, Figure 6 · p. 8
• Moving λresid to the other residual connection or replacing its vector with a scalar does not fix the large-scale issue and marginally hurts small-scale performance.The authors suggest the negative result is caused by scale rather than interaction with NormFormer.Source: §4.2 · p. 8

Removing any of the added operations hurts performance

• In 125M-parameter language modeling, removing any proposed operation worsens perplexity; removing HeadScale is most damaging, while removing post-attention LayerNorm is least damaging.Adding three more LayerNorm operations leaves perplexity unchanged at a fixed number of updates and reduces training speed by 5%.Source: §5 · p. 8, Table 4 · p. 9

Other experiments

• Replacing FFN LayerNorm with FFNGeGlu degrades performance in the tested 125M CLM setting; a LayerNorm variant performs equally well with fewer parameters but is about 2x slower.The slower variant removes bias and mean subtraction; the authors therefore do not adopt it.Source: §5 · p. 9
• Related stabilization techniques for text-to-image models do not help in this setting, apart from post-attention LayerNorm.Source: §5 · p. 9
• NormFormer outperforms the baseline in all seven reported hyperparameter configurations.The configurations include longer warmup, GPT-3 settings, and gradient clipping.Source: §5 · p. 9

• Pre-LN builds on work addressing Post-LN gradients, and NormFormer extends Pre-LN to improve stability and efficiency.Source: §6 · p. 9
• Press et al. (2020a) place attention before feed-forward layers, increasing the number of late FFN parameters rather than their importance and gradient norm.The passage says this arrangement does not impact stability.Source: §6 · p. 9
• Chen et al. (2021) use low HeadScale γ parameters to prune attention heads; NormFormer uses γ parameters to improve pretraining performance.Source: §6 · p. 9
• GradInit (Zhu et al., 2021) and Admin (Liu et al., 2020) use variance heuristics for initialization or profiling; their work also uses variants of ResScale.The authors find ResScale helpful at small scale and harmful at large scale.Source: §6 · p. 9
• Unlike initialization-focused approaches, NormFormer adds learnable parameters without variance heuristics, extra stages, or initialization changes.Related examples include ReZero (Bachlechner et al., 2020), FixUp, LookLinear, and DALL-E (Ramesh et al., 2021).Source: §6 · p. 9

• NormFormer’s three added operations alleviate Pre-LN gradient mismatch and improve validation perplexity and downstream performance for causal and masked language models.Source: §7 · p. 10
• Earlier Pre-LN layers receive larger gradients, while optimal residual scaling is larger in earlier layers than later ones.Source: §7 · p. 10
• Removing any added operation degrades performance toward the baseline, and the tested additional normalization does not improve performance.Source: §7 · p. 10
• The authors suggest testing whether all three operations are needed at every layer and whether FFN LayerNorm can be fused to reduce computational overhead.Mean centering and bias do not appear to improve pretraining perplexity.Source: §7 · p. 10

• NormFormer outperforms the baseline in all seven settings reported in the hyperparameter table.The settings include longer warmup, GPT-3 settings, and gradient clipping.Source: Table 5 · p. 11
• On Wikitext-103, NormFormer reaches the baseline perplexity in 70% as many steps and improves convergence perplexity, but shows no improvement in the last 30% of training.The experiments add the three operations to the architecture and hyperparameters of Baevski & Auli (2019).Source: §8 · p. 11, Table 6 · p. 11
• The appendix defines Steps to Final PPL as the percentage of 280K steps needed to reach 18.70 perplexity, and PPL as best perplexity.The associated hyperparameter-table budget lets the baseline run for 100,000 updates.Source: Table 6 · p. 11, Table 7 · p. 11
• The gradient comparison reports average norms from steps 950–1000, normalized across losses; NormFormer raises gradients to later fully connected parameters and lowers attention-parameter gradients at all layers.A y-axis value of 2.0 means a parameter’s gradient is twice the baseline average.Source: Figure 7 · p. 12