直接扩展到无限长，谷歌Infini-Transformer终结上下文长度之争

发布时间：2024-04-13

点击次数：

不知 Gemini 1.5 Pro 是否用到了这项技术。

谷歌又放大招了，发布下一代 transformer 模型 infini-transformer。

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

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

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

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

该方法使 Transformer LLM 在有限内存的情况下扩展到无限长上下文，并以流的方式处理极长的输入进行计算。

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

本文贡献总结如下：

引入了一种实用且强大的注意力机制 Infini-attention—— 具有长期压缩内存和局部因果注意力，可用于有效地建模长期和短期上下文依赖关系；
Infini-attention 对标准缩放点积注意力（ standard scaled dot-product attention）进行了最小的改变，并通过设计支持即插即用的持续预训练和长上下文自适应；
该方法使 Transformer LLM 能够通过流的方式处理极长的输入，在有限的内存和计算资源下扩展到无限长的上下文。

论文链接：https://arxiv.org/pdf/2404.07143.pdf
论文标题：Le*e No Context Behind: Efficient Infinite Context Transformers with Infini-attention

方法介绍

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 上下文长度的任务。