上下文输入无限制，谷歌发布Infini-Transformer

去年，百川智能发布号称全球最长的上下文窗口大模型Baichuan2-192K，一次性可输入35万字，超越GPT-4。

今年3月，Kimi智能助手宣布在上下文窗口技术上突破200万字。

紧追其后，国内各大互联网巨头纷纷布局升级自家大模型产品，360智脑正式内测500万长文本处理能力功能，阿里宣布通义千问开放1000万字长文本能力，百度也即将开放200-400万的长文本能力。

而近日，谷歌推出Infini-Transformer框架，彻底结束LLM上下文长度之争，支持无限长的输入。

4月10日，谷歌发布论文《Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention》，提出一种名为"Infini-attention（无限注意力）"的新型注意力机制，可以将基于 Transformer 的LLM 扩展到无限长输入，而不增加内存和计算需求。

论文链接：https://arxiv.org/pdf/2404.07143.pdf

使用该技术，研究者成功将一个 1B 的模型上下文长度提高到 100 万；应用到 8B 模型上，模型能处理 500K 的书籍摘要任务。

不同于传统的 Transformer 使用局部注意力丢弃旧片段，为新片段释放内存空间。Infini-attention将压缩记忆整合进标准的点积注意力机制，并在单个Transformer块内同时实现了掩码局部注意力和长期线性注意力机制。

该方法使现有LLM能够通过持续预训练和fine-tuning自然地扩展到处理无限长上下文，并以流的方式处理极长的输入进行计算。Infini-attention复用了标准注意力计算的所有键(Key)、值(Value)和查询(Query)状态,将旧的KV状态存储在压缩记忆中,而不是像标准注意力机制那样丢弃它们。在处理后续序列时,Infini-attention通过使用注意力查询状态来从记忆中提取值。为计算最终的上下文输出,Infini-attention会聚合长期记忆提取的值和局部注意力上下文。

Infini-attention的核心思想是将压缩记忆整合进标准点积注意力机制中，如图所示。

Infini-Transformer与Transformer-XL的操作方式类似,都是在一个个序列段上进行计算。在每个段内计算标准因果点积注意力上下文。不同的是,Transformer-XL在处理下一个段时会丢弃前一个段的注意力状态,而Infini-Transformer复用旧的 KV 注意力状态，以通过压缩存储来维护整个上下文历史。因此，Infini-Transformer 的每个注意力层都具有全局压缩状态和局部细粒度状态。

Infini-attention的具体实现如图所示：

它与标准的多头注意力(MHA)类似,每个注意力层都维护H个parallel的压缩记忆。Infini-attention首先计算标准的点积注意力上下文A_dot,然后从压缩记忆中检索出长期记忆上下文A_mem。最后，它通过一个学习的门控scalar β来结合这两部分上下文，得到最终的注意力输出O

Infini-attention的记忆更新和检索机制借鉴了先前工作,采用了简单高效的线性注意力形式。具体地,记忆检索通过将查询Q与之前存储的键值对M进行线性注意力计算得到;而记忆更新则是将新的键值对以联想绑定的方式累加到M中。研究人员还引入了delta规则来进一步优化记忆更新,尝试有选择地只更新那些不存在于记忆中的新信息。

与标准（多头注意力）MHA相比,Infini-attention仅引入了极少量的额外参数(每头一个标量值),就能够学习长短期上下文信息的最佳平衡。这种设计不仅计算高效,而且便于将Infini-attention无缝集成到现有的Transformer LLM中，支持即插即用的长上下文适应。

在具体实验中,首先在长上下文语言建模基准测试(PG19和Arxiv-math)上评估了小型Infini-Transformer模型。表2结果显示，Infini-Transformer不仅优于Transformer-XL和Memorizing Transformers等基线模型,而且存储参数减少了114倍。研究人员还发现，当训练序列长度增加到100K时,模型的困惑度进一步降低。

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