Unlimiformer: Long-Range Transformers with Unlimited Length Input

Transformer-based models typically have a predefined bound to their input length, because of their need to potentially attend to every token in the input. In this work, we propose Unlimiformer: a general approach that can wrap any existing pretrained encoder-decoder transformer, and offload the attention computation across all layers to a single $k$-nearest-neighbor index; this index can be kept on either the GPU or CPU memory and queried in sub-linear time. This way, we can index extremely long input sequences, while every attention head in every decoder layer retrieves its top-$k$ keys, instead of attending to every key. We demonstrate Unlimiformers's efficacy on several long-document and multi-document summarization benchmarks, showing that it can summarize even 350k token-long inputs from the BookSum dataset, without any input truncation at test time. 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 this https URL .

Transformer-based models typically have a predefined bound to their input length, because of their need to potentially attend to every token in the input. In this work, we propose Unlimiformer: a general approach that can wrap any existing pretrained encoder-decoder transformer, and offload the attention computation across all layers to a single $k$-nearest-neighbor index; this index can be kept on either the GPU or CPU memory and queried in sub-linear time. This way, we can index extremely long input sequences, while every attention head in every decoder layer retrieves its top-$k$ keys, inst

related reading

• A Mathematical Framework for Transformer Circuitstransformer-circuits.pub
• The Annotated Transformernlp.seas.harvard.edu
• Extending Context is Hard | kaiokendevkaiokendev.github.io
• Transformer Explainer: LLM Transformer Model Visually Explainedpoloclub.github.io
• Transformers from Scratche2eml.school
• Transformers Inference Optimization Toolset | AstraBlogastralord.github.io
• The Transformer Family Version 2.0 | Lil'Loglilianweng.github.io
• On the Tradeoffs of SSMs and Transformers | Goomba Labgoombalab.github.io
• The Annotated Transformernlp.seas.harvard.edu
• 1706.03762arxiv.org
• Getting Caught Up to Modern LLM Research | Samarth Goeldev.samarthgoel.com
• Overleaf Examplearxiv.org