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 https://github.com/abertsch72/unlimiformer .

翻译：基于Transformer的模型通常对输入长度有预设限制，这是因为它们需要潜在地关注输入中的每个词元。本文提出Unlimiformer：一种通用方法，可封装任何现有的预训练编码器-解码器Transformer，并将所有层的注意力计算卸载到单个$k$-近邻索引中；该索引可存放在GPU或CPU内存中，并以次线性时间进行查询。通过这种方式，我们能够索引极长的输入序列，同时每个解码器层的注意力头仅检索其前$k$个键，而非关注所有键。我们在多个长文档和多文档摘要基准上验证了Unlimiformer的有效性，表明它能在测试时无需截断输入的情况下，摘要BookSum数据集中长达35万个词元的输入。Unlimiformer无需额外学习权重或修改代码，即可扩展BART和Longformer等预训练模型至无限输入。我们的代码和模型已公开发布于https://github.com/abertsch72/unlimiformer。