Summary
The paper proposes a method to increase context lengths of encoder-decoder transformers to very long input sequences. The idea is to essentially encode all of the tokens of the entire input (on overlapping context-length chunks) and create an index of the input tokens. Just prior to decoding, k-nearest neighbor (kNN) encoded tokens are selected from the index to maximize $QK^T$ dot product of the cross-attention phase, between the encoder and decoder hidden states, during decoding. One of the key contributions in this paper is that they re-order the computation of the matrix products in the cross-attention ($QK^T$) so that creating a single index of the encoder tokens is sufficient, and kNN look-up can be performed efficiently to compute ($QK_{kNN}^T \approx QK_{best}^T$). Additionally, their method has the advantage of being applied to existing pre-trained architectures without the need for additional weights or tuning. It can be applied directly during inference at test time. Applying it during training gives an additional performance boost.
Strengths
* The idea is good and the experimental evidence is strong.
* The evaluation comparisons using “low-cost” tuning methods, i.e. applying this idea just during validation and testing is also clever.
* The proposed method has interesting advantages,
1. While the idea of using kNN is not new, it has been done nicely and efficiently,
2. The method can also be applied to existing pre-trained encoder-decoder models during inference directly.
3. There's further boost in performance when applied during training.
Weaknesses
No major weaknesses in the overall quality of the contribution. There are minor weaknesses which the authors can address e.g. the clarity of the writing can be improved, results in some tables are incomplete etc. I have framed those as questions below for the authors to address.
Questions
1. The Fig. 2 shows encoder encoding disjoint chunks, whereas text lines L73-74 say it is overlapping chunks. Can you clarify which of these is correct, and change the figure to represent that more accurately? The figure, overall, can be improved.
2. For the k nearest neighbor indexing and search, you mention the use of the FAISS library. Do you use approximate nearest neighbors, or the exact nearest neighbors? Considering most contexts are fairly small, if you’re not using exact nearest neighbors, can you provide a small comparison to using exact nearest neighbors, and share whether or not it affects results (and/or time).
3. In the Comparison to Wu et. al. section, it would be good to highlight that Memorizing Transformer paper applies their approach to decoder-only models. Your approach is applied to the encoder-decoder model.
4. Train chunked +test Unlimiformer (L137-141) requires more details. e.g., I don’t understand what treating each chunk as its own training example entails. Is this what you do: For an example that’s longer than the context length, do you chunk it into say C non-overlapping chunks, treat each (chunk, output) pair as a training example (to create C examples), but allow retrieval over the tokens of all the C chunks during decoding? (After reading retrieval training, *I guess there is actually no retrieval happening over training data* in the Train chunked +test Unlimiformer)
* 4.1 For Train chunked +test Unlimiformer, how does the duplication of some training examples because of the C (chunk, output) pairs figure into evaluation? Do you try to combine the C outputs to a single one, or take the average, or ignore the duplication and just treat them as more example pairs than what other methods would have? (Just clarify this detail in the paper)
5. Is Sec. 3.2 also doing fine-tuning? I was quite confused by the statement “We also consider training Unlimiformer directly”. Perhaps a different way to separate out section 3.1 and 3.2 is to say that in 3.1. you are applying kNN look up during model evaluation (while evaluating the models after each period of training), Vs 3.2 where you are applying kNN strategy on the training data examples (where the choice of the kNN values directly affects the model gradients during tuning).
6. During retrieval training, for the training examples that are longer than the context length, do you just have 1 copy of the (input, output) pair example in training – unlike the train chunked case where you could have multiple copies? Is this true even for training examples that are longer than 16k tokens?
7. In Table 4 when you say “Unlimiformer (this work)” do you mean the “Retrieval training” regime?
8. In Table 5 for BookSum, why isn’t there a line reporting results for BART-base Unlimiformer (Retrieval Training)?
9. Table 4 doesn’t directly include PRIMERA base training (standard finetuning) results, but worth repeating that result line from Table 3, so that this sentence (L206-207) is easier to follow: “in Table 4 highlight two important points: first, Unlimiformer+BARTbase performs better than the base PRIMERA across all metrics and datasets”
10. In the Comparison to SLED Ivgi et. al. The statements in L274-276
*“…This in practice limits SLED to only about 16k token-long inputs on a single GPU; in contrast, instead of attending to all input tokens, Unlimiformer attends only to the top-k input tokens for every attention head, and thus can process unlimited inputs in practice”*
Contradicts your statements in Sec. 3.2 Retrieval Training L149-150
*“When inputs are longer than 16k tokens, we truncated the input to 16k tokens at training time due to GPU memory requirements.”*
So, it seems like atleast in practice, in the experiments reported in this paper the input length was restricted to max 16k tokens due to practical limitations. Is that correct to say?
11. [Comment] In Table 2 (Sec 4.1. Datasets), for the distribution plots it might be worth adding a vertical line to indicate where the average length (avg. number of input tokens) falls in the distribution curve.
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Limitations
Perhaps worth noting that this is not applied to decoder-only models.
Also worth noting the 16k limitation to fit into GPU during training.
---
**Post rebuttal notes**
The authors rebuttal and subsequent discussion addressed all my questions and concerns.