# Test-Time Training on Nearest Neighbors for Large Language Models

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

## Key facts

- Cited by 88 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 34 references that are in the map.
- First posted on arXiv 29 May 2023.

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

- [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)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Don't Stop Pretraining: Adapt Language Models to Domains and Tasks](https://mlsanity.com/paper/2004.10964/) (2020)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Atlas: Few-shot Learning with Retrieval Augmented Language Models](https://mlsanity.com/paper/2208.03299/) (2022)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [Test-Time Training with Masked Autoencoders](https://mlsanity.com/paper/2209.07522/) (2022)
- [DEMix Layers: Disentangling Domains for Modular Language Modeling](https://mlsanity.com/paper/2108.05036/) (2021)
- [Learning by Transduction](https://mlsanity.com/paper/1301.7375/) (2013)
- [Efficient Nearest Neighbor Language Models](https://mlsanity.com/paper/2109.04212/) (2021)
- [Training Data is More Valuable than You Think: A Simple and Effective Method by Retrieving from Training Data](https://mlsanity.com/paper/2203.08773/) (2022)
- [Neuro-Symbolic Language Modeling with Automaton-augmented Retrieval](https://mlsanity.com/paper/2201.12431/) (2022)
- [Dynamic Evaluation of Neural Sequence Models](https://mlsanity.com/paper/1709.07432/) (2017)
- [Dynamic Evaluation of Transformer Language Models](https://mlsanity.com/paper/1904.08378/) (2019)

## Built on by (most cited first, 30 of 88)

- [Learning to (Learn at Test Time): RNNs with Expressive Hidden States](https://mlsanity.com/paper/2407.04620/) (2024)
- [Reinforcement Learning via Self-Distillation](https://mlsanity.com/paper/2601.20802/) (2026)
- [TTRL: Test-Time Reinforcement Learning](https://mlsanity.com/paper/2504.16084/) (2025)
- [The Surprising Effectiveness of Test-Time Training for Few-Shot Learning](https://mlsanity.com/paper/2411.07279/) (2024)
- [Retaining by Doing: The Role of On-Policy Data in Mitigating Forgetting](https://mlsanity.com/paper/2510.18874/) (2025)
- [End-to-End Test-Time Training for Long Context](https://mlsanity.com/paper/2512.23675/) (2025)
- [Test-Time Learning for Large Language Models](https://mlsanity.com/paper/2505.20633/) (2025)
- [Seek in the Dark: Reasoning via Test-Time Instance-Level Policy Gradient in Latent Space](https://mlsanity.com/paper/2505.13308/) (2025)
- [A Survey of Test-Time Compute: From Intuitive Inference to Deliberate Reasoning](https://mlsanity.com/paper/2501.02497/) (2025)
- [Test-Time Training on Video Streams](https://mlsanity.com/paper/2307.05014/) (2023)
- [Frustratingly Easy Test-Time Adaptation of Vision-Language Models](https://mlsanity.com/paper/2405.18330/) (2024)
- [Self-Improving LLM Agents at Test-Time](https://mlsanity.com/paper/2510.07841/) (2025)
- [Cognitive Memory in Large Language Models](https://mlsanity.com/paper/2504.02441/) (2025)
- [How Benchmark Prediction from Fewer Data Misses the Mark](https://mlsanity.com/paper/2506.07673/) (2025)
- [Test3R: Learning to Reconstruct 3D at Test Time](https://mlsanity.com/paper/2506.13750/) (2025)
- [Let's (not) just put things in Context: Test-Time Training for Long-Context LLMs](https://mlsanity.com/paper/2512.13898/) (2025)
- [Learning to (Learn at Test Time)](https://mlsanity.com/paper/2310.13807/) (2023)
- [Puzzle: Distillation-Based NAS for Inference-Optimized LLMs](https://mlsanity.com/paper/2411.19146/) (2024)
- [Test-Time Training Provably Improves Transformers as In-context Learners](https://mlsanity.com/paper/2503.11842/) (2025)
- [Latent learning: episodic memory complements parametric learning by enabling flexible reuse of experiences](https://mlsanity.com/paper/2509.16189/) (2025)
- [Continuous Self-Improvement of Large Language Models by Test-time Training with Verifier-Driven Sample Selection](https://mlsanity.com/paper/2505.19475/) (2025)
- [Latent Context Compilation: Distilling Long Context into Compact Portable Memory](https://mlsanity.com/paper/2602.21221/) (2026)
- [DISCOVER: Automated Curricula for Sparse-Reward Reinforcement Learning](https://mlsanity.com/paper/2505.19850/) (2025)
- [Adaptive Defense against Harmful Fine-Tuning for Large Language Models via Bayesian Data Scheduler](https://mlsanity.com/paper/2510.27172/) (2025)
- [Tuning-Free Accountable Intervention for LLM Deployment -- A Metacognitive Approach](https://mlsanity.com/paper/2403.05636/) (2024)
- [Model-free Test Time Adaptation for Out-Of-Distribution Detection](https://mlsanity.com/paper/2311.16420/) (2023)
- [TeleRAG: Efficient Retrieval-Augmented Generation Inference with Lookahead Retrieval](https://mlsanity.com/paper/2502.20969/) (2025)
- [Local Mixtures of Experts: Essentially Free Test-Time Training via Model Merging](https://mlsanity.com/paper/2505.14136/) (2025)
- [Learning with Calibration: Exploring Test-Time Computing of Spatio-Temporal Forecasting](https://mlsanity.com/paper/2506.00635/) (2025)
- [Learning on the Job: Test-Time Curricula for Targeted Reinforcement Learning](https://mlsanity.com/paper/2510.04786/) (2025)

## Abstract (arXiv)

Many recent efforts augment language models with retrieval, by adding retrieved data to the input context. For this approach to succeed, the retrieved data must be added at both training and test time. Moreover, as input length grows linearly with the size of retrieved data, cost in computation and memory grows quadratically for modern Transformers. To avoid these complications, we simply fine-tune the model on retrieved data at test time, using its standard training setup. We build a large-scale distributed index based on text embeddings of the Pile dataset. For each test input, our system retrieves its neighbors and fine-tunes the model on their text. Surprisingly, retrieving and training on as few as 20 neighbors, each for only one gradient iteration, drastically improves performance across more than 20 language modeling tasks in the Pile. For example, test-time training with nearest neighbors significantly narrows the performance gap between a small GPT-2 and a GPT-Neo model more than 10 times larger. Sufficient index quality and size, however, are necessary. Our work establishes a first baseline of test-time training for language modeling.