# Small-scale proxies for large-scale Transformer training instabilities

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

## Key facts

- Cited by 183 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 56 references that are in the map.
- First posted on arXiv 25 Sep 2023.

## Builds on (most cited first, 30 of 41)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [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)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Sigmoid Loss for Language Image Pre-Training](https://mlsanity.com/paper/2303.15343/) (2023)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing](https://mlsanity.com/paper/1808.06226/) (2018)
- [An Empirical Study of Training Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.02057/) (2021)
- [LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale](https://mlsanity.com/paper/2208.07339/) (2022)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [On the Variance of the Adaptive Learning Rate and Beyond](https://mlsanity.com/paper/1908.03265/) (2019)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Generating Wikipedia by Summarizing Long Sequences](https://mlsanity.com/paper/1801.10198/) (2018)
- [Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer](https://mlsanity.com/paper/2203.03466/) (2022)
- [Using the Output Embedding to Improve Language Models](https://mlsanity.com/paper/1608.05859/) (2016)
- [Transformer Quality in Linear Time](https://mlsanity.com/paper/2202.10447/) (2022)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability](https://mlsanity.com/paper/2209.15594/) (2022)
- [Learning-Rate-Free Learning by D-Adaptation](https://mlsanity.com/paper/2301.07733/) (2023)

## Built on by (most cited first, 30 of 183)

- [Gemma 3 Technical Report](https://mlsanity.com/paper/2503.19786/) (2025)
- [Chameleon: Mixed-Modal Early-Fusion Foundation Models](https://mlsanity.com/paper/2405.09818/) (2024)
- [Cosmos World Foundation Model Platform for Physical AI](https://mlsanity.com/paper/2501.03575/) (2025)
- [Show-o: One Single Transformer to Unify Multimodal Understanding and Generation](https://mlsanity.com/paper/2408.12528/) (2024)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [Kimi K2: Open Agentic Intelligence](https://mlsanity.com/paper/2507.20534/) (2025)
- [Generative Verifiers: Reward Modeling as Next-Token Prediction](https://mlsanity.com/paper/2408.15240/) (2024)
- [Massive Activations in Large Language Models](https://mlsanity.com/paper/2402.17762/) (2024)
- [MM1: Methods, Analysis & Insights from Multimodal LLM Pre-training](https://mlsanity.com/paper/2403.09611/) (2024)
- [When Attention Sink Emerges in Language Models: An Empirical View](https://mlsanity.com/paper/2410.10781/) (2024)
- [TokenFlow: Unified Image Tokenizer for Multimodal Understanding and Generation](https://mlsanity.com/paper/2412.03069/) (2024)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations](https://mlsanity.com/paper/2405.18392/) (2024)
- [Why Transformers Need Adam: A Hessian Perspective](https://mlsanity.com/paper/2402.16788/) (2024)
- [Language models scale reliably with over-training and on downstream tasks](https://mlsanity.com/paper/2403.08540/) (2024)
- [Scaling Laws for Precision](https://mlsanity.com/paper/2411.04330/) (2024)
- [Stable LM 2 1.6B Technical Report](https://mlsanity.com/paper/2402.17834/) (2024)
- [Mamba-3: Improved Sequence Modeling using State Space Principles](https://mlsanity.com/paper/2603.15569/) (2026)
- [Resolving Discrepancies in Compute-Optimal Scaling of Language Models](https://mlsanity.com/paper/2406.19146/) (2024)
- [Disentangling the Causes of Plasticity Loss in Neural Networks](https://mlsanity.com/paper/2402.18762/) (2024)
- [Hyper-Connections](https://mlsanity.com/paper/2409.19606/) (2024)
- [Don't be lazy: CompleteP enables compute-efficient deep transformers](https://mlsanity.com/paper/2505.01618/) (2025)
- [Test-time regression: a unifying framework for designing sequence models with associative memory](https://mlsanity.com/paper/2501.12352/) (2025)
- [How Does Critical Batch Size Scale in Pre-training?](https://mlsanity.com/paper/2410.21676/) (2024)
- [Distillation Scaling Laws](https://mlsanity.com/paper/2502.08606/) (2025)
- [Proteina: Scaling Flow-based Protein Structure Generative Models](https://mlsanity.com/paper/2503.00710/) (2025)
- [Normalization and effective learning rates in reinforcement learning](https://mlsanity.com/paper/2407.01800/) (2024)
- [Spike No More: Stabilizing the Pre-training of Large Language Models](https://mlsanity.com/paper/2312.16903/) (2023)
- [Parameters vs FLOPs: Scaling Laws for Optimal Sparsity for Mixture-of-Experts Language Models](https://mlsanity.com/paper/2501.12370/) (2025)
- [Unified Multimodal Discrete Diffusion](https://mlsanity.com/paper/2503.20853/) (2025)

## Abstract (arXiv)

Teams that have trained large Transformer-based models have reported training instabilities at large scale that did not appear when training with the same hyperparameters at smaller scales. Although the causes of such instabilities are of scientific interest, the amount of resources required to reproduce them has made investigation difficult. In this work, we seek ways to reproduce and study training stability and instability at smaller scales. First, we focus on two sources of training instability described in previous work: the growth of logits in attention layers (Dehghani et al., 2023) and divergence of the output logits from the log probabilities (Chowdhery et al., 2022). By measuring the relationship between learning rate and loss across scales, we show that these instabilities also appear in small models when training at high learning rates, and that mitigations previously employed at large scales are equally effective in this regime. This prompts us to investigate the extent to which other known optimizer and model interventions influence the sensitivity of the final loss to changes in the learning rate. To this end, we study methods such as warm-up, weight decay, and the μParam (Yang et al., 2022), and combine techniques to train small models that achieve similar losses across orders of magnitude of learning rate variation. Finally, to conclude our exploration we study two cases where instabilities can be predicted before they emerge by examining the scaling behavior of model activation and gradient norms.