Unlimiformer：支持无限长度输入的长程 Transformer
计算与语言
2023-11-01 v3
摘要
自 Transformer 提出以来，由于需要对输入中每个 token 进行注意力计算，这类模型一直受限于有界的输入长度。本文提出 Unlimiformer：一种通用方法，可封装任意已有的预训练编码器-解码器 Transformer，并将交叉注意力计算卸载到单一的 k-最近邻（kNN）索引上，而返回的 kNN 距离即为注意力的点积分数。该 kNN 索引可保存在 GPU 或 CPU 内存中，并以亚线性时间查询；由此，我们几乎可以索引无限长的输入序列，同时每个解码器层中每个注意力头检索其 top-k 个键，而非关注每一个键。我们在多个长文档和书籍摘要基准上评估 Unlimiformer，表明其能够处理 BookSum 数据集中长达 50 万 token 的输入，且在测试时无需任何输入截断。我们证明，Unlimiformer 通过扩展 BART 和 Longformer 等预训练模型以支持无限输入，无需额外可学习权重，也无需修改其代码。我们的代码和模型已在 https://github.com/abertsch72/unlimiformer 公开。
引用
@article{arxiv.2305.01625,
title = {Unlimiformer: Long-Range Transformers with Unlimited Length Input},
author = {Amanda Bertsch and Uri Alon and Graham Neubig and Matthew R. Gormley},
journal= {arXiv preprint arXiv:2305.01625},
year = {2023}
}
备注
NeurIPS 2023