Methods of improving LLM training stability

TL;DR

This paper investigates training instabilities in large language models, focusing on the growth of linear layer outputs, and proposes normalization techniques to improve training stability and perplexity.

Contribution

It extends previous work by analyzing all linear layer outputs in Transformers and introduces normalization methods that enable higher learning rates and better model performance.

Findings

Applying layer normalization to additional layers increases stable learning rates.

Layer normalization after QK layers allows 1.5x higher learning rates without divergence.

Proposed methods significantly improve perplexity over baseline models.

Abstract

Training stability of large language models(LLMs) is an important research topic. Reproducing training instabilities can be costly, so we use a small language model with 830M parameters and experiment with higher learning rates to force models to diverge. One of the sources of training instability is the growth of logits in attention layers. We extend the focus of the previous work and look not only at the magnitude of the logits but at all outputs of linear layers in the Transformer block. We observe that with a high learning rate the L2 norm of all linear layer outputs can grow with each training step and the model diverges. Specifically we observe that QKV, Proj and FC2 layers have the largest growth of the output magnitude. This prompts us to explore several options: 1) apply layer normalization not only after QK layers but also after Proj and FC2 layers too; 2) apply layer normalization after the QKV layer (and remove pre normalization). 3) apply QK layer normalization together with softmax capping. We show that with the last two methods we can increase learning rate by 1.5x (without model divergence) in comparison to an approach based on QK layer normalization only. Also we observe significant perplexity improvements for all three methods in comparison to the baseline model.