Unlimiformer: Long-Range Transformers with Unlimited Length Input

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

Keywords: retrieval augmentation, summarization, long-context, generation, long-input, encoder-decoder, transformers, language models, natural language generation, natural language processing, deep learning, neural networks

TL;DR: Unlimiformer modifies pretrained encoder-decoder models with retrieval to allow unlimited input at test time.

Abstract: 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 ($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 and Longformer by extending them to unlimited inputs without additional learned weights and without modifying their code. Our code and models are publicly available at https://github.com/abertsch72/unlimiformer , and support LLaMA-2 as well.

Submission Number: 3720