Test-Time Training on Nearest Neighbors for Large Language Models
Authors: Moritz Hardt
Yu Sun
Venue: ICLR 2024 poster (2024)
Keywords: Test-time training
Primary area: general machine learning (i.e., none of the above)

Abstract: 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.

Published as conference paper at ICLR 2024
TEST-TIME TRAINING ON NEAREST NEIGHBORS FOR LARGE LANGUAGE MODELS
Moritz Hardt
Yu Sun
Max Planck Institute for Intelligent Systems, Tübingen Tübingen Al Center, University of Tübingen
Stanford University
ABSTRACT
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, 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 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
Machine learning traditionally separates training and testing. Once trained, a model remains frozen during evaluation. But nearly as old as machine learning the intriguing idea to update the model test time with data relevant to each individual test instance. Variants of this idea have existed for almost fifty years, including locally weighted regression (Stone, 1977; Cleveland, 1979), local learning (Bottou Vapnik, 1992), and SVM-KNN (Zhang et al., 2006). Recently, test-time training has grown again in popularity with the rise of deep learning, exposing a large space of possible heuristics (Krause et al., 2018; 2019; Sun et al., 2020).
We investigate simple, yet powerful, heuristic in this space, called test-time training nearest neighbors (TTT-NN), for the task of language modeling. For each test instance, we retrieve its nearest neighbors from huge database, and fine-tune the model on those neighbors before applying it to the test instance.
Fine-tuning language models is a well known practice to boost its performance on specific domains or tasks. However, there seems to be no consensus on what exactly defines a domain task (Gururangan et al., 2021). and conventional boundaries between previously defined domain or tasks have become increasingly unclear with the recent advances in general-purye purppose language models (Brown al., 2020; Bubeck et al., 2023). From the perspective of test-time training, each test instance defines its own "domain". Our hope is that sufficiently large database will contain enough data relevant each "domain" induced by test instance, such that fine-tuning improves performance locally.
1.1 OUR CONTRIBUTIONS
We present a system for test-time training on nearest neighbors, based on a large-scale index as illustrated in Figure 1. Our distributed index can serve each nearest neighbor query to approximately 200 million vectors and 1TB of data in approximately one second on standard hardware. The vectors represent text embeddings of all training sequences in the Pile dataset (Gao et al., 2020).
'Code, index files, and model checkpoint: https: popp:///iee hta.hha....mooool..fooaaatttoeee

---

Published conference paper at ICLR 2024
text
emhheddings
Server
Server
Server 3
Server 180
input
oulput
TTT-NN
Figure 1: System architecture for test-time training with nearest neighbors (TTT-NN).
We evaluate method all 22 tasks for language modeling from the Pile benchmark. We focus three causal language models of increasing size: a small GPT-2 model with 117M parameters, large GPT-2 model with 774M parameters, and GPT-Neo model with 1.3B parameters. We find that training for only gradient iteration on as few as 50 neighbors reduces normalized perplexity measurement, the bits per byte metric, by 20%. In fact, most of the gain is already realized after only 20 neighbors. When it comes some tasks, especially code generation as in the pile-gi hub task, the bits per byte metric goes down by more than 60% during time training. See Figure
github
pile-cc
Bits por hytn hytz
1.0-
0.5
0.5
Bits
Bits
0.0
0.0
9pd2
9pi2-1t
gpt2- te2-tarre
gpt2. (i2-targettt tuz
gpizen
opizen-iti
uptz opt2 tre 9m12 larrte ge12 harggrreet ttE
opinen
spiuwo-1tt I1!
mpe2
apt2 itt
gp12 large
upt2 larts--tt ttt
uptnon uption-tt
Figure 2: Comparison of different models before and after TTT-NN with 50 neighbors. Left: All Pile tasks.
Center: Best performing Pile task (pile le github). Right: Largest Pile task
The GPT-Neo model was specifically trained to convergence on the entire Pile dataset. In contrast, GPT-2 trained only some subsets of the dataset, roughly pile.pile-cc, but not others such pile.github. Therefore, we can see how TTT-NN on a model that was not pre-trained on a specific Pile task stacks up model that unseen during training, such as pi le www.gitt..b...... can be dramatic. TTT-NN still helps on seen tasks, such as pile-cc, but the improvements
We that the improvements due test-time training on tasks
moderate.
Our results suggest that test-time training can increase the effective capacity of a model. This comes at the obvious cost of increased inference time, like other retrieval augmented techniques. While this cost may be prohibitive for some real-time applications, test-time training can be stopped at any time and still yield significant improvement (Figure 6). It is also less of a concern in other applications, such offline code generation and synthesis of sequences for scientific purposes. From this perspective, it is encouraging, and worth investigating further, that test-time training performs best on code generation within the Pile.
2
RELATED WORK
2.1 LOCAL LEARNING
As discussed at the beginning. the idea of using local data to train parametric model has long history and many variants, depending how locality is defined and what the model is. The oldest literature in nonparametric linear regression (Stone, 1977) and locally weighted regression (Cleveland, 1979; Cleveland & Devlin, 1988; Ruppert al., 1995; Atkeson et al., 1997) uses linear model and engineered distance
2

---

Published as conference paper at ICLR 2024
metrics, e.g. kernels, to assign more weights to training data closer to each test instance. Because initialization does not affect the optimal solution of linear regression, each test instance trains a model from scratch. SVM-KNN (Zhang et al., 2006) applies the same idea to support vector machines.
The most relevant paper is Bottou & Vapnik (1992). Their method, called local learning, first trains convolutional neural network on images of handwritten digits. The local training data here are chosen to be the nearest neighbors to the test instance, in terms of Euclidean distance in pixel space. Because it is slow to train the entire network for many epochs on each subset of neighbor, they instead only train the last linear layer of the network, as a trade-off between computational cost and performance gain. comparison, our method can afford to fine-tune the entire network, using its standard training setup, because each neighbor only needs one gradient iteration.
2.2 LANGUAGE MODELS USING RETRIEVAL
Many recent advances in language models retrieve nearest neighbors from the training set, but use them differently from our approach. The most successful line of work uses those neighbors as additional context, also known as feature augmentation, for the test instance. Some prominent examples are Guu et al. (2020), Lewis et al. (2020), Borgeaud et al. (2022), Wang et al. (2022b), Wang et al. (2022b), and Izacard et al. (2022). The success of these methods at test time requires having the retrieved neighbors additional context also during training, mirroring the procedure at test time. In comparison, our method requires no special training; simply downloaded pre-trained models from HuggingFace. In Subsection 5.1, evaluate against baseline that uses neighbors as context at test time without special training.
At test time, our method requires one forward and backward pass through each of the neighbors individually. Using the neighbors as additional context requires only one forward pass through the neighbors, but the time complexity of self-attention quadratic in the context length, consequently the number of neighbors. With 50 neighbors, the quadratic cost can be significant. Our method does not have this concern.
Another related line of work is KNN-LM (Khandelwal al., 2019; He et al., 2021; Alon et al., 2022). The basic idea is that autoregressive language modeling maps each context window of tokens to a predicted distribution over the next token. This can alternatively be achieved with a nonparametric approach: find all tokens in the training data whose context window is similar to the test instance, and use their (weighted) histogram as the predicted distribution. Khandelwal et al. (2019) simply uses a convex combination of the predicted distributions from the two approaches as the final prediction. In Subsection 5.1, we evaluate against a variant of KNN-LM that is computationally feasible for the Pile.
2.3 TEST-TIME TRAINING AND DYNAMIC EVALUATION
Beside the ones mentioned in Subsection 2.1, many papers explore the broader idea of training different model specific to each unlabeled test instance. An early attempt is transductive learning, beginning in the 1990s (Gammerman et al., 1998). Vapnik (2013) states the principle of transduction as: "Try to get the answer that you really need but not a more general one." This principle has been most commonly applied on SVMs (Vapnik, 2013; Collobert et al., 2006; Joachims, 2002) by using the test data to specify additional constraints on the margin of the decision boundary.
More recently, the idea of test-time training has been applied deep learning, for generalization under distribution shifts (Sun et al., 2020; Gandelsman et al., 2022; Wang et al., 2022a; Sun et al., 2023). It produces different model for every single test input through self-supervised learning. Another line of work called dynamic evaluation (Krause et al., 2018; 2019) fine-tunes the language model on the context of the test instance. In Subsection 5.1, we compare with dynamic evaluation baseline.
3 NEAREST NEIGHBOR INDEX
Our nearest neighbor index is based text embeddings of the Pile training set. The training set is split across 30 files, each approximately 7M sequences and 43GB of text. The entire dataset has 210M sequences and size 1.3TB. In addition, the Pile dataset has a validation set and test set that we do not include in the index.
3

---

Published as conference paper at ICLR 2024
Distances nearest neighbor
Distances 0.0001
Distances between random queries
400
40
300
400
300
20
200
100
100
2
3
4
2
0
Figure 3: Left: Histogram of distances to nearest neighbor for 10, 000 random queries from the validation set. Center: Focusing on the 400 smallest distances. Right: Distances between 10, 000 pairs of random queries.
distance nearest neighbor
Mean distance to nearest neighbor
Index size
2.0
Fuclidean disamnon
20%
1.8
1.7
40%
1.6
1.6
1.4
1.54
60%
70%
1.2
14-
80%
90%
1.0-
100%
0.8
20
60
100
20
40
60
80
100
Rank of nearest neighbor
Rank of nearest neighbor
Figure 4: Left: Mean distances to the nearest neighbor as index size grows. Right: Mean distances with standard deviation shown as shaded region for the full size index.
The text embedding model we use is based on large Roberta model (Liu al., 2019) with 355M parameters. Starting from a pre-trained Roberta model, we train the model on the Pile for 1.7M iterations, with the standard batch size of 128. corresponding to approximately one pass over the training set. The embedding dimension is 1024. Training the embedding model on the entire training set proved to be important for index quality. To build the index, we apply the embedding model to each sequence in the training set. For simplicity, we naively truncate long sequences to the maximum sequence length of the embedding model. In other words, on long sequences we select neighbors based on only a prefix of the sequence.
We store the text embeddings in nearest neighbor data structure. As proof of concept, we simply use the FAISS Flat L2 index (Johnson et al., 2019). The entire index 810GB, adding up to TTB for data and vectors combined. To speed up query time across the data structure, we build distributed client-server architecture, illustrated in Figure 1, on top of the FAISS index. The system splits the index across 180 servers each running on one chunk of the index and dataset. A client sends out nearest neighbor queries to each server and receives the respective neighbors in each chunk. The client then builds a local nearest neighbor structure out of the received results and queries the local structure for the final results.
The distributed index can answer query to the entire Pile dataset in around one second, including time to compute the query embedding, on standard hardware and network access without any GPU acceleration on the server side, which would further accelerate query time. Specifically, we observed a mean query time of 1.35 seconds with standard deviation of 0. 12 seconds, computed across 1000 random queries form the Pile training set. The distributed architecture gives a flexible way to split the index across different hardware footprints. Increasing the number of servers decreases the local query processing time on each server, but increases the time for network traffic. We find that 180 servers is a reasonable trade-off for the cluster we use.
Figure takes a closer look at nearest neighbor distances. We pick 10, 000 random points from the validation set and retrieve their neighbors from the index. The index does not include the validation or test set. However, there is a small cluster of about 4% of the points that have a near exact match in the index. In general, distances to the neighbors the right panel of the figure shows. Figure shows that distances to nearest neighbors continue to decline as increase the index size. This reflects the benefits of using a larger index.
substantially smaller than distances between the embeddings of two random queries,

---

Published as conference paper at ICLR 2024
4 TEST-TIME TRAINING ON NEAREST NEIGHBORS
Although the idea of test-time training on nearest neighbors is simple and intuitive, making it work involves few important design choice.
The first is how exactly to train on the neighbors. Our best practice trains on each neighbor sequentially in increasing distance, starting from the very nearest neighbor. The opposite order of going from farthest to nearest performs significantly worse. This might seem counter-intuitive. A priori, it could have been plausible that we should train on the nearest last, because it contains the most relevant information that the model should not forget. But the opposite turns out be the case. One interpretation of this empirical phenomenon is that, it could be advattgeouus to take the largest gradient step at the beginning on the best available data point, and therefore ground the fine-tuning process in a better region of the loss landscape. We also experiment with much more costly variant: put all neighbors into a batch, and take gradient steps on this entire batch of data repeatedly. This is much slower than the sequential method, but does not lead significant improvements.
The second design choice considers the many sequences in the Pile dataset that exceed the maximum sequence length of the base models. such cases, we split the long sequences into chunks, of length equal to the model's 'ss maximum sequence length. This means that single retrieved neighbor can result in several gradient updates. We again iterate over these chunks sequentially, one gradient update per chunk.
After fine-tuning the model on the current test instance, reset its trainable parameters to their original states - those of the pre-trained base model. We could potentially let changes to the model accumulate during the entire evaluation run, but intuitively this would only make sense under some additional assumptions temporal smoothness, that the next test instance is similar to the current one in some sense.
Hyper-parameeers- - or lack thereof. Beyond these design choices, the method requires no tuning. A remarkable aspect is that we can simply reuse the default hyper- param--rrsmtterss for the model and the optimizer available for each model in the HuggingFace library. This in contrast with other test-time augmentation methods that often invest labor and computational resources in tuning the augmentations.
5 RESULTS
We focus on language modeling tasks with causal language models, specifically, models in the GPT family, When it comes to language modeling, the most widely studied metric has been perplexity. But there are actually several different ways researchers have computed and reported perplexity. This is, in part, due to the computational burden of evaluating perplexities exactly across large corpora. To ease reproducibility of our numbers, we use Eleuther-Al's 1m-eva weavaluuaiino aaaallaatiinhhh on-ha rness library (Gao et al., 2021) for evaluating perplexities on the Pile. We report the bits per byte metric as recommended by Gao et al. (2020). This metric proportional the negative likeliiol loss normalized by a dataset-specific constant. The Pile test has 214, 584 sequences. For efficiency, we evaluate on 20% of the test set, corresponding to 42, 916 sequences. We defer bootstrap error bars to the supplementary materials as they are generally small and do not affect the interpretation of the results.
We first evaluate on a small GPT-2 model with 117M parameters the default HuggingFace gpt model from the trans f formers library. See Figure 5. Figure 6 showcases the decrease in three different perplexity measures as we increase the number of nearest neighbors to train on. We can see that using 20 neighbors already achieves most of the decrease in perplexity, computationally costing less than half than using all the neighbors. Using additional neighbors continues to decrease perplexity. Figure focuses on the six largest Pile tasks in increasing order. These six tasks combined make up more than 70% of the Pile benchmark.
Figure summarizes our results for a large GPT-2 model with 774M parameters, specifically, the Huggingface gpt pprglarge model. Finally, evaluated our method on the GPT-Neo model with 1.3B parameters by Eleuther-Al, that aims to replicate a small GPT3 model. Note this model is specifically trained on the Pile, so it gives an important insight into how test-time training compares with pre-training on the entire dataset. What we can see from Figure 2 and Figure 8 is that remarkably. TTT-NN on the 2-T-r- model can recover most of the gains of pre-training model twice as large to convergence on the full dataset.
5

---

Published conference paper at ICLR 2024
atxiv
bookcorpus2
books3
dm-mathematics
%
1.0
1.5
20
1.5
1.0
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
before
before
europarl
freelaw
github
gutenberg
hackemews
nih-exporter
1.5
1.0
1.0
86
1.0
1.0
0.5
0.5
0.5
0.0
atter
after
opensubtitles
openwebtext2
philpapers
pile-cc
pubmed-abstracts
pubmed-central
82 %
34 %
87 %
94 %
1.0
1.0
86 %
1.0
0.5
0.5
0.5
after
after
stackexchange
ubuntu-ire
uspto
wikipedia
youtubesubuitles
80 %
#8 %
0.75
1.0
88
1.0
0.50
0.5
0.25
before
alter
before
añtr
before
after
before
alter
before
after
Figure 5: Bits per byte results on all Pile tasks for small GPT-2 model (117M parameters) before and after test-time training on 50 nearest neighbors.
Splitting sequences to avoid retrieval-evaluation overlap. There is an important subtlety. In our standard protocol, use the same sequence for retrieving nearest neighbors that we are training on sequences that were chosen with knowledge of the target of prediction. We address this concern by splitting the evaluation sequence into two chunks equal length, using the first for retrieval and the second for evaluation. It turns out that doing changes our results very little, for each task by typically less than 1% and rarely up 2%. On the entire pile all task, the bits per byte number is 83%. up from 82% for the variant without splits. The supplementary materials contain the detailed breakdown of numbers the splitting protocol. Since the split is somewhat arbitrary, we present the results with no split in the main body of the paper.
evaluate on. This raises the concern
Training costs. Training costs vary greatly depending the length of the retrieved sequences. Some tasks have book-length sequences, while others consist of very short sequences. Recall that long sequences are This leads to highly heterogeneous training cost profile across
split into multiple chunks during different tasks as illustrated in Figure 9. Even within tasks there is significant variation. The figure shows training cost in seconds per neighbor
single NVIDIA A100 GPU.
Needless to say, test-time training significantly increases the compute footprint at inference time. For some applications, running all 50 iterations may be prohibitive. However, shown in Figure those gradient iterations can be stopped at any time and still yield a significant improvement, with the first few steps being the most helpful, thanks to the sequential order discussed the beginning of this section. In addition, there are many applications of language modeling that are not real-time, such as generating sequences for scientific applications, or offline code synthesis.
5.1 COMPARISONS WITH BASELINES
We compare TTT-NN with three natural baselines: 1) in-context neighbors, 2) interpolation with the distribution of tokens among the neighbors, and 3) dynamic evaluation. We found that method outperforms all three.
6

---

Published conference paper at ICLR 2024
github
arxiv
openwebtext2
books
pubmed-centre
pile-co
hits per.per byte
1.10
1.075
1.5
1.16
1.0
1.2
1.0
1.14
1.05
1.025
25
50
25 neighbors
50
25
50
25 neighbors
50
0
25
50
25
50
neighbors
neighbors
neighbors
neighbors
github
openwebtext2
books3
nubmed-centryy
pile-cc
bytr, perplextry
3
2.4
2.15
2.25
2.0
2.10
2.20
1.9
2.05
25
25
50
0 neighbors
25
50
25 neighhors
noighbors
neighbers
neighbors
neighbors
github
arxiv
openwebtext2
books3
pubmed-central
pile-cc
word perplextii
120
1000
250
250000
110
80
90
100
70
neighbors
neighhors
neighbors
neighbers
neighbors
neighbors
Figure 6: How different measurements of perplexity, provided by Gao et al. (2020), decrease with the number of neighbors on the six largest Pile tasks in ascending order, for GPT-2-Small, Top: Bits per byte. Center: Byte perplexity. Bottom: Word perplexity.
github
arxiv
openwebtext2
books3
pubmed-central
pile-cc
1.0
87 %
1.0
93
1.0
87 %
1.0
93
0.5
0.5
0.5
0.5
26 %
0.0
before
after
hefore
after
before
after
before
after
before:
after
before
after
github
arxiv
openwebtext2
books3
pubmed-centrr
pile-ce
bits per bbytc
22
0.95
0.90
1.02
0.85
0.90
0.80
25
25
50
255
50
25
50
25
25
neighbors
neighbors
neighbors
neighbors
neighbors
neighbors
Figure 7: Results for GPT-2-Large. Top: Before and after TTT-NN with 50 neighbors on the top six tasks. Bottom: How the bits per byte metric decreases with additional neighbors.
Table shows detailed results on the large GPT-2 model. Section 2 discusses those baselines the context of related work. The next few paragraphs discuss the implementation of those baselines.
Many methods use retrieved neighbors as additional context for the test instance, without test-time training. Unlike models for those methods need be trained also with retrieval. Due the prohibitive cost of training with retricval, we experiment with a baseline that simply uses the neighbors in-context at test time. Specifically, concatenate the neighbors in increasing distance, and include much of the concatenated text as possible into the context window the test instance, in addition to its original context. This baseline improves performance in few cases, e.g. pile.enreen but does help much overall.

---

Published conference paper at ICLR 2024
github
arxiv
openwebtext2
books3
pubmed-central
pile-ce
38 %
97 %
99 %
0.6
97 %
0.4
%
0.6
0.75
0.4
0.30
0.2
0.25
0.2
0.25
0.0
0.0
0.00
0.0
0.0
0.00
before
before
before
before
after
github
arxiv
openwebtext2
books3
pubmed-central
pile-co
bits per byte
0.86
0.67
0.584
0,40
0.83
0.910
0.85
0.583
0.35
0.66
0.82
0.905
0.84
0.582
25
25
50
25
50
50
25
25
neighbars
neighbors
neighhors
noig hbors
neighbors
neighbors
Figure 8: Results for GPT-Neo (1.3B parameters). Top: Before and after TTT-NN with 50 neighbors on the top six tasks. Bottom: How the bits per byte metric decreases with additional neighbors.
Training cost (sec)
101
109
pile all
pite bookcorpus?
pile argzv
pile dms-mathematis
pile books3
pile earon
pile europari
pile freelaw
pile gutenberg
githab
pile hackernows
pile nlb-exparter
pile opensubtitles
pile philpapers
pile puhmed abstracts
pile pile-oc
pile nubmed-camte cantral
pile. stackerchangs
pile abuntu-ire
pile uspto
pile wikepedia
pile joutubesabtttles
pile upenwablexi
Figure 9: Training costs in seconds per neighbor each task.
Also discussed in Subsection 2.2, the KNN-LM line of work interpolates between the distribution of the token among the neighbors, and the distribution predicted by the language model. However, the cost of directly implementing their method on large datasets like the Pile prohibitive. While our method creates database entry for each document, theirs creates one for each token in the training data; also, retrieval needs to be performed for each token in the test instance.
To make the comparison feasible, we propose modification. For each next token in the test instance, instead of using a different distribution by querying with the previous tokens as context, can make use of our document-level datahase by querying with all tokens in the test instance. The rest remains the same: We interpolate this distribution with the one predicted by the language model, and choose the best interpolation hyper-parameter pile-all (2% NN, 98% LM). In Table 1, this modified baseline is called Interpolate.
Finally, we experiment with a dynamic evaluation baseline, discussed in Subsection 2.3. Because dynamic evaluation directly trains the text that would otherwise be the target of prediction, this baseline only makes sense in settings where the text for test-time training and evaluation do not overlap. The last paragraph of Scction 5 describes such setting for method to avoid retrivval-evaluation overlap: the sequence of tokens in the test instance are split in two halves, the first half for retrieval, and the second for evaluation.
Applying dynamic evaluation to this setting, instead the first half for test-time training, allowing for many as 50 iterations. However, dynamic evaluation starts overfitting after 7 iterations. We therefore report the results with 7 iterations. We observe that this baseline leads to meaningful improvements in many tasks,

---

Published conference paper at ICLR 2024
Table 1: Comparison with baselines. All results to those labeled after in Figure 7, and for Base Only correspond to before. The baselines are detailed in Subsection 5.1. In-Context simply adds the retrieved neighbors the context. Interpolate is modified version of KNN-LM (Khandelwal et al., 2019). Dyn. Eval. is dynamic evaluation (Krause et al., 2018).
GPT-2-Large. Results for 7TT-NN (Ours) correspond
Task
In- In-Context
Interpolate
Dyn. Eval.
TTT-NN
Base Only
pile all
1.06
0,99
1.03
0.85
1.07
pile_arxiv
1.14
1.13
1.13
0.99
1.14
pile.bookcorpus2
0,95
0,96
0.91
1.00
0.96
pile.books3
1.03
1.03
1.09
1.01
1.04
pile.dm--athhemtiica
1.42
1.42
1.30
0.99
1.43
pile.enron
1.28
1.16
1.38
0.72
1.44
pile_europazz
1,96
1.92
2.13
1.15
1.96
file.freelaw
0.88
0.85
0.85
0.72
0.89
pile.github
1.88
1.06
1.29
0.51
1.95
pile.gutenberg
1.20
1.17
1.14
1.18
1.21
pile.hackernews
1,28
1,26
1.16
1.09
1.30
pile.nih-exporter
0.75
0.78
0.86
0.63
0.78
pile.opensubtitles
1.15
1.15
1.18
0.75
1.15
pile.openwebtext2
0.95
0.95
1.00
0.88
0.95
pile.philpapers
1.20
1.36
1.38
1.07
1.22
pile.pile-cc
0.94
0.95
0.99
0.88
0.96
pile-pubhee-abbsrraetss
0.77
0,81
0.92
0.74
0.82
pile pubmed-cent
0.90
0.90
0.89
0.79
0.90
sile.stackexchage
1.07
1.03
1.05
0.88
1.09
pile.ubuntu-ire
1.27
1.27
1.30
1.13
1.27
pile.uspto
0.73
0.75
0.78
0.57
0.75
pile.wikipedia
0,92
0.95
1.03
0.81
0.98
pile.youtubeesutittee.
1.55
1.41
1.34
1.42
1.51
and in three cases better than ours. It is especially helpful, although not much ours, for pile-gi sele-iithul thub. Note that the same hcuristic can also be added to our method: Our first iteration can simply be taken on the query text the test-time training data for dynamic evaluation, with the remaining iterations the neighbors.
6 LIMITATIONS AND FUTURE WORK
Our work shows that test-time training with nearest neighbors boost the effective capacity of model in terms of perplexity for language modeling tasks. But it does not show if test-time training also exhibits some of the emergent phenomena associated with larger models. Broader evaluation other non-perplexity tasks is an important direction for future work. Code generation seems to be particularly promising direction due to the strong results of TTT-NN on pile-github. code synthesis tasks in this direction. Furthermore, we would like to understand if this strength extends to life sciences applications, such as those involving large-scale medical or protein databases.
would be interesting to conduct studies of end-to-end
Also worth investigating is the use of test-time training in the context of fairness and safety. Our approach might help mitigate biases against under-represeen represeepde groups in the data source that large language models are trained on. As long as the database contains some high-quality data from particular under-represented group, TTT-NN could potentially benefit this group at least as much as the larger groups Our experiments on the Pile give some support for this intuition, the smaller tasks with less than 1% size, e.g. pile_enron and gileeeurropai,, see much more significant improvements than the larger ones, c.g. pile seteepplleccc and pile- www.e--ubbedeetteeeeee adversarial behaviors, including data poisoning attacks, by superimposing data at test time from a trusted data source. We hope
with more than 30% size together. Test-time training might also help mitigate
future work in this direction.

---

Published conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We are grateful to Sasha Rush for tremendously helpful conversations at an early stage of the project. We also thank Carsten Eickhoff and Tatsunori Hashimoto for helpful feedback and suggestions. Yu Sun would like to thank his other PhD) advisor, Alexei A. Efros, his internship manager Meta, Xinlei Chen, and his friend, Beidi Chen, for their help in the general idea of this project. Yu Sun is supported in part by Oracle Cloud credits and related resources provided by the Oracle for Research program.
REFERENCES
Uni Alon, Frank Xu, Junxian He, Sudipta Sengupta, Dan Roth, and Graham Neubig. Neuro-symbolic language modeling with antomatoo-augmentee retrieval. In International Conference Machine Learning, pp. 468 485. PMLR, 2022.
Christopher G Atkeson, Andrew W Moore, and Stefan Schaal. Locally weighted learning. Lazy learning, pp. 11-73, 1997.
Sehastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, fean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improv- ing language models by retrieving from trillions tokens. In International conference on machine learning, pp. 2206-2240. PMLR, 2022.
Léon Bottou and Vladimir Vapnik. Local learning algorithms. Neural computation, 4(6)) 1992.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 13:187711 2020.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303. 12712, 2023.
William S Cleveland. Robust locally weighted regression and smoothing scatterplots. Journal of the American statistical association, 74(368):829-836 1979.
William Cleveland and Susan Devlin. Locally weighted regression: an approach to regression analysis by local fitting, Journal of the American statistical association, 83(403):596-610 1988.
Ronan Collobert, Fabian Sinz, Jason Weston, Léon Bottou, and Thorsten Joachims. Large scale transductive syms. Journal of Machine Learning Research, 7(8), 2006.
Alexander Gammerman, Volodya Vovk, and Vladimir Vapnik. Learning by transduction. In Proceedings of the Fourteenth conference on Uncertainty in artificial intelligence, pp. 148-155, Morgan Kaufmann Publishers Inc., 1998.
Yossi Gandelsman, Yu Sun, Xinlei Chen, and Alexei A. Efros. Test-time training with masked autoencoders. Advances in Neural Information Processing Systems, 2022.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding. Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprinl S00ii:22011000022. 2020.
Leo Gao, Jonathan Tow, Stella Biderman. Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang. Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou, A framework for few-shot language model evaluation. Zenodo, September 2021. doi: 10.528 2e2.1///enooo.3711111 5371628. URL https: // doi. org /10. 5281/ zenodo. 5371628.
Suchin Gururangan, Mike Lewis, Ari Holtzman, Noah A Smith, and Luke Zettlemoyer. Demix layers: Disentangling domains for modular language modeling, arXiv preprint arXiv:2108.055::: 2021.
10

---

Published as conference paper at ICLR 2024
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training. In International conference on machine learning, pp. 3929-3938. PMLR, 2020.
Junxian He, Graham Neubig, and Taylor King-Kikpaatric. Efficient nearest neighbor language models. arXiv preprint arXiv:2 109.0421,, 2021.
Gautier Izacard, Patrick Lewis. Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi- Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Few-shot learning with retrieval augmented language models. arXiv preprint (at::2208.032999 2022.
Thorsten Joachims. Learning classify text using support vector machines, volume 668. Springer Science & Business Media, 2002.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, 7(3):535-3) 547, 2019.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. arXiv preprint 1XXi:1911.00172. 2019.
Ben Krause, Emmanuel Kahembwe, lain Murray, and Steve Renals. Dynamic evaluation of neural sequence models. International Conference on Machine Learning, pp. 2766-2775. PMIR, 2018.
Ben Krause, Emmanuel Kahembwe, lain Murray, and Steve Renals. Dynamic evaluation of transformer language models. arXiv preprint arXiv:1904. 08378, 2019.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmeened generation for knowledge- intensive nlp tasks. Advances in Neural Information Processing Systems, 33:9459 33:9459-9477, 2020.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692 2019.
David Ruppert, Simon J Sheather, and Matthew P Wand. An effective bandwidth selector for local least squares regression. Journal of the American Statistical Association, 90(432):1257-1227) 1995.
Charles Stone. Consistent nonparametric regression. The annals of statistics, pp. 595-620, 1977.
Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with self-supervision for generalization under distribution shifts. In International Conference Machine Learning, pp. 9229-9248. PMLR, 2020.
Yu Sun, Xinhao Li, Karan Dalal, Chloe Hsu, Sanmi Koyejo, Carlos Guestrin, Xiaolong Wang. Tatsunori Hashimoto. and Xinlei Chen. Learning (leam at test time). arXiv preprint arXhv:2310.13807 2023.
Vladimir Vapnik. The nature of statistical learning theory. Springer science & business media, 2013.
Renhao Wang, Yu Sun, Yossi Gandelsman, Xinlei Chen, Alexei A Efros, and Xiaolong Wang. Test-time training video streams. 2022a.
Shuohang Wang, Yichong Xu, Yuwei Fang, Yang Liu, Siqi Sun, Ruochen Xu, Chenguang Zhu, and Michael Zeng. Training data is more valuable than you think: A simple and effective method by retrieving from training data. arXiv preprint 2200:22003.07773 2022b.
Hao Zhang, Alexander Berg, Michael Maire, and Jitendra Malik. Svm-knn: Discriminative nearest neighbor classification for visual category recognition. In 2006 IEEE Computer Society Conference Computer Vision and Pattern Recognition (CVPR' (C))) volume 2, pp. 2126-2136 IEEE, 2006.
11

---

Published conference paper at ICLR 2024
A RESULTS FOR GPT-2-SMALI ON SPLIT SEQUENCES
Here show the results where we split each sequence in half, using the first part for retrieval and the second part for evaluation. If the sequence is more than twice the maximum sequence length, then our first part has length equal to the maximum sequence length, and the second part the rest. For these results, we focus on the small GPT-2 model.
The model available on HuggingFace at https:/ https::/huggigggiigtet/ hww.gggggggggg.eee/ee/a co/gpt2. We use learning rate of 2e-5 for the Adam optimizer with value 1e-08 The maximum sequence length of the model is 1048 tokens.
arxiv
bookcorpus2
books3
Al--mathmaatics
enron
1.0
1.5
34 %
102 %
1.0
97
1.5
70
1.5
65 %
1.0
1.0
0.5
0.5
0.5
0.5
0.5
0.5
before after
before
after
before
after
before
atter
before
after
before
after
europarl
freelaw
github
gutenberg
hackernews
nih nih-exporter
1.0
85 %
1.5
91 %
1.5
92
1.0
1.0
1
0.5
37 %
0.5
0.5
0.5
0.0
0.0
0.0
0.0
before after
before
after
before
after philpapers
before
after
before
after
before
after
opensubtitles
openwebtext2
pile-cc
pubmed-abstracts
pubmed-central
82 %
94 %
1.5
88 %
1.0 95 %
1.0
94 %
1.0
85
1.0
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
0.0
0.0
before
after
before
after
belore
after
before
after
before
after
before
after
stackexchange
ubuntu-ire
uspto
wikipedia
youtubesubtitles
1.5
85 %
87
1.0
1.0
0.5
0.5
before
after
before
after
before
after
before
after
before
after
github
arxiv
openwebtext2
books3
pubmed-central
pile-cc
2.0
1.3
1.150
1.14
1.12
1.5
1.0
1.10
1.2
1.125
1.0
1.12
1.100
0.9
25
25
50
25
50
25
25
50
50
noighbors
neighbors
noighbors
neighhors
neighbors
nnighhors
Figure 10: Bits per byte results on all Pile tasks for GPT-2-Small (117M parameters) before and after training 50 nearest neighbors, where split the sequence into non-overlapping halves for query and evaluation.
12

---

Published conference paper at ICLR 2024
B ALL RESULTS FOR GPT-2-LARGE,
The model is available on HuggingFace htttp::/hugginggeee co/gpt2-large We use learn- ing rate of 2e-5 for the Adam optimizer with value 1e-08. The maximum sequence length of the model is 1048 tokens.
all
arxiv
bookcorpus2
books3
dm-mathematics
enron
0.5 1.0
1.0
1.0
105 %
98 %.
1.5
1.5
70 %
50 %.
0.5
0.5
before after
before after
before after
before
after
befors
after
before
after
europarl
freelaw
github
gutenberg
hackernews
nih-exporter
1.0
98
1.0
1.0
84
0.75
61 %
0.5
0.50
26
0.5
0.0
0.0
0.0
0.00
opensubtitles
openwebtex12
philpapers
pile-cc
pubmed-al abstracts
pubmed central
1.0
1.0
93
1.0
87 %
93
0.75
B7 %
0.5
0.5
0.50
0.5
0.5
0.25
0.0
0.0
0.0
0.00
0.0
before
after
before
after
before after
before
after
before
after
before
after
stackexchange
ubuntu-irc
uspto
wikipedia
youtubesubtitles
0.75
1.5
94
1.0
1.0
76
83 %
0.50
1.0
0.5
0.5
0.5
before
after
before afer
before
after
before
after
before
alter
github
arxiv
openwebtext2
books3
pubmed-central
pile-cc
0.95
0.90
1.02
0.85
bils
0.90
0.90
0.80
1.00
25
50
25
50
25
50
25
25
50
25
50
neighbors
neighbors
neighbors
neighhors
neighbors
ntighhors
Figure 11: Bits per byte results on all Pile tasks for GPT-2-Large (771M parameters) before and after training on 50 nearest neighbors.
13

---

Published conference paper at ICLR 2024
C ALL RESULTS FOR GPT-NEO
The model available on HuggingFace at https: heattps:/huggiiiii www.ggggggggg.eeeeeee/aaaaagaeee how.e//eeoeetttee///ttteeeee 3B. We used a learning rate of 5e-6 for the Adam optimizer with value 1e-08. The maximum sequence length of the model 2048 tokens.
all
arxiv
bookcorpus2
books3
dm-mathematics
96 %
%
1.0
1.0
0.6
0.75
'? %
0.50
0.5
0.5
0.4
0.25
0.2
0.25
0.2
0.00
0.0
0.0
0.0
0.00
0.0
after
after
europarl
freelaw
github
gutenberg
backernews
nih- nih-exporter
0.75
101 %
75 %
1.0
103
1.0
100 %
0.6
87 %
0.4
0.2
0.5
0.2
0.2
0.00
0.0
0.0
0.0
0.0
0.0
before after
before after
before
after
before
after
before
after
befors after
opensubtitles
openwebtext2
philpapers
pile-cc
pubmed-abstrenn
pubmed-central
1.0
72 %
%
1.0
97 %
0.75
97 %
0.6
100
0.5
0.50
0.5
0.50
0.25
0.25
0.2
before after
before after
before after
before after
before after
before after
stackexchange
ubuntu-irc
uspto
wikipedia
youtubesubtitles
1.0
0.6
%
0.75
96 %
1.0
98 %
0.4
0.5
0.5
0.2
0.2
0.25
0.0
0.0
0.0
0.00
0.0
github
arxiv
openwebtext2
books
pubmed-central
pile-cc
byte
0.84
0.86
0.67
0.40
0.910
0.85
0.35
0.66
0.82
0.905
25
50
0
25
25
50
25
50
50
25
50
neighbors
neighbors
meighbors
neighbors
neighbors
neighbors
Figure 12: Bits per byte results on all Pile tasks for GPT Neo (1.3B parameters) before and after training 50 nearest neighbors.
14

---

Published conference paper at ICLR 2024
D BOOTSTRAP ERROR BARS
In this section, we report empirical bootstrap enror bars for our bits per byte evaluation metrics. Computing valid error bars for perplexity metrics is non-trivial for two reasons. First, perplexities can assumee large values and are generally unbounded the top. Second, the metric is an exponential of sum values, rather than sum of values for individual data points as would be the case with metric like accuracy. For these reasons, we follow established practices in statistics and report empirical bootstrap error bars.
For given task consisting of sequences, bootstrap sample will sample the individual negative log likelihoods on each sequence i.i.d. with replacement n times, For each bootstrap sample we compute the bits per byte metric. Error bars then refer to quantiles of the bits per byte values. We choose 1000 bootstrap samples and report the (0.1, 0.9) quantiles as lower and upper error bars, respectively.
all
bookcorpus2
books3
dn-mathematics
1.5
1.0
1.5
1.5
I
1.0
1.0
1.0
0.5
0.5
0.5
0.5
before after
before after
before after
before after
before after
before after
europarl
freelaw
github
gutenberg
hackernews
nih-exporter
1.5
1.0
1.0
1.0
1.0
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
opensubtitles
openwebtext2
philpapers
pile-cc
pubmed-ahee abstracts
pubmed-ceme central
1.0
I
1.0
0.5
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
0.0
before after
before after
before after
before after
before after
before after
stackexchange
ubuntu-irc
uspto
wikipedia
youtubesubtitles
1.5
0.75
1.0
1.0
0.50
0.5
0.5
0.25
0.5
0.0
0.0
0.00
before
alter
before
afer
before
after
before
after
before
aller
Figure 13: Bits per byte results on all Pile tasks for GPT-2-Small (117M parameters) with error bars indicating 0.1 and 0.9 quantiles across 1000 bootstrap samples. This means only 10% of bootstrap samples fall outside the error bars the top and bottom, respectively.
The error bars for the larger models are only smaller, since the perplexities are generally smaller and have fewer large values. We therefore omit the corresponding plots for GPT-2-Large and GPT-Neo.
15

---

Published conference paper at ICLR 2024
E TRAINING COSTS
For completeness, we compare all training costs below across all tasks and the three models we evaluated on.
Training cost (sec)
10a
109
10
pile all
pile arsie
Inle. bookeorpuss
pile tta-mathematics
pile books]
pile euron
pile. europari
pile freelaw
pile githab
pile gutenberg
pile hackerneew
pile nihk ait--rrorterr
pile opensubtitles
pile. opanwabtezt
pile philpapors
pile puhmed abotraris
pile pile-cc
pile wwwppuu--cccc www.-nhdddccntralll -antrrl
pile stackexchange
pile abuntu-ire
pile wikipedia
pile. youtuueeuubttlee
uppto
Training rost (soc)
102
fp
109
R
pile all
pilo arziv
pile bookeurpus2
pile an--mathonatics
pile booke3
pile ourn
pile europar
pile frenlaw
Pill github
pile putenbery
pile hackemews
pile nih expartar
pile opensuhtali tale
pile openwebtexi2
pile phipapers
pile pubmed ahstracts
pila pile-ar
pile pwwupdddddooe puumeed-cantral contral
pile stackecccaa
pile ubantu-ire
pile uppto
pile. wilipodia
pile yww.yobbbbsbbtileee
Training (are)
10?
10t
109
1
pile all
pile bookcorpue2
pile arzir
plle An-matthematics
pile booke
pile enmn
pale euruport
pile freelow
pile mithab
pile gutmbary
pile tackemnos
pile aib-nsparinr
pile openauhotlos
pile openwehlextz
pile phipapera
pile pulmaed abbttacta
pilo pile-or
pile pwwpadd-aaaa cantral contral
mile stackexchange
pile ubuntu-ire
pile uspto
pile. wikipadiia
pile youtubesubttlle
Figure 14: Training costs in seconds neighbor. Top: GPT-2-Smal Middle: GPT-2, Bottom: GPT-Neo.
16


Reviews:
[{"review_id": "SCnC8jMbIy", "reviewer": "Reviewer_kVHL", "summary": "This paper proposes a method of training at test time", "questions": "See weaknesses above.\n\n* Also, will the code/data/datastore be released so others can reproduce these studies?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The method is not too complicated, and could likely be reproduced.\n* In some ways, the evaluation was very impressive. Quite large scale, showing benefits with an index that spans the whole Pile across many domains.\n* The baselines of kNN and in-context prompting also seemed relevant/strong.", "weaknesses": "There were some weaknesses. I think this paper still could have value, but I would be more confident in recommending that the paper be accepted if the following could be addressed:\n\n1. There are some clarity issues with the paper. For instance, it was not very clear to me if retrieval is done after every token or at some other cadence.\n\n2. There is a discussion of inference speed, but it is not very concrete. Could inference throughput be added to table 1?\n\n3. While bits/byte based LLM evaluation is good, it would also be really nice to see results on extrinsic tasks as well.\n\n4. This is not so much a weakness as a missed reference, but this paper is very relevant: https://arxiv.org/abs/1609.06490\nLi, Xiaoqing, Jiajun Zhang, and Chengqing Zong. \"One sentence one model for neural machine translation.\" arXiv preprint arXiv:1609.06490 (2016).\nI think it should definitely be cited, but I think even if a similar idea has been proposed before in the neural (conditional) LM space, the modernized evaluation of the current paper has significant value.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DIKuA6ll0p", "reviewer": "Reviewer_Cq76", "summary": "This paper addresses the problem of improving the language model perplexities by using the training data during inference. The core idea is to find the sequences similar to the test sequence from the indexed training data and finetune the base model with these nearest neighbor sequences. \n\nThe paper evaluates this on pile benchmark, where the training data is indexed using representations obtained from the Roberta model. For each test sequence from various pile categories, the nearest neighbors are picked to finetune the model and later the test sequence is evaluated for perplexity measure. The empirical results show usefulness of the approach as it improves the LM perplexities.", "questions": "1. It is not clear to me how indexes handle larger sequences? Bigger sequences are chunked [chunk1, chunk2, chunk3, ..] and if the nearest neighbor match happens at chunk2, what is the process? \n\n2. While it is neat that this method doesn't require hyper-paramter tuning? What happens when one tries that? (I agree it is prohibitively expensive, but could be done for few test sequences)\n\n3. How does KNN-LM work with document level index? For the original work, it was context -> next_word, how do we get token probabilities with document level index.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow. The idea is clearly described and empirically validated.\n\n2. The empirical results on various pile benchmark show the usefulness of the approach.", "weaknesses": "1. While the idea is neat and simple to implement, as shown in Figure 9, the training costs for each neighbor is expensive, thus limiting the usefulness in real-time applications. \n\n2. While the results on the LM perplexity are useful, it would be interesting to see how this compares in an end-to-end task such as code generation, etc. Few-shot prompt tuning (with or without retrieval augmented learning) are popular paradigms that are used in bigger LLMs. It would be interesting to see the comparison with such methods (in offline evaluation settings).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gY2dNvWWYY", "reviewer": "Reviewer_zh5b", "summary": "This paper investigates test-time training on nearest neighbors (TTT-NN) in the context of large language models, specifically transformer models. The authors create a large-scale distributed nearest neighbor index based on text embeddings of the Pile dataset. For each test instance, the system retrieves nearest neighbors from this index and fine-tunes the model on these neighbors before applying it to the test instance. The method is evaluated on 22 language modeling tasks from the Pile benchmark, using three causal language models of increasing size (small GPT2, large GPT2, and GPTNeo).\n\nThe results show that training for only one gradient iteration on as few as 50 neighbors can reduce a normalized perplexity measure (bits per byte metric) by 20%. Test-time training narrows the performance gap between a small GPT2 model and a GPTNeo model, which was specifically trained to convergence on the Pile. The improvements due to TTT-NN are more dramatic on unseen tasks, while still helpful on seen tasks. Test-time training can increase the effective capacity of a model, though at the cost of increased inference time. The authors conclude that their work establishes a valuable baseline for implementing test-time training in large language models and opens the door to further research in this area.", "questions": "see above", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The organization of this paper is well-structured, making it easy to read and comprehend.\n2. This paper presents a simple test-time training approach on nearest neighbors (TTT-NN), which significantly improves performance across more than twenty language modeling tasks in the Pile benchmark with minimal fine-tuning.\n3. Test-time training effectively increases the capacity of a model, showcasing its potential to narrow the performance gap between smaller and larger models, and offering a valuable  baseline for implementing test-time training in the context of large language models.\n4. The large-scale distributed nearest neighbor index built on text embeddings of the Pile dataset enables efficient retrieval of relevant data for test-time training, serving queries to approximately 200 million vectors and 1TB of data in just one second.", "weaknesses": "1. Why not use the PQ (Product Quantization) Index, which can significantly reduce storage overhead and thus avoid the cost of distributed retrieval? Although the vectors after PQ are approximations of the original vectors, recent works such as ”**[KNN-MT](https://openreview.net/forum?id=7wCBOfJ8hJM)“** have demonstrated better performance using this approach.\n2. Retrieval plus k*seq_len gradient updates may reduce inference speed. How much of a difference is there between the inference speed of the proposed method and the original model?\n3. How is the database used by TTT-NN constructed? Is it built using the training data from Pile?\n4. How does the baseline \"interpolation with the distribution of tokens among the neighbors\" work? What is the key used for retrieval when predicting the next token? Also, for KNN-LM, directly constructing the database is indeed very costly. Dai's work ”**[SK-MT/SK-LM](https://openreview.net/forum?id=uu1GBD9SlLe)“** provides an efficient construction method for their KNN-LM, i.e., first using BM25 to retrieve similar N documents, and then using these N documents to build a token-level database for interpolation and prediction. Considering that this work retrieves k documents, how are the k documents and the original model's predicted probability distribution interpolated? Is it similar to Dai's work mentioned above?\n5. For Section 4.1, \"Splitting sequences to avoid retrieval-evaluation overlap,\" suppose the test sequence is $x_t = ABCDEFGH$. Do the authors mean that there might be a sentence $x_d = ABCDEFGH$ in the database? If so, does this introduce test data leakage? Moreover, even if we split $x_t$ into $x_t^{'}=ABCD$, according to the description in Section 3, when using the prefix for retrieval, can we still retrieve $x_d$ and thus cause test data leakage?\n6. What does \"plain\" refer to in Table 1? No specific definition was found.\n7. Some missing related works\n    - [REALM: Retrieval-Augmented Language Model Pre-Training](https://arxiv.org/abs/2002.08909)\n    - [Training Language Models with Memory Augmentation.](https://aclanthology.org/2022.emnlp-main.382.pdf) The work is to retrieve similar k neighbors to aid training\n    - To some extent, this paper can be considered as an explicit extrapolation of KNN-LM/KNN-MT. Specifically, Gao's work“[Nearest Neighbor Machine Translation is Meta-Optimizer on Output Projection Layer](https://arxiv.org/abs/2305.13034)” demonstrates that the working mechanism of KNN-LM/KNN-MT is to perform implicit gradient updates using the retrieved k nearest neighbors. In contrast, this paper explicitly uses the retrieved k nearest neighbors for explicit gradient updates.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sy81WjOvep", "reviewer": "Reviewer_y6wT", "summary": "The paper proposes a retrieval-augmented LM recipe with these steps:\n\n1) Given a query context, retrieve N sequences.\n2) Sequentially finetune the LM on the N sequences.\n3) Use the finetuned LM to predict the next tokens. (Although results in the main text use the same text for retrieval and testing)\n\nThe paper shows improve perplexity on multiple datasets of the pile, and uses the pile for retrieval (done with a distributed neighbor index). The main weaknesses of the paper are lack of analysis to show why retrieval is helpful (perhaps it only helps when finding exact matches in retrieval?) and weakly implemented retrieval-based baselines.", "questions": "Q1: What if the retrieved contexts are not relevant to the query context?\n\nQ2: How would GPU acceleration further improve the speed? Isn't the data much too large for any GPU?\n\nQ3: Is there any plan to release the code for the distributed server? Is it meant to be a novel contribution of this work? Is there existing work that achieves anything similar? My impression is that there are multiple options for this type of distributed neighbor index used in industry.\n\nQ4: \"largest gradient step\" I am confused what is meant by this. I assume it is meant to correspond to learning rate scheduling, but in general I would think the furthest neighbors may have a larger gradient magnitude if they are more different from the existing context. Also, can you simply measure the gradients and see if this is true or not---whether the first step is the largest?\n\nSuggestions for kNN-LM\n\n* Use GPT2 for retrieval. This has been shown to work well with kNN-LM.\n* Use only a single dataset, e.g. github, for retrieval and encode every token for all or contiguous subset of the data.\n* Retrieve at every token, and follow the recipe from kNN-LM.\n* Alternatively, simply report perplexity on a dataset that kNN-LM was already evaluated for.\n\n\nMinor notes\n\n* Fig 7: What are the top tasks? I assume they are the largest.\n\nOther Related Work\n\n* Basu et al: This paper does test-time training on other tasks and also presents a modern theoretical view on the value of test-time training.\n* Drozdov et al: This paper improves upon the kNN-LM by adapting the interpolation to the quality of retrieval.\n* Ram et al: An effective application of the in-context approach.\n\nBasu et al. A Statistical Perspective on Retrieval-Based Models\n\nDrozdov et al: When and how to rely on retrieval in the kNN-LM\n\nRam et al: In-Context Retrieval-Augmented Language Models", "limitations": "", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The approach is simple and effective.\n\n2. There is evaluation using a large amount of data for retrieval, and across many datasets. Although, lack of analysis is concerning and makes it hard to understand how significant the improvements are.\n\n3. The approach relies on a distributed neighbor index. In general, it will be helpful for the community working on retrieval-enhanced ML to see papers that operate with such large retrieval. Although there are not many details about how the distributed index is implemented besides the server count and amount of data used. \"We find that 180 servers is a reasonable trade-off for the cluster we use.\" is not backed up by any statistics. Also, sometimes retrieval is done on this large scale with sparse retrieval, e.g. BM25, and dense retrieval is done here but not not with a model designed for retrieval such as DPR or others.", "weaknesses": "1. The paper is lacking in analysis. Consider this statement from the intro \"In this paper, our hope is that a sufficiently large database will contain data relevant enough to each “domain”, induced by each test instance, such that fine-tuning improves local performance.\". Perhaps it is worth measuring if the data found was from the matching domain? Although we see improvements in perplexity, we do not have a sense of why these improvements are happening. In addition, there is mention that 4 percent of retrieval is nearly an exact match---if the perplexity improvements is only due to exact match from the training data then it is not clear how useful this is.\n\n2. The retrieval baselines are very weak and not well configured for this task. Also, \"Unlike ours, models for those methods need to be trained also with retrieval.\" Except many of the models do not need to be trained with retrieval, including the ones used as baselines. Is the implication that models trained for retrieval are better than test time training?\n\n2a. The kNN-LM comparison is far from fair. By looking at the interpolation param, the value is clearly much worse than expected. The param is only 0.02 but in the kNN-LM paper it is much higher 0.25 when the retrieved and test data are from the same domain, and 0.6 when they are from different domains. This is a huge difference and either the kNN-LM results should be excluded or amended. I make suggestions for how to amend the results in the questions section.\n\n2b. The \"in-context\" baseline gives almost the same result as the base model. This suggests the in-context approach is almost not doing anything, probably for two reasons: 1) the context length is very limiting and 2) LMs often do not attend well to retrieved information in such large contexts (Sun et al and Liu et al). In contrast, Shi et al shows that when done properly the in-context approach can be very effective both for language modeling and QA. To be done properly, probably chunks should be retrieved and the query shortened. If there is still a context limit issue then Shi et al proposes an ensembling technique that enables scaling up the number of retrievals more efficiently.\n\n3. Despite the clarification in 4.1 about retrieval-evaluation overlap, it seems inappropriate to include the results in Figure 5 as the model is using the text it is meant to predict for retrieval. Fortunately, this should be an easy fix as the fair comparison is already in the appendix. \n\n4. In general, there are not many insights about how effective and influential retrieval is. It would be helpful to include an alternative retrieval, such as BM25. Similarly, it may be helpful to include an alternative dense retriever---since in this setup the LM is disjoint from the retriever, then it makes sense to include a model designed specifically for retrieval.\n\nSun et al. Do Long-Range Language Models Actually Use Long-Range Context?\n\nLiu et al. Lost in the Middle: How Language Models Use Long Contexts\n\nShi et al. REPLUG: Retrieval-Augmented Black-Box Language Models", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)