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

172 阅读5分钟

去年,百川智能发布号称全球最长的上下文窗口大模型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 扩展到无限长输入,而不增加内存和计算需求。


图片论文链接:arxiv.org/pdf/2404.07…

使用该技术,研究者成功将一个 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 β来结合这两部分上下文,得到最终的注意力输出OInfini-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 上下文长度的任务。

图片

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

图片

随着算法不断突破,新技术不断涌现,势必将造就智能算力的爆发式增长。英智公司智能调度全球高端算力,助力企业抓住人工智能发展趋势,满足政企在模型训练、应用部署中的需要。 英智公司希望每个政企都能抓住AI浪潮的风口,实现数字化转型和高质量发展。