谷歌重塑Transformer：无限记忆力，无限长输入，LLM基础研究重大突破

生成式人工智能的重大突破和现在我们熟知的几乎所有大模型都来源于谷歌的Transformer架构。

长期以来，记忆一直被视为智能的基石，因为它使得针对特定上下文进行高效计算成为可能。然而，Transformer及其衍生的大型语言模型(LLM)由于注意力机制的固有特性，其上下文相关记忆能力受到限制。

标准Transformer架构中的注意力机制在内存占用和计算时间复杂度上呈现二次增长。以一个500B模型为例，如果批大小为512，上下文长度为2048，其注意力Key-Value(KV)状态的内存占用就高达3TB。可见，要使LLM扩展到更长的序列(如1M个tokens)，单凭标准Transformer架构是非常具有挑战性的，且部署这样的长上下文模型也会带来巨大的经济成本。

相比之下，压缩记忆系统在处理极长序列方面具有更好的可扩展性和效率。压缩记忆主要维护一个固定数量的参数来存储和调用信息，从而实现有限的存储和计算开销。在压缩记忆中，新信息通过改变其参数来添加，并可以在之后被取回。然而，目前的LLM还未见到一种既简单又高质量的压缩记忆技术。

近日谷歌发表论文：Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention 重塑了Transformer。

Google提出了一种名为"Infini-attention（无限注意力）"的新型注意力机制，旨在使基于Transformer的大型语言模型(LLM)能够有效地处理无限长的输入序列，同时保持有限的内存占用和计算资源。

该方法的关键是将压缩记忆整合进标准的点积注意力机制，并在单个Transformer块内同时实现了掩码局部注意力和长期线性注意力机制。

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

主要见解和贡献

1.通过持续的预训练和微调，Infini-attention 可使现有语言模型即插即用地适应超长语境；

2. Infini-Transformers在长语境语言建模任务中的表现优于强基线，在内存占用方面实现了超过 100 倍的压缩率；

3. 参数为 1B 的 Infini-Transformer 解决了输入长度高达 100 万个词组的密钥检索任务，即使对更短的序列进行微调，也能显示出很强的长度泛化能力；

4. 经过持续的预训练和微调，8B Infini-Transformer 在长度为 500K 的书籍摘要任务中取得了最先进的性能。对所学门控得分的可视化揭示了专注于局部或长期上下文的专门注意头的出现，以及两者兼而有之的混合注意头的出现。

其他细节：

1.利用稳定的训练技术，将压缩记忆更新和检索过程作为线性注意机制；

2. 将训练序列长度从 32K 增加到 100K，进一步提高了语言建模的复杂性；

3. Infini-Transformers利用固定大小的内存循环处理输入片段，从而实现快速的流推理；

4. 消融实验比较了线性和 "线性 + 三角形 "内存更新规则，两者的性能相当；

5. 本文比较了各种段级内存转换器模型的内存占用和有效上下文长度。

技术详解和实验结果

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

Infini-Transformer与Transformer-XL的操作方式类似，都是在一个个序列段上进行计算。在每个段内计算标准因果点积注意力上下文。不同的是，Transformer-XL在处理下一个段时会丢弃前一个段的注意力状态，而Infini-Transformer则建议保留这些旧的KV注意力状态，利用压缩记忆来维持整个上下文历史。

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模型。结果显示，Infini-Transformer不仅优于Transformer-XL和Memorizing Transformers等基线模型，而且内存占用仅为后者的1/114。研究人员还发现，当训练序列长度增加到100K时，模型的困惑度进一步降低。

研究人员还在1M长度密钥检索任务和500K长度图书摘要任务上验证了Infini-Transformer的性能。对于密钥检索任务，通过持续预训练1B LLM并注入Infini-attention，在仅fine-tuning 5K长度输入的情况下，就能解决1M长输入的问题。而对于图书摘要任务，8B Infini-Transformer模型在持续预训练和fine-tuning后，在500K长输入上达到了新的SOTA结果。

结语

Infini-attention 机制为 Transformer 语言模型处理超长上下文提供了一种高效而强大的方法，同时不会对内存或计算量造成过多的增加。该方法在具有挑战性的长语境任务中释放出了新的能力，对于调整现有模型以适应长输入也非常实用。实验证明，与之前的方法相比，该方法在性能、压缩和泛化方面都有很强的优势。总之，这项工作为提高长语境语言建模的可扩展性和有效性做出了重大贡献。

来源：AI寒武纪