# Stable Language Model Pre-training by Reducing Embedding Variability

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 53 references that are in the map.
- First posted on arXiv 12 Sep 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [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)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Scaling Language Models: Methods, Analysis & Insights from Training Gopher](https://mlsanity.com/paper/2112.11446/) (2021)
- [The LAMBADA dataset: Word prediction requiring a broad discourse context](https://mlsanity.com/paper/1606.06031/) (2016)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [On the Variance of the Adaptive Learning Rate and Beyond](https://mlsanity.com/paper/1908.03265/) (2019)
- [Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning](https://mlsanity.com/paper/2012.13255/) (2020)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [8-bit Optimizers via Block-wise Quantization](https://mlsanity.com/paper/2110.02861/) (2021)
- [One Billion Word Benchmark for Measuring Progress in Statistical Language Modeling](https://mlsanity.com/paper/1312.3005/) (2013)
- [PyHessian: Neural Networks Through the Lens of the Hessian](https://mlsanity.com/paper/1912.07145/) (2019)
- [Speeding up Convolutional Neural Networks with Low Rank Expansions](https://mlsanity.com/paper/1405.3866/) (2014)
- [Using the Output Embedding to Improve Language Models](https://mlsanity.com/paper/1608.05859/) (2016)
- [Mitigating Neural Network Overconfidence with Logit Normalization](https://mlsanity.com/paper/2205.09310/) (2022)
- [ReLoRA: High-Rank Training Through Low-Rank Updates](https://mlsanity.com/paper/2307.05695/) (2023)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [The Shattered Gradients Problem: If resnets are the answer, then what is the question?](https://mlsanity.com/paper/1702.08591/) (2017)
- [On Identifiability in Transformers](https://mlsanity.com/paper/1908.04211/) (2019)

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

- [PolyPythias: Stability and Outliers across Fifty Language Model Pre-Training Runs](https://mlsanity.com/paper/2503.09543/) (2025)
- [Peri-LN: Revisiting Normalization Layer in the Transformer Architecture](https://mlsanity.com/paper/2502.02732/) (2025)
- [Cross-lingual Transfer of Reward Models in Multilingual Alignment](https://mlsanity.com/paper/2410.18027/) (2024)
- [Exploiting Vocabulary Frequency Imbalance in Language Model Pre-training](https://mlsanity.com/paper/2508.15390/) (2025)
- [On the Spectral Flattening of Quantized Embeddings](https://mlsanity.com/paper/2602.00969/) (2026)
- [Continuity and Ordinality Matter: Constraining Time Series Tokens for Effective Time Series Analysis with Large Language Models](https://mlsanity.com/paper/2605.28866/) (2026)

## Abstract (arXiv)

Stable pre-training is essential for achieving better-performing language models. However, tracking pre-training stability by calculating gradient variance at every step is impractical due to the significant computational costs. We explore Token Embedding Variability (TEV) as a simple and efficient proxy for assessing pre-training stability in language models with pre-layer normalization, given that shallower layers are more prone to gradient explosion (section 2.2). Moreover, we propose Multi-head Low-Rank Attention (MLRA) as an architecture to alleviate such instability by limiting the exponential growth of output embedding variance, thereby preventing the gradient explosion (section 3.2). Empirical results on GPT-2 with MLRA demonstrate increased stability and lower perplexity, particularly in deeper models.