Reproducibility Report: Test-Time Training on Nearest Neighbors for Large Language Models

TL;DR

This paper reproduces and validates the effectiveness of test-time training using nearest neighbors to adapt large language models, showing significant performance improvements across diverse datasets and models, with practical implementation insights.

Contribution

It confirms the benefits of nearest-neighbor test-time training for large language models and introduces memory-efficient retrieval methods for large-scale deployment.

Findings

Test-time training reduces perplexity across datasets.

Models benefit more if not pretrained on the target data.

Memory-efficient retrieval enables large-scale adaptation.

Abstract

We reproduce the central claims of Test-Time Training on Nearest Neighbors for Large Language Models (Hardt and Sun, 2024), which proposes adapting a language model at inference time by fine-tuning on retrieved nearest-neighbor sequences. Using pretrained RoBERTa embeddings indexed with Faiss, we retrieve 20 neighbors per test input and apply one gradient update per neighbor across GPT-2 (117M, 774M), GPT-Neo (1.3B), and R1-Distilled-Qwen2.5-1.5B. Our experiments confirm that test-time training significantly reduces perplexity and bits-per-byte metrics across diverse domains from The Pile, with the largest improvements in structured or specialized datasets such as GitHub and EuroParl. We further validate that models not pretrained on The Pile benefit more from this adaptation than models already trained on similar data, allowing smaller models to approach the performance of larger ones. Due to infrastructure limitations, we introduce a memory-efficient retrieval implementation that loads only required line offsets rather than entire files, reducing RAM requirements from over 128 GB per server to 32 GB. We also extend the original study by evaluating R1-Distilled-Qwen2.5-1.5B, showing that test-time training yields consistent gains even for modern reasoning-optimized architectures. Overall, our results support the robustness and generality of nearest-neighbor test-time training while highlighting practical considerations for reproducing large-scale retrieval-augmented adaptation.