Google | 提出Infini-Transformer架构，可让LLMs处理无限长上下文，内存节约114倍！
Google | 提出Infini-Transformer架构，可让LLMs处理无限长上下文，内存节约114倍！
引言
为解决大模型（LLMs）在处理超长输入序列时遇到的内存限制问题，本文作者提出了一种新型架构：Infini-Transformer，它可以在有限内存条件下，让基于Transformer的大语言模型（LLMs）高效处理无限长的输入序列。实验结果表明：Infini-Transformer在长上下文语言建模任务上超越了基线模型，内存最高可节约114倍。
https://arxiv.org/pdf/2404.07143.pdf
背景介绍
对于人工智能(AI)来说，内存资源是神经网络模型进行高效计算的必要条件。然而，由于Transformers 中的注意力机制在内存占用和计算时间上都表现出二次复杂度，所以基于Transformer的大模型就更加依赖内存资源。
例如，对于批量大小为 512、上下文长度为 2048 的 500B 模型，注意力键值 (KV) 状态的内存占用为 3TB。这意味着对于标准的Transformer架构大模型，要想支持越长的上下文，那么需要的内存成本也将越高。为了降低内存成本，且要让LLMs能够支持更长的上下文，这无疑是一项非常具有挑战的任务。
面对超长序列，相比注意力机制，内存压缩技术更具扩展性。内存压缩不使用随输入序列长度而增长的数组，而是在有限的内存资源上，维护固定数量的参数来进行信息的存储和回调。然而，目前的LLMs尚未有一种有效、实用的内存压缩技术，可以在简单性与质量之间取得平衡。
基于以上背景，本文作者提出了一种新架构：Infini-Transformer，能够让基于Transformer的大模型在有限内存、计算资源的条件下，处理无限长的上下文输入。
Infini-Transformer模型
下图是本文Infini-Transformer模型、Transformer-XL模型的对比图。与Transformer-XL类似，Infini-Transformer处理的是一系列片段。即在每个片段内计算 standard causal 点积attention context（注意力上下文）。因此，点积注意力计算在某种意义上是局部的，它覆盖了索引为 S 的当前片段的总共 N 个标记。
然而，局部注意力在处理下一个片段时会丢弃前一个片段的注意力状态。在Infini-Transformer中，并没有忽略旧的键值（KV）注意力状态，而是通过内存压缩技术重新使用它们来保持整个上下文历史。因此，Infini-Transformer的每个注意力层都具有全局压缩和局部细粒度状态，这就是前面提到的无限注意力（Infini-attention）。
Infini-attention
Infini-Transformer模型的关键组成部分为：Infini-attention，这是一种新型Attention技术，如下图所示。它将计算局部和全局上下文状态，并将它们组合到一块作为输出。与多头注意力（MHA）类似，除了点积注意力之外，它还为每个注意力层维护 H 个并行压缩内存（H 是注意力头的数量）。
其中：
「缩放点积Attention」 在Infini-attention机制中发挥了重要作用。Scaled Dot-product Attention是多头注意力（Multi-head Attention, MHA）的一个变种，它在处理序列数据时能够有效地捕捉序列中的依赖关系。对于MHA，本文并行地为序列中的每个元素计算H个注意力上下文向量，将它们沿着第二维拼接起来，最后将拼接后的向量投影到模型空间以获得注意力输出。
「Compressive Memory」 它通过固定数量的参数关联矩阵来存储和检索信息，当处理新的输入序列时，模型会更新关联矩阵中的参数，以存储当前序列的信息。更新过程通常涉及到将新的键值对（KV）与记忆矩阵中的现有参数进行某种形式的结合。在处理后续序列时，模型会利用查询向量（Query）来检索关联矩阵中的信息，从而获取之前序列的上下文信息。
实验结果
如下图所示，在PG19和Arxiv-math数据集上，Infini-attention在长上下文语言建模任务上取得了优于基线模型的性能，同时在内存大小上实现了114倍的压缩比。
如下图所示，在1M序列长度的密钥检索任务中，Infini-attention在仅使用5K长度输入进行微调后，成功解决了任务，展示了其在处理极长输入序列时的能力。
如下图，在500K长度的书籍摘要任务中，Infini-attention通过持续的预训练和任务微调，达到了SOTA，证明了其在长文本摘要任务上的性能。
