# Why do Nearest Neighbor Language Models Work?

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

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- Placed in Generation and context integration, one of 15 directions in Retrieval-augmented generation.
- Builds on 33 of its 44 references that are in the map.
- First posted on arXiv 7 Jan 2023.

## Placed in

- Retrieval-augmented generation / [Generation and context integration](https://mlsanity.com/rag/generation-and-context-integration/)

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

- [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)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [Regularizing Neural Networks by Penalizing Confident Output Distributions](https://mlsanity.com/paper/1701.06548/) (2017)
- [Adaptive Input Representations for Neural Language Modeling](https://mlsanity.com/paper/1809.10853/) (2018)
- [Nearest Neighbor Machine Translation](https://mlsanity.com/paper/2010.00710/) (2020)
- [Regularizing and Optimizing LSTM Language Models](https://mlsanity.com/paper/1708.02182/) (2017)
- [Breaking the Softmax Bottleneck: A High-Rank RNN Language Model](https://mlsanity.com/paper/1711.03953/) (2017)
- [On NMT Search Errors and Model Errors: Cat Got Your Tongue?](https://mlsanity.com/paper/1908.10090/) (2019)
- [Training Language Models with Memory Augmentation](https://mlsanity.com/paper/2205.12674/) (2022)
- [Efficient Nearest Neighbor Language Models](https://mlsanity.com/paper/2109.04212/) (2021)
- [Generating Sentences by Editing Prototypes](https://mlsanity.com/paper/1709.08878/) (2017)
- [Efficient softmax approximation for GPUs](https://mlsanity.com/paper/1609.04309/) (2016)
- [Neuro-Symbolic Language Modeling with Automaton-augmented Retrieval](https://mlsanity.com/paper/2201.12431/) (2022)
- [Generalized Entropy Regularization or: There's Nothing Special about Label Smoothing](https://mlsanity.com/paper/2005.00820/) (2020)
- [Best-First Beam Search](https://mlsanity.com/paper/2007.03909/) (2020)
- [Unbounded cache model for online language modeling with open vocabulary](https://mlsanity.com/paper/1711.02604/) (2017)
- [Efficient Cluster-Based k-Nearest-Neighbor Machine Translation](https://mlsanity.com/paper/2204.06175/) (2022)
- [Stolen Probability: A Structural Weakness of Neural Language Models](https://mlsanity.com/paper/2005.02433/) (2020)
- [You can't pick your neighbors, or can you? When and how to rely on retrieval in the kNN-LM](https://mlsanity.com/paper/2210.15859/) (2022)

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

- [Retrieval-Augmented Generation for Natural Language Processing: A Survey](https://mlsanity.com/paper/2407.13193/) (2024)
- [Privacy Implications of Retrieval-Based Language Models](https://mlsanity.com/paper/2305.14888/) (2023)
- [Chameleon: a Heterogeneous and Disaggregated Accelerator System for Retrieval-Augmented Language Models](https://mlsanity.com/paper/2310.09949/) (2023)
- [CAGRA: Highly Parallel Graph Construction and Approximate Nearest Neighbor Search for GPUs](https://mlsanity.com/paper/2308.15136/) (2023)
- [Non-Exchangeable Conformal Language Generation with Nearest Neighbors](https://mlsanity.com/paper/2402.00707/) (2024)
- [Retrofitting Large Language Models with Dynamic Tokenization](https://mlsanity.com/paper/2411.18553/) (2024)
- [Memory Decoder: A Pretrained, Plug-and-Play Memory for Large Language Models](https://mlsanity.com/paper/2508.09874/) (2025)
- [KNN-LM Does Not Improve Open-ended Text Generation](https://mlsanity.com/paper/2305.14625/) (2023)
- [FT2Ra: A Fine-Tuning-Inspired Approach to Retrieval-Augmented Code Completion](https://mlsanity.com/paper/2404.01554/) (2024)
- [Surface-Based Retrieval Reduces Perplexity of Retrieval-Augmented Language Models](https://mlsanity.com/paper/2305.16243/) (2023)
- [On the Predictive Power of Representation Dispersion in Language Models](https://mlsanity.com/paper/2506.24106/) (2025)
- [NNOSE: Nearest Neighbor Occupational Skill Extraction](https://mlsanity.com/paper/2401.17092/) (2024)
- [On the Generalization Ability of Retrieval-Enhanced Transformers](https://mlsanity.com/paper/2302.12128/) (2023)
- [Incorporating Precedents for Legal Judgement Prediction on European Court of Human Rights Cases](https://mlsanity.com/paper/2409.18644/) (2024)
- [Languages Transferred Within the Encoder: On Representation Transfer in Zero-Shot Multilingual Translation](https://mlsanity.com/paper/2406.08092/) (2024)
- [MemSFT: Mitigating Alignment Tax with an External Parametric Memory](https://mlsanity.com/paper/2607.25614/) (2026)
- [On Retrieval Augmentation and the Limitations of Language Model Training](https://mlsanity.com/paper/2311.09615/) (2023)
- [Local Topology Measures of Contextual Language Model Latent Spaces With Applications to Dialogue Term Extraction](https://mlsanity.com/paper/2408.03706/) (2024)
- [Nearest Neighbor Machine Translation is Meta-Optimizer on Output Projection Layer](https://mlsanity.com/paper/2305.13034/) (2023)
- [Efficient Nearest Neighbor based Uncertainty Estimation for Natural Language Processing Tasks](https://mlsanity.com/paper/2407.02138/) (2024)
- [Long-Tail Crisis in Nearest Neighbor Language Models](https://mlsanity.com/paper/2503.22426/) (2025)
- [Studying the Role of Input-Neighbor Overlap in Retrieval-Augmented Language Models Training Efficiency](https://mlsanity.com/paper/2505.14309/) (2025)
- [Memorizing Documents with Guidance in Large Language Models](https://mlsanity.com/paper/2406.15996/) (2024)
- [Experimental comparison of graph-based approximate nearest neighbor search algorithms on edge devices](https://mlsanity.com/paper/2411.14006/) (2024)
- [Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts](https://mlsanity.com/paper/2601.02144/) (2026)
- [Adaptive kNN graph model](https://mlsanity.com/paper/2601.16509/) (2026)
- [Nearest-Neighbor Radii under Dependent Sampling](https://mlsanity.com/paper/2605.14343/) (2026)
- [One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context](https://mlsanity.com/paper/2609.01311/) (2026)
- [Forking: Sudden Overfitting Under Replay](https://mlsanity.com/paper/2610.00394/) (2026)

## Abstract (arXiv)

Language models (LMs) compute the probability of a text by sequentially computing a representation of an already-seen context and using this representation to predict the next word. Currently, most LMs calculate these representations through a neural network consuming the immediate previous context. However recently, retrieval-augmented LMs have shown to improve over standard neural LMs, by accessing information retrieved from a large datastore, in addition to their standard, parametric, next-word prediction. In this paper, we set out to understand why retrieval-augmented language models, and specifically why k-nearest neighbor language models (kNN-LMs) perform better than standard parametric LMs, even when the k-nearest neighbor component retrieves examples from the same training set that the LM was originally trained on. To this end, we perform a careful analysis of the various dimensions over which kNN-LM diverges from standard LMs, and investigate these dimensions one by one. Empirically, we identify three main reasons why kNN-LM performs better than standard LMs: using a different input representation for predicting the next tokens, approximate kNN search, and the importance of softmax temperature for the kNN distribution. Further, we incorporate these insights into the model architecture or the training procedure of the standard parametric LM, improving its results without the need for an explicit retrieval component. The code is available at https://github.com/frankxu2004/knnlm-why.