Unlimiformer: Long-Range Transformers with Unlimited Length Input — part 2/4

Cleaned-text code points [15866, 31839). Bibliography omitted here; source.md includes it.

SLED ( Ivgi et al., 2022 ) extends encoder-decoder models for longer contexts by applying fusion in-decoder ( Izacard and Grave, 2021 ) : 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 a transformer with a trainable attention gate that moderates between the standard cross-attention and attention over retrieved keys from a datastore. Since their public implementation 4 4 4 https://github.com/google-research/meliad is “not officially supported” and is not fully reproducible, we approximated it by using attention over the index in only a single decoder layer; this is equivalent to their setting with the learned interpolation parameter g set to 1. 5 5 5 Wu et al. (2022) note that in their experiments that most heads learned a value for g such that they attended “almost exclusively” to the external memory. 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-tuning; further, Wu 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 .

## 5 Results

### 5.1 Long Document Summarization

#### Low-cost training

Table 3 shows the results in the long-document summarization datasets. First, we can see that applying Unlimiformer on an existing checkpoint without any training ( +test Unlimiformer ) improves BART base {}_{\text{base}} by, for example, 1.8 ROUGE-1 points on both datasets, 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.

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-1 and 3 BERTScore points improvements, while training is as computationally cheap as standard finetuning.

#### Long-range training

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

The PRIMERA experiments in Table 4 highlight two important points: first, Unlimiformer+BART base {}_{\text{base}} performs better than the base PRIMERA across all metrics and datasets, even though PRIMERA is larger and was pretrained on much more data, using a pretraining objective that was designed for summarization; second, not only can Unlimiformer outperform Longformer-based models such as PRIMERA, Unlimiformer can also be applied on top of existing long-range transformers and further improve them: Unlimiformer+PRIMERA improves over PRIMERA across all metrics and datasets. Additional results on the validation set are provided in Appendix E .

### 5.2 Book Summarization

Table 5 shows the result on BookSum. As shown, Unlimiformer improves both base models BART base {}_{\text{base}} and PRIMERA, in both low-cost training approaches such as Early stop w/ Unlimiformer , as well as in the long-range training approaches. Random-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 5 ), the base BART ( Standard finetuning , which truncates the input to the first 1024 tokens) shows competitive ROUGE and BERTScore metrics. This is strongly counterintuitive for book summarization, where the book’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 F for an example with analysis). However, this is not reflected in n-gram-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 11.6 to 25.5 in Unlimiformer+PRIMERA.

## 6 Analysis

#### Is the long input really needed?

As found in various recent papers ( Shaham et al., 2022 ; Kedzie et al., 2018 ) , 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 4 shows the performance of Unlimiformer in BookSum: EntMent increases almost monotonically with input length, suggesting Unlimiformer exploits the longer inputs to 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 not necessarily the beginning. We observed this trend in WikiSum, a multi-document summarization dataset where the inputs are all references of a Wikipedia article and the output summary is the intro paragraph of the article ( Liu* et al., 2018 ) 6 6 6 A full copy of WikiSum is not available online; details of our scraped copy are in Appendix B . . As a strong baseline, we followed Liu* et al. (2018) , and ranked the input paragraphs according to TF-IDF. Unlimiformer did not improve over a 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 dataset 7 7 7 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. .

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 to the input length in Figure 4 . When all inputs are restricted to 1,024 tokens, Unlimiformer requires a small additional time overhead relative to the baseline for indexing and search. However, the benefits of Unlimiformer are clear as input length increases: the total GPU-time required increases sublinearly with input length 8 8 8 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. . 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 et al., 2021 ) , a question-answering dataset over NLP papers; Contract NLI ( Koreeda and Manning, 2021 ) , a natural language inference dataset over legal contracts; QMSum ( Zhong et al., 2021 ) , a query-based summarization dataset over meeting transcripts; and NarrativeQA ( Kočiský et al., 2018 ) , a reading comprehension dataset over narratives 9 9 9 We report results on the validation split from SCROLLS Shaham et al. (2022) . .

Table 6 shows the performance of BART-Unlimiformer (with alternating training) relative to base BART. On three of the four datasets, applying Unlimiformer improves over the base model.

.

What is attended to?

We plotted the frequency of retrieval for keys across the full decoding process for the test set of BookSum, the dataset with the longest inputs. The average number of input embeddings retrieved at least once varied by method, from 43.5% of all tokens for the test-time-only Unlimiformer to 64.5% of all tokens for the alternating-training model 10 10 10 Note that this is not the percentage of the input that influenced the output; we retrieve encoded, contextualized hidden states, so even vectors that were not directly retrieved by the decoder impacted the final output. .

Figure 5 shows the retrieval locations for the alternating-training model. We 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.

## 7 Related Work

Long-range transformers Previous long-range transformers change the transformer architecture to reduce its space or time requirements ( Tay et al., 2020 ) . Most solutions achieve this reduction through sparsifying the attention mechanism ( Child et al., 2019 ; Kitaev et al., 2020 ; Beltagy et al., 2020 ; Roy et al., 2020 ; Ainslie et al., 2020 ; Zaheer et al., 2020 ) . Other works approximate or replace the attention mechanism entirely ( Wang et al., 2020 ; Katharopoulos et al., 2020 ; Choromanski et al., 2020 ; Lee-Thorp et al., 2021 ) . All these approaches change the standard transformer architecture or its training objective ( Zhong et al., 2022 ) , and thus require pretraining the model from scratch, which does not allow to leverage existing pretrained models. In contrast, Unlimiformer is generic , can be injected into any encoder-decoder transformer, and improve it either without training or with merely fine-tuning. This way, Unlimiformer can leverage any already-pretrained 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, and due to memory constraints can thus apply their approach only to a single decoder layer. In contrast, thanks to our attention reformulation ( Section 2.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 5 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 ( Ivgi 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 a single GPU; in contrast, instead of attending to all input tokens, Unlimiformer attends only to the top- k 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 to a k k NN index, to allow for unlimited length input. Instead of attending to all keys, this k k NN index allows every cross-attention head in every decoder layer to retrieve and attend only to its top- k k keys. We evaluate Unlimiformer on several long-document and book-summarization 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-based models, Unlimiformer can be applied on top of Longformer-based models and further improve them.

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 already-pretrained models to process longer inputs, we aim to 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://github.com/abertsch72/unlimiformer . Our code is based on HuggingFace Transformers ( Wolf et al., 2020 ) , without changing any individual architecture’s code, and thus can be injected into any encoder-decoder model, and supports decoder models such as LLaMA-2 as well.

## 9 Limitations

In our experiments, we have only considered English-language datasets. While we have no reason to believe the method would suffer from the use of a different high-resourced language, the quality of the nearest-neighbors 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 a concern when using smaller GPUs or larger models.

## Acknowledgments

We are grateful to Sireesh Gururaja for useful feedback on a 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 et 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. DGE2140739. 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.

## Appendix A Training details