无限上下文方案Leave No Context Behind

Infini-attention是一种创新的注意力机制,它将压缩记忆与局部causal注意力和长期线性注意力相结合,极大地增强了Transformer语言模型处理超长文本的能力。局部注意力机制负责处理当前的文本片段,而长期注意力则从压缩记忆中检索出与当前上下文最相关的历史信息。Infini-attention巧妙地复用点积注意力层已有的query,key和value状态来更新压缩记忆,避免了额外的计算和存储开销。压缩记忆通过一个关联矩阵以线性时空复杂度存储和检索历史上下文的key-value绑定,使得模型可以在保持记忆大小不变的情况下处理任意长度的输入序列。

将Infini-attention集成到Transformer结构中,就得到了一种新的语言模型架构——Infini-Transformer。Infini-Transformer在多个长文本建模任务上展现出了卓越的性能。例如,在长距离语言建模基准测试中,它以114倍的记忆体积压缩率超越了之前的最佳结果;在百万长度的passkey上下文检索任务中,仅需要在5000长度的序列上进行短暂的微调就轻松解决;在50万token的超长书籍摘要任务上,也刷新了SOTA成绩。

与之前的长文本Transformer模型相比,Infini-Transformer展现出了更优的性能和效率。相比于Transformer-XL和Memorizing Transformers等基于片段级别显式存储上下文的方法,Infini-Transformer以较小的参数开销实现了更大幅度的上下文压缩,并在准确性上更胜一筹。相比于RMT和AutoCompressor等基于输入压缩的方法,Infini-attention采用了一种即插即用的设计,只需要对预训练的语言模型进行轻量级的增量训练,就能灵活高效地适应超长文本推理,在工程实践中更具优势。

Infini-attention为Transformer语言模型打开了高效处理超长文本的大门,有望在对话、问答、摘要、检索等领域大幅拓展语言模型的应用范围和性能上限。它简洁优雅的设计也为后续研究指明了可能的方向,即进一步探索基于神经记忆(而非显式存储)的高效注意力机制。