Google又出大绝招了，发表下一代 Transformer 模型 Infini-Transformer。?

Infini-Transformer 引入了一种有效的方法，可以将基于 Transformer 的大型语言模型 (LLM) 拓展到无限长输入，而不增加记忆体和运算需求。使用该技术，研究者成功将一个 1B 的模型上下文长度提高到 100 万；应用到 8B 模型上，模型能处理 500K 的书籍摘要任务。?

自 2017 年开创性研究论文《Attention is All You Need》问世以来，Transformer 架构就一直主导着生成式人工智慧领域。而Google对 Transformer 的优化设计最近比较频繁，几天前，他们更新了 Transformer 架构，发表 Mixture-of-Depths（MoD），改变了以往 Transformer 运算模式。没过几天，Google又放出了这项新研究。

专注 AI 领域的研究者都了解记忆体的重要性，它是智慧的基石，可以为 LLM 提供高效的运算。然而，Transformer 和基于 Transformer 的 LLM 由于注意力机制的固有特性，即 Transformer 中的注意力机制在记忆体佔用和运算时间上都表现出二次複杂性。例如，对于大小为 512、上下文长度为 2048 的 500B 模型，注意力键 - 值 (KV) 状态的记忆体佔用为 3TB。但事实上，标準 Transformer 架构有时需要将 LLM 拓展到更长的序列（如 100 万 token），这就带来巨大的记忆体开销，并且随着上下文长度的增加，部署成本也在增加。?

基于此，Google引入了一种有效的方法，其关键组成部分是一种称为 Infini-attention（无限注意力）的新注意力技术。不同于传统的 Transformer 使用局部注意力丢弃旧片段，为新片段释放记忆体空间。Infini-attention 增加了压缩记忆体（compressive memory），可以将使用后的旧片段储存到压缩记忆体中，输出时会聚合当前上下文资讯以及压缩记忆体中的资讯，因而模型可以检索完整的上下文历史。?

该方法使 Transformer LLM 在有限记忆体的情况下拓展到无限长上下文，并以串流的方式处理极长的输入进行运算。?

实验表明，该方法在长上下文语言建模基準测试中的性能优于基线，同时记忆体参数减少了 100 倍以上。当使用 100K 序列长度进行训练时，该模型实现了更好的困惑度。此外该研究发现，1B 模型在 5K 序列长度的金钥实例上进行了微调，解决了 1M 长度的问题。最后，论文展示了具有 Infini-attention 的 8B 模型经过持续的预训练和任务微调，在 500K 长度的书籍摘要任务上达到了新的 SOTA 结果。?

本文贡献总结如下：?

方法介绍?

Infini-attention 使 Transformer LLM 能够透过有限的记忆体佔用和运算有效地处理无限长的输入。如下图 1 所示，Infini-attention 将压缩记忆融入到普通的注意力机制中，并在单个 Transformer 块中构建了遮罩局部注意力和长期线性注意力机制。

对 Transformer 注意力层进行这种微妙但关键的修改可以透过持续的预训练和微调将现有 LLM 的上下文视窗拓展到无限长。?

Infini-attention 採用标準注意力运算的所有键、值和查询状态，以进行长期记忆巩固（memory consolidation）和检索，并将注意力的旧 KV 状态储存在压缩记忆体中，而不是像标準注意力机制那样丢弃它们。在处理后续序列时，Infini-attention 使用注意查询状态从记忆体中检索值。为了运算最终的上下文输出，Infini-attention 聚合了长期记忆检索值和局部注意力上下文。?

如下图 2 所示，研究团队比较了基于 Infini-attention 的 Infini-Transformer 和 Transformer-XL。与 Transformer-XL 类似，Infini-Transformer 对 segment 序列进行操作，并运算每个 segment 中的标準因果点积注意力上下文。因此，点积注意力运算在某种意义上是局部的。

然而，局部注意力在处理下一个 segment 时会丢弃前一个 segment 的注意力状态，但 Infini-Transformer 複用旧的 KV 注意力状态，以透过压缩储存来维护整个上下文历史。因此，Infini-Transformer 的每个注意力层都具有全域压缩状态和局部细细微性状态。?

与多头注意力（MHA）类似，除了点积注意力之外，Infini-attention 还为每个注意力层维护 H 个并行压缩记忆体（H 是注意力头的数量）。

?

下表 1 列出了几种模型根据模型参数和输入 segment 长度，定义的上下文记忆体佔用和有效上下文长度。Infini-Transformer 支援具有有限记忆体佔用的无限上下文视窗。

实验?

该研究在长上下文语言建模、长度为 1M 的金钥上下文块检索和 500K 长度的书籍摘要任务上评估了 Infini-Transformer 模型，这些任务具有极长的输入序列。对于语言建模，研究者选择从头开始训练模型，而对于金钥和书籍摘要任务，研究者採用不断预训练 LLM 的方式，以证明 Infini-attention 随插即用的长上下文适应能力。?

长上下文语言建模。表 2 结果表明 Infini-Transformer 优于 Transformer-XL 和 Memorizing Transformers 基线，并且与 Memorizing Transformer 模型相比，储存参数减少了 114 倍。

金钥任务。表 3 为 Infini-Transformer 在 5K 长度输入上进行微调后，解决了高达 1M 上下文长度的金钥任务。实验中输入 token 的範围从 32K 到 1M，对于每个测试子集，研究者控制金钥的位置，使其位于输入序列的开头、中间或结尾附近。实验报告了零样本準确率和微调準确率。在对 5K 长度输入进行 400 个步骤的微调后，Infini-Transformer 解决了高达 1M 上下文长度的任务。

摘要任务。表 4 将 Infini-Transformer 与专门为摘要任务构建的编码器 - 解码器模型进行了比较。结果表明 Infini-Transformer 超越了之前最佳结果，并且通过处理书中的整个文本在 BookSum 上实现了新的 SOTA。

研究者还在图 4 中绘製了 BookSum 资料验证分割的总体 Rouge 分数。根据折线趋势表明，随着输入长度的增加，Infini-Transformers 提高了摘要性能指标。