# Methods of improving LLM training stability

arXiv: https://arxiv.org/abs/2410.16682 · HTML: https://mlsanity.com/paper/2410.16682/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 21 references that are in the map.
- First posted on arXiv 22 Oct 2024.

## Builds on (most cited first, 16 of 16)

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing](https://mlsanity.com/paper/1808.06226/) (2018)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Going deeper with Image Transformers](https://mlsanity.com/paper/2103.17239/) (2021)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [Neural Combinatorial Optimization with Reinforcement Learning](https://mlsanity.com/paper/1611.09940/) (2016)
- [Query-Key Normalization for Transformers](https://mlsanity.com/paper/2010.04245/) (2020)
- [Small-scale proxies for large-scale Transformer training instabilities](https://mlsanity.com/paper/2309.14322/) (2023)
- [Primer: Searching for Efficient Transformers for Language Modeling](https://mlsanity.com/paper/2109.08668/) (2021)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing](https://mlsanity.com/paper/2306.12929/) (2023)
- [Stable and low-precision training for large-scale vision-language models](https://mlsanity.com/paper/2304.13013/) (2023)
- [A Theory on Adam Instability in Large-Scale Machine Learning](https://mlsanity.com/paper/2304.09871/) (2023)
- [Transformer Normalisation Layers and the Independence of Semantic Subspaces](https://mlsanity.com/paper/2406.17837/) (2024)

## Built on by (most cited first, 13 of 13)

- [Rope to Nope and Back Again: A New Hybrid Attention Strategy](https://mlsanity.com/paper/2501.18795/) (2025)
- [HybridNorm: Towards Stable and Efficient Transformer Training via Hybrid Normalization](https://mlsanity.com/paper/2503.04598/) (2025)
- [Scaling Behavior Cloning Improves Causal Reasoning: An Open Model for Real-Time Video Game Playing](https://mlsanity.com/paper/2601.04575/) (2026)
- [YuLan-Mini: An Open Data-efficient Language Model](https://mlsanity.com/paper/2412.17743/) (2024)
- [LIBMoE: A Library for comprehensive benchmarking Mixture of Experts in Large Language Models](https://mlsanity.com/paper/2411.00918/) (2024)
- [Pion: A Spectrum-Preserving Optimizer via Orthogonal Equivalence Transformation](https://mlsanity.com/paper/2605.12492/) (2026)
- [Instability in Downstream Task Performance During LLM Pretraining](https://mlsanity.com/paper/2510.04848/) (2025)
- [Controlling changes to attention logits](https://mlsanity.com/paper/2511.21377/) (2025)
- [Spectral Alignment as Predictor of Loss Explosion in Neural Network Training](https://mlsanity.com/paper/2510.04202/) (2025)
- [GeoNorm: Unify Pre-Norm and Post-Norm with Geodesic Optimization](https://mlsanity.com/paper/2601.22095/) (2026)
- [Stability as a Liability:Systematic Breakdown of Linguistic Structure in LLMs](https://mlsanity.com/paper/2601.18588/) (2026)
- [A Unifying View of Attention Sinks: From Mechanisms to Architectural Interventions](https://mlsanity.com/paper/2606.08105/) (2026)
- [One QK Channel, Many Sources: Tracing Low-Precision Attention Collapse](https://mlsanity.com/paper/2608.02091/) (2026)

## Abstract (arXiv)

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.