SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The PRIMERA experiments in Table 4 highlight two important points: first, Unlimformer+BART<sub>base</sub> 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 Unlimformer outperform Longformer-based models such as PRIMERA, Unlimformer can also be applied on top of existing long-range transformers and further improve them: Unlimformer+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, Unlimformer improves both base models BART<sub>base</sub> and PRIMERA, in both low-cost training approaches such as Early stop w/ Unlimformer, 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 Unlimformer 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
Figure 3: As the maximum datastore size increases, the entity recall generally increases. At all datastore sizes, Unlimformer outperforms the baseline (BART, in red).
Figure 4: As the maximum datastore size increases, the inference cost increases sublinearly. The plot shows total wall-clock inference time per example.
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 Unlimformer to achieve significantly better Entity Mention Recall (EntMent): the Unlimformer models exhibit far higher EntMent, and even adding Unlimformer only at test time without costly training (Early stop w/ Unlimformer) doubles the entity recall compared to the base model. Further, Unlimformer improves EntMent in the base PRIMERA from 11.6 to 25.5 in Unlimformer+PRIMERA.
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 Unlimformer really utilizes long inputs, we experimented with limiting the input length in BookSum. Figure 3 shows the performance of Unlimformer in BookSum: EntMent increases almost monotonically with input length, suggesting Unlimformer 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)<sup>7</sup>. As a strong baseline, we followed Liu* et al. (2018), and ranked the input paragraphs according to TF-IDF. Unlimformer 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<sup>8</sup>.
Computational cost Although Unlimformer 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, Unlimformer requires a small additional time overhead relative to the baseline for indexing and search. However, the benefits of
<sup>7</sup>A full copy of WikiSum is not available online; details of our scraped copy are in Appendix B.
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.
Base model	Training method	QASPER F1	Contract NLI Exact Match	QMSum ROUGE 1 / 2 / L	Narrative QA F1
BARTbase	Standard finetuning	22.0	77.5	30.8 / 8.7 / 20.8	15.5
BARTbase	Unlimiformer	27.5	77.7	30.9 / 8.0 / 19.9	18.5
Table 6: Results on question answering, query-based summarization, and NLI datasets.
Unlimformer are clear as input length increases: the total GPU-time required increases sublinearly with input length<sup>9</sup>. Additional GPU-time measurements are reported in in Appendix D.
Performance on other tasks We measure the performance of Unlimformer 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<sup>10</sup>.
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 Unlim-iformer to  $64.5\%$  of all tokens for the alternatingtraining model11.
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.
Figure 5: Histogram of location of retrieved embeddings in the original document (averaged over the BookSum test set). There is a slight bump at the beginning (first  $10\%$  of the book), but otherwise no strong trend, with tokens retrieved quite uniformly from the entire inputs.
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,
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.
<sup>10</sup>We report results on the validation split from SCROLLS Shaham et al. (2022).
<sup>11</sup>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.
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) Unlimformer 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 Unlimformer 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$  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.
We present Unlimformer, an approach for augmenting pretrained encoder-decoders and offloading the cross-attention computation to a  $k\mathrm{NN}$  index, to allow for unlimited length input. Instead of attending to all keys, this  $k\mathrm{NN}$  index allows every cross-attention head in every decoder layer to retrieve and attend only to its top- $k$  keys. We evaluate Unlimformer on several long-document and book-summarization benchmarks having inputs of up to 500K tokens, and show that Unlimformer improves existing models, even without further training. When training with Unlimformer, not only that Unlimformer makes smaller models such as BART perform better than larger Longformer-based models, Unlimformer 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.
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, Unlimformer can process the longest inputs when the index is offloaded to the CPU memory. In this case, Unlimformer 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.
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.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumont, Clement Delangue, Anthony Moi, Pierrick 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: State-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 Papers), pages 5245-5263, Dublin, Ireland. Association for Computational Linguistics.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. 2020. Big bird: Transformers for longer sequences.
Longxiang Zhang, Renato Negrinho, Arindam Ghosh, Vasudevan Jagannathan, Hamid Reza Hassanzadeh, Thomas SchAAF, and Matthew R. Gormley. 2021. Leveraging pretrained models for automatic summarization of doctor-patient conversations. In Findings of the Association for Computational Linguistics: EMNLP 2021, pages 3693-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<sup>n</sup>: A multi-stage summarization framework for long input dialogues and documents. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), 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 multi-domain meeting summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: 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).
At training time, we must backpropagate through the operations described above. Thus, the input length is bounded more strictly – the number of tokens in the full input must fit in GPU memory while the model is loaded. For the computationally expensive methods, we 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 (Ivgi 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 on 2,000 randomly sampled examples (1,000 validation, 1,000 test), maintaining the same sample across all experiments. When sampling, we respect the original WikiSum train/validation/test split. We release the subset we 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, we evaluate using facebook/bart-large-mnli instead. This model has context size 1024. For BookSum, we experimented with using allenai/longformer-large-4096 (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  $\mathrm{SpaCy}^{12}$  to tag all named entities in the gold summary and collected a set of unique entities. We then tagged each candidate summary and computed the percentage of entities present in this summary, that is, recall of unique entities. For the named entity recognition in EntMent, we used SpaCy's en_core_web-lg model.
D Computational Cost
We estimate the total GPU time for results presented in this paper did not exceed approximately 116 days of time on a single 48-GB A6000. The longest-training models, SLED and retrieval training for GovReport, took approximately 10 days to train.
GPU-time Table 7 shows the relative cost for each method. The Unlimformer training methodologies are 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 a CPU datastore is many times slower than a GPU datastore because of slower search and the need to transfer retrieved embeddings to the GPU. In our experiments, we were able to use a GPU datastore for input examples exceeding 500k tokens (on GPUs no larger than 48 GBs), but this may be a concern when using smaller GPUs or even larger inputs. Additionally, CPU indices are 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.
12https://spacy.io
Method	Relative GPU-time
Baseline 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-encoded 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 3 runs on BookSum using a 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 Unlimformer at different numbers of layers in BART.
Memory usage Table 8 shows the memory required to apply Unlimformer on varying numbers of layers in a BART model. Using Unlimformer requires more memory than the base model for two reasons: index construction and the additional input processed. There is a slight overhead for constructing and storing the index. But more crucially, the base BART ("normal inference") is truncating the input to the first 1024 tokens. When Unlimformer is used, we process the full inputs, some of which are  $>500,000$  tokens, and so much of the additional memory cost comes from storing the additional hidden states. However, the GPU memory consumption remains constant even when we use Unlimformer in more layers, highlighting the scalability of Unlimformer compared to 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 metrics for GovReport and SummScreen.
F Sample Outputs
These outputs from BookSum are summaries of The Brothers Karamazov, an elaborate novel about a 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 in the original. The hallucinations in the Unlimiformer output are more constrained; for instance, it incorrectly describes Dmitri as a "nobleman" instead of a landowner and says he has been sentenced to death instead of jail. This summary features more of the novel'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 a family tragedy centered around a 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 a greedy landowner, a bawdy lecher, and a neglectful father. Hence, the Karamazov brothers end up growing into young men under the care of various
Base model	Training method	ROUGE 1 / 2 / L / BERTScore
GovReport	SummScreen
Low-cost training methods:
BARTbase	Standard finetuning	47.7 / 18.5 / 22.3 / 64.0	30.0 / 6.5 / 17.7 / 56.7
BARTbase	+test SLED	46.0 / 16.3 / 20.3 / 62.8	28.4 / 5.9 / 17.0 / 56.0
BARTbase	+test Unlimiformer	49.5 / 19.6 / 21.9 / 64.8	31.8 / 7.1 / 18.6 / 57.8
BARTbase	+early stop w/ Unlimiformer	51.0 / 20.6 / 21.6 / 65.9	32.5 / 7.2 / 19.9 / 57.9
BARTbase	Train chunked	48.3 / 18.1 / 22.3 / 63.8	29.4 / 6.3 / 17.6 / 56.8
BARTbase	+test Unlimiformer	52.9 / 22.2 / 22.4 / 65.8	29.4 / 6.3 / 17.6 / 56.8
Long-range training methods:
BARTbase	SLED (Ivgi et al., 2022)	55.5 / 24.8 / 25.8 / 66.9	34.2 / 8.2 / 19.2 / 58.8
BARTbase	Memorizing Transformers	55.8 / 25.6 / 26.9 / 67.7	32.8 / 7.6 / 19.3 / 57.7
BARTbase	Unlimiformer	57.4 / 26.4 / 27.9 / 68.2	35.0 / 8.3 / 19.6 / 58.4
Low-cost training methods:
PRIMERA	Standard finetuning	55.0 / 23.6 / 25.9 / 66.9	33.0 / 7.8 / 18.8 / 57.4
PRIMERA	+test Unlimiformer	56.4 / 24.7 / 26.4 / 67.6	33.1 / 7.9 / 18.7 / 57.4
PRIMERA	+early stop w/ Unlimiformer	56.4 / 25.0 / 26.4 / 67.6	33.5 / 8.2 / 19.3 / 57.7
Long-range training methods:
PRIMERA	Memorizing transformers	57.0 / 25.6 / 26.8 / 67.8	32.9 / 7.7 / 18.5 / 57.5
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.