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

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

# Unlimiformer: Long-Range Transformers with Unlimited Length Input

###### Abstract

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 k -nearest-neighbor ( k k NN) index, while the returned k k NN distances are the attention dot-product scores. This k 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 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 1 1 1 https://github.com/abertsch72/unlimiformer .

## 1 Introduction

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 ( Narayan et al., 2018 , XSum; ) ( Nallapati et al., 2016 , CNN/DM; ) . 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 ( Kryściń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 resources. Other architectures such as Longformer-Encoder-Decoder ( Beltagy et al., 2020 , LED; ) 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 k -nearest-neighbor ( k k NN) index over the hidden states of all input tokens. Then, every standard cross-attention head in every decoder layer queries the k k NN index, such that the k k NN distances are the attention dot-product scores, and attends only to the top- k k input tokens. In preliminary experiments, we found that the top- k k attention keys cover more than 99% of the attention mass, and thus attending only to the top- k 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 Unlimiformer 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 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 a k k NN search into each decoder layer: prior to cross-attention, the model performs a nearest-neighbor search in a k k NN 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 k k NN index, using a library such as Faiss ( 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 k first tokens in the input sequence.

Instead of attending only to this k k -token prefix of the input, we retrieve the top- k k hidden states from the k k NN index for each cross-attention head, and attend only to these top- k 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 k k NN index; then, the index of encoded hidden states is queried at each decoding step. The k k NN 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 𝒉 d {\bm{h}}_{d} be the decoder hidden state and 𝒉 e {\bm{h}}_{e} be an encoder’s last layer hidden state. The standard cross-attention computation for a single head in a transformer is: Attn ⁡ ( Q , K , V ) = softmax ⁡ ( Q ​ K T d k ) ​ V \mathrm{Attn}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^{T}}{\sqrt{d_{k}}}\right)V (1) where Q = 𝒉 d ​ W q Q={{\bm{h}}_{d}}W_{q} is the product of the decoder states 𝒉 d {\bm{h}}_{d} and the query weight matrix W q W_{q} ; the keys K = 𝒉 e ​ W k K={{\bm{h}}_{e}}W_{k} are the product of the last encoder hidden states 𝒉 e {\bm{h}}_{e} with the key weight matrix W k W_{k} ; and V = 𝒉 e ​ W v V={{\bm{h}}_{e}}W_{v} is similarly the product of 𝒉 e {\bm{h}}_{e} with the value weight matrix W v W_{v} . Our goal is to retrieve a set of keys K b ​ e ​ s ​ t K_{best} that maximize Q ​ K b ​ e ​ s ​ t T QK_{best}^{T} , with the size of K b ​ e ​ s ​ t K_{best} fixed to the size of the model’s context window, and then compute the standard attention over K b ​ e ​ s ​ t K_{best} only.

Note that the linear layers W q W_{q} , W k W_{k} , and W v W_{v} are layer-specific and head-specific. Thus, naïvely creating an index from the keys K = 𝒉 e ​ W k K={{\bm{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 × L × H 2\times L\times H indexes, where L L is the number of decoder layers and H H is the number of attention heads. In fact, this exact naïve approach was taken by Memorizing Transformers ( Wu et al., 2022 ) , who pioneered the use of a k k NN index for previously encoded inputs. 2 2 2 See Memorizing Transformers’ official implementation at https://github.com/google-research/meliad/blob/main/transformer/memory_factory.py#L78-L79 and https://github.com/google-research/meliad/blob/main/transformer/memory_layer.py#L334-L339 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 dot-product part of the transformer’s attention computation can be rewritten as follows: 3 3 3 For brevity, we omit the linear layers’ bias term, as softmax is invariant to constants added to all inputs. Q ​ K T = \displaystyle QK^{T}= ( 𝒉 d ​ W q ) ​ ( 𝒉 e ​ W k ) ⊤ \displaystyle\left({{\bm{h}}_{d}}W_{q}\right)\left({{\bm{h}}_{e}}W_{k}\right)^{\top} (2) = \displaystyle= ( 𝒉 d ​ W q ) ​ W k ⊤ ​ 𝒉 e ⊤ \displaystyle\left({{\bm{h}}_{d}}W_{q}\right)W_{k}^{\top}{{\bm{h}}_{e}}^{\top} = \displaystyle= ( 𝒉 d ​ W q ​ W k ⊤ ) ​ 𝒉 e ⊤ \displaystyle\left({{\bm{h}}_{d}}W_{q}W_{k}^{\top}\right){{\bm{h}}_{e}}^{\top} Thus, the retrieval step can be formulated as choosing the encoder hidden states 𝒉 e {\bm{h}}_{e} that maximize ( 𝒉 d ​ W q ​ W k ⊤ ) ​ 𝒉 e ⊤ \left({{\bm{h}}_{d}}W_{q}W_{k}^{\top}\right){{\bm{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 𝒉 e {\bm{h}}_{e} only, and just project the queries to 𝒉 d ​ W q ​ W k ⊤ {{\bm{h}}_{d}}W_{q}W_{k}^{\top} using head-specific and layer-specific W q W_{q} and W k W_{k} ; second, the values can be calculated trivially given 𝒉 e {\bm{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 × L × H 2\times L\times H indexes and retrieving from all indexes during each decoding step, we construct a single index from 𝒉 e {\bm{h}}_{e} and retrieve from it by just projecting the decoder hidden states to per-head per-layer 𝒉 d ​ W q ​ W k ⊤ {{\bm{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, 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, as we show later in Table 3 . Next, we turn our focus to training approaches to further improve the performance of Unlimiformer. Table 1 summarizes and contrasts the training approaches described below, and Appendix A contains further implementation details.

### 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 as compared to the standard finetuning regime. +test Unlimiformer: As the simplest case, we use a standard fine-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.

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.

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 Unlimiformer 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 Unlimiformer 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 Unlimiformer 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.

At each training step, the keys for each decoder head and layer are selected using a k k NN 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.

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 k keys as well.

## 4 Experimental Settings

### 4.1 Datasets

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 ( Kryściń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 .

### 4.2 Baselines

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.

PRIMERA ( Xiao et al., 2022 ) is a Longformer-Encoder-Decoder ( Beltagy et al., 2020 , LED large {}_{\texttt{large}} ; ) (447M parameters), pretrained specifically for multi-document summarization, with maximum input length of 4096 tokens.