直接扩展到无限长,谷歌Infini-Transformer终结上下文长度之争
作者:LR •更新时间:2025-08-07 17:56:47•阅读 0
不知 Gemini 1.5 Pro 是否用到了这项技术。
谷歌又放大招了,发布下一代 Tran**ormer 模型 Infini-Tran**ormer。
Infini-Tran**ormer 引入了一种有效的方法,可以将基于 Tran**ormer 的大型语言模型 (LLM) 扩展到**长输入,而不增加内存和计算需求。使用该技术,研究者成功将一个 1B 的模型上下文长度提高到 100 万;应用到 8B 模型上,模型能处理 500K 的书籍摘要任务。自 2017 年开创性研究论文《Attention is All You Need》问世以来,Tran**ormer 架构就一直主导着生成式人工智能领域。而谷歌对 Tran**ormer 的优化设计最近比较**,几天前,他们更新了 Tran**ormer 架构,发布 Mixture-of-Depths(MoD),改变了以往 Tran**ormer 计算模式。没过几天,谷歌又放出了这项新研究。
专注 AI 领域的研究者都了解内存的重要性,它是智能的基石,可以为 LLM 提供**的计算。然而,Tran**ormer 和基于 Tran**ormer 的 LLM 由于注意力机制的固有特性,即 Tran**ormer 中的注意力机制在内存占用和计算时间上都表现出二次复杂性。例如,对于批大小为 512、上下文长度为 2048 的 500B 模型,注意力键 - 值 (KV) 状态的内存占用为 3TB。但事实上,标准 Tran**ormer 架构有时需要将 LLM 扩展到更长的序列(如 100 万 token),这就带来巨大的内存开销,并且随着上下文长度的增加,部署成本也在增加。
基于此,谷歌引入了一种有效的方法,其关键组成部分是一种称为 Infini-attention(**注意力)的新注意力技术。不同于传统的 Tran**ormer 使用局部注意力丢弃旧片段,为新片段释放内存空间。Infini-attention 增加了压缩内存(compressive memory),可以将使用后的旧片段存储到压缩内存中,输出时会聚合当前上下文信息以及压缩内存中的信息,因而模型可以检索完整的上下文历史。
该方法使 Tran**ormer LLM 在有限内存的情况下扩展到**长上下文,并以流的方式处理极长的输入进行计算。
实验表明,该方法在长上下文语言建模基准测试中的性能优于基线,同时内存参数减少了 100 倍以上。当使用 100K 序列长度进行训练时,该模型实现了更好的困惑度。此外该研究发现,1B 模型在 5K 序列长度的密钥实例上进行了微调,解决了 1M 长度的问题。**,论文展示了具有 Infini-attention 的 8B 模型经过持续的预训练和任务微调,在 500K 长度的书籍摘要任务上达到了新的 SOTA 结果。
本文贡献总结如下:
- 引入了一种实用且强大的注意力机制 Infini-attention—— 具有长期压缩内存和局部因果注意力,可用于有效地建模长期和短期上下文依赖关系;
- Infini-attention 对标准缩放点积注意力( standard scaled dot-product attention)进行了**的改变,并通过设计支持即插即用的持续预训练和长上下文自适应;
- 该方法使 Tran**ormer LLM 能够通过流的方式处理极长的输入,在有限的内存和计算资源下扩展到**长的上下文。

- 论文链接:https://arxiv.org/pdf/2404.07143.pdf
- 论文标题:Leave No Context Behind: Efficient Infinite Context Tran**ormers with Infini-attention
方法介绍
Infini-attention 使 Tran**ormer LLM 能够通过有限的内存占用和计算有效地处理**长的输入。如下图 1 所示,Infini-attention 将压缩记忆融入到普通的注意力机制中,并在单个 Tran**ormer 块中构建了掩码局部注意力和长期线性注意力机制。

Infini-attention 采用标准注意力计算的所有键、值和查询状态,以进行长期记忆巩固(memory consolidation)和检索,并将注意力的旧 KV 状态存储在压缩内存中,而不是像标准注意力机制那样丢弃它们。在处理后续序列时,Infini-attention 使用注意查询状态从内存中检索值。为了计算**的上下文输出,Infini-attention 聚合了长期记忆检索值和局部注意力上下文。
如下图 2 所示,研究团队比较了基于 Infini-attention 的 Infini-Tran**ormer 和 Tran**ormer-XL。与 Tran**ormer-XL 类似,Infini-Tran**ormer 对 segment 序列进行操作,并计算每个 segment 中的标准因果点积注意力上下文。因此,点积注意力计算在某种意义上是局部的。

与多头注意力(MHA)类似,除了点积注意力之外,Infini-attention 还为每个注意力层维护 H 个并行压缩内存(H 是注意力头的数量)。


该研究在长上下文语言建模、长度为 1M 的密钥上下文块检索和 500K 长度的书籍摘要任务上评估了 Infini-Tran**ormer 模型,这些任务具有极长的输入序列。对于语言建模,研究者选择从头开始训练模型,而对于密钥和书籍摘要任务,研究者采用不断预训练 LLM 的方式,以证明 Infini-attention 即插即用的长上下文适应能力。
长上下文语言建模。表 2 结果表明 Infini-Tran**ormer 优于 Tran**ormer-XL 和 Memorizing Tran**ormers 基线,并且与 Memorizing Tran**ormer 模型相比,存储参数减少了 114 倍。


