Test-Time Training on Nearest Neighbors for Large Language Models

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.

2. 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?

2a. 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.

2b. 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.

3. 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.

4. 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.

Sun et al. Do Long-Range Language Models Actually Use Long-Range Context?

Liu et al. Lost in the Middle: How Language Models Use Long Contexts

Shi et al. REPLUG: Retrieval-Augmented Black-Box Language Models

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“ have demonstrated better performance using this approach.
2. 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?
3. How is the database used by TTT-NN constructed? Is it built using the training data from Pile?
4. 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“ 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?
5. 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?
6. What does "plain" refer to in Table 1? No specific definition was found.
7. Some missing related works
   • REALM: Retrieval-Augmented Language Model Pre-Training
   • Training Language Models with Memory Augmentation. 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” 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.

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.

2. 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).

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:

1. 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.

2. There is a discussion of inference speed, but it is not very concrete. Could inference throughput be added to table 1?

3. While bits/byte based LLM evaluation is good, it would also be really nice to see results on extrinsic tasks as well.

4. 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.