NormFormer: Improved Transformer Pretraining with Extra Normalization

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 .

References in corpus (12)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Gaussian Error Linear Units (GELUs)
• Zero-Shot Text-to-Image Generation
• GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding
• CogView: Mastering Text-to-Image Generation via Transformers
• fairseq: A Fast, Extensible Toolkit for Sequence Modeling
• ReZero is All You Need: Fast Convergence at Large Depth
• The Shattered Gradients Problem: If resnets are the answer, then what is the question?
• Understanding the Difficulty of Training Transformers
• GLU Variants Improve Transformer
• Primer: Searching for Efficient Transformers for Language Modeling
• GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training

Cited by in corpus (6)

• MetaFormer Baselines for Vision
• Findings of the BabyLM Challenge: Sample-Efficient Pretraining on Developmentally Plausible Corpora
• PC-JeDi: Diffusion for Particle Cloud Generation in High Energy Physics
• Masked Particle Modeling on Sets: Towards Self-Supervised High Energy Physics Foundation Models
• Machine Learning in High Energy Physics: A review of heavy-flavor jet tagging at the LHC
• Learning to Optimise Wind Farms with Graph Transformers