mlsanityData to 7 Oct 2026

Why do Nearest Neighbor Language Models Work?

arXiv 2301.02828 ↗First posted 7 Jan 2023cs.CL cs.LGcited by 29

On the map: Retrieval-augmented generation → Generation and context integration

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.

Key facts. Cited by 29 papers in the map (data through 7 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.

Builds on

33 of 44 references.

References · 33

  1. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
  2. 2020Language Models are Few-Shot Learnerscited by 32,199
  3. 2019RoBERTa: A Robustly Optimized BERT Pretraining Approachcited by 13,266
  4. 2015Distilling the Knowledge in a Neural Networkcited by 10,076
  5. 2016Layer Normalizationcited by 5,579
  6. 2015Rethinking the Inception Architecture for Computer Visioncited by 5,451
  7. 2019DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lightercited by 4,243
  8. 2016Pointer Sentinel Mixture Modelscited by 2,992
  9. 2015Neural Machine Translation of Rare Words with Subword Unitscited by 2,898
  10. 2019The Curious Case of Neural Text Degenerationcited by 2,881
  11. 2017Billion-scale similarity search with GPUscited by 2,477
  12. 2021Improving language models by retrieving from trillions of tokenscited by 1,278
21 more
  1. 2019Generalization through Memorization: Nearest Neighbor Language Modelscited by 788
  2. 2017Regularizing Neural Networks by Penalizing Confident Output Distributionscited by 461
  3. 2018Adaptive Input Representations for Neural Language Modelingcited by 226
  4. 2020Nearest Neighbor Machine Translationcited by 209
  5. 2017Regularizing and Optimizing LSTM Language Modelscited by 205
  6. 2017Breaking the Softmax Bottleneck: A High-Rank RNN Language Modelcited by 147
  7. 2019On NMT Search Errors and Model Errors: Cat Got Your Tongue?cited by 114
  8. 2022Training Language Models with Memory Augmentationcited by 112
  9. 2021Efficient Nearest Neighbor Language Modelscited by 93
  10. 2017Generating Sentences by Editing Prototypescited by 88
  11. 2016Efficient softmax approximation for GPUscited by 67
  12. 2022Neuro-Symbolic Language Modeling with Automaton-augmented Retrievalcited by 60
  13. 2020Generalized Entropy Regularization or: There's Nothing Special about Label Smoothingcited by 40
  14. 2020Best-First Beam Searchcited by 40
  15. 2017Unbounded cache model for online language modeling with open vocabularycited by 37
  16. 2022Efficient Cluster-Based k-Nearest-Neighbor Machine Translationcited by 24
  17. 2020Stolen Probability: A Structural Weakness of Neural Language Modelscited by 21
  18. 2022You can't pick your neighbors, or can you? When and how to rely on retrieval in the kNN-LMcited by 18
  19. 2020Learning Sparse Prototypes for Text Generationcited by 15
  20. 2022Low-Rank Softmax Can Have Unargmaxable Classes in Theory but Rarely in Practicecited by 14
  21. 2022Nearest Neighbor Knowledge Distillation for Neural Machine Translationcited by 10

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2023: 8; 2024: 11; 2025: 4; 2026: 6.

29 papers.

Sort:

By area:

  1. 2024Retrieval-Augmented Generation for Natural Language Processing: A Surveycited by 70
  2. 2023Privacy Implications of Retrieval-Based Language Modelscited by 31
  3. 2023Chameleon: a Heterogeneous and Disaggregated Accelerator System for Retrieval-Augmented Language Modelscited by 23
  4. 2023CAGRA: Highly Parallel Graph Construction and Approximate Nearest Neighbor Search for GPUscited by 22
  5. 2024Non-Exchangeable Conformal Language Generation with Nearest Neighborscited by 21
  6. 2025Memory Decoder: A Pretrained, Plug-and-Play Memory for Large Language Modelscited by 16
  7. 2024Retrofitting Large Language Models with Dynamic Tokenizationcited by 15
  8. 2023KNN-LM Does Not Improve Open-ended Text Generationcited by 10
  9. 2024FT2Ra: A Fine-Tuning-Inspired Approach to Retrieval-Augmented Code Completioncited by 9
  10. 2023Surface-Based Retrieval Reduces Perplexity of Retrieval-Augmented Language Modelscited by 7
  11. 2025On the Predictive Power of Representation Dispersion in Language Modelscited by 6
  12. 2024NNOSE: Nearest Neighbor Occupational Skill Extractioncited by 5
17 more
  1. 2023On the Generalization Ability of Retrieval-Enhanced Transformerscited by 4
  2. 2024Incorporating Precedents for Legal Judgement Prediction on European Court of Human Rights Casescited by 4
  3. 2024Languages Transferred Within the Encoder: On Representation Transfer in Zero-Shot Multilingual Translationcited by 3
  4. 2023On Retrieval Augmentation and the Limitations of Language Model Trainingcited by 2
  5. 2024Local Topology Measures of Contextual Language Model Latent Spaces With Applications to Dialogue Term Extractioncited by 2
  6. 2026MemSFT: Mitigating Alignment Tax with an External Parametric Memorycited by 2
  7. 2023Nearest Neighbor Machine Translation is Meta-Optimizer on Output Projection Layercited by 1
  8. 2024Efficient Nearest Neighbor based Uncertainty Estimation for Natural Language Processing Taskscited by 1
  9. 2025Long-Tail Crisis in Nearest Neighbor Language Modelscited by 1
  10. 2025Studying the Role of Input-Neighbor Overlap in Retrieval-Augmented Language Models Training Efficiencycited by 1
  11. 2024Memorizing Documents with Guidance in Large Language Modelscited by 0
  12. 2024Experimental comparison of graph-based approximate nearest neighbor search algorithms on edge devicescited by 0
  13. 2026Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Expertscited by 0
  14. 2026Adaptive kNN graph modelcited by 0
  15. 2026Nearest-Neighbor Radii under Dependent Samplingcited by 0
  16. 2026One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Contextcited by 0
  17. 2026Forking: Sudden Overfitting Under Replaycited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2024Great Memory, Shallow Reasoning: Limits of kNN-LMs
  2. 2022kNN-Prompt: Nearest Neighbor Zero-Shot Inference
  3. 2024Reliable, Adaptable, and Attributable Language Models with Retrieval
  4. 2024Nearest Neighbor Speculative Decoding for LLM Generation and Attribution
  5. 2025Reproducibility Report: Test-Time Training on Nearest Neighbors for Large Language Models
  6. 2022When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories
  7. 2024More Room for Language: Investigating the Effect of Retrieval on Language Models
  8. 2024Let your LLM generate a few tokens and you will reduce the need for retrieval