不知 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
论文标题：Leave 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 上下文长度的任意。
择要任意。表 4 将 Infini-Transformer 与专门为择要任意构建的编码器 – 解码器模型举行了比较。结果表明 Infini-Transformer 超越了之前最佳结果，并且通过处置惩罚书中的整个文本在 BookSum 上实现了新的 SOTA。
研究者还在图 4 中绘制了 BookSum 数据验证分割的总体 Rouge 分数。根据折线趋势表明，随着输出长度的增加，Infini-Transformers 提高了择要性能指标。
