PaperlayerView exact revision on arXiv

Source-linked AI summary

Test-Time Training on Nearest Neighbors for Large Language Models

Moritz Hardt, Yu Sun

arXiv:2305.18466v3cs.CLcs.LG

TL;DR

Conventional domain boundaries are increasingly unclear for language-model fine-tuning, motivating a test instance’s nearest neighbors as a local domain. The paper retrieves those neighbors, fine-tunes the model on them at test time, and finds that this can recover much of the benefit of substantially larger pretraining while increasing inference cost.

  • Problem

    Conventional boundaries between domains and tasks are increasingly unclear for general-purpose language models, while each test instance can define its own local domain.

    Source: §1 · p. 1

  • Method

    For each test instance, TTT-NN retrieves nearest neighbors from a large text-embedding index and fine-tunes the model on them before prediction.

    Source: §1 · p. 1, §1.1 · p. 1

  • Results

    Using 20 nearest neighbors achieves most of the perplexity decrease at less than half the computational cost of using all neighbors, while GPT-2-large recovers most gains from pretraining a model twice as large.

    Source: §5 · p. 5

  • Takeaways & Limitations

    Test-time training can increase effective model capacity, but its usefulness must be weighed against increased inference time.

    Source: §1.1 · p. 2

    The method’s demonstrated benefits are limited mainly to perplexity-based language modeling, leaving emergent phenomena and broader non-perplexity tasks untested.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper introduces test-time training on nearest neighbors (TTT-NN), retrieving relevant data for each test instance and fine-tuning the model before evaluation. Across Pile language-modeling tasks, this approach improves performance with few neighbors and can narrow differences between models of different sizes.

1 INTRODUCTION

  • TTT-NN retrieves each test instance’s nearest neighbors and fine-tunes the model on them before applying it to that instance.The method uses a large-scale index of Pile text embeddings.Source: §1 · p. 1, §1.1 · p. 1

1.1 OUR CONTRIBUTIONS

  • 20% lower bits per byte is achieved after one gradient iteration on as few as 50 neighbors across all 22 Pile language-modeling tasks.Most of the gain is already realized after 20 neighbors.Source: §1.1 · p. 2
  • More than 60% lower bits per byte occurs on some tasks, especially pile github, during test-time training.The comparison covers a small GPT-2, a large GPT-2, and GPT-Neo.Source: §1.1 · p. 2
  • TTT-NN produces dramatic gains on tasks unseen during GPT-2 training, while improvements on seen tasks such as pile-cc are more moderate.GPT-Neo was trained to convergence on the entire Pile, whereas GPT-2 saw only some subsets.Source: §1.1 · p. 2
  • TTT-NN can recover most of the effective-capacity gains associated with larger models, but it increases inference time and may be unsuitable for some real-time applications.The paper identifies offline code generation and scientific sequence synthesis as settings where this cost is less concerning.Source: §1.1 · p. 2

2 RELATED WORK

The related work connects TTT-NN to local learning, retrieval-augmented language modeling, KNN-LM, transductive learning, and dynamic evaluation. Its distinction is to fine-tune the pretrained model on retrieved neighbors rather than append them as context or train a separate local predictor.

2.1 LOCAL LEARNING

  • Local-learning methods train models on data selected near each test instance, with earlier approaches using distance-weighted regression, SVMs, or partial neural-network updates.Bottou & Vapnik (1992) trained only the last linear layer, whereas TTT-NN fine-tunes the entire network for one gradient iteration per neighbor.Source: §2.1 · pp. 2–3

2.2 LANGUAGE MODELS USING RETRIEVAL

  • Retrieval-augmented language models use neighbors as additional context and therefore require matching retrieval-aware training, unlike TTT-NN’s use of standard pretrained models.Relevant examples include Guu et al. (2020), Lewis et al. (2020), Borgeaud et al. (2022), Wang et al. (2022b), and Izacard et al. (2022).Source: §2.2 · p. 3
  • KNN-LM predicts the next token from similar training contexts, while dynamic evaluation fine-tunes on the test instance’s context; both are evaluated as baselines against TTT-NN.Dynamic evaluation is part of broader test-time-training work that produces a different model for each test input.Source: §2.2 · p. 3, §2.3 · p. 3

3 NEAREST NEIGHBOR INDEX

The paper builds a distributed nearest-neighbor index over Pile training sequences and finds that larger, higher-quality indexes yield closer neighbors. The system serves full-dataset queries in about one second on standard hardware.

3 NEAREST NEIGHBOR INDEX

  • The index contains text embeddings for the Pile training set, while validation and test sequences are excluded from retrieval.The training set contains 210M sequences and 1.3TB of text.Source: §3 · p. 3
  • The embedding model is a 355M-parameter RoBERTa model trained on the Pile, producing 1024-dimensional representations for indexing.Long sequences are truncated to the embedding model’s maximum length, so neighbor selection uses only their prefixes.Source: §3 · p. 4
  • A distributed FAISS Flat L2 index partitions the 810GB vector store across 180 servers and combines shard results to return final neighbors.Data and vectors together occupy 2.1TB.Source: §3 · p. 4
  • 1.35 seconds mean query time, with 0.12 seconds standard deviation, is achieved across 1,000 random Pile training queries using 180 servers.The reported time includes computing the query embedding and uses no server-side GPU acceleration.Source: §3 · p. 4
  • Nearest-neighbor distances decline as index size increases, indicating that larger indexes provide closer retrieved examples.For validation queries, neighbor distances are generally smaller than distances between random query pairs, and about 4% have near-exact matches.Source: §3 · p. 4, Figure 4 · p. 4

4 TEST-TIME TRAINING ON NEAREST NEIGHBORS

TTT-NN fine-tunes the model on retrieved neighbors before evaluating each test instance, using simple sequential updates and resetting parameters afterward. The method requires no additional hyper-parameter tuning beyond each model’s defaults.

4 TEST-TIME TRAINING ON NEAREST NEIGHBORS

  • Sequentially training on neighbors from nearest to farthest performs best, while reversing the order performs significantly worse.Batching all neighbors and repeatedly updating is much slower without significant improvement.Source: §4 · p. 5
  • TTT-NN splits overlength neighbors into model-length chunks and applies one sequential gradient update per chunk.A single retrieved sequence can therefore produce several gradient updates.Source: §4 · p. 5
  • After each test instance, the model’s trainable parameters are reset to the pretrained base state rather than accumulated across evaluation examples.Accumulation would require assumptions such as temporal smoothness between successive test instances.Source: §4 · p. 5
  • The method reuses each model’s default HuggingFace optimizer and hyper-parameters without additional tuning.This contrasts with test-time augmentation methods that tune their augmentations using extra computation and labor.Source: §4 · p. 5

5 RESULTS

Test-time training on retrieved neighbors improves language-model evaluation across Pile tasks, with most gains appearing after relatively few neighbors. The approach outperforms several retrieval and adaptation baselines, while incurring heterogeneous and substantially higher inference-time compute.

5 RESULTS

  • 20 neighbors achieve most of the perplexity decrease for GPT-2-Small, while additional neighbors continue to improve results across the six largest Pile tasks.These tasks together represent more than 70% of the Pile benchmark.Source: §5 · p. 5, Figure 6 · p. 7
  • TTT-NN on GPT-2-Large recovers most of the gains from pre-training a model more than twice as large to convergence on the full dataset.The comparison uses GPT-2-Large and GPT-Neo results shown in the paper's figures.Source: §5 · p. 5, Figure 7 · p. 7
  • Splitting each evaluation sequence to separate retrieval from evaluation changes results by typically less than 1% and rarely more than 2%.On pile all, bits per byte rises from 82% without splitting to 83% with splitting.Source: §5 · p. 6
  • Inference-time training substantially increases compute, with costs varying across tasks because retrieved sequences differ greatly in length.Training cost is measured in seconds per neighbor on a single NVIDIA A100 GPU.Source: §5 · p. 6

5.1 COMPARISONS WITH BASELINES

  • TTT-NN outperforms in-context neighbors, interpolation, and dynamic evaluation in the reported baseline comparison.Dynamic evaluation improves many tasks but begins overfitting after seven iterations, while TTT-NN is especially stronger on pile github.Source: §5.1 · pp. 6, 8, 9, Table 1 · p. 9
  • In-context neighbors improve performance in a few cases, such as pile enron, but provide little overall benefit.The baseline concatenates retrieved neighbors into the test context without test-time training.Source: §5.1 · p. 7

6 LIMITATIONS AND FUTURE WORK

The paper identifies unresolved questions about whether nearest-neighbor test-time training extends beyond perplexity language modeling and to fairness, safety, and specialized applications.

Limitations

  • Test-time training is shown to improve effective model capacity for perplexity-based language modeling, but broader non-perplexity evaluation remains needed.The authors specifically identify code generation and life-sciences applications as future directions.Source: §6 · p. 9

Future work

  • The method’s potential benefits for fairness and safety remain topics for future investigation rather than established findings.The paper suggests possible benefits for under-represented groups and trusted-data defenses against poisoning, while framing them as hypotheses.Source: §6 · p. 9
  • Smaller Pile tasks such as pile enron and pile europarl reportedly improve more than larger tasks such as pile-cc and pile pubmed-central.The paper presents this pattern as support for investigating fairness-related effects, not as a demonstrated fairness result.Source: §6 · p. 9

A RESULTS FOR GPT-2-SMALL ON SPLIT SEQUENCES

These results evaluate GPT-2-Small after splitting each sequence into retrieval and evaluation halves, with bits per byte results shown before and after training on 50 nearest neighbors.

Protocol

  • GPT-2-Small results use the first sequence half for retrieval and the second half for evaluation.For sequences longer than twice the maximum length, the first half is capped at the model’s maximum sequence length.Source: A RESULTS FOR GPT-2-SMALL ON SPLIT… · p. 12

Implementation

Results

B ALL RESULTS FOR GPT-2-LARGE

These results report GPT-2-Large performance across Pile tasks before and after training on 50 nearest neighbors, using bits per byte as the metric.

Implementation

Results

  • GPT-2-Large results compare bits per byte across all Pile tasks before and after training on 50 nearest neighbors.The supplied figure caption defines the comparison scope but does not state individual task outcomes.Source: Figure 11 · p. 13

Model

C ALL RESULTS FOR GPT-NEO

These results report GPT-Neo performance across Pile tasks before and after training on 50 nearest neighbors, using bits per byte as the metric.

Implementation

Results

  • GPT-Neo results compare bits per byte across all Pile tasks before and after training on 50 nearest neighbors.The supplied figure caption defines the comparison scope but does not state individual task outcomes.Source: Figure 12 · p. 14

Model

D BOOTSTRAP ERROR BARS

The paper uses empirical bootstrap error bars for bits per byte because perplexity-based metrics are unbounded and nonlinear. It reports 0.1 and 0.9 quantiles from 1000 bootstrap samples, with smaller error bars for larger models.

D BOOTSTRAP ERROR BARS

  • Perplexity metrics motivate bootstrap intervals because they can be unbounded and are exponentials of sums rather than sums over individual data points.Source: D BOOTSTRAP ERROR BARS · p. 15
  • 1000 bootstrap samples produce 0.1 and 0.9 quantiles as lower and upper error bars for bits per byte.Each sample resamples sequence-level negative log likelihoods with replacement and recomputes the metric.Source: D BOOTSTRAP ERROR BARS · p. 15
  • Figure 13 reports bits per byte results for GPT-2-Small across all Pile tasks with bootstrap error bars.The bars span the 0.1 and 0.9 quantiles, leaving 10% of bootstrap samples outside each end.Source: Figure 13 · p. 15, D BOOTSTRAP ERROR BARS · p. 15
  • The larger models have smaller error bars, so the corresponding GPT-2-Large and GPT-Neo plots are omitted.Source: D BOOTSTRAP ERROR BARS · p. 15

E TRAINING COSTS

The paper compares training costs across all evaluated tasks and three models. Figure 14 organizes costs in seconds per neighbor for GPT-2-Small, GPT-2, and GPT-Neo.

E TRAINING COSTS

  • Training costs are compared across all tasks and the three evaluated models.Source: E TRAINING COSTS · p. 16
  • Figure 14 reports training costs in seconds per neighbor for GPT-2-Small, GPT-2, and GPT-Neo.The figure places GPT-2-Small at the top, GPT-2 in the middle, and GPT-Neo at the bottom.Source: Figure 14 · p. 16
Loading 2305.18466v3…