Paper: Test-Time Training on Nearest Neighbors for Large Language Models (ICLR-2024)

Reviewer comment (Evaluation):
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.

Author response:
Our motivation is actually that each test instance is its own “domain”. This so-called “domain” (in quotes) is not defined as, for example, a specific task in the Pile like github, but the test instance itself. Therefore, Figure 3 and 4 already serve as analysis for this motivation. The improvements in perplexity, from our point of view, happens because there are sufficiently close training instances to this test instance. When we manually inspect the 4% of data that have a nearly exact match, the matches are not the same text as the queries. They are just different texts that seem to have the same style. Different passages from the same novel, for example, could easily exhibit this pattern. This is not a case of test set leakage.

Reviewer comment (Experiments):
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?

Author response:
The implication is that models trained with retrieved neighbors require a huge amount of extra resources for training. This is why we choose not to compare with them.

Reviewer comment (Experiments):
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.

Author response:
Generally speaking, KNN-LM has substantial improvements for medium size datasets like Wikitext-103, but is extremely resource intensive, to the point of almost impractical, due to its token-level retrieval and database. Our method is practical for very large datasets. In fact, the larger the dataset, the more relevant our retrieved neighbors, and the better our method should be. Our baseline is meant to be a practical implementation on the Pile, which would otherwise be too large for KNN-LM. Following the recipe of KNN-LM would use hundreds of times more memory and compute. Downgrading to a smaller dataset like Wiktext-103 would force us into a regime where our method is not meant to be advantageous.

Reviewer comment (Experiments):
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.

Author response:
The in-context baseline is an approximation to retrieval-augmented LMs if they are not trained with retrieval. The point is to put their training cost and ours into the same regime. On the other hand, we agree that the context length of Roberta and GPT-2 is limiting the performance of this baseline.

Reviewer comment (Experiments):
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.

Author response:
Thank you for the suggestion. We will highlight the results of the split evaluation.

Reviewer comment (Experiments):
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.

Author response:
Our choice of Roberta for retrieval (instead of GPT-2) is exactly motivated by the consideration that Roberta is more suitable for retrieval, since each output embedding takes advantage of the global context. We agree that other alternatives for building the index would be interesting. We chose to use the simplest possible option, to demonstrate the potential of this method if extended with more sophisticated techniques.

Reviewer comment (Experiments):
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.

Author response:
There are a lot of alternatives for building the index. We chose to use the simplest possible option, to demonstrate the potential of this method if extended with more sophisticated techniques.

Reviewer comment (Evaluation):
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?

Author response:
Please see the table in Appendix E for details on inference speed. We're happy to highlight these numbers in the main body.

Reviewer comment (Experiments):
How is the database used by TTT-NN constructed? Is it built using the training data from Pile?

Author response:
Yes, using the training data from the Pile.

Reviewer comment (Reproducibility):
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?

Author response:
We simply use the retrieved documents to get a distribution over tokens. The same distribution is used for all predictions of the next token in the same instance. The first paragraph on page 8 contains a detailed description.

Reviewer comment (Experiments):
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?

Author response:
The split evaluation is not for test data leakage in the training set. The purpose is to eliminate the possibility of using the target tokens as parts of the query for retrieval. Since the point of language modeling is to predict the next tokens as target, it could potentially be viewed as unprincipled to use those tokens for retrieval. The split evaluation avoids that.

Reviewer comment (Writing):
What does "plain" refer to in Table 1? No specific definition was found.

Author response:
In Table 1, plain means the baseline before TTT, i.e. direct inference with the pre-trained language model.

Reviewer comment (Novelty):
Some missing related works  
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://arxiv.org/abs/2002.08909)  
- [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  
- 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.

Author response:
We have already cited Guu et al. (2020) for REALM. Thank you for the other two references, we’ll incorporate them.

Reviewer comment (Evaluation):
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.

Author response:
While our method is more expensive than direct inference without retrieval, it is less expensive than retrieval-augmented LMs (see in Section 2.2), which is one of the baselines.

Reviewer comment (Experiments):
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).

Author response:
We agree that few-shot evaluations would be nice.

Reviewer comment (Reproducibility):
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?

Author response:
For long sequences, we only use chunk1 for retrieval, as explained in the first paragraph on page 4. Specifically, chunk1 contains the maximum number of tokens that can fit into the context window. Using a more elaborate scheme would probably improve performance even further.

Reviewer comment (Experiments):
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)

Author response:
If you find our current results convincing, then the lack of hyper-parameter tuning means that they could be even better. So the method might have still more potential than it appears in the results.

Reviewer comment (Reproducibility):
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.

Author response:
We simply use the retrieved documents to get a distribution over tokens. The same distribution is used for all predictions of the next token in the same instance.

Reviewer comment (Writing):
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.

Author response:
Retrieval is done once based on a prefix of the query. There is no retrieval after each token.

Reviewer comment (Evaluation):
There is a discussion of inference speed, but it is not very concrete. Could inference throughput be added to table 1?

Author response:
Please see the table in Appendix E for details inference speed. We're happy to highlight these numbers in the main body.

Reviewer comment (Experiments):
While bits/byte based LLM evaluation is good, it would also be really nice to see results on extrinsic tasks as well.

Author response:
We agree that additional evaluation is a good step for future work.

Reviewer comment (Novelty):
This is not so much a weakness as a missed reference, but this paper is very relevant: https://arxiv.org/abs/1609.06490 Li, Xiaoqing, Jiajun Zhang, and Chengqing Zong. "One sentence one model for neural machine translation." arXiv preprint arXiv:1609.06490 (2016). I 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.

Author response:
Thanks for the pointer. We'll incorporate the reference.

Reviewer comment (Reproducibility):
Will the code/data/datastore be released so others can reproduce these studies?

Author response:
All of our code and, in fact, the complete index have already been made publicly available. The URL was omitted to preserve anonymity during the reviewing process.