Research map / Architectures and efficiency / Attention and transformers
Test-Time Training on Nearest Neighbors for Large Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Fine-tunes a causal LM on each test input's nearest neighbors retrieved from a 210M-sequence Pile index (one gradient step per neighbor), reducing bits-per-byte across Pile tasks.
Problem
Retrieval-augmented language models need retrieved data in context at both training and test time, and long contexts make Transformer cost grow quadratically. The paper asks whether test-time fine-tuning on nearest neighbors can improve language modeling without such retrieval-aware training.
Approach
For each test instance, retrieve its nearest neighbors from a large distributed embedding index of the Pile, fine-tune the pretrained model on them sequentially (nearest last-to-first ordering, one gradient step per chunk), then reset parameters; no hyperparameter tuning beyond HuggingFace defaults.
Key results
- On GPT-2-Large, TTT-NN on pile_all reduces bits per byte from 1.07 to 0.85 with 50 neighbors.
- On pile_github, bits per byte drops from 1.95 to 0.51 for GPT-2-Large under TTT-NN, outperforming in-context, interpolation and dynamic evaluation baselines.
- Most gains are realized with about 20 neighbors, and test-time training narrows the gap between GPT-2 and the 10x larger GPT-Neo 1.3B.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| test-time training on nearest neighbors (TTT-NN, 50 neighbors, one gradient step each) | GPT-2-Large · Pile all tasks | bits per byte: 1.07 -> 0.85 | Helped |
| test-time training on nearest neighbors | GPT-2-Large · pile_github | bits per byte: 1.95 -> 0.51 | Helped |
| test-time training on nearest neighbors | GPT-2-Large · pile_gutenberg, pile_youtubesubtitles | bits per byte: pile_gutenberg 1.21 -> 1.18 (improves slightly); pile_youtubesubtitles 1.51 -> 1.42 but dyn. eval. better | Hurt |
| interpolating neighbor-token distribution with LM (Interpolate, KNN-LM style) | GPT-2-Large · Pile all tasks | bits per byte: 1.07 -> 0.99 | Helped |
| dynamic evaluation (7 iterations on first half of sequence) | GPT-2-Large · Pile all tasks | bits per byte: 1.07 -> 1.03 | Helped |
| adding retrieved neighbors in-context at test time without special training | GPT-2-Large · Pile all tasks | bits per byte: 1.07 -> 1.06 | No clear effect |
| training on neighbors in increasing distance order (nearest last) vs. nearest first | TTT-NN · Pile | bits per byte | No clear effect |
| training on all neighbors in a single batch repeatedly instead of sequential updates | TTT-NN · Pile | bits per byte | No clear effect |
| splitting evaluation sequence into retrieval half and evaluation half (no overlap) (ablation) | TTT-NN GPT-2-Small · Pile all tasks | bits per byte: 82% -> 83% | No clear effect |
| using only ~20 neighbors instead of 50 (ablation) | GPT-2-Small · Pile | bits per byte | No clear effect |
Surprising findings
- Training on neighbors in increasing-distance order (nearest last) performs worse than nearest-first; the nearest-first order is used.
- Dynamic evaluation overfits after about 7 iterations, and in three tasks it beats TTT-NN.
Limitations
- Substantially increases inference compute, since each test instance requires fine-tuning on many neighbors.
- Evaluated only on perplexity metrics; no evidence for emergent or non-perplexity downstream tasks.
- Uses a naive FAISS flat index and truncates long sequences to the embedding model's maximum length.
Experimental setup
Datasets: The Pile
Benchmarks: The Pile (22 language modeling tasks)
Metrics: bits per byte byte perplexity word perplexity
Baselines: In-Context neighbors Interpolate (KNN-LM variant) Dynamic evaluation Base model only
Builds on: dynamic evaluation local learning KNN-LM RoBERTa
Closest papers by meaning
- Efficiently Learning at Test-Time: Active Fine-Tuning of LLMs (ICLR 2025 · similarity 0.89)
- DEPT: Decoupled Embeddings for Pre-training Language Models (ICLR 2025 · similarity 0.89)
- Test-Time Learning for Large Language Models (ICML 2025 · similarity 0.88)
- Compact Language Models via Pruning and Knowledge Distillation (NeurIPS 2024 · similarity 0.88)
- Pruning Large Language Models with Semi-Structural Adaptive Sparse Training (AAAI 2025 · similarity 0.88)
- Scaling Embedding Layers in Language Models (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.