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

TL;DR

The paper investigates whether test-time fine-tuning on retrieved nearest-neighbor sequences (TTT-NN) can substantially reduce perplexity and close the gap between smaller and larger language models. It reproduces the approach using $k=20$ neighbors with a single gradient update per neighbor, leveraging RoBERTa embeddings, Faiss indexing over The Pile, and evaluating across GPT-2 variants, GPT-Neo, and R1-Distilled-Qwen-1.5B. The study confirms the core claims, showing robust reductions in bits per byte across diverse domains, with larger gains on specialized datasets, and demonstrates that smaller models not pre-trained on The Pile can approach larger-model performance after TT-TN-NN, though exact numbers differ due to the reduced neighbor count. It also introduces memory-efficient indexing to enable reproduction under limited RAM and extends the evaluation to modern architectures, underscoring the method's robustness and practical considerations for large-scale retrieval-augmented inference.

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.

Figures (6)

• Figure 1: Training cost in seconds per neighbor using GPT2
• Figure 2: Result comparison between before and after applying TTT-NN with 20 nearest neighbors on GPT-2 on 7 sub-datasets from the Pile, including wikipedia, arxiv, books3, enron, dm-mathematics, europarl and github. The metric used here is bits per byte.
• Figure 3: Model performance steadily improves as we increase the number of neighbors used in TTT-NN on GPT-2. The first row shows perplexity, the second row shows bits per bytes, and the third row shows training loss.
• Figure 4: Performance comparison (bits per byte) before and after applying TTT-NN on the EuroParl (left) and Github (right) datasets using 20 nearest neighbors across models of increasing size.
• Figure 5: Performance comparison (bits per byte) before and after applying TTT-NN on the DeepMind Mathematics dataset using 20 nearest neighbors across models of increasing size.