# Efficient Training of Language Models with Compact and Consistent Next Token Distributions

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 30 of its 51 references that are in the map.
- First posted on arXiv 3 Jul 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Distributed Representations of Words and Phrases and their Compositionality](https://mlsanity.com/paper/1310.4546/) (2013)
- [QLoRA: Efficient Finetuning of Quantized LLMs](https://mlsanity.com/paper/2305.14314/) (2023)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems](https://mlsanity.com/paper/1905.00537/) (2019)
- [What Disease does this Patient Have? A Large-scale Open Domain Question Answering Dataset from Medical Exams](https://mlsanity.com/paper/2009.13081/) (2020)
- [PubMedQA: A Dataset for Biomedical Research Question Answering](https://mlsanity.com/paper/1909.06146/) (2019)
- [When Does Label Smoothing Help?](https://mlsanity.com/paper/1906.02629/) (2019)
- [Neural Network Acceptability Judgments](https://mlsanity.com/paper/1805.12471/) (2018)
- [GPT-NeoX-20B: An Open-Source Autoregressive Language Model](https://mlsanity.com/paper/2204.06745/) (2022)
- [OLMo: Accelerating the Science of Language Models](https://mlsanity.com/paper/2402.00838/) (2024)
- [BioGPT: Generative Pre-trained Transformer for Biomedical Text Generation and Mining](https://mlsanity.com/paper/2210.10341/) (2022)
- [Neural Text Generation with Unlikelihood Training](https://mlsanity.com/paper/1908.04319/) (2019)
- [MAUVE: Measuring the Gap Between Neural Text and Human Text using Divergence Frontiers](https://mlsanity.com/paper/2102.01454/) (2021)
- [One Billion Word Benchmark for Measuring Progress in Statistical Language Modeling](https://mlsanity.com/paper/1312.3005/) (2013)
- [A Contrastive Framework for Neural Text Generation](https://mlsanity.com/paper/2202.06417/) (2022)
- [Exploring the Limits of Language Modeling](https://mlsanity.com/paper/1602.02410/) (2016)
- [The MiniPile Challenge for Data-Efficient Language Models](https://mlsanity.com/paper/2304.08442/) (2023)
- [Contrastive Search Is What You Need For Neural Text Generation](https://mlsanity.com/paper/2210.14140/) (2022)
- [ContraCLM: Contrastive Learning For Causal Language Model](https://mlsanity.com/paper/2210.01185/) (2022)
- [Generalizing and Hybridizing Count-based and Neural Language Models](https://mlsanity.com/paper/1606.00499/) (2016)
- [Language Modelling via Learning to Rank](https://mlsanity.com/paper/2110.06961/) (2021)
- [BPTree: an ℓ2 heavy hitters algorithm using constant memory](https://mlsanity.com/paper/1603.00759/) (2016)
- [New Algorithms for Heavy Hitters in Data Streams](https://mlsanity.com/paper/1603.01733/) (2016)

## Abstract (arXiv)

Maximizing the likelihood of the next token is an established, statistically sound objective for pre-training language models. In this paper we show that we can train better models faster by pre-aggregating the corpus with a collapsed n-gram distribution. Previous studies have proposed corpus-level n-gram statistics as a regularizer; however, the construction and querying of such n-grams, if done naively, prove to be costly and significantly impede training speed, thereby limiting their application in modern large language model pre-training. We introduce an alternative compact representation of the next token distribution that, in expectation, aligns with the complete n-gram distribution while markedly reducing variance across mini-batches compared to the standard next-token loss. Empirically, we demonstrate that both the n-gram regularized model and our approximation yield substantial improvements in model quality and convergence rate compared to existing methods. Furthermore, our approximation facilitates scalability of gains to larger datasets and models compared to the straightforward n-gram regularization method.