Unlimiformer: Long-Range Transformers with Unlimited Length Input
Authors: Amanda Bertsch
Uri Alon
Graham Neubig
Matthew R. Gormley
Venue: NeurIPS 2023 poster (2023)
Keywords: retrieval augmentation
summarization
long-context
generation
long-input
encoder-decoder
transformers
language models
natural language generation
natural language processing
deep learning
neural networks
TL;DR: Unlimiformer modifies pretrained encoder-decoder models with retrieval to allow unlimited input at test time.

Abstract: Since the proposal of transformers, these models have been limited to bounded input lengths, because of their need to attend to every token in the input. In this work, we propose Unlimiformer: a general approach that wraps any existing pretrained encoder-decoder transformer, and offloads the cross-attention computation to a single $k$-nearest-neighbor ($k$NN) index, while the returned $k$NN distances are the attention dot-product scores. This $k$NN index can be kept on either the GPU or CPU memory and queried in sub-linear time; this way, we can index practically unlimited input sequences, while every attention head in every decoder layer retrieves its top-$k$ keys, instead of attending to every key.  We evaluate Unlimiformer on several long-document and book-summarization benchmarks, showing that it can process even **500k** token-long inputs from the BookSum dataset, without any input truncation at test time. We demonstrate that Unlimiformer improves pretrained models such as BART and Longformer by extending them to unlimited inputs without additional learned weights and without modifying their code. Our code and models are publicly available at https://github.com/abertsch72/unlimiformer , and support LLaMA-2 as well.

Unliniformer: Long-Range Transformers with Unlimited Length Input
Amanda Bertsch
Uri Alon"
Graham Neubig
Matthew R. Gormley
Carnegie Mellon University, USA
(abertsch, ualon, gneubig mgorml ey Bcs. cmu.cmu........mmmmmmmmm. cmw.cmu
Abstract
Since the proposal of transformers (Vaswani et al. (017), these models have been limited to bounded input lengths, because of their need to attend every token in the input. In this work, we propose Unlimiformer: general approach that wraps any existing pretrained encoder-decoder transformer, and offloads the cross-attention computation returned &NN distances are the attention dot-product scores. This kNN index can be kept either the GPU or CPU memory and queried in sub-linear time; this way, we can index practically unlimited input sequences, while every attention head in every decoder layer retrieves its top-k keys, instead of attending every key. We evaluate Unlimiformer on several long-document and hook-summarizztion benchmarks, showing that it process even 500k token-long inputs from the BookSum dataset, without any input truncation at test time. We demonstrate that Unlimiformer improves pretrained models such BART (Lewis et al. (222a)) and Longformer (Beltagy et al. 2020) by extending them to unlimited inputs without additional learned weights and without modifying their code. Our code and models are publicly available, and support LLaMA-2 as wel
single k-nnarest-neighhor (kNN) index, while the
Introduction
Transformers (Vaswani al. 2017) have risen the dominant sequence-to-sequence architecture. Pretrained transformers generally have a context window of 512 (e.g. BERT (Devlin et al.) 2019). T5
(Raffel al. ) or 1024 tokens (e.g. BART (Lewis al. (020b)), which sufficient lengths
for many current conditional generation datasets (XSum; Narayan et al. 2018) (CNN/DM; Nallapati al. 2016). To address inputs between 1024 and 16,384 tokens, specialized long-contexl models sparsity or approximate attention (e.g. Longformer (Beltagy et al. 2020). Performers (Choromanski
al. 2020)), allowing the maximum input length quadruple while remaining computationally feasible. Most long-document document summarization and question question-answering datasets, such as SCROLLS
(Shaham al. 2022). are included in this range.
Yet tasks that involve long narratives, such as book summarization (Kryšciáski et al. 2021). can con- tain inputs exceeding 500k tokens. Figure [ecuree[showr shows the input lengths of several popular summarization and inestion-answering datasets, plotted against common context window lengths; the longest inputs are more than 34 times longer than Longformer's context window.
In these sutremely-longinpue cases, vanilla transformers cannot be simply scaled, as naïve self- attention has quadratic complexity. Long-input transformers usually modify the base architecture, and thus necessitate re-pre-training the model from scratch, which requires significant computational
Now at Google DeepMind
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Index one long input
Retrieved hidden states
Encoder
ANN Search
Decoder Layer
Encode chunks
query
Cross attention
c
d
C
Input:
d
Figure 2: In this example, given LM's encoder's maximum input length is tokens. A 6-token input is encoded in chunks and indexed in index. We inject Unlimiformer into each decoder layer prior to cross-attention. In Unlimiformer, we perform kNN search to select 2-token context for each attention head from the index. This makes cross-attention attend to tokens from the entire input sequence, without adding parameters and without changing the given LM's architecture.
(LED; Beltagy et al. 2020) can leverage pretrained models, but they still need to further train new position embeddings or global attention weights, which is computationally and environmentally costly.
resources. Other architectures such as
We introduce Unlimiformer, a retrieval-based approach to augment pretrained language mod- els to accept inputs of unbounded length at test time. Given a long input sequence, Unlimi- former constructs d--neares--neighbor (kNN) index over the hidden states of all input tokens. Then, every standard cross-attention head in ery decoder layer queries the kNN index, such that the ANN distances are the attention dot- product scores, and attends only to the top-k input tokens.
16384 tokers
1024 tokonn
105
Input tokens
10a
103
preliminary experiments, we found that the top-k attention keys cover more than 99% of the attention mass, and thus attend- ing only to the top-k keys is an accurate approx- imation of the full, exact, attention. Unlimi- former can be injected into any existing encoder- decoder transformer to permit unbounded inputs. The index can be stored in either GPU or CPU memory, needs to hold only a single vector per input token, and can be queried in sublinear time. Unlimiformer is illustrated in Figure 2
XSum (Avg)
CNI/DM (Avg)
Arxiv (Avg)
GovReport (Avg)
WikiSum (Avg)
Narrative (Avi))
BBok (Avg)
NarrativeQA (Max)
BookSum (Max)
WilkiSum (Max)
Figure 1: Long-range transformers can avoid input truncation in some datasets; however, there are datasets with inputs many times longer than these models' maximum input length. The dotted lines represent three common maximum input lengths for models; the bars are the average maximum input length in each dataset, as indicated. Averages for datasets from Koh al. (2022).
Unlimiformer is generic approach: it can be applied to trained models and improve existing checkpoints without adding weights and without further training. When finetuning Unliniformer, performance is even further improved: across a variety of long-range datasets, not only that Unlimiformer performs better than strong long-range -range transformers such LED (Beltagy et al. 2020), PRIMERA (Xiao et al.) 2022), SLED (Ivgi et al. (2022) and Memorizing Transtormers (Wu et al. 2022). but Unlimiformer can be applied on top of such models to further improve them.
2
Unlimiformer
Given a trained encoder-decoder transformer, Unlimiformer allows each cross-attention head to choose separate keys to attend to from the full-length input, at each decoding step. We inject ANN

---

search into each decoder layer: prior to cross- -osssaattnnioon the model performs a nearest -eeihst search KNN index choose a set of per-decoder-layee per-attention-hert head tokens attend to.
2.1 Encoding
To encode an input sequence that longer than the model's context window, we use the given model's encoder to encode overlapping chunks of the input, following Ivgi et al. (2022). We keep only the middle half of the encoded vectors from each chunk, to ensure that the encodings have sufficient context on both sides. Finally, we index the encoded inputs in a kNN index, using a library such as Faiss (Johnson et al., 2019), using dot-product as the index's nearest-neighhor similarity metric.
2.2 Retrieval-augmentte Cross-Attention
In standard wross-attention, a transformer decoder attends to the encoder's top-layer hidden states, where the encoder usually truncates the input and encodes only the k: first tokens in the input sequence.
Instead of attending only this k-token prefix of the input, we retrieve the top-k hidden states from the kNN index for each cross-attention head, and attend only to these top-k. This allows retrieval from the entire input sequence instead of truncating. Our approach is also cheaper, in computation and GPU-memory, than attending to all input tokens; and because softmax is dominated by the largest values, retrieving the aost-dttended tokens preserves the vast majority of attention mass.
Figure 2 illustrates our generic changes to any full input encoded using the encoder in chunks and indexed in kNN index; then, the index of encoded hidden states is queried at each decoding step. The ANN search step is non-parametric and can be injected into any pretrained seq2seq transformer. The search step reformulates attention for space efficiency, as detailed below.
transformer's ' architecture. The
2.3
A Attention reformulation
Let h be the decoder hidden state and he be an encoder's s last layer hidden state. The standard cross-attention computation for single head in transformer is:
( V
Attn(Q, K) softmax
(1)
where Q = h,W, is the product of the decoder states h2 and the query weight matrix Wg; the keys the product of the last encoder hidden states h. with the key weight matrix WR; and
K = W, V = h2 W, is similarly the product of h. with the value weight matrix W,, Our goal is to retrieve set of keys Krest that maximize QK with the size of Kbest fixed the size of the model's context window, and then compute the standard attention over Kbest only.
Note that the linear layers Wg. Wx. and W. are layer-specific and head-specific. Thus, naïvely creating an index from the keys = h W3 and querying this index using the query vectors will require constructing separate indexes for the keys and values each layer and each head, for total of 2 x L x H indexes, where L is the number of decoder layers and H is the number of attention heads. In fact, this exact naïve approach was taken by Memorizing Transformers (Wu al. 2022), who pioneered the use of a ANN index for previously encoded inputs P A separate index for each attention head in each decoder layer is both time-intensive to create and inace-intensive store. So, not surprisingly, Wu et al. (2022) apply their memory layer to only a single decoder layer.
Instead, present different order of computing the well-known transformer attention formula, which allows us to store a single index across all attention heads and all decoder layers, without changing the mathematical definition of the transformer's 's standard dot- dro-prr attention. The
*See Memorizing Transformers' official implementation at bbob/maii/rraasffommrmmmmmmm factory. L79 | and https//gttbbboooooooooooeeeeceea transformer/memorry layer.py#H3344-333

---

dot- cot-product part of the transformer's attention computation can be rewritten as follows
QK" = (haW (k.)
(2)
= (haWg) W he
(haW, h.
Thus, the retrieval step can be formulated as choosing the encoder hidden states h. that maximize (haWWW, he This rewriting has two major advantages: first, there need index the keys for each head and layer separately: we can create a single index of the hidden states h, only, and just project the queries W using head-specific and layer-specific Wg and Wa; second, the values can be calculated trivially given he. so there is no need to store the values in separate index from the keys before decoding. Thus, instead of constructing 2 L x H indexes and retrieving from all indexes during each decoding step, we construct single index from he and retrieve from it by just projecting the decoder hidden states to per-head per-layer h.W,W.
Using our reformulation, the index stores only a single vector per input token. Using 16-bit floats and hidden states of size 1024, this requires only 2GB of memory for 1, 1,000,000 input tokens. Since indexes can be offloaded to the CPU memory, Unlimiformer's input length is practically unlimited.
3 Training Unlimiformer
Unlimiformer can be used, at test time, with an already-trained model, and lead to gains without further training, show later in Table Next, turn focus to training approaches to further and contrasts the training approaches
improve the performance of Unlimiformer. described below, and Appendix contains further implementation details.
total tokens example seen at training time
Validation input (carly stopping)
Method name
Training input
Test input
Bascline
1024
1024
1024
1024
+test Unlimiformer
1024
1024
1024
unlimited
+early stop w/ Unlimiformer
1024
1024
unlimited
unlimited
Train chunked +test Unlimiformer
1024
all
unlimited
unlimited
SLED ([vgi et al. [202
16k
T6k
Tōk
16k
Longformer (Beltagy al. 2020
16k
16k
16k
16k
Sandom-encoden traning
8-16k
8-16k
unlimited
unlimited
Retrieval training
8-16k
8-16k
unlimited unlimited
unlimited
Alternating training
8-16k
8-16k
unlimited
Table 1: A comparison the training approaches using BART (context window size 1024) as a running example. The dashed line separates methods that approximately the same training-time cost as the baseline, from those that require significant additional compute.
3.1 Low (additional-) Cost Training Methods: Applying Unlimiformer at validation or test-time only
We first consider training approaches that do not require significant additional compute compared to the standard finetuning regime.
+test Unlimiformer:
As the simplest case, we use a standard fine-tume tie-tuning regime, where the input is truncated during training. At inference time only, we inject Unlimiformer into the trained model to process full-length inputs.
+early stop w/ Unlimiformer: We train without Unlimiformer, but when we evaluate the model for early stopping, we use Unlimiformer for generation on the validation set. This results in choosing a slightly different checkpoint to stop training at; the additional computational cost here is minor, and comes only from the application of Unlimiformer over the validation set.
For brevity, omit the linear layers' bias term, softmax invariant to constants added to all inputs.

---

Train chunked +est Unlimiformer: As data augmentation approach, we split each training ex ample into non-overlapping chunks of the context-window size, and treat each chunk as its own training example. Then, we finetune the model as normal, with this augmented set of examples as the training data. This is orthogonal to the Unlimiformer model, but has the advantage that all tokens from the full-length training example are observed during training instead of truncated-
across several examples. We apply early stopping with Unlimiformer the validation set; when validating, we do not chunk inputs.
3.2 Long-range Training Methods: Applying Unlimiformer at training time
We also consider training Unlimiformer directly, which introduces additional computational cost.
Random-encoded training: At each training step, the full (longer-than-conter context-windows training example is encoded in chunks; then, the keys for each decoder layer chosen randomly from the encoded hidden states. This weakly simulates nearest-neighhors search, but is computationally cheaper.
Retrieval training: At each training step, the keys for each decoder head and layer are selected using kNN search. When inputs are longer than 16k tokens, truncated the input 16k tokens at training time due to GPU memory requirements. This training approach is the closest to the test-time computation.
Alternating training: In this approach we alternate batches of Random-encoder training and Re- trieval training. Retrieval training is identical to the test-time setting, while Random- encoded introduces regularization that makes the model attend to non-top-k keys as well.
4 Experimental Settings
4.1 Datasets
Avg # tokens
Dataset
Domain
# examples
Input
Output
Input length distribution
GovReport
Government
19,402
9,616
597
74
303192
SummScreen
TV shows
4,348
8.987
137
22635
BookSum
Literature
436
143,301
1294
8388
642376
Table 2: Dataset statistics. The last column is visualization of the distribution of input example lengths in each dataset; the histogram is binned by powers of 2, with the minimum and maximum input size displayed either end. The dotted line indicates the mean length.
We experiment with two long-document- and one somm-auumaiizttion datasets from varying domains. Table 2 summarizes statistics for each dataset. GovReport and SummScreen were taken from the SCROLLS benchmark (Shaham et al. 2022). GovReport (Huang et al. 2021) a long-document summarization dataset where the task to write the executive summary of US government report. SummScreen (Chen et al. 2022) is a long-document summarization dataset where the task is to write the recap of TV show episode (such as "Friends"), given the transcript of the entire episode. BookSum (Kryšcinski et al. 202I) is suomaaaamaiizationn dataset of entire books. BookSum has setting, where
paragraph, chapter, and book- settings; consider the hardest BOOKSUM-BE the task is to generate a book-level summary given the full text of the novel as input.
Metrics We report ROUGE 1/2/L (Lin, 2004) and BERTScore Fl (Zhang al. 2019). Following Zhang et al. (202I). in BookSum we also used Entity Mention Recall ("EntMent" ") as proxy for the informativeness of the candidate summaries. EntMent measured the fraction of gold entities mentioned in the candidate summary. Additional evaluation details are provided in Appendix C
4.2 Baselines
BART (base) (Lewis et al. 2020b a pretrained seq2seq model (139M parameters), commonly used for summarization tasks. Its maximum input sequence length is 1024 tokens.
5

---

Base model
Training method
ROUGE L/ BERTScore
GovReport
SummScreen
BARTbase
Standard finetuning
48.7 19.2/ 22.8 64.3
29.7/62/
BaRTpase
+test SLED lvgi et al. 2022
45.8 16.1 20.2 62.7
27.5
16.7
BARTbase
+lest Unlimi ormer
49.7 19.6/9.6/2200 22.0/ 64.8
30.9 /0.9/6.55 /
BARThase
+early stop w/ Unlimiformer
51.0/20.5/155 65.1
32.1 6.8 18.6 57.6
BARThase
Train chunked
46.2/17.8/21.7
/ 17.0/
BARThase
+lest Unlimiformer
$3.4/22.55 22.5 66.0
29.3 6.6/ 17.6 57.0
PRIMERA
Standard finetuning
/3.9 25.9 67.0
/ 18.3 57.1 33.3/7.77 19.1 57.6
PRIMERA
+lest Unlimiformer
56.5 24.8/ / 26.3 / 67.7
Table 3: Results long-document summarization, low-cost training methods: the training costs no higher than standard finetuning that truncates the inputs to the model's 's max input size. The best metric in every training category marked in bold. PRIMERA (Xiao et al., 2022) is a Longformer- Srooder-Decoder (Beltagy et al. 2020) with additional commarization-speeliie pretraining.
PRIMERA (Xiao et al. 22022) is Longformer-Encodee Encoder-Decoder (LED largei Beltagy al. 2020) (447M parameters), pretrained specifically for multi-document summarization, with maximum imput length of 4096 tokens.
SLED (Igi et al. 2022 extends encoder-decoder models for longer contexts by applying fusion in-decoder (Izacard and Grave, 202I): the long input is encoded in chunks, and the decoder then attends to all input tokens. This allows the use of pretrained models, albeit with expensive fine-tuning. The input sequence length is eventually memory bounded.
Memorizing Transformers (Wu et al., 2022) is the most similar work to ours; they propose extending transformer with trainable attention gate that moderates between the standard cross-attention and attention over retrieved keys from a datastore. Since their public implementation is "not officially supported" and is not fully reproducible, a single decoder layer; this is equivalent to their setting with the learned interpolation parameter g set to 1 Our work differs from Memorizing Transformers in several key ways: Wu et al. (2022) added additional weights, and thus cannot easily leverage pretrained LMs, while Unlimiformer is fully non-parametric and can improve performance without fine-tunin:: further, Wwu et al. (2022) applies retrieval attention to only a single layer because of computational constraints, while our attention reformulation enables the use of Unlimiformer in every decoder layer with individualized retrieval per-head, while still being more efficient than Memorizing Transformers, as we detail in Section 2.3
approximated it by using attention over the index in only
5 Results
5.1
Long Document Summarization
Low-cost training Table 3 shows the results the long-document summarization datasets. First, existing checkpoint without any training (+test Unlimi- former) improves BART by, for example, 1.8 ROUGE-1 points on both datascts, and improves PRIMERA by 1-1.4 ROUGE-1 points. In contrast, without additional training, SLED decreases performance. Thus, Unlimiformer is the only model that can provide benefits without further training.
can see that applying Unlimiformer
Early stop w/ Unlimiformer further improves the base model without any special training: it provides, for example, 3.3 ROUGE-1 points gain on GovReport, while the training computational cost is identical to standard finetuning. Train chunked does not provide benefits on its own; however injecting Unlimiformer applied at test time results in the most significant gains: 7.2 ROUGE- and BERTScore points improvements, while training is as computationally cheap as standard finetuning.
https///itti.oooooooogoeeeeeeete research/meliad
Wu et al. (0022) note that in their experiments that most heads learned value for such that they attended "almost exclusively" to the external memory.

---

Base model
Training method
ROUGE L/2/ BERTScore GovReport
SummScreen
DARTpaae
Standard finetuning
48.7 22.8 64.3
BART baRe
SLED (Ivgi et al., 2022)
54.7 /4.4 /2..2/222/////11 25.4
32.7 19.1//99911111
BARTpase
Memorizing transtormers
55.2 222/255///20 /6.4//26.667.
19.4/99//555 3 19.9 58.5
BART BARTaase
Unlimiformer (this work)
56.6 / 27.6/27.6/00
PRIMERA
Standard finetuning
55.1/ /
32.3/ 18..11//133//3
PRIMERA
Memorizing transformers
57.0 25.3 26.5 67.7
33.0/ /.30/733/ 18.3 33.3 7.6 18.9 57.7
PRIMERA
Unlimiformer (this work)
57.4 / 26.2 28.0 / 68.1
Table 4: Test results on long- oonu-eecmennt datasets, when allowing compute-costly, long training methods, using different base models. The best metric in every dataset and every training category marked bold. The Unlimiformer results in this table are from using the alternating training strategy.
Base model
Training method
ROUGE 1 /2/L
EntMent
BART
Hierarchical (Kryšciáski et al. 2021
30.0/6.0/ 11.0
BART
Standard finetuning
36.4/7.6/ 15.3
10.0
BART
+est Unlimiformer
35.5/7.7 15.4
21.9
BART
+early stop w/ Unlimiformer
35.5 15.4
21.9
BART base
Memorizing Transformers
35.6/6.4/ 14.6
10.1
BARTBESe
Unlimiformer (retrieval training)
36.8/8.3/15.7
20.3
BART base
Unlimiformer (random-encoded training)
37.3/6.7/15.2
20.8
BART
Unlimiformer (alternating training)
36.7 15.5
20.3
PRIMERA
Standard finetuning
38,6/7.2 15.6
11.6
PRIMERA
Hest Unlimiformer
38.3 15.9
18.9
PRIMERA
+early stop w/ Unlimiformer
39.5 9..5/733/ 15.8
22.2
PRIMERA
Unlimiformer (retrieval training)
37.9 8.22/ 16.3
25.5
PRIMERA
Unlimiformer (random encoded training)
39.5 1/5/7.1/550
19.7
PRIMERA
Unlimiformer (alternating training)
38.2 0
23.4
Table 5: Results BookSum (average imput length ss 143k tokens). EntMent is entity recall. Hierar- chical summarization a baseline reported by Kryšcinski et al. (202I). where chapter summaries are condensed to form a book summary. The best metric in every dataset is marked in bold.
shows results when allowing computationally expensive training approaches. As shown, in almost all metrics and datasets, Unlimiformer outperforms the SLED and Memorizing Transformers baselines when using the same base model.
Long-range training Table 4
The PRIMERA experiments in Table hhightiht two important points: first. Unlimiformer+ +BAttorrmmetBBtBBTBT base performs better than the base PRIMERA across all metrics and datasets, even though PRIMERA is larger and was pretrained much more data, using pretraining objective that designed for summarization; second, not only can Unlimiformer outperform Longformer-bant PRIMERA, Unlimiformer improve them: Molimiforme:+PRIMEE
models such as also be applied top of existing long-range transformers and further
improves over PRIMERA across all metrics and datasets. Additional results the validation are provided in Appendix
5.2 Book Summarization
Table shows the result BookSum. As shown, Unlimiformer improves both base models BART buse and PRIMERA, in both low-cost training approaches such as Early stop w/ Unlimiformer, as well as in the rangee training approaches. Random- Random-encoded, encoded- Retrieval-, and Alternating- training show competitive performance, with the best method varying across datasets and models.
We found that although Unlimiformer outperforms all base models on BookSum (Table the base BART (Standard finetuning, which truncates the input to the first 1024 tokens) shows competitive ROUGE and BERT Score metrics. This is strongly counterintuitive for book summarization, where

---

25
Relative time per example
Unlimiformer
Entity Mention Recall
BART base (truncates to 1024)
base
20
15
4
10
Unlimiformer
BART base
2
8
16
32
64
100
350
0 10
20
30
40
50
60
70
80
90
100
Max input size (in thousands tokens)
Max input size (in thousands of tokens)
As the maximum datastore size in- creases, the entity recall generally increases. At all datastore sizes, Unlimiformer outperforms the baseline (BART, in red).
Figure 3:
Figure 4: creases, the inference cost increases sublinearly. The plot shows total wall-clock inference time per example.
As the maximum datastore size in-
the book's s plot should not be apparent from reading only the first pages. In the outputs from this base model, we observe limited coherence and a high rate of hallucination (see Appendix For example with analysis). However, this is not reflected in n--ram-based overlaps, and BERTScore does not strongly distinguish between any of the BookSum models.
Nonetheless, the ability to attend to unlimited inputs at test time allows Unlimiformer to achieve significantly better Entity Mention Recall (EntMent): the Unlimiformer models exhibit far higher EntMent, and even adding Unlimiformer only at test time without costly training (Early stop w/ Unlimiformer) doubles the entity recall compared to the base model. Further, Unlimiformer improves EntMent in the base PRIMERA from 116 to 25.5 in Unlimifomme++PRIMRAA
6 Analysis
Is the long input really needed? As found in various recent papers (Shaham et al. 2022 Kedzic al.) (0188), many text generation datasets do not require long-range modeling, since most of the needed information is concentrated at the beginning of the input. To evaluate whether Unlimiformer really utilizes long inputs. we experimented with limiting the input length in BookSum. Figure shows the performance of Unlimiformer in BookSum: EntMent increases almost monotonically with input length, suggesting Unlimiformer exploits the longer inputs generate better outputs.
Other work (Jiang and Bansal, 2019) has found that in some datasets, the needed information is concentrated in only part of the input, which is necessarily the beginning. We observed this trend in WikiSum, a multi-document summarization dataset where the inputs are all references of Wikipedia article and the output summary the intro paragraph of the article (Liu* al. 2018 As strong baseline, we followed Liu et al. (2018), and ranked the input paragraphs according to TF-IDF. Unlimiformer did not improve over baseline that uses only the first 1024 tokens of this sorted input, suggesting that the full input is not necessary to produce the summary on this datase
Computational cost Although Unlimiformer does not introduce additional trained parameters, the encoding of the full input, index construction, and index search increase the processing time during both training and inference. We plot the computational cost of inference with respect the input length in Figure 4 When all inputs are restricted to 1,024 tokens, Unlimiformer requires a small additional time overhead relative the baseline for indexing and search. However, the benefits of
A full copy of WikiSum is not available online; details of scraped copy Appendix 'It is also possible that the baseline's strong performance on this task is due to some contamination from pretraining data: BART was trained on books and Wikipedia data, which likely includes some articles used in WikiSum. Any contamination from pretraining here strengthens the baseline performance.

---

QASPER
Contract NLI Exact Match
QMSum
Narrative QA
Base model
Training method
ROUGE /
F1
BARTbase
Standard finetuning
22.0
77.5
30.8/ 20.8
BARTbase
Unlimiformer
27.5
77.7
30.9 8.0 19.9
18.5
Table Results on question answering, query-based summarization, and NLI datasets.
Unlimiformer are clear as input length increases: the total GPU-time required increases sublinearly with imput length | Additional GPU-time measurements are reported in in Appendix D
Performance on other tasks We measure the performance of Unlimiformer relative to the base model on 4 additional datasets: QASPER (Dasigi al. 20211. question-answering dataset over NLP papers; Contract NLI (Koreeda and Manning/ 2021). a natural language inference dataset legal contracts; QMSum (Zhong et al.] 2001). query- summarization dataset meeting transcripts; and NarrativeQA Kočiský et al. 2018). a reading comprehension dataset over narratives
Table shows the performance BART-Unlimiforne (with alternating training) relative base BART. On three of the four datasets, applying Unlimiformer improves over the basc model.
What is attended to?
We plotted the frequency of retrieval for keys across the full decoding process for the test of BookSum, the dataset with the longest in- puts. The average number of input embeddings retrieved at least once varied by method, from 43.5% of all tokens for the test-time-only Unlim- iformer to 64.5% of all tokens for the alternating- training mode
% totel retneval ocouring acauin in this region
10.0%
6.0%
4.0%
Figure 5 shows the retrieval locations for the alternating- training model. We found no specifie skew or pattern in the retrieved keys, and keys from the entire input were used by the model; for all models, the median location of a retrieved key was between 49.73% and 49.87% of the way through the input document.
2.0%
0.0
0.2
0.4
0.6
0.8 1.0
locaaion decument
Figure 5: Histogram of location of retrieved em- beddings in the original document (averaged over the BookSum test set). There is a slight bump at the beginning (first 10% of the book), but other- wise no strong trend, with tokens retrieved quite uniformly from the entire inputs.
Related Work
Long-range transformers Previous long-range
transformers change the transformer architecture to reduce its space or time requirements TTy et al.
2020 Most solutions achieve this reduction through sparsifying the attention mechanism (Child
et al. 2019 Kitaev al. 2020 [Beltagy et al., 20200 Roy et al. 2020 Ainslie et al.) 2020; Zaheen
al. 2020). Other works approximate or replace the attention mechanism entirely (Wang et al.) 22220]
Katharopoulos et al. 2020 Choromanski et al.) 2020 Lee-Thorp et al. p. All these approaches change the standard transformer architecture or its training objective (Zhong al., 2022), and thus require pretraining the model from scratch, which does not allow to leverage existing pretrained models. In contrast, Unlimiformer generic, can injected into any encoder-decoder transformer, and improve it either without training with merely fine-tuning. This way, Unlimiformer leverage any preaad--pretraied model.
Comparison to Wu et al. (2022) The closest work to ours is Memorizing Transformers (Wu et al. 2022). Memorizing Transformers construct two datastores for each attention head in each layer,
"The time to encode the input increases linearly and the time to decode increases sublinearly; because the encoding is only a small part of the total time during inference, this results in an overall sublinear trend.
JeWe report results on validation split from SCROLLS Shaham al. (2022).
Note that this is not the percentage of the input that influenced the output, we retrieve encoded, contextualized hidden states, vectors that were not directly retrieved by the decoder impacted the final output.

---

and due to memory constraints can thus apply their approach only to single decoder layer. In contrast, thanks our attention reformulation (Section (.3) Unlimiformer can use a single index for all decoder layers, and thus allow all cross-attention heads in all decoder layers retrieve from the long context. As we show in Section 5] this results in significant empirical gains over retrieving only at a single layer. Further, Memorizing Transformers introduce additional learned weights, thus they must be trained to incorporate their memory, and thus cannot easily leverage pretrained models; as we show in Section Unlimiformer can improve existing models without any training, and thus can be applied to any existing transformer. Additionally, Memorizing Transformers focused on decoder-only models; while our approach could also be applied to decoder-only models (and would provide a space efficiency boost there as well). we focus on encoder-decoder models in this work.
Comparison to Ivgi et al. (2022) Another related work to ours is SLED (Ivi et al. 2022). SLED encodes long inputs in chunks, similarly to Unlimiformer, but the decoder in SLED attends to all inputs at the same time. This in practice limits SLED to only about 16k token-long inputs on 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, while preserving more than 99% of the attention mass. Further, SLED requires computationally costly training, while Unlimiformer can provide benefits without any training.
Additional related work is discussed in Appendix G
8 Conclusions
We present Unlimiformer, an approach for augmenting pretrained encoder-decoders and offloading the cross-attention computation attending to all keys, this ANN index allows every cross-attention head in every decoder layer to retrieve and attend only to its top-k keys. We evaluate Unlimiformer on several long-document and book-summarizztion benchmarks having inputs of up to 500K tokens, and show that Unlimiformer improves existing models, even without further training. When training with Unlimiformer, not only that Unlimiformer makes smaller models such as BART perform better than larger Longformer-basee models, Unlimiformer can be applied on top of Longformer- based models and further improve them.
ANN index, allow for unlimited length input. Instead of
Many real-world NLP tasks require processing large amounts of data or text. Yet pretraining large models incurs substantial carbon costs (Strubell et al., 2019), which increase with the length of the context window; by choosing instead to modify inputs, we aim gain the benefits of long contexts with less computational cost. We hope that our approach will allow the democratization of long-range transformers, especially for researchers and practitioners with low-compute resources. Toward this end, we release our code at https: cetthub.coo/abberssde /ebbertcch//mmmmmfoooneee
prretrriee models to process longer
Our code is based on HuggingFace Transformers (Wolf ct al.) (020)) without changing any individual architecture's 's code, and thus can be injected into any encoder-decoder model, and supports decoder models such as LLaMA-2 as well.
Limitations
In our experiments, we have only considered tnglish-language datasets. While have reason to believe the method would suffer from the use of a different high-resourced language, the quality of the nearest-neighhors search depends on the quality of the indexed keys.
The length of inputs that can be used at training time is limited by the GPU memory, as the embeddings and their computational graph must be stored for backpropagation. Multi-GPU training would allow longer inputs at training time.
At inference time, Unlimiformer can process the longest inputs when the index is offloaded to the CPU memory. In this case, Unlimiformer requires to index only a single vector per input token, which practically means unlimited inputs for any modern server and even small machines during inference. However, offloading the index to the CPU results in higher test-time latency compared to storing the encoded hidden states and the index on the GPU. In our experiments, we were able to use a GPU index for input examples exceeding 500k tokens (on GPUs no larger than 48 GBs), but this may be concern when using smaller GPUs larger models.
10

---

Acknowledgments
We are grateful to Sireesh Gururaja for useful feedback draft of this paper. We also thank Maor Ivgi for the help in reproducing results from SLED (Ivgi et al. 2022) and for sharing code and models, and Uri Shaham for the discussions about the SCROLLS benchmark (Shaham al. 2022). We are also grateful to the anonymous reviewers for their useful comments and suggestions.
This work was supported in part by grants from 3M- M*Modal and from the National Science Foundation Graduate Research Fellowship Program under Grant No. DGE2140733 Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the sponsors.
References
Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. 2020. Etc: Encoding long and Structured inputs in transformers)
Uri Alon, Frank F. Xu, Junxian He, Sudipta Sengupta, Dan Roth, and Graham Neubig. 2022. Neuro symbolic language modeling with cettomatoo--ugmmentin retrieval
Ahsaas Bajaj, Pavitra Dangati, Kalpesh Krishna, Pradhiksha Ashok Kumar, Rheeya Uppaal, Bradford Windsor, Eliot Brenner, Dominic Dotterrer, Rajarshi Das, and Andrew McCallum. 2021. Long
document summarization in a low resource setting using pretrained language models] In Proceed- ings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing: Student Research Workshop, pages 71-80, Online, Association for Computational Linguistics.
Iz Beltagy, Matthew E Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. arXiv preprint w0xv::0004005/111
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean- Baptttttttt Lespiau, Bogdan Damoc, Aidan Clark, et al. 2022. Improving language models by retrieving from trillions of tokens. In International conference on machine learning, pages 2206-2240 PMLR.
Mingda Chen, Zewei Chu, Sam Wiseman, and Kevin Gimpel. 2022. SummScreen: A dataset for abstractive screenplay summarization] In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 8602-8615, Dublin, Ireland. Association for Computational Linguistics.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. 2019. Generating long sequences with sparse transformers)
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. 2020. Rethinking attention with performers)
Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A. Smith, and Matt Gardner. 2021. A dataset of nnformation-seekinn questions and answers anchored in research papers) In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 4610, Online. Association for Computa- tional Linguistics.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language understanding In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota. Association for Computational Linguistics.
Andrew Drozdov, Shufan Wang, Razieh Rahimi, Andrew McCallum, Hamed Zamani, and Mohit lyyer. 2022. You can't 't pick your neighbors, or can you? when and how to rely on retrieval in the KNN-LM In Findings of the Association for Computational Linguistics: EMNLP 2022, pages 2997-3007, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics.
11

---

Quentin Grail, Julien Perez, and Eric Gaussier. 2021. Globalizing BERT-basee transformer architec tures for long document summarization In Proceedings the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume, pages 1792-1810, Online. Association for Computational Linguistics.
Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu Wang. 2021. Efficient attentions
for long document summarization, In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1436, Online. Association for Computational Linguistics.
Maor Ivgi. Uri Shaham, and Jonathan Berant. 2022. Efficient long-text understanding with short-text models
Gautier Izacard and Edouard Grave. 2021. Leveraging passage retrieval with generative models
for open domain question answering In Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume, pages 874-880, Online. Association for Computational Linguistics.
Yichen Jiang and Mohit Bansal. 2019, Avoiding reasoning shortcuts: Adversarial evaluation, training and model development for multi-hop QA In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2726-27366 Florence, Italy, Association for Computational Linguisties.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2019. Billion-scale similarity search with GPUs.
IEEE Transactions on Big Data, 7(3):535-547.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Flcuret. 2020. Transformers are mns: Fast autoregressive transformers with lincar attention
Chris Kedzie, Kathleen McKeown, and Hal Daumé III. 2018. Content selection in deep learning models of summarization/ In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 1818-1828, Brussels, Belgium. Association for Computational Linguistics.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2019, General ization through memorization: Nearest neighbor language models
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. 2020. Reformer: The efficient transformer
Tomá's Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. 2018. The NarrativeQA reading comprehension challenge Transactions of the Association for Computational Linguistics, 6:317-328.
Huan Yee Koh, Jiaxin Ju, Ming Liu, and Shirui Pan. 2022. An empirical survey on long document Summarization: Datasets, models, and metrics ACM Comput. Surv., 55(8).
Yuta Koreeda and Christopher Manning. 2021. ContractNLI: A dataset for document-level natural language inference for contracts In Findings of the Association for Computational Linguistics: EMNLP 2021, pages 1907-1919, Punta Cana, Dominican Republic. Association for Computational Linguistics.
Wojciech Kryšciáski, Nazneen Rajani, Divyansh Agarwal, Caiming Xiong, and Dragomir Radev, 2021. Booksum: collection of datasets for long-form narrative summarization
Juho Lee, Yoonho Lee, Jungtack Kim, Adam R. Kosiorek, Seungjin Choi, and Yee Whye Teh. 2019. Set transtormer: A framework for attention-based permutation-invarian neural networks
James ThorpT Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. 2021. with fourier transforms
Fnet: Mixing tokens
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. 2020a. Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 7871-7880.
12

---

Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov. and Luke Zettlemoyer. 2020b. BART: Denoising sequence to-eueence pre-training for natural language generation, translation, and comprehension In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 7871-7880, Online. Association for Computational Linguistics.
Chin-Yew Lin. 2004. ROUGE: A package for automatic evaluation of summaries In Text Summa- rization Branches Out, pages 74-81, Barcelona, Spain. Association for Computational Linguistics.
Peter J. Liu*, Mohammad Saleh*, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. 2018. Generating wikipedia by summarizing long sequences) In International Conference on Learning Representations.
Yang Liu and Mirella Lapata. 2019. Hierarchical transformers for multi-document summarization, In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 5070-5081, Florence, Italy. Association for Computational Linguistics.
Ramesh Nallapati, Bowen Zhou, Cicero dos Santos, Çağlar Gulçehre, and Bing Xiang. 2016. Ab stractive text summarization using sequence-to -equence-to-eqquence RNNs and beyond In Proceedings of the 20th SIGNLL Conference on Computational Natural Language Learning, pages 280-290, Berlin, Germany, Association for Computational Linguistics.
Shashi Narayan, Shay Cohen, and Mirella Lapata. 2018. Don't give me the details, just the summary! ! topic-aware convolutional neural networks for extreme summarization In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 1797-1807, Brussels, Belgium. Association for Computational Linguistics.
Guanghui Qin and Benjamin Van Durme. 2023. Nugget: Neural agglomerative embeddings of text International Conference Machine Learning.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2020. Exploring the limits of transfer learning with unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5885-551.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. 2020. Efficient content-based sparse attention with routing transformers
Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. 2022. SScrolls: Standardized comparison over long language sequences
Emma Strubell, Ananya Ganesh, and Andrew McCallum. 2019. Energy and policy considerations for deep learning in NLP In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 3645-3650, Florence, Italy. Association for Computational Linguisties.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. 2020. Efficient transformers: A survey
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, t ukasz Kaiser, and Illia Polosukhin. 2017. Attention all you need Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
Apoorv Vyas, Angelos Katharopoulos, and François Fleuret. 2020. Fast transformers with clustered attention
Shuohang Wang, Luowei Zhou, Zhe Gan, Yen-Chun Chen, Yuwei Fang, Siqi Sun, Yu Cheng, and Jingjing Liu. 2021. Cluster-former: Clustering- sparse transformer for question answering In Findings of the Association for Computationall Linguistics: ACI-IJCNLLP 2021, pages 1958-3968, Online. Association for Computational Linguistics.
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. 2020. Linformer: Self-attention with linear complexity
13

---

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick Von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. 2020. [Transformers: Ofate-of-the-art natural language processing In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38-45, Online. Association for Computational Linguistics.
Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. 2022. Memorizing transformers In International Conference on Learning Representations.
Wen Xiao, Iz Beltagy, Giuseppe Carenini, and Arman Cohan. 2022. PRIMERA: Pyramid-based masked sentence pre-training for multi-document summarization In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics ( (Volume 1: Long Pupers), pages 5245-5263, Dublin, Ireland. Association for Computational Linguistics.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago On- tanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. 2020. BBg bird: Transformers for longer sequences
Longxiang Zhang, Renato Negrinho, Arindam Ghosh, Vasudevan Jagannathan, Hamid Reza Has- sanzadeh, Thomas Schaaf, and Matthew R. Gormley. 2021. Leveraging pretrained models for aattmattic summarization of doctor doctor-patient patient conversations In Findings of the Association for Computational Linguistics: EMNLP 2021, pages (6993-3712 Punta Cana, Dominican Republic. Association for Computational Linguistics.
Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q Weinberger, and Yoav Artzi. 2019. Bertscore: Evaluating text generation with bert
Yusen Zhang, Ansong Ni, Ziming Mao, Chen Henry Wu, Chenguang Zhu, Budhaditya Deb, Ahmed Awadallah. Dragomir Radev, and Rui Zhang. 2022. (Summ": A multi-stage summarization frramwork for long input dialogues and documents) In Proceedings of the 60th Annual Meeting of The Association for Computational Linguistics ( (Volume I: LLong Pupers), pages 1592-1604, Dublin, Ireland. Association for Computational Linguistics.
Ming Zhong, Da Yin, Tao Yu, Ahmad Zaidi, Mutethia Mutuma, Rahul Jha, Ahmed Hassan Awadallah, Asli Celikyilmaz, Yang Liu, Xipeng Qiu, and Dragomir Radev. 2021. QMSum: A new benchmark for query-based mnulti-domal meeting summarization In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, pages 5905-5921, Online. Association for Computational Linguistics.
Zexuan Zhong, Tao Lei, and Danqi Chen. 2022. Training language models with memory augmentation. In Empirical Methods in Natural Language Processing (EMNLP).
14

---

A Training details
At training time, we must backpropagate through the operations described above. Thus, the input length bounded more strictly the number tokens the full input must in GPU memory while the model is loaded. For the computationally expensive methods, train using batch size 1 and truncate the longest inputs (generally, to 16k tokens). At test time, we use the full input without truncation. We train one model per setting, using the hyperparameter settings from SLED (lvgi et al. 2022) and early stopping.
B WikiSum scraping
We rescraped the dataset, following the same preprocessing steps as the original authors. We observe that many inputs in the scraped dataset are shorter than reported, likely due to changes in availability of the data since 2017; as a preprocessing step, we remove all inputs that are less than 1457 words, which is the 40th percentile of citation size for the original dataset. We trained on 10,000 randomly selected examples from this version of WikiSum and evaluate 2,000 randomly sampled examples (1,000 validation, 1,000 test). maintaining the same sample across all experiments. When sampling, respect the original WikiSum train/validen veinn/alidation/test split. We release the subset trained on as well as our modified version of the scraping code.
C Evaluation details
Vanilla BERTScore is only well-defined up to 512 tokens; for GovReport and ScriptSumm, evaluate using Eaceboos/paart-aargeeeeere
instead. This model has context size 1024. For BookSum, experimented with using allenai/ woogooormeeeeeeeeeeee -areeer--aa--0000 6 (context size 4096), as many references are longer than 1024 tokens; however, we found that this approach had no distinguishing power between model outputs, ranking all models tested within 0.3 points of each other despite observing significant differences with ROUGE, EntMent, and manual inspection.
For computing Entity Mention Recall (EntMent), we used SpaCy to tag all named entities in gold summary and collected a set of unique entities. We then tagged each candidate summary and computed the percentage entities present in this summary, that is, recall of unique entities. For the named entity recognition in EntMent, used SpaCy's en_core web 1g model.
12
D Computational Cost
We estimate the total GPU time results presented this paper did not exceed approximately 116 days of time single 48-GB A6000. The oonget-traiiingg models, SLED and retrieval training for GovReport, took approximately 10 days to train.
GPU-time Table shoows the relative cost for each method. The Unlimiformer training methodolo- gies higher cost than the base training: however, the largest difference occurs during inference, where the full input (in Booksum, an average of 112,885 tokens) must be encoded, instead of the 1,024 tokens encoded in the baseline approach.
Using CPU datastore many times slower than GPU datastore because of slower search and the need to transfer retrieved embeddings to the GPU. In our experiments, we were able to use GPU datastore for input examples exceeding 500k tokens (on GPUs no larger than 48 GBs), but this may be concern when using smaller GPUs or even larger inputs. Additionally, CPU indices necessary for models with context windows larger than 2048 tokens, as the Faiss GPU index implementation does not support retrieving more than 2048 nearest neighbors; however, the datastore can still be stored on GPU.
https:///ppay.io
15

---

Method
Relative GPU-time
Bascline training
1.00 0.00
Chunked training
1.02 = 0.02
+early stop w/ Unlimiformer
1.00 0.00
Retrieval training
1.89 = 0.06
Random-enee enccoedd training
2.87 = 0.28
Baseline inference
1.00 = 0.00
Unlimiformer inference
4.48 0.56
Table 7: Computational effort per epoch for different training methodologies, relative to the baseline of standard finetuning and inference. All are averaged over runs on BookSum using single 48 GB A6000 GPU, 32 GB RAM, and 16 CPUs.
Number of layers using Unlimiformer
Memory consumption (GB)
0 (normal inference)
1.61
1
7.33
2
7.32
3
7.36
4
7.32
5
7.33
6 (all)
7.35
Table 8: Memory consumption for applying Unlimiformer at different numbers of layers in BART.
Memory usage Table showws the memory required to apply Unlimiformer on varying numbers of layers in a BART model. Using Unlimiformer requires more memory than the base model for two reasons: index construction and the additional input processed. There a slight overhead for constructing and storing the index. But more erucially, the base BART ("normal inference") is truncating the input to the first 1024 tokens. When Unlimiformer is used, we process the full inputs, some of which are >500,000 tokens, and much of the additional memory cost comes from storing the additional hidden states. However, the GPU memory consumption remains constant even when use Unlimiformer in more layers, highlighting the scalability of Unlimiformer compared other approaches such as Memorizing Transformers, which need to allocate more memory with every layer and every attention head.
E
Validation Results
Table 9 shows the validation metries for GovReport and SummScreen.
F
Sample Outputs
These outputs from BookSum are summaries of The Brothers Karamazov, an elaborate novel about Russian family. Neither summary is fully factually correct, but the summary from the input-truncated model hallucinates several plotlines (e.g. a lover from the Congo, the many deaths of Pavel) which are not present the original. The hallucinations in the Unlimiformer output are more constrained; for instance, it incorrectly describes Dmitri as "nobleman" instead of landowner and says he has been sentenced death instead of jail. This summary features more of the novel's s characters and identifies plot details from the later parts of the book, such as Dmitri's trial.
Gold (reference) summary:
The Brothers Karamazov is family tragedy centered around father and his sons. Fyodor, the eldest Karamazov, has three sons: Dmitri, Ivan, and Alyosha. Ivan and Alyosha have the same mother, but Dmitri, the oldest, has a different mother. Fyodor is greedy landowner, bawdy lecher, and a neglectful father. Hence, the Karamazov brothers end up growing into young men under the care of various
16

---

Base model
Training method
ROUGE 1/2/L/ L/ BERTScore GovReport
SummScreen
Low-cost training methods:
BARTEase
Standard finetuning
47.7 / 18.5
30.0/6.5 17.7 56.7
BARTesse
+estt SLED
46.0 / 16.3 20.3
28.4/5.9/11.0/
BART base
+eett Unlimiformer
49.5 / 19.6/ 21.9/
31.8/7.1/ 57.8
BART BARTease base
+early stop w/ Unlimiformer
51.0 / 20.6/21.6/65)
32.5/7.2/199/ 57.9
BART base
Train chunked
48.3 / /18.1/22.3// /3.88
/.4/6.//7.6//5
BART base
+test Unlimiformer
52.9 / 22.2
29.4/ /
Long-range training methods:
BART base
SLED (Ivgi et al. 20222)
55.5 / 24.8 25.8 66.9
19.2/ 1.2/8/./19/2/88
BARTesse
Memorizing Transformers
55.8 / 25.6 26.9 67.7
32.8 /../ /9.3/57.7
BART base
Unlimiformer
57.4 26.4 /7.9 68.2
35.0 / 8.3 19.6 58.4
Low-cost training methods:
PRIMERA Standard finetuning
55.0 23.6
33.0 / 7.8 /
PRIMERA
+test Unlimiformer
56.4 / 24.7 / 26,4 67.6
33.1 1
PRIMERA
+early stop w/ Unlimiformer
56.4 / 25.0/ /6.4/ 67.6
33.5 8.2 1933
Long-range training methods:
PRIMERA
Memorizing transformers
57.0/ / 25.6 / 26.8 67.8
32.9 / 7.7 /
PRIMERA Unlimiformer
58.0 / 26.5 / 28.6 / 68.3
34.1 / 7.9 19.0 57.8
Table 9: Validation results on long-document datasets (average input length between 4k to 16k tokens). The best metric in every dataset and every training category is marked in bold.
other people. But they all have returned home visit their father, and it is the first time they all have been together for quite some time. Dmitri has a dispute with Fyodor over his inheritance, and Alyosha, who is living in monastery, suggests that they see Father Zossima, Alyosha's 's mentor. Alyosha believes that the wise old man can settle the dispute peacefully. Father Zossima is patient and kind, but Fyodor and Dmitri end up quarreling anyway. After Fyodor drives the men to frustration, they leave the monastery separately, and Alyosha worries about their family's 's future. Alyosha talks to Dmitri, who confesses his complicated situation with women and money. Dmitri promised to marry girl named Katerina, and she lent him 3,000 rubles. Instead of paying it back, he spent it on another girl named Grushenka. He wants to run away with Grushenka, but he feels that needs to pay Katerina back before he can do so. This is why he is so interested in getting the money from Fyodor. Back at Fyodor's house, Smerdyakov is talking to the Karamazovs. Smerdyakov is an epileptic servant who Marfa, Fyodor's other servants. He was born to a woman named Lizaveta who died in childbirth. She was the town idiot, and she lived off charity from the other townspeople. Everyone called her "Stinking Lizaveta," " and when the town found out she was pregnant, they were furious whoever could do such thing to a helpless girl. They decided Fyodor must have been the culprit. Grigory and Marfa gave birth deformed child, and when they buried the child, they found Lizaveta, who had just given birth to Smerdyakov. They adopted the child immediately, and Fyodor named him. Father Zossima is dying, and Alyosha is distraught. Instead of asking Alyosha to stay with him during his last days, however, Father Zossima tells Alyosha he should leave the monastery to be with his family. His life gets even more complicated when a young crippled girl named Lise expresses that she has feelings for him. Alyosha visits Katerina, the girl who is engaged marry Dmitri. Ivan is in love with her, but he feels that Dmitri is a better match for her. Frustrated and disgusted with his family's 's situation, Ivan says he going leave town. Alyosha sees boy being picked on by his schoolmates, and he tries to talk to the boy, but he bites Alyosha's hand and runs away. Later, when Alyosha is bringing money to a man named Captain Snegiryov, who has been beaten by Dmitri, he recognizes the man's son. is Ilusha, the boy who bit his hand. The family is poor, but Captain Snegiryov refuses to take the money because he feels that he needs
adopted by Grigory and
17

---

to earn his son's respect after being humiliated by Dmitri-and accepting charity, especially from Karamazov, out of the question. When Alyosha goes back see Katerina, he finds Lise, Madame Hohlakov's 'ss daughter. The two realize that they love each other, and they decide to get married. Alyosha goes to visit Ivan, and he finds him in a restaurant. Ivan has gone there to get away from his father, and Alyosha sits down with him to have an intimate talk. Ivan tells his brother about his thoughts regarding God and the world. He recites to Alyosha a poem he has written called "The Great Inquisitor." The poem describes Christ returning to earth in the sixteenth century. The Church throws him in jail, and The Great Inquisitor explains to him that his presence is problematic for the world. The Church has spent years trying to replace the sense of freedom Christ gave man with security. He talks about how cruel the world is, especially to innocent children. After their meal, Alyosha and Ivan part ways, feeling closer than ever. Ivan sees Smerdyakov when he goes back to his father's ' house, and Smerdyakov tells him he is worried about Fyodor. He is worried Dmitri will come kill him and the old man will be helpless to save himself. Ivan goes to sleep very troubled. Father Zossima is on his deathbed, and Alyosha goes to visit him. The Elder tells those around him how much Alyosha reminds him of his older brother, a boy who died when he was a youth. He talks about being a profligate youth in the army. One day, he challenged another man to a duel because of girl. Before the duel, however, he had a change of heart. He did not shoot and, after the duel, he retired from the army and joined a monastery. He talks about how much the Bible has affected him and says that everyone should embrace the world and the people in it. He dies. Many predicted that miracle would happen upon Father Zossima's 's death, but his body begins to putrefy, filling the monastery with an awful smell. This fills the other monks with doubt that Father Zossima was the saintly man they thought he was. Alyosha is shaken by the news. He goes to see Grushenka, who has sent for him, and she admits to wanting to "ruin" him. When he tells her that Father Zossima has died, however, she becomes contrite about her callousness. She says she thinks she is wicked person, and the two comfort each other. When Alyosha leaves, he has a renewed faith in Father Zossima and his teachings because Alyosha feels how wonderful it is to love and be loved in return. Meanwhile, Dmitri has become desperate. He wants to be with Grushenka, but he wants to pay Katerina back first. He goes
odyssey, hoping that he can depend the charity of others. He visits a man named Samsanov, a man who used to pursue Grushenka, and he hates Dmitri. He sends Karamazov to see surly drunk, tricking Dmitri into thinking this man may be helpful. The man is practically incoherent, however, and Dmitri goes find Madame Hohlakov. She tells Dmitri that the only way he will find 3,000 rubles is in the gold mines. In confusion, Dmitri concludes that Grushenka has gone to visit his father, and he goes to his father's house in a rage. carrying a brass pestle. When he arrives, he does not find Grushenka, but as he is leaving, Grigory, his father's s servant, thinks he has come to murder Fyodor. The two scuffle, and Dmitri hits Grigory on the head with the pestle. After determining that the man is not dead, Dmitri flees the scene and looks for Grushenka. She is with Kalganov, a former lover who had treated her poorly. Dmitri decides that he will not end up with Grushenka and decides to kill himself after seeing her one more time. He crashes her party and sits down with her gentleman friend and some other men. The situation becomes tense, and after the gentlemen make some disparaging remarks about Russians and Dmitri, Grushenka decides she does not want to be with such an insulting and vicious man. She decides that she loves Dmitri, and as the two are coming to terms with their love, the police come to arrest him for the murder of Fyodor. As the police question Dmitri, it becomes clear that the facts all support the conclusion that he did indeed murder his father, even though he did not commit the crime. He was at the scene of the crime, wielding a weapon, the night of the murder. He had said he would kill his father on several occasions. He publicly announced he was looking for 3,000 rubles and was desperate to find them, and Fyodor reportedly had envelope with 3,000 rubles that was stolen the night of the murder. Dmitri is carried away, and very few people believe that he is innocent of Fyodor's murder. Mcanwhile, Alyosha is visiting Ilusba, the boy who bit his
18

---

hand, in the hospital. The boy has fallen quite ill, and Alyosha has gotten to know many of the boy's friends, who are also visiting him. One boy, Kolya Krassotkin, is leader among the boys. He and Ilusha were friends, but they had a falling out because Ilusha fed a pin to a dog, and Kolya did not approve of his cruelty. When Alyosha comes to visit, he and Kolya talk for quite some time. The boy looks up to this wise man about which he has heard so much from the other boys, and he wants to impress him. The two become friends, and Alyosha treats all the boys as equals. When Kolya goes in to see Ilusha, he gives him a dog as a present. He reveals that the dog is none other but the dog Ilusha gave the piece of bread with a pin in it. Kolya has nursed the dog back to health and has fully trained him as a gesture of friendship to llusha. The mood is dampened, however, when the doctors go in to see Ilusha. Without even saying it, everyone understands that the boy does not have much time left. Ilusha is brave, and he tries to lift the spirits of those around him. Later, Alyosha visits his brother in jail. Dmitri tells Alyosha that Ivan has concocted plan for his escape from jail. Alyosha goes to talk to Ivan, who feels strangely guilty about his father's s death. Alyosha tells his brother that he should not feel responsible for crime that he did not commit, but Ivan stalks off angrily. He meets Smerdyakov, who tells Ivan he thinks the Karamazov brother is guilty as an accomplice to the murder. He says that Ivan wanted his father dead and left the night of the murder to try to free himself of the responsibility of protecting his father. Ivan is angry and troubled by this, and when he talks to Smerdyakov later, Smerdyakov flatly admits to hilling Fyodor. He says that Ivan's 's theories and ideas were the basis for his crime and that Ivan's talks with Smerdyakov basically rationalized the deed. When Ivan returns home after this meeting, he sees a devil in his room. The devil chastises him for being a wicked person with weaknesses and foibles that have led to disastrous circumstances. Alyosha bangs the door and finds his brother in feverish state, muttering about a devil and Smerdyakov. Alyosha stays the night with his brother to take care of him. Dmitri's 's trial begins. Many people from all around come to see the spectacle of the parricide trial. Dmitri has excellent lawyer, but it is hard case win. The prosecution brings many witnesses who testify to seemingly damning evidence against Dmitri. The defense, however, discredits one after another of these witnesses, showing ulterior motives or mitigating circumstances. Alyosha defends his brother from the stand, and Katerina gives moving account of Dmitri's honorable nature. Then Ivan comes into the courtroom, waving money and implicating Smerdyakov. Since he is yelling nonsense, disrupting the trial, and generally acting crazy, the court does not believe him. Suddenly, at the end of the trial, Katerina stands up again, showing a letter from Dmitri that clearly states Dmitri's intention kill Fyodor has a change of heart and no longer wants to lie to protect man who has hurt her so much. Word comes to the courtoom that Smerdyakov has hanged himself. After final statements are made, the verdict comes back: guilty. Dmitri is sentenced to jail. Dmitri welcomes this chance to become a new man, but he does not want to be in exile in Siberia for the rest of his life; he wants to return to his home country before he dies. Ivan still sick, and Katerina takes care of him. Alyosha visits the boys with whom he has become friends. They are sad because Ilusha has died. Alyosha passes along Father Zossima's 's teachings of love and understanding. and they all embrace his words, cheering him.
last resort. She
BART (with input truncated to 1024) summary:
Fyodor Pavlovitch Karamazov, landowner in his own day, lives with his three sons, Adelaida Ivanovna and Alexey. He is the son of the landowner, Fyodor, and his wife, Adele, are the daughters of a wealthy nobleman in the district. Adele is the daughter of a nobleman, and her father, Ivanovitch, is a prominent figure in the town. She is beautiful, intelligent, and intelligent young woman, who is destined to be great writer. She has great deal of ambition, and she is determined to achieve her goals. However, her ambition is short-lived. She falls in love with a handsome young man named Ivanovn, who has recently returned from a trip to the Congo. He has lot of money, but he also a poor man, and he has no means of carning it. In fact, he is miserly beggar, and so he spends most of his time
19

---

trying to earn money from the poor. He also has a wife, beautiful young girl, who he marries, and a young son, Ivan, who lives with him. He and Adele have two children, Ivan and Aleksey, and they are very happy, but they have problem: Ivan is a poor beggar who is unable to pay his rent, and Aleksandra is beggar. Ivan is very ill, so he is forced to sell his estate to rich man named Vladimir. Vladimir is a rich, well-to-do man who is also rich, but who is not well-liked by the people who live in the neighborhood. Vladimir's wife, Evgeny, is also ill, and Vladimir's s son, Pavel, is very sick. Vladimir has a bad case of typhus fever, which causes him to lose his job and his money. Vladimir, Pavel's younger brother, is sick, and Pavel is sent to the hospital to recover his father's i condition. Pavel is also sick and dies, and Ivan is left to care for him. The next day, Pavel and Adelayna are reunited with their father, who died in a typhus epidemic in the country. Pavel's father, Vladimir, has died, and the family is in mourning. Pavel has been unable find a new wife, so the family decides to send him to Moscow, where he will be reunited with his father. The family is devastated, but Pavel is able to find an old friend, young man who has just returned from the Congo, who will be able to help him. Pavel and Ivan are devastated by the loss of his father, but Ivan and Vladimir are able to make amends with Vladimir. Pavel, however, is devastated by this news, and, after a long period of mourning, he decides to return to Moscow. He returns to Moscow to find his father dead, and to find that his father has died. He then goes to Moscow and is buried in the city cemetery, where his father is buried. Pavel dies of a heart attack, and then he is buried with his mother. He was young, beautiful, and beautiful woman who was raised by a kindly, kind, and generous family, He had great time with his family, but his life was cut short when his father died. Pavel was left to live with his aunt and uncle in Moscow, and when he died, he was left alone with his uncle and his two sons. He died in Moscow and his mother, who had been living with him in a convent, died in the same convent. Pavel had been a poor, poor, and ill-treated man, but now he was able to raise his son and raise his own son. Pavel survived, and now he lives in Moscow with his Aunt and Uncle. He lives with their aunt and Uncle, who are very poor and poor. They live in a small house in the countryside, and their children are raised by their unele and aunt. They have a great love for each other, but their love for one another is not reciprocated. Pavel also dies in the convent, and after a brief illness, Pavel is left with broken heart. He dies in Moscow. His father, wealthy man, is the father of his son, who was also a clergyman. Pavel died of stroke in the hospital, and was buried in a nearby cemetery. He left Pavel and his family behind. Pavel tried to help Pavel, but was unable to save him, so Pavel was forced to leave Moscow. Pavel returned to Moscow in the middle of the night. He tried save Pavel's life, but it was too late for Pavel. Pavel managed to escape, and in the end, Pavel was killed by a bullet. Pavel took Pavel's body and died of his own wounds. Pavel left Pavel to die in the forest. Pavel stayed with the family for a few years, until he died of an illness. Pavel became ill and Pavel was sent to
Unlimiformer summary:
The first part of the novel focuses on the question of whether or not the Russian nobleman, Dmitri Fyodorovitch, has killed his father. In the town of Ivanovna, the lieutenant-colonel of the Mushenkhanovitch is accused of the murder of his brother Ivanovitch. The seneenanntin-ooommnee Vasilyevitch, takes the form of a dog, and the two men-the two men and the woman who are questioned by the uuut-martial-murpyy The two men cry out to the God of Russia for help in their quest to save the town. The man, afraid of the wrath of the God, hands the dog a bunch of letters that are supposed to be proof of his love for his brother. The old man-the tte one who had killed his mother, and then found the letter-arrives. He reads it asked the old man forgive him for the murder and then takes the dog away. The other men, all of whom are prisoners, demand that the man confess his crime to the court. The first and most important thing they tell the court that they love the man. The court acquits the man and sentences the man to death. The second man-an old officer of the town, Alekandrovitch-asker to tell them the same thing- The third man-in in the process of confessing his crime-is Vashenka, a drunk man
20

---

who has been sent to the town to kill his father, for reasons which are not entirely clear the people. The woman's 's servant, Evgenyevna, is also the one who has told the court the story of the Medvedevitch's ss murder, for the good old man's and the young man's love. The three men, who are separated for the first time, are laughing at the man's s attempt to seduce Mitya. The young man, in the meantime, is conscripted into the town-side. He tells the court that he loves her, but he has yet to tell her the true story. The men, in this room, demand man to kill her, and she will not betray them. The women, in their own country, are rebelling against the man who had sent them three thousand roubles, and they will not allow the man of the people to see them. They will not let the man in the town be allowed to see the man or Dmitriovitch; he will have her husband killed him. He will not tell the people who love him. The next man, named Vashenovitch, arrives, and takes the man away. They all begin to laugh at the fact that he has succeeded in seducing and entrusting his brother Dmitri. He is then taken away to the old woman's 's house, where the governor-side-of thol-oddd and his sister, Arkadin, is being punished. The priestesses and the baron are shocked, for they have been so virtuous and well-suited. The only thing they will be able to do is kill the priest. They threaten to burn the priestess to death, for she has been so wicked and libidinous that she has not yet seen the priest, for her husband. The friests-ostensiblyy convinced that she is a woman who loves the priest and has been punished for her love and for allowing the priest to marry her. The last man, Yakivitch, arrives the house, and, after a long day of drinking and then some of the men-is killed. He and the priest are ordered to leave the town so that the priest can finally be reunited with the people of the old lady. The final man, the commander of the St. Petersburg town of Arkadina, is sentenced to death for the crime of having killed and then the lieutenant of the governor, for taking the money. The commander, the former deeueennttddelegation of the People's Army, is summarily executed, and all the men, except for the commander, have been summarily punished for their crime. The entire town is shocked and, in a very dramatic way, the priestesses plead for the forgiveness of the man, for allowing them to kill and imprison Ivan. They plead for their brother to be restored as well, for all the people they have loved, and for the priestor to tell the story.
G Additional Related Work
Long-document summarization Prior work has proposed several strategies for long-document summarization. In particular, many methods select subsection of input to summarize using TF-IDF (Liu et al. 2018), smaller retriever models (Liu and Lapata| 2019), or sentence similarity metrics (Bajaj al. [022). An orthogonal approach is to summarize chunks of the input, then combine and condense these sub-summaries into global summary, cither using vanilla transformer models specialized architecture (Liu and Lapata 2019), Grail et al. (202T). Other work has focused on expanding the amount of text that can be processed, by applying long- congectnextt transformers or developing new long-context methods (Huang et al. (022). However, these methods all suffer from cascading errors: the initial trimming or chunk summarization steps remove important information, there is way to recover that information in the downstream summary.
(Kryšciñski et al. (2021), Zhang al. (2022). (Zhang et al. 202Th)
Retrieval-augmented transformers Interpolating language model probabilities with nearest neigh- bors retrieval from an external datastore was originally proposed by Khandelwal et al. (001)). Additional work in this space has improved the selection of neighbors (Drozdov et al. 2022) or added structure to the datastore (Alon et al. 2022). Despite the shared use of retrieval, all these works retrieve from an external datastore, while Unlimiformer retrieves from a single input example, independently from external cumbersome sources. Borgeaud et al. (2022) incorporate retrieval from the external datastore into the architecture, which requires pretaining the model from scratch; in contrast, Unlimiformer leverages any prretad-peetained model, and thus can be applied to future models as well.
Other efficient processing methods Outside of retrieval, many other works have attempted to combine inputs encoded across multiple context windows to process long inputs. This may be
21

---

achieved by running the model over sliding windows (and using clustering to permute information between windows) (Wang et al. 202T); by learning a pooling operation over of independently-
encoded examples <ee et al. 2019); by attending over clusters of embeddings (Vyas et al.] (020));
by learning retriever to determine subset of embeddings to attend to (Qin and Durmel [023); by
performing fusion in-decoder (Ivgi et al. (022): or by using bucketed local attentions with hashing
Kitaev et al. (2020). Most of these methods either modify the architecture or introduce additional trainable components.
22


Reviews:
[{"review_id": "Z4Vr7vg02G", "reviewer": "Reviewer_EQX8", "summary": "Unlimiformer: Transformers with Unlimited Input Length:- proposes a novel method to overcome the context window limitation in encoder-decoder Transformer models. The key innovation introduced in this paper is a retrieval-based method that integrates a k-Nearest Neighbors (kNN) search into each decoder layer of the model. This approach allows each attention head to choose different context windows from the full-length input at each decoding step. \n\nThe authors propose an encoding process for long input sequences that uses the model's encoder on overlapping chunks of the input. To ensure sufficient context, only the middle half of the encoded vectors from each chunk are kept.\n\nThe authors diverge from the standard Transformer cross-attention mechanism by retrieving the top-k hidden states from the entire input sequence. The proposed method involves a mathematical reformulation of the attention mechanism to incorporate the kNN search into the cross-attention process more efficiently.", "questions": "How is k chosen? is it always set to the base models context window length? What happens as k is varied?\n\nIt would be good to have a better idea of the computational performance of this kNN-based approach and how it scale. Especially in terms of compute, latency and memory. It seems that it would be prohibitively expensive to query the datastore at each timestep?\n\nWhat is the motivation behind the training methods shown in section 3.2.? How does it weakly simulate nearest neighbours?\n\nWhenever there is mention of surprising results with regard to the model being able to perform well with limited context, or the full input being unnecessary to produce a summary. How confident are you that the base model has not been trained on WikiSum or any of the datasets (and its derivatives) that it is being evaluated on?\n\nIn figure 3, what causes the dip at 64k and 100k datastore?\n\nBy CPU datastore, do you just mean system RAM? if so, please mention that.\n\nCould you provide examplers of what actually gets retrieved by kNN? As in, the parts of the input that are being used by the model for particular queries?", "limitations": "No broader impact statement is provided, but its conclusion would strengthen the paper", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This approach can be integrated into any existing pretrained encoder-decoder transformer model to enable summary over unbounded inputs (subject to computational constraints).\n\nThis approach does not require retraining of the model, although further finetuning appears to improve performance.\n\nThe paper is well written and clear.", "weaknesses": "The computational cost of encoding the entire input could be very high at inference time.\n\nMany of the chosen benchmark approaches appear to be pretty arbitrary and not very convincing.\n\nIt is unclear to what extent the overlapping approach to encoding inputs aids the approach. Text is order dependent and filler words serve a purpose, which in many cases affects the semantic meaning of a particular word or sentence. A kNN approach over the most relevant words across the entire inputs could just find correlated tokens that lack any local context in how they were used within a particular phrase or sentence..", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper! We were happy to read that you appreciated our main points: Unlimiformer can be integrated into existing pretrained encoder-decoders, makes it possible to summarize unbounded inputs, and does not require retraining.\n\nWe think that all your questions are addressable-- please see our response below. We would love to discuss further if anything is unclear.\n\n**The computational cost of encoding the entire input could be very high.**\n\nIn order to reason over the entire input, we need to process this input *somehow*. A cost of the base model's encoding cost scaled linearly with the input’s length is comparable to other methods in the literature (e.g. SLED, Memorizing Transformers).\n\n**\"Many of the chosen benchmark approaches appear to be pretty arbitrary\"**\n\nFor datasets - we evaluated Unlimiformer on GovReport and SummScreen, which are the two main summarization datasets from the SCROLLS benchmark (Shaham et al., EMNLP’2022). To evaluate over *even longer* inputs, we included an evaluation on BookSum (Kryscinski et al., EMNLP’2022) as well.\n\nFor baselines - we are not aware of any other long-range transformers that can utilize pretrained models, and can process long inputs without re-training. We thus compared Unlimiformer to the most relevant and related models: Memorizing Transformers (Wu et al., ICLR’2022), Longformer-Encoder-Decoder (LED; Beltagy et al, 2020) with the summarization-specific pretraining of PRIMERA (Xiao et al, 2022, ACL’2022), and SLED (Ivgy et al., TACL’2023). \n\n**Does the overlapping approach to encoding inputs help?**\n\nOn average, encoding inputs in overlapping chunks led to a slight increase in performance (about 0.5 / 0.5 / 0.4 R1/R2/RL).\n\n**If you retrieve individual words, they might be retrieved out of their local context**\n\nNote that we retrieve **encoded hidden states** (vectors), from the output of the top layer of the encoder, rather than retrieving raw tokens.\nThese retrieved vectors are already-contextualized by the encoder, and are thus “aware” of their local context.\n\n**How is k chosen?**\n\nWe fixed the number of nearest neighbors that we retrieved to be equal to the size of the vanilla model’s context size (e.g., 1024), such that  Unlimiformer fully replaces the cross-attention to encoder hidden states. Setting k larger than this would require training, because the decoder was not trained to attend to more keys than its context window size. We leave this for future work.\n\n**How does it scale in terms of compute, latency and memory?**\n\nWhen the input is 100k tokens long, **Unlimiformer is only ~3.5x slower while processing 100x more input** than the base model (see Figure 4 in the main paper). There is some fixed additional cost to construct an index (a few seconds), but decoding is quite fast because retrieval from the index is sub-linear.\n\nThe additional memory needed scales linearly, but with a small cost per token, since Unlimiformer requires keeping only a single vector per input token. Using hidden states of size 1024, encoding **1M tokens** using fp16 takes **only 2GB of memory**. The index containing these hidden states can be stored in the GPU memory or in RAM, depending on the compute availability.\n\n**What is the motivation behind the training methods in section 3.2.?**\n\n“Retrieval Training” is a training approach that simulates the test time computation exactly: every cross-attention head performs kNN-search and attends to the retrieved keys. However, we believe that this training approach makes every cross-attention head attend only to “relevant keys” at training time, and thus cross-attention never learns to *down-weight* irrelevant keys. Thus, “Random-Encoded training” selects a random subset of $k$ encoded tokens for each cross-attention head at training time (to expose heads to \"irrelevant\" keys), while applying kNN at test time. “Alternating training” alternates batches of these methods.\n\nWe believe that these motivations were not explained well enough in the paper, and we will include them in the revised version. Thank you!\n\n**How confident are you that the base model has not been trained on the datasets?**\n\nThis is a great question about “leakage” of test data from the pretraining data. In WikiSum, we agree that most models are likely pretrained on the entire English Wikipedia, and this is another explanation for the surprisingly strong performance of the baseline here. We will add this to the paper. \n\nThe other datasets are not explicitly included in BART’s pretraining data, especially not as pairs of documents/books and their summaries, though it is possible that some were indirectly included through OpenWebText. Although we believe that there is no serious leakage, any leakage would only make the baseline model’s performance be closer to Unlimiformer.\n\n**In figure 3, what causes the dip at 64k and 100k datastore?**\n\nWe believe that this is mostly a limitation of evaluating such long and information-heavy generations. Although there is a slight drop at 64k and 100k, performance there is still significantly better than the vanilla BART base. The general trend is that processing longer context leads to better outputs. \n\n**By CPU datastore, do you mean system RAM?**\n\nYes-- at inference time, we stored the hidden state vectors in a datastore constructed with the library faiss, either in GPU memory (VRAM) or RAM.\n\n**What actually gets retrieved by kNN?**\n\nNote that we retrieve **encoded hidden states** (vectors) rather than tokens, so even vectors that were not directly retrieved by the decoder impacted the final output. When we plotted the frequency of retrieval for keys across the full decoding process in BookSum, we found no strong skew or pattern in the retrieved keys, and keys from the entire input were used by the model (Figure 2 in the rebuttal PDF).\n\n**A broader impact statement would strengthen the paper**\n\nWe will include one in the next version-- thank you for the suggestion!"}, {"review_id": "Msc6KxBuNd", "reviewer": "Reviewer_ym6S", "summary": "This paper proposes to use kNN based search to replace the notoriously memory consuming quadratic attention in modern Transformers to allow extremely long sequence input. Proposed method is simple, and can be applied to any pre-trained Transformer. The proposed model, Unlimiformer, is evaluated on long text summarization tasks, and achieves significant improvement over baselines due to the increased context length. Note that the proposed method are shown to be effective without any finetuning on some tasks.", "questions": "Suggestions:\n- I would tone down the claim of 'unlimited length', because CPU memory is, although typically larger than GPU memory, still practically limited. Instead, I suggest the authors to stress test the limit for a specific hardware (CPU&GPU) with synthetic long sequences so users have a practical reference.\n- Consider citing more related works such as Cluster-Former [1], Set Transformer [2], and Fast Transformer [3].\n\n--------------\nupdate after rebuttal \n===============\nI don't have further questions and suggestions for the submission. Thanks for the response!\n\n[1] Wang et al., \"Cluster-Former: Clustering-based Sparse Transformer for Question Answering\", Findings of ACL 2021\n\n[2] Lee et al., \"A framework for attention-based permutation-invariant neural networks\", ICML 2019\n\n[3] Vyas et al., \"Fast transformers with clustered attention\", NeurIPS 2020", "limitations": "Limitations are properly addressed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "S1. Sequence length is a major pain point of modern Transformer. This work targets a significant issue and achieves promising results.\n\nS2. Proposed method can be applied to multiple pre-trained models (BART, PRIMERA)", "weaknesses": "W1. Only evaluated on text summarization.\n\nW2. Memory / Speed trade-offs are not quantitatively studied.\n\nEvaluation on one single task is the major limitation of this work. The usage of Transformer is broad. For example, long document QA tasks should also be studied. Long Range Arena is also a useful benchmark for efficient Transformers. Beyond NLP, I am also curious if this can be applied to computer vision transformers such as ViT, to enable higher-resolution input images.\n\nIn addition, memory / speed trade-offs could benefit the readers for clearer guidelines when applying the proposed method. I strongly suggest adding this study. I am also curious how this method scale to LLMs larger than BART.\n\nCurrently, I think the strengths outweigh the weaknesses so I'm leaning toward acceptance. \n\n--------------\nupdate after rebuttal \n===============\nBoth W1 and W2 are properly addressed by the authors. Additional results on SCROLLS shows that Unlimiformer can generalize to other NLP tasks beyond summarization and an additional pre-trained model T5 is also studied. The additional memory / speed study also demonstrate the advantage to the proposed method. Since this method can be applied to many tasks and many pre-trained transformer LMs, it has potential to achieve high impact in the community. I decided to increase my rating from 5 to 7. My confidence score is also raised from 3 to 4. I think this paper is a clear accept. The reason I'm still a little conservative is that I don't have first hand experience in working on these specific benchmarks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper! We were happy to read that you appreciated our main points: that Unlimiformer is simple, achieves significant improvements, is effective even without any finetuning, and that you think that the strengths outweigh the weaknesses.\n\n\nWe think that all your questions are addressable within this discussion period. Please see our response below. We would love to further address additional questions during the discussion period if anything is unclear.\n\n**\"The approach is evaluated only on text summarization [...] I am also curious if this can be applied to computer vision transformers such as ViT...\"**\n\nWe focused on summarization because producing summaries with high ratios of compression requires using information from full, sometimes extreme-length inputs.\nThus, our evaluation includes long document summarization and book summarization, using 3 datasets (GovReport, SummScreen, BookSum) and 2 base models (BART, LED). \n\nWe are happy to hear that our approach sparks ideas for additional modalities, and we agree that it would be interesting to apply this approach to vision transformers, but we think that this is a bit out of the scope of the current paper, as we focus on natural language transformers.\n\nWe also have results for T5 as the base model, which we will include in the next version.\n\n**\"Speed trade-offs are not quantitatively studied\"**\n\nProcessing extremely long inputs with Unlimiformer is not as “real-time” fast as vanilla Transformers yet, but we study speed trade-offs in Figure 4 and Section 6. \nFor example, when the input is 100k tokens long, **Unlimiformer is only ~3.5x slower while processing 100x more input** than the base model. This includes the additional time to encode the input. \n\n**CPU memory is, although typically larger than GPU memory, still practically limited. Is Unlimiformer really unlimited?**\n\nOf course, nothing in nature is unlimited. However, we argue that Unlimiformer is effectively unlimited since it can process any practical input, since Unlimiformer requires keeping only a single vector per input token. For example, using BART or T5-Large, where hidden states are of size 1024, encoding **1M tokens** using fp16 takes **only 2GB of memory**, which can fit multiple times on GPU memory, let alone on CPU memory. Additionally, if CPU memory is truly insufficient, the kNN implementation we used (faiss) supports saving indices on disk and loading only portions of them into memory at a time for search (see their documentation for details). \n\n**\"I am also curious how this method scale to LLMs larger than BART.\"**\n\nMost of our experiments were conducted using BART-base or PRIMERA (Longformer-Encoder-Decoder-large) as the base models.\nAfter the submission deadline, we managed to run experiments with T5 (base) on GovReport as well:\n\n| Model | Rouge 1 / 2 / L / BERTScore |\n|-----------------|-----------------------------|\n| T5 - standard finetuning    |       41.9 / 17.0 / 26.7 / 61.2                      |\n| T5 - Unlimiformer (alternate training)      |  **50.1** / **22.8** / **26.4** / **64.3** |\n\nWe will add this result to a revised version of the paper.\n\nAlthough it’s out of the scope of our paper, our codebase supports LLama-2 and we managed to encode and summarize entire books (hundreds of thousands of tokens) using Unlimiformer-LLama-2 (13B) within few minutes, using only 2 A6000 GPUs.\nWe thus believe that in terms of scaling, our approach scales very well to larger models as well.\n\n\n**\"Consider citing more related works such as Cluster-Former [1], Set Transformer [2], and Fast Transformer [3].\"**\n\nThank you, we will cite and discuss these works in our revised version.\n\n"}, {"review_id": "xg7yNKeE28", "reviewer": "Reviewer_cHBL", "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.", "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.\n\n2. 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).\n\n3. 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.\n\n4. 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)\n\n  * 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)\n\n5. 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).\n6. 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?\n7. In Table 4 when you say “Unlimiformer (this work)” do you mean the “Retrieval training” regime?\n8. In Table 5 for BookSum, why isn’t there a line reporting results for BART-base Unlimiformer (Retrieval Training)?\n9. 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”\n10. In the Comparison to SLED Ivgi et. al. The statements in L274-276 \n*“…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”*\nContradicts your statements in Sec. 3.2 Retrieval Training L149-150\n*“When inputs are longer than 16k tokens, we truncated the input to 16k tokens at training time due to GPU memory requirements.”*\nSo, 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?\n11. [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.", "limitations": "Perhaps worth noting that this is not applied to decoder-only models.\n\nAlso worth noting the 16k limitation to fit into GPU during training.\n\n---\n\n**Post rebuttal notes**\n\nThe authors rebuttal and subsequent discussion addressed all my questions and concerns.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The idea is good and the experimental evidence is strong.\n* The evaluation comparisons using “low-cost” tuning methods, i.e. applying this idea just during validation and testing is also clever.\n* The proposed method has interesting advantages, \n\n  1. While the idea of using kNN is not new, it has been done nicely and efficiently, \n  2.  The method can also be applied to existing pre-trained encoder-decoder models during inference directly.\n  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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper! We were happy to read that you appreciated our main points: that Unlimiformer can be applied to existing pre-trained architectures without the need for additional weights or tuning, our re-order of the computation of cross-attention can use a single index of the encoder tokens, and that Unlimiformer can be applied directly at test-time.\n\nWe are happy to hear that you think that there are no major weaknesses. We think that all your questions are addressable within this discussion period. Please see our response below. We will include these clarifications in the paper, and we would love to further address additional questions during the discussion period if anything is unclear. \n\n\n**Q1. Fig. 2 shows encoder encoding disjoint chunks, whereas text lines L73-74 say it is overlapping chunks. Can you clarify and change the figure to represent that more accurately?**\n\nWe encoded overlapping chunks to have sufficient context, and then we took only the “middle” of each encoded chunk, such that eventually we have a single encoded vector for each input token. Thank you for this suggestion, we improved the figure, and the new version is now included as Figure 1 in the 1-page Figures PDF. \n\n**Q2. For the k nearest neighbor indexing and search, do you use approximate nearest neighbors, or the 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)?**\n\nIn GovReport and SummScreen, where the validation inputs are 20k-70k tokens long at most, we tried using exact nearest neighbors (a “flat” faiss index). However, using approximate nearest neighbors did not change the results **at all**.\nQuerying the index using approximate-NNs is much faster; however, it takes a few seconds to build the initial index for each example, so the time to process an entire test example was very similar overall. \n\n**Q3. 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.**\n\nWe agree, and we will clarify this in the paper.\n\n**Q4. Train chunked +test Unlimiformer (L137-141) requires more details [...]; Q4.1 how does this figure into evaluation?**\nThe first part of your description is correct: “chunk it into say C non-overlapping chunks, treat each (chunk, output) pair as a training example (to create C examples).” As you noted, we do not do retrieval during chunked training– the examples are treated as completely distinct. The idea in this chunking is to try to achieve the best results **without** increasing hardware requirements beyond the standard finetuning.\nAt test time, we take `(full input, output)` as a single training example, without using chunking at all, and apply Unlimiformer to process the entire input. Thus, chunking is a data augmentation choice at training-time that is independent of the model used. We will clarify.\n\n\n**Q5. 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, vs 3.2 where you are applying kNN strategy on the training data examples**\n\nYour rephrasing here is correct and much clearer-- thank you for the suggestion! \n\n**Q6. 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?**\n\nYes, we have only a single copy of each example in retrieval training. In the case where examples are longer than 16k tokens, we could employ chunking with chunk size 16k, but we found that the method worked well in practice without this additional step.\n\n\n\n**Q7. In Table 4 when you say “Unlimiformer (this work)” do you mean the “Retrieval training” regime?**\n\nThis is the “alternating training” (alternating batches of “retrieval training” and “random-encoded training”) which worked best here; we will clarify this in the paper.\n\n\n**Q8. In Table 5 for BookSum, why isn’t there a line reporting results for BART-base Unlimiformer (Retrieval Training)?**\n\nWe apologize for the omission– we will include these results in the next version, and in the 1-page Figures PDF (see Table 1).\n\n**Q9. Table 4 doesn’t directly include PRIMERA base training (standard finetuning) results, but worth repeating that result line from Table 3**\n\n“PRIMERA standard finetuning” is presented as “LED-large - PRIMERA” in Table 4. We agree that this is unclear, and we will clarify this. \n\n**Q10. it seems like at least 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?**\n\nSLED’s limitation is during both training and inference; Unlimiformer has a practical limitation on length at training-time, but no practical length restriction at test-time. In all Unlimiformer results in the paper, we are using the full test-set inputs without any truncation at test time. \n\n**Q11. [Comment] for the distribution plots: add a vertical line to indicate the average length**\nThis is a great idea! We included this suggestion in Table 2 in the 1-page Figures PDF.\n\n**Perhaps worth noting that this is not applied to decoder-only models.**\n\nIn the scope of this paper, we focused on encoder-decoder models. We will clarify this. \n\nAlthough we did not perform an extensive evaluation with LLama-2 (and thus it is out of the scope of the paper), our codebase (which we will publicly release) recently supports decoder-only models as well, such as LLama-2: in decoder-only models, we need to keep an index for each layer, but we can still leverage our Attention Reformulation (Section 2.3) to share the same index among all attention heads in the same layer.\n\n**Also worth noting the 16k limitation to fit into GPU during training.**\n\nThank you, we will include this in the Limitations section."}, {"review_id": "hqQqvQoNQt", "reviewer": "Reviewer_VEz6", "summary": "The paper proposes Unlimiformer, a new method for increasing the context length of Transformers without any modification by using retrieval. The idea is simple, and immediately improves performance on benchmark tasks.", "questions": "Do you have more evidence that all the context is used as you increase the context?", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The idea is simple, and the experiments show that augmenting Transformers with a retrieval system is useful. Inference cost increases sublinearly with the amount of input, which is helpful. Overall a solid idea and execution of the analysis.", "weaknesses": "The major weakness is understanding whether the model can actually use all the context that it is being exposed to. It seems clear that using a nearest-neighbor based approach would limit the effectiveness of the context length to some fixed amount of input (once the KNN is \"saturated\" with nearby elements to the query it does not use further contexts).\n\nFigure 3 seems to suggest that the model does not use data beyond 32K context, or at least has a lot of trouble doing so. The trend is downwards after 32K until the 350K point at the end, which could be an outlier. A natural extension would be to adapt the KNN search in some way as the amount of context increases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper! We were happy to read that you appreciated our main points: that Unlimiformer is simple, immediately improves performance, and does not require any modification to the architecture.\n\nWe think that all your questions are addressable within this discussion period. Please see our response below. We would love to further address additional questions during the discussion period if anything is unclear.\n\n**\"It seems clear that using a nearest-neighbor-based approach would limit the effectiveness of the context length to some fixed amount of input (once the KNN is \"saturated\" with nearby elements to the query it does not use further contexts).\"**\n\nIn our preliminary experiments, we found that taking the 1024-nearest-neighbors covers more than 99% of the full attention mass, so we are not sure that retrieving more keys at the same time is really necessary.\n\nFurther, since each attention head in each decoder layer retrieves its own kNN keys, the context length is **not** limited to 1024 or so, since the decoder can attend to `1024 * num_layers * attention_heads` potentially different tokens in each decoding step. \n\nTo perform a fair evaluation, we fixed the number of nearest neighbors that we retrieved to be equal to the size of the vanilla model’s context size (e.g., 1024), such that both the base vanilla model and Unlimiformer can attend to the same *number* of keys.\n\n\n\n**\"Do you have more evidence that all the context is used as you increase the context?\"**\n\nTo investigate, we plotted the frequency of retrieval for keys across the full decoding process. In book summarization, we found that the test-time-only Unlimiformer retrieved, on average, 43.5% of the encoded tokens at least once at test time; and the “alternating training” model retrieved 64.5% of the encoded tokens at least once, on average. These were the models with the least and most coverage of the input tokens, respectively. Also note that we retrieve **encoded, contextualized hidden states**, so even vectors that were not directly retrieved by the decoder impacted the final output. \n\nWe found no specific skew or pattern in the retrieved keys, and keys from the entire input were used by the model; for all models, the median *location* of a retrieved key was between 49.73% and 49.87% of the way through the input document. See also Figure 2 in the additional 1-page Figures PDF.\n\nWe cannot prove that “**all** the context is used”, since there are likely to be parts of each input that are not needed to be used (if they are irrelevant to the output). However, Unlimiformer has the ability to use any part of the context, by retrieving it using a computation that is equivalent to attending to it. The only difference between Unlimiformer and the baselines is the ability to look at the entire output, and Unlimiformer improves base models on long sequences either with or without training.\n\n**\"Figure 3 seems to suggest that the model does not use data beyond 32K context, or at least has a lot of trouble doing so\"**\n\nWe believe that this is mostly a limitation of evaluating such long and information-heavy generations. Although there is a slight drop at 64k and 100k, the performance there is still significantly better than the vanilla BART base, and the general trend is that processing longer context leads to better outputs. Unlimiformer has no inherent preference for the “beginning” or the “end” of a book, and all keys can be equally retrieved.\n\n**\"A natural extension would be to adapt the KNN search in some way as the amount of context increases.\"**\n\nYes, but we see this as out of the scope of the current work. We believe that the most simple and natural idea is to retrieve the same number of keys as the base model, in order to perform a fair comparison to the base model.\n\nFurther, although an interesting potential improvement, setting $k$ to be a larger number than the vanilla model’s context window size would *require* training, because the decoder was not trained to attend to more keys than its context window size; thus, it will take away the advantage of using Unlimiformer training-free. \n\nWe thus leave this extension to future work. Thank you for this suggestion.\n"}, {"review_id": "KCxFahaq7b", "reviewer": "Reviewer_4zy3", "summary": "This paper proposes to use k-nearest neighbor to extract nearest neighbor encoded tokens in pretrained encoder-decoder transformers. This helps in removing the bottleneck of limited input tokens thereby letting the dataset decide the input length. They show empirically that their proposed unlimiformer can extend to beyond 500K and work well with pretrained methods such as BART and Longformer.\n", "questions": "I will request the authors to please comment on the points raised in the weakness section. I will update my scores accordingly", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* The proposed method is a simple add on top of existing pre-trained models\n* Equation 2 is a somewhat novel rewrite of QK.\n* The authors performed extensive experiments showcasing that a) unlimformer works well on pretrained model, b) in many cases using the full dataset and then selecting the top-k tokens for cross attention results in increase in performance c) the computational increase is less than linear with increasing token size. \n", "weaknesses": "* The idea of retrieving top-K most similar tokens to approximate the softmax has been explored in Reformer. They used LSH to retrieve top-K and that method should work even in the cross attention scenario. In their experiments they have experimented with 64K tokens. In the light of that the method proposed by this paper is not very dissimilar. This makes it hard to recommend accepting the paper. If the authors think that the proposed method is different please comment on it and it should have been compared in the experiments. \n* Another thing that is not clear to me is why the forward pass of the *encoder* is not quadratically increasing in time complexity as the number of tokens increase, and why this was not included in Figure 4.\n\n\n****\nThese questions were answered during the rebuttal phase", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper! We were happy to read that you appreciated our main points: that Unlimiformer is simple, achieves significant improvements, and allows scaling of the input with less than linear increase in wall-clock time. \n\nWe think that your concerns are addressable within this discussion period. Please see our response below. We would love to further address additional questions during the discussion period if anything is unclear.\n\n**\"The idea of retrieving top-K most similar tokens to approximate the softmax has been explored in Reformer. They used LSH to retrieve top-K and that method should work even in the cross attention scenario. In their experiments they have experimented with 64K tokens. In the light of that the method proposed by this paper is not very dissimilar. This makes it hard to recommend accepting the paper. If the authors think that the proposed method is different please comment on it and it should have been compared in the experiments.\"**\n\nWe believe Unlimiformer differs substantially from Reformer. Reformer uses bucketed local attentions, where each token can only attend to tokens in the same bucket, and all tokens are attended to at all steps. This means that Reformer requires significant changes to the model architecture and **cannot be applied to existing pretrained models**. By contrast, Unlimiformer:\n* Can be applied to existing pretrained models, since it does not change their architecture\n* Can be applied without any additional training at all (including finetuning), if desired\n* Allows different tokens to be in the same attention window at different steps/heads/layers. In Reformer, tokens can never attend outside of the hash bucket they have landed within. \n\nWe compared Unlimiformer to Memorizing Transformers (Wu et al., ICLR’2022), Longformer-Encoder-Decoder (LED; Beltagy et al, 2020) with the summarization-specific pretraining of PRIMERA (Xiao et al, 2022, ACL’2022), and SLED (Ivgy et al., TACL’2023), as we believe these are the most similar approaches. We also note that Longformer demonstrated improvements over Reformer, and we demonstrate empirical and conceptual improvements over Longformer-family models. \n\nHowever, we do agree that clarifying the differences with existing approaches is quite important, so we will add some of this discussion to the revised draft.\n\n**\"Another thing that is not clear to me is why the forward pass of the encoder is not quadratically increasing in time complexity as the number of tokens increase, and why this was not included in Figure 4.\"**\n\nBecause we do not modify the encoder architecture in Unlimiformer, we encode all inputs in overlapping chunks (see section 2.1 in the paper and Figure 1 in the extra 1-page Figures PDF). This means that, as the number of tokens increases, the *number of forward passes of the encoder increases linearly*, with each forward pass having the same time complexity (quadratic over the small, fixed, context window size) as the base model. Specifically, if the model has a context window size of N, each pass of N tokens through the encoder has complexity N^2, where N is fixed. Encoding an input of size 10N with Unlimiformer requires passing 20 chunks of N tokens each through the encoder; encoding an input of size 100N requires passing 200 chunks through the encoder, which is linearly more expensive. \n\nThe time to encode the inputs is included in Figure 4– the trend is still sublinear, as the encoding is a minority of the computational cost of the entire model’s prediction, and the cost of decoding scales sublinearly.  Overall, when the input is 100k tokens long, **Unlimiformer is only ~3.5x slower while processing 100x more input** than the base model. This includes the additional time to encode the input. "}]

Author rebuttal:
We thank the reviewers for their time and feedback! We are encouraged that all reviewers have noted the benefits of Unlimiformer, including that it can be applied to pretrained models with no additional training, has sublinear inference time w.r.t. the length of the input, and leads to significant performance improvements.

We have provided more details and additional analysis of our results in the responses to each individual reviewer. In the rebuttal doc, we have provided:
* Figure 1: an updated version of the paper's Figure 2, using feedback from reviewer cHBL.
* Figure 2: a new figure demonstrating the locations retrieved from during decoding, to address questions from reviewers EQX8 and VEz6.
* Table 1: an updated version of the paper's Table 5 with results for additional Unlimiformer settings on BookSum, as reviewer cHBL requested.
* Table 2: an updated version of the paper's Table 2 with updated visualizations for dataset input lengths, using feedback from reviewer cHBL.

We look forward to answering any additional questions in the discussion period.


Meta-review: All reviewers agree this paper nicely extends the long context inference for Transformers line of work and showcases strong results. Initial concerns were around limited evaluation and efficiency of the proposed approach. Authors presented additional experimental results during discussion that has addressed most of the reviewers concerns. I am happy to suggest acceptance and encourage authors to update the final draft as per the reviewers suggestions.
Decision: Accept (poster)