# Output Embedding Centering for Stable LLM Pretraining

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 31 references that are in the map.
- First posted on arXiv 5 Jan 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Chameleon: Mixed-Modal Early-Fusion Foundation Models](https://mlsanity.com/paper/2405.09818/) (2024)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [Using the Output Embedding to Improve Language Models](https://mlsanity.com/paper/1608.05859/) (2016)
- [Representation Degeneration Problem in Training Natural Language Generation Models](https://mlsanity.com/paper/1907.12009/) (2019)
- [All-but-the-Top: Simple and Effective Postprocessing for Word Representations](https://mlsanity.com/paper/1702.01417/) (2017)
- [Small-scale proxies for large-scale Transformer training instabilities](https://mlsanity.com/paper/2309.14322/) (2023)
- [What Language Model Architecture and Pretraining Objective Work Best for Zero-Shot Generalization?](https://mlsanity.com/paper/2204.05832/) (2022)
- [Resolving Discrepancies in Compute-Optimal Scaling of Language Models](https://mlsanity.com/paper/2406.19146/) (2024)
- [Spike No More: Stabilizing the Pre-training of Large Language Models](https://mlsanity.com/paper/2312.16903/) (2023)
- [Better Embeddings with Coupled Adam](https://mlsanity.com/paper/2502.08441/) (2025)

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

- [Grokking or Glitching? How Low-Precision Drives Slingshot Loss Spikes](https://mlsanity.com/paper/2605.06152/) (2026)
- [Sparse Readout Prism: Explaining Logit-Lens Scores in Features Instead of Tokens](https://mlsanity.com/paper/2609.01936/) (2026)
- [Softmax Reparameterization for Output-Head Quantization](https://mlsanity.com/paper/2609.31291/) (2026)
- [Why Adaptive Optimizers Underestimate Rare Tokens](https://mlsanity.com/paper/2609.37535/) (2026)

## Abstract (arXiv)

Pretraining of large language models is not only expensive but also prone to certain training instabilities. A specific instability that often occurs at the end of training is output logit divergence. The most widely used mitigation strategies, z-loss and logit soft-capping, merely address the symptoms rather than the underlying cause of the problem. In this paper, we analyze the instability from the perspective of the output embeddings' geometry and identify anisotropic embeddings as its source. Based on this, we propose output embedding centering (OEC) as a new mitigation strategy, and demonstrate that it suppresses output logit divergence. OEC can be implemented in two different ways: as a deterministic operation called μ-centering, or a regularization method called μ-loss. Our experiments show that both variants outperform z-loss in terms of training stability, while being on par with logit soft-capping. This holds true both in the presence and the absence of weight tying. As a secondary result, we find that μ-loss is significantly less sensitive to regularization hyperparameter tuning than z-loss.