NormFormer: Improved Transformer Pretraining with Extra Normalization

By Sam Shleifer, Jason Aaron Edward Weston, Myle Ott. 2021

97 citations recorded by the source corpus.

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 .

Subfields: Artificial Intelligence, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Furu Wei
2. Jianfeng Gao
3. Mohit Bansal
4. Quoc V. Le
5. Luke S. Zettlemoyer
6. Li Dong
7. Jiashi Feng
8. Graham Neubig

In Model Architectures

1. Jiashi Feng
2. Irwan Bello
3. Quoc V. Le
4. Furu Wei
5. Shuicheng Yan
6. Barret Zoph
7. Weihao Yu
8. Xinchao Wang

Citations in the corpus

• GLM-130B: An Open Bilingual Pre-trained Model (2022, background): “As a remedy, on top of Pre-LN, CogView (later in Normformer (Shleifer et al., 2021)) develops Sandwich-LN which appends extra normalization to the end of each residual branch.”
• OFA: Unifying Architectures, Tasks, and Modalities Through a Simple Sequence-to-Sequence Learning Framework (2022, methodology): “To stabilize training and accelerate convergence, we add head scaling to self attention, a post-attention layer normalization (LN) [65], and an LN following the first layer of FFN [66].”
• Learning to (Learn at Test Time): RNNs with Expressive Hidden States (2024, unknown): “[67].”
• Dual PatchNorm (2023, methodology): “Our observation also applies to other alternate LayerNorm placements: NormFormer (Shleifer et al., 2021) and Sub-LN (Wang et al., 2022), which in isolation, do not improve over strong ViT classification models.”
• Designing Effective Sparse Expert Models (2022, result): “In line with our findings, Shleifer et al. (2021) found adding a learned multiplicative scalar to the residual connection in Transformers made them much more unstable.”
• Stabilizing Transformer Training by Preventing Attention Entropy Collapse (2023, background): “Besides the popular post-LN and pre-LN configurations, other variants have been proposed (Wang et al., 2022; Shleifer et al., 2021; Liu et al., 2020a).”
• A generalist vision–language foundation model for diverse biomedical tasks (2023, unknown): “62.”
• A Comprehensive Overview of Large Language Models (2023, unknown): “DeepNorm: While pre-LN has certain benefits over postLN training, pre-LN training has an unwanted effect on the gradients [59].”
• ST-MoE: Designing Stable and Transferable Sparse Expert Models (2022, result): “In line with our findings, Shleifer et al. (2021) found adding a learned multiplicative scalar to the residual connection in Transformers made them much more unstable.”
• Foundation Transformers (2022, methodology): “We compare M AGNETO with vanilla Pre-LN Transformer and Normformer (Shleifer et al., 2021).”
• DeepNet: Scaling Transformers to 1,000 Layers (2022, methodology): “In order to alleviate the above issue, there have been efforts on improving the optimization of deep Transformer by means of better initialization [16], [17], [18], or better architecture [15], [19], [20], [21].”
• MetaFormer Baselines for Vision (2022, methodology): “Instead, we introduce a new activation StarReLU and other two modifications [60, 55, 10] to improve MetaFormer.”