SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Unlimformer: Long-Range Transformers with Unlimited Length Input
Amanda Bertsch
Uri Alon*
Graham Neubig
Matthew R. Gormley
Carnegie Mellon University, USA
{abertsch,ualon,gneubig,mgormley}@cs.cmu.edu
Since the proposal of transformers (Vaswani et al., 2017), 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 (Lewis et al., 2020a) 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 well $^{2}$ .
Transformers (Vaswani et al., 2017) have risen as the dominant sequence-to-sequence architecture. Pretrained transformers generally have a context window of 512 (e.g. BERT (Devlin et al., 2019), T5 (Raffel et al., 2020)) or 1024 tokens (e.g. BART (Lewis et al., 2020b)), which are sufficient lengths for many current conditional generation datasets (XSum; Narayan et al., 2018) (CNN/DM; Nallapati et al., 2016). To address inputs between 1024 and 16,384 tokens, specialized long-context models sparsify or approximate attention (e.g. Longformer (Beltagy et al., 2020), Performers (Choromanski et al., 2020)), allowing the maximum input length to quadruple while remaining computationally feasible. Most long-document summarization and question-answering datasets, such as SCROLLS (Shaham et al., 2022), are included in this range.
Yet tasks that involve long narratives, such as book summarization (Krysciński et al., 2021), can contain inputs exceeding  $500k$  tokens. Figure 1 shows the input lengths of several popular summarization and question-answering datasets, plotted against common context window lengths; the longest inputs are more than 34 times longer than Longformer's context window.
In these extremely-long-input 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
$^{2}$ https://github.com/abertsch72/unlimformer
Figure 2: In this example, a given LM's encoder's maximum input length is 2 tokens. A 6-token input is encoded in chunks and indexed in an index. We inject Unlimiformer into each decoder layer prior to cross-attention. In Unlimiformer, we perform  $k$ NN search to select a 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.
resources. Other architectures such as Longformer-Encoder-Decoder (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.
We introduce Unlimiformer, a retrieval-based approach to augment pretrained language models to accept inputs of unbounded length at test time. Given a long input sequence, Unlimiformer constructs a  $k$ -nearest-neighbor ( $k\mathrm{NN}$ ) index over the hidden states of all input tokens. Then, every standard cross-attention head in every decoder layer queries the  $k\mathrm{NN}$  index, such that the  $k\mathrm{NN}$  distances are the attention dot-product scores, and attends only to the top- $k$  input tokens. In preliminary experiments, we found that the top- $k$  attention keys cover more than  $99\%$  of the attention mass, and thus attending only to the top- $k$  keys is an accurate approximation of the full, exact, attention. Unlimiformer 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.
Unlimiformer is a generic approach: it can be applied to trained models and improve existing checkpoints without adding weights and without further training. When finetuning Unlimiformer, performance is even further improved: across a variety of long-range datasets, not only that
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 or maximum input length in each dataset, as indicated. Averages for datasets from Koh et al. (2022).
Unlimformer performs better than strong long-range transformers such as LED (Beltagy et al., 2020), PRIMERA (Xiao et al., 2022), SLED (Ivgi et al., 2022) and Memorizing Transformers (Wu et al., 2022), but Unlimformer can be applied on top of such models to further improve them.
2 Unlimiformer
Given a trained encoder-decoder transformer, Unlimformer allows each cross-attention head to choose separate keys to attend to from the full-length input, at each decoding step. We inject a  $k\mathrm{NN}$
search into each decoder layer: prior to cross-attention, the model performs a nearest-neighbor search in a kNN index to choose a set of per-decoder-layer per-attention-head tokens to attend to.
2.1 Encoding
To encode an input sequence that is 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 Faisss (Johnson et al., 2019), using dot-product as the index's nearest-neighbor similarity metric.
2.2 Retrieval-augmented Cross-Attention
In standard cross-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 to this  $k$ -token prefix of the input, we retrieve the top- $k$  hidden states from the  $k$ NN 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 most-attended tokens preserves the vast majority of attention mass.
Figure 2 illustrates our generic changes to any sequence-to-sequence transformer's architecture. The full input is encoded using the encoder in chunks and indexed in a kNN index; then, the index of encoded hidden states is queried at each decoding step. The kNN 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.
2.3 Attention reformulation
Let  $h_d$  be the decoder hidden state and  $h_e$  be an encoder's last layer hidden state. The standard cross-attention computation for a single head in a transformer is:
$$
\operatorname {A t t n} (Q, K, V) = \operatorname {s o f t m a x} \left(\frac {Q K ^ {T}}{\sqrt {d _ {k}}}\right) V \tag {1}
$$
where  $Q = \pmb{h}_dW_q$  is the product of the decoder states  $\pmb{h}_d$  and the query weight matrix  $W_{q}$ ; the keys  $K = \pmb{h}_{e}W_{k}$  are the product of the last encoder hidden states  $\pmb{h}_{e}$  with the key weight matrix  $W_{k}$ ; and  $V = \pmb{h}_{e}W_{v}$  is similarly the product of  $\pmb{h}_{e}$  with the value weight matrix  $W_{v}$ . Our goal is to retrieve a set of keys  $K_{best}$  that maximize  $QK_{best}^{T}$ , with the size of  $K_{best}$  fixed to the size of the model's context window, and then compute the standard attention over  $K_{best}$  only.
Note that the linear layers  $W_{q}$ ,  $W_{k}$ , and  $W_{v}$  are layer-specific and head-specific. Thus, naively creating an index from the keys  $K = h_{e}W_{k}$  and querying this index using the query vectors will require constructing separate indexes for the keys and values at each layer and each head, for a total of  $2 \times L \times H$  indexes, where  $L$  is the number of decoder layers and  $H$  is the number of attention heads. In fact, this exact naive approach was taken by Memorizing Transformers (Wu et al., 2022), who pioneered the use of a kNN index for previously encoded inputs. A separate index for each attention head in each decoder layer is both time-intensive to create and space-intensive to store. So, not surprisingly, Wu et al. (2022) apply their memory layer to only a single decoder layer.
Instead, we present a 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 standard dot-product attention. The
<sup>3</sup>See Memorizing Transformers' official implementation at https://github.com/google-research/meliad/blob/main/transformer/memoryfactory.py#L78-L79 and https://github.com/google-research/meliad/blob/main/transformer/memory_layer.py#L334-L339
dot-product part of the transformer's attention computation can be rewritten as follows:4
$$
\begin{array}{l} Q K ^ {T} = \left(\boldsymbol {h} _ {d} W _ {q}\right) \left(\boldsymbol {h} _ {e} W _ {k}\right) ^ {\top} \tag {2} \\ = \left(\boldsymbol {h} _ {d} W _ {q}\right) W _ {k} ^ {\top} \boldsymbol {h} _ {e} ^ {\top} \\ = \left(\boldsymbol {h} _ {d} W _ {q} W _ {k} ^ {\top}\right) \boldsymbol {h} _ {e} ^ {\top} \\ \end{array}
$$
Thus, the retrieval step can be formulated as choosing the encoder hidden states  $\pmb{h}_{e}$  that maximize  $\left(\pmb{h}_{d}W_{q}W_{k}^{\top}\right)\pmb{h}_{e}^{\top}$ . This rewriting has two major advantages: first, there is no need to index the keys for each head and layer separately: we can create a single index of the hidden states  $\pmb{h}_{e}$  only, and just project the queries to  $\pmb{h}_{d}W_{q}W_{k}^{\top}$  using head-specific and layer-specific  $W_{q}$  and  $W_{k}$ ; second, the values can be calculated trivially given  $\pmb{h}_{e}$ , so there is no need to store the values in a separate index from the keys before decoding. Thus, instead of constructing  $2\times L\times H$  indexes and retrieving from all indexes during each decoding step, we construct a single index from  $\pmb{h}_{e}$  and retrieve from it by just projecting the decoder hidden states to per-head per-layer  $\pmb{h}_{d}W_{q}W_{k}^{\top}$ .
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,000,000 input tokens. Since indexes can be offloaded to the CPU memory, Unlimformer's input length is practically unlimited.
3 Training Unlimiformer
Unlimformer can be used, at test time, with an already-trained model, and lead to gains without further training, as we show later in Table 3. Next, we turn our focus to training approaches to further improve the performance of Unlimformer. Table 1 summarizes and contrasts the training approaches described below, and Appendix A contains further implementation details.
Method name	Training input	total # tokens in example seen at training time	Validation input (early stopping)	Test input
Baseline	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 (Ivgi et al., 2022)	16k	16k	16k	16k
Longformer (Beltagy et al., 2020)	16k	16k	16k	16k
Random-encoded training	8-16k	8-16k	unlimited	unlimited
Retrieval training	8-16k	8-16k	unlimited	unlimited
Alternating training	8-16k	8-16k	unlimited	unlimited
Table 1: A comparison of the training approaches using BART (context window size 1024) as a running example. The dashed line separates methods that are approximately the same training-time cost as the baseline, from those that require significant additional compute.
3.1 Low (additional-) Cost Training Methods: Applying Unlimformer at validation or test-time only
We first consider training approaches that do not require significant additional compute as compared to the standard finetuning regime.
+test Unlimformer: As the simplest case, we use a standard fine-tuning regime, where the input is truncated during training. At inference time only, we inject Unlimformer into the trained model to process full-length inputs.
+early stop w/ Unlimformer: We train without Unlimformer, but when we evaluate the model for early stopping, we use Unlimformer 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 Unlimformer over the validation set.
For brevity, we omit the linear layers' bias term, as softmax is invariant to constants added to all inputs.
Train chunked +test Unlimformer: As a data augmentation approach, we split each training example 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 Unlimformer model, but has the advantage that all tokens from the full-length training example are observed during training instead of truncated—albeit across several examples. We apply early stopping with Unlimformer on 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 Unlimformer directly, which introduces additional computational cost.
Random-encoded training: At each training step, the full (longer-than-context-window) training example is encoded in chunks; then, the keys for each decoder layer are chosen randomly from the encoded hidden states. This weakly simulates a nearest-neighbors search, but is computationally cheaper.
Retrieval training: At each training step, the keys for each decoder head and layer are selected using a kNN search. When inputs are longer than 16k tokens, we truncated the input to 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-encoded training and Retrieval 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.
Dataset	Avg # tokens	Input length distribution
Domain	# examples	Input	Output
GovReport	Government	19,402	9,616	597	74	303192
SummScreen	TV shows	4,348	8,987	137	2365	22635
BookSum	Literature	436	143,301	1294	8388	642376
Table 2: Dataset statistics. The last column is a 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 on either end. The dotted line indicates the mean length.
We experiment with two long-document- and one book-summarization 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) is a long-document summarization dataset where the task is to write the executive summary of a US government report. SummScreen (Chen et al., 2022) is a long-document summarization dataset where the task is to write the recap of a TV show episode (such as "Friends"), given the transcript of the entire episode. BookSum (Krysciński et al., 2021) is a book-summarization dataset of entire books. BookSum has paragraph, chapter, and book-level settings; we consider the hardest BOOKSUM-Book setting, where 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 F1 (Zhang et al., 2019). Following Zhang et al. (2021), in BookSum we also used Entity Mention Recall ("EntMent") as a 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.
BART (base) (Lewis et al., 2020b) is a pretrained seq2seq model (139M parameters), commonly used for summarization tasks. Its maximum input sequence length is 1024 tokens.
Base model	Training method	ROUGE 1 / 2 / L / BERTScore
GovReport	SummScreen
BARTbase	Standard finetuning	48.7 / 19.2 / 22.8 / 64.3	29.7 / 6.2 / 17.7 / 56.3
BARTbase	+test SLED (Ivgi et al., 2022)	45.8 / 16.1 / 20.2 / 62.7	27.5 / 5.5 / 16.7 / 55.9
BARTbase	+test Unlimiformer	49.7 / 19.6 / 22.0 / 64.8	30.9 / 6.5 / 18.2 / 57.5
BARTbase	+early stop w/ Unlimiformer	51.0 / 20.5 / 21.5 / 65.1	32.1 / 6.8 / 18.6 / 57.6
BARTbase	Train chunked	46.2 / 17.8 / 21.7 / 63.3	28.1 / 5.6 / 17.0 / 55.6
BARTbase	+test Unlimiformer	53.4 / 22.5 / 22.5 / 66.0	29.3 / 6.6 / 17.6 / 57.0
PRIMERA	Standard finetuning	55.1 / 23.9 / 25.9 / 67.0	32.3 / 7.1 / 18.3 / 57.1
PRIMERA	+test Unlimiformer	56.5 / 24.8 / 26.3 / 67.7	33.3 / 7.7 / 19.1 / 57.6
Table 3: Results on long-document summarization, low-cost training methods: the training costs are no higher than standard finetuning that truncates the inputs to the model's max input size. The best metric in every training category is marked in **bold**. PRIMERA (Xiao et al., 2022) is a Longformer-Encoder-Decoder (Beltagy et al., 2020) with additional summarization-specific pretraining.
PRIMERA (Xiao et al., 2022) is a Longformer-Encoder-Decoder  $(\mathrm{LED}_{\text {large }}$  ; Beltagy et al., 2020) (447M parameters), pretrained specifically for multi-document summarization, with maximum input length of 4096 tokens.
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<sup>5</sup> 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.<sup>6</sup> 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 Unlimformer 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 Unlimformer 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.1 Long Document Summarization
Low-cost training Table 3 shows the results in the long-document summarization datasets. First, we can see that applying Unlimformer on an existing checkpoint without any training (+test Unlimformer) improves  $\mathrm{BART_{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, Unlimformer is the only model that can provide benefits without further training.
Early stop  $w/$  Unlimformer 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 Unlimformer 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.
<sup>5</sup>https://github.com/google-research/meliad
<sup>6</sup>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.
Base model	Training method	ROUGE 1 / 2 / L / BERTScore
GovReport	SummScreen
BARTbase	Standard finetuning	48.7 / 19.2 / 22.8 / 64.3	29.7 / 6.2 / 17.7 / 56.3
BARTbase	SLED (Ivgi et al., 2022)	54.7 / 24.4 / 25.4 / 67.0	32.7 / 7.9 / 19.1 / 58.4
BARTbase	Memorizing transformers	55.2 / 25.1 / 26.4 / 67.5	32.7 / 7.4 / 19.2 / 57.4
BARTbase	Unlimiformer (this work)	56.6 / 26.3 / 27.6 / 68.2	34.7 / 8.5 / 19.9 / 58.5
PRIMERA	Standard finetuning	55.1 / 23.9 / 25.9 / 67.0	32.3 / 7.1 / 18.3 / 57.1
PRIMERA	Memorizing transformers	57.0 / 25.3 / 26.5 / 67.7	33.0 / 7.3 / 18.4 / 57.3
PRIMERA	Unlimiformer (this work)	57.4 / 26.2 / 28.0 / 68.1	33.3 / 7.6 / 18.9 / 57.7
Table 4: Test results on long-document datasets, when allowing compute-costly, long-range training methods, using different base models. The best metric in every dataset and every training category is marked in bold. The Unlimformer results in this table are from using the alternating training strategy.
Base model	Training method	ROUGE 1 / 2 / L	EntMent
BARTbase	Hierarchical (Krysciński et al., 2021)	30.0 / 6.0 / 11.0	-
BARTbase	Standard finetuning	36.4 / 7.6 / 15.3	10.0
BARTbase	+test Unlimiformer	35.5 / 7.7 / 15.4	21.9
BARTbase	+early stop w/ Unlimiformer	35.5 / 7.7 / 15.4	21.9
BARTbase	Memorizing Transformers	35.6 / 6.4 / 14.6	10.1
BARTbase	Unlimiformer (retrieval training)	36.8 / 8.3 / 15.7	20.3
BARTbase	Unlimiformer (random-encoded training)	37.3 / 6.7 / 15.2	20.8
BARTbase	Unlimiformer (alternating training)	36.7 / 7.3 / 15.5	20.3
PRIMERA	Standard finetuning	38.6 / 7.2 / 15.6	11.6
PRIMERA	+test Unlimiformer	38.3 / 7.5 / 15.9	18.9
PRIMERA	+early stop w/ Unlimiformer	39.5 / 7.3 / 15.8	22.2
PRIMERA	Unlimiformer (retrieval training)	37.9 / 8.2 / 16.3	25.5
PRIMERA	Unlimiformer (random-encoded training)	39.5 / 7.1 / 15.9	19.7
PRIMERA	Unlimiformer (alternating training)	38.2 / 7.1 / 16.0	23.4
Table 5: Results on BookSum (average input length ≈ 143k tokens). EntMent is entity recall. Hierarchical summarization is a baseline reported by Krysciński et al. (2021), where chapter summaries are condensed to form a book summary. The best metric in every dataset is marked in bold.
Long-range training Table 4 shows results when allowing computationally expensive training approaches. As shown, in almost all metrics and datasets, Unlimformer outperforms the SLED and Memorizing Transformers baselines when using the same base model.