PaperlayerView exact revision on arXiv

Source-linked AI summary

Unlimiformer: Long-Range Transformers with Unlimited Length Input

Amanda Bertsch, Uri Alon, Graham Neubig, Matthew R. Gormley

arXiv:2305.01625v3cs.CL

TL;DR

Pretrained encoder-decoder transformers struggle with inputs far beyond their bounded context windows, especially in long-document tasks. Unlimiformer uses a shared kNN index to replace full cross-attention with top-k retrieval, extending existing models to unbounded inputs. It processes inputs up to 500k tokens and improves pretrained and long-range models without additional learned weights or code changes.

  • Problem

    Long-document tasks can contain inputs exceeding 500k tokens, beyond the bounded context windows of pretrained transformers and many long-context models.

  • Method

    Unlimiformer indexes encoder hidden states and lets each decoder cross-attention head retrieve its top-k keys from the full input through kNN search.

  • Results

    Unlimiformer processes inputs up to 500k tokens and improves existing models, including without further training.

  • Takeaways & Limitations

    The approach extends pretrained encoder-decoder models to unlimited inputs without adding learned weights or modifying individual architecture code.

  • Takeaways & Limitations

    Training input length remains limited by GPU memory, while CPU index offloading enables longer inference inputs at higher test-time latency.

Abstract

from arXiv · show

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 (kNN) index, while the returned kNN distances are the attention dot-product scores. This kNN 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. We make our code and models publicly available at https://github.com/abertsch72/unlimiformer .

1 Introduction

Pretrained transformers are bounded by context windows, while some narrative and long-document tasks contain inputs far beyond those limits. Unlimiformer addresses this gap by extending pretrained encoder-decoder models to unbounded inputs through retrieval-based cross-attention without adding weights or changing the architecture.

  • Book summarization inputs can exceed 500k tokens, with the longest datasets exceeding Longformer’s context window by more than 34 times.
  • Naïve self-attention scales quadratically, while long-input architectures often require architectural changes and costly re-pretraining or additional training.
  • Unlimiformer indexes hidden states for all input tokens and uses kNN retrieval so every decoder cross-attention head attends only to its top-k keys.
  • Preliminary experiments found that the top-k attention keys cover more than 99% of the attention mass.
  • Unlimiformer can improve trained checkpoints without adding weights or further training, and can also improve strong long-range models when applied on top of them.

2 Unlimiformer

Unlimiformer encodes long inputs in chunks, indexes their hidden states, and replaces full cross-attention with per-head retrieval from the entire sequence. An attention reformulation permits one shared index across layers and heads while preserving standard dot-product attention.

  • Long inputs are encoded in overlapping chunks, retained hidden vectors are indexed with dot-product similarity, and the index is queried before each decoder cross-attention operation.
  • Each cross-attention head retrieves top-k hidden states from the full input instead of attending only to a truncated prefix.
  • Attention reformulation: Reordering the attention computation allows a single index to serve all decoder layers and attention heads without changing standard dot-product attention.
  • Attention reformulation: The reformulation stores one vector per input token; with 16-bit floats and hidden size 1024, one million tokens require 2GB of memory.
  • The index may be offloaded to CPU memory, making input length practically unlimited while trading lower memory pressure for slower retrieval.

3 Training Unlimiformer

The paper evaluates ways to apply Unlimiformer during training, ranging from inference-only use with standard fine-tuning to retrieval-based training. These approaches differ in whether they preserve baseline training cost or add substantial computation.

  • Test-time Unlimiformer applies retrieval only during inference to a conventionally fine-tuned model whose training inputs were truncated.
  • Chunked training exposes the model to all tokens by splitting each long example into non-overlapping context-window-sized examples.
  • Random-encoded training samples keys from chunked encodings, while retrieval training selects keys with kNN search at each decoder head and layer.
  • Retrieval training truncates inputs to 16k tokens during training because of GPU memory requirements.
  • Alternating training combines random-encoded and retrieval training, with random encoding regularizing attention to include non-top-k keys.

4 Experimental Settings

The experiments cover two long-document datasets and the hardest book-level BookSum setting, using standard summarization metrics and entity recall. Baselines include pretrained BART, PRIMERA, SLED, and Memorizing Transformers, while Table 3 emphasizes training costs no higher than truncated standard fine-tuning.

  • Datasets: The evaluation uses GovReport, SummScreen, and the hardest BOOKSUM-Book setting, spanning government reports, TV transcripts, and entire novels.BOOKSUM-Book generates a book-level summary from the full novel text.
  • Metrics: Performance is measured with ROUGE-1/2/L and BERTScore F1, plus Entity Mention Recall as a BookSum informativeness proxy.Entity Mention Recall measures the fraction of gold entities mentioned in a candidate summary.
  • Baselines: BARTbase is a 139M-parameter pretrained sequence-to-sequence baseline with a 1024-token maximum input length.
  • Baselines: PRIMERA is a 447M-parameter Longformer-Encoder-Decoder pretrained for multi-document summarization, with a 4096-token maximum input length.
  • Baselines: SLED encodes long inputs in chunks and lets the decoder attend to all tokens, but its input length remains memory bounded after expensive fine-tuning.
  • Baselines: Memorizing Transformers add a trainable attention gate, whereas Unlimiformer is non-parametric and can retrieve in every decoder layer with individualized per-head retrieval.The comparison passage notes that Unlimiformer can improve pretrained language models without fine-tuning.

5 Results

Unlimiformer improves long-document and book-summarization results across low-cost and long-range training settings, including gains over strong long-range baselines. On BookSum, its larger entity recall gains expose limitations of overlap-based metrics for evaluating full-book summaries.

  • Low-cost training: 1.8 ROUGE-1 points improve BARTbase on both long-document datasets without additional training, while PRIMERA gains 1–1.4 ROUGE-1 points.
  • Low-cost training: 3.3 ROUGE-1 points improve GovReport with Early stop w/ Unlimiformer at identical training cost to standard fine-tuning.
  • Low-cost training: 7.2 ROUGE-1 and 3 BERTScore points improve results when Unlimiformer is injected after chunked training, with standard fine-tuning-level computational cost.
  • Long-range training: In almost all metrics and datasets, Unlimiformer outperforms SLED and Memorizing Transformers when using the same base model.
  • Long-range training: Unlimiformer+BARTbase beats base PRIMERA across all metrics and datasets, while Unlimiformer+PRIMERA further improves PRIMERA across all metrics and datasets.
  • Book summarization: Unlimiformer improves both BARTbase and PRIMERA on BookSum across low-cost and long-range training approaches.
  • Book summarization: BookSum’s truncated BART baseline has competitive ROUGE and BERTScore despite limited coherence and frequent hallucination, which overlap metrics do not reflect strongly.
  • Book summarization: Unlimiformer’s entity recall generally rises with datastore size while inference time increases sublinearly, and it outperforms BART at all datastore sizes.

6 Analysis

Unlimiformer benefits from longer inputs on BookSum, but not every dataset requires the full context. Its retrieval and encoding strategy also introduces computational costs that grow sublinearly overall.

  • EntMent increases almost monotonically with BookSum input length, suggesting Unlimiformer exploits longer inputs to generate better outputs.
  • Unlimiformer did not improve over the first 1024 TF-IDF-ranked tokens on WikiSum, suggesting the full input is unnecessary for that dataset.
  • Full-input encoding, index construction, and search increase processing time, while total GPU-time grows sublinearly with input length.
  • On three of four additional datasets, BART-Unlimiformer with alternating training improves over base BART.
  • 43.5% of BookSum input tokens were retrieved at least once by test-time-only Unlimiformer, versus 64.5% with alternating training.
  • Retrieved keys came from across BookSum inputs, with median locations between 49.73% and 49.87% of each document.

7 Related Work

Unlimiformer differs from prior long-range approaches by using a shared kNN index for cross-attention across all decoder layers and heads. This design supports broader retrieval than comparable methods while avoiding some of their architectural costs.

  • Most prior long-range transformers reduce resource requirements by sparsifying, approximating, or replacing the attention mechanism.
  • Memorizing Transformers use separate datastores and, because of memory constraints, apply their approach only to a single decoder layer.
  • Unlimiformer uses one index for all decoder layers, allowing every cross-attention head to retrieve from the long context.
  • SLED attends to all encoded inputs simultaneously and is practically limited to about 16k tokens on a single GPU.
  • Unlike SLED, Unlimiformer attends only to the top-k input tokens for every attention head, enabling unlimited inputs in practice.

8 Conclusions

Unlimiformer augments pretrained encoder-decoder models with kNN-based cross-attention for unlimited-length inputs. Across benchmarks, it improves existing models, including Longformer-based systems, with or without further training.

  • Unlimiformer offloads cross-attention to a kNN index so every decoder-layer head retrieves only its top-k keys from the full input.
  • The method is evaluated on long-document and book-summarization benchmarks with inputs up to 500K tokens.
  • Unlimiformer improves existing models even without further training.
  • With Unlimiformer training, smaller BART models outperform larger Longformer-based models, while Longformer-based models can also improve further.
  • The approach aims to obtain long-context benefits by modifying already-pretrained models with less computational cost than pretraining large models.

9 Limitations

The experiments leave important scope and resource boundaries: training length is constrained by GPU memory, inference latency depends on index placement, and evaluations cover only English datasets.

  • Experiments consider only English-language datasets, and nearest-neighbor quality depends on the quality of indexed keys for other languages.
  • Training input length is limited by GPU memory because embeddings and their computational graph must remain stored for backpropagation.
  • CPU index offloading supports the longest inference inputs but increases test-time latency relative to GPU storage.
  • At test time, full inputs are used without truncation, but computationally expensive training methods generally truncate the longest inputs to 16k tokens.
  • The WikiSum evaluation uses a rescraped dataset with shorter-than-reported inputs and excludes inputs below 1457 words.

C Evaluation details

The evaluation uses context-adjusted BERTScore variants and Entity Mention Recall to assess generated summaries, with metric choices determined by input length and entity coverage.

  • BERTScore is evaluated with context size 1024 for GovReport and ScriptSumm because vanilla BERTScore is defined only up to 512 tokens.
  • For BookSum, Longformer-large with a 4096-token context window failed to distinguish model outputs, ranking all tested models within 0.3 points.
  • Entity Mention Recall measures the percentage of unique named entities from the gold summary that appear in each candidate summary.
  • Named entities for Entity Mention Recall are identified with SpaCy’s en_core_web_lg model.

D Computational Cost

Unlimiformer increases computational and memory demands because it encodes full inputs and builds/searches an index, while maintaining GPU-memory scalability across decoder layers.

  • Inference costs rise most because BookSum inputs average 112,885 tokens and must be fully encoded instead of truncating to 1,024 tokens.
  • CPU datastores are many times slower than GPU datastores because of slower search and transferring retrieved embeddings to the GPU.
  • Table 7 reports relative computational effort per epoch across training methodologies, averaged over three BookSum runs on one 48-GB A6000.
  • Table 8 reports memory consumption when applying Unlimiformer to different numbers of BART layers.
  • GPU memory remains constant as Unlimiformer is applied to more BART layers, unlike Memorizing Transformers, which allocate more memory per layer and head.

F Sample Outputs

Sample outputs compare full-input Unlimiformer summaries with truncated BART summaries on BookSum, showing that full-input access produces more constrained hallucinations and broader story coverage.

  • BookSum examples: Unlimiformer’s Brothers Karamazov summary has more constrained hallucinations than the input-truncated model’s summary.The truncated output invents plotlines involving a Congo lover and Pavel’s deaths, whereas Unlimiformer makes narrower factual errors.
  • Reference: The reference summary presents The Brothers Karamazov as a family tragedy centered on a father and his sons.
  • Reference: The reference describes Dmitri’s inheritance dispute, Alyosha’s mediation attempt, and Dmitri’s debts involving Katerina and Grushenka.
  • BookSum examples: The truncated BART output introduces unsupported characters, places, illnesses, and deaths while inaccurately describing the Karamazov family.
  • BookSum examples: The Unlimiformer output identifies the novel’s central question as whether Dmitri Fyodorovitch killed his father.
  • Related approaches: Prior long-document summarization methods select subsections, summarize chunks for later fusion, or use other efficient cross-window processing strategies.
Loading 2305.01625v3…