Transformer革新：Infini-Transformer在长文本处理中的突破

 人工智能咨询培训老师叶梓 转载标明出处

在当今信息爆炸的时代，大型语言模型（LLMs）在处理长文本数据方面的需求日益增长。无论是科学研究、法律分析还是医学诊断，长文本的处理能力都显得尤为重要。然而，现有的基于Transformer的模型在处理这类数据时遇到了重大挑战。由于其注意力机制的二次方复杂度，随着输入序列的增长，所需的内存和计算资源呈指数级增加，这不仅限制了模型处理长文本的能力，也带来了巨大的经济成本。

本文针对这一挑战，提出了一种创新的解决方案——Infini-attention技术。这一技术通过引入压缩记忆（compressive memory）的概念，对传统的注意力机制进行了根本性的改进。压缩记忆能够在保持固定数量参数的同时，存储和检索与输入序列长度无关的信息，从而显著降低了存储和计算成本。Infini-attention技术不仅在理论上具有吸引力，更在实际应用中展现出了卓越的性能。

Figure 1显示Infini-attention如何通过增加压缩记忆和线性注意力机制来处理无限长的上下文。图中展示了{KV}s−1和{KV}s分别代表当前和之前输入片段的注意力键和值，Qs代表注意力查询。PE代表位置嵌入

模型对长序列数据有着出色处理能力。在1M序列长度的passkey检索任务中，模型能够准确地从大量干扰信息中检索出关键信息。在500K长度的书籍摘要任务中，模型生成的摘要质量也达到了前所未有的高度。

图1Infini-attention的核心机制，即如何在Transformer模型中引入压缩记忆（compressive memory）来处理无限长的上下文。它通过结合局部压缩记忆和线性注意力来处理无限长的上下文信息。这种机制的关键创新在于它能够在处理当前输入片段的同时，有效地利用之前处理过的信息。

• 压缩记忆（Compressive memory）：这是Infini-attention中新增的一个组件，它允许模型存储和检索先前处理过的信息。这种记忆机制与传统的注意力机制不同，因为它不随输入序列长度的增加而增长，从而实现了对内存使用的优化。

• 线性注意力（Linear attention）：与标准的点积注意力相比，线性注意力是一种更高效的计算方法，它可以减少计算复杂度，同时保持对上下文信息的敏感性。

• 注意力键值对（{KV}s−1 和 {KV}s）：这里，{KV}s−1 表示上一个输入片段的键值对，而 {KV}s 表示当前输入片段的键值对。在传统的Transformer模型中，这些键值对仅在当前片段的处理中使用。然而，在Infini-attention中，它们被存储在压缩记忆中，以便在后续处理中重复使用。

• 注意力查询（Qs）：这些是当前输入片段的注意力查询，用于从压缩记忆中检索相关信息。

• 位置嵌入（PE）：位置嵌入是一种机制，用于向模型提供关于单词在序列中位置的信息。这有助于模型理解单词的顺序和上下文关系。

Infini-attention的工作流程

1. 局部注意力计算：首先，模型使用当前输入片段的键值对和查询来计算局部注意力，这与标准的Transformer操作相似。

2. 压缩记忆更新：在局部注意力计算之后，模型更新压缩记忆中的内容。这包括将当前片段的键值对添加到记忆中，并根据需要调整已有的键值对。

3. 记忆检索：使用当前片段的查询从压缩记忆中检索信息。这一步骤允许模型访问和整合过去的上下文信息。

4. 上下文整合：最后，模型将局部注意力的结果和从压缩记忆中检索到的信息结合起来，形成一个综合的上下文表示，这个表示将被用于后续的模型层或任务特定的处理。

通过这种设计，Infini-attention能够在保持计算效率的同时，处理非常长的上下文信息，这对于大型语言模型在复杂任务上的表现至关重要。

Infini-Transformer通过引入Infini-attention机制，能够在每个片段上计算标准的因果点积注意力上下文，同时保持内存占用的恒定。这种设计允许模型在处理每个新片段时，不仅考虑当前片段的信息，还能通过压缩记忆机制有效地利用之前片段的信息。

上图比较了Infini-Transformer和Transformer-XL在处理序列时对上下文历史的记忆方式。Infini-Transformer保留了整个上下文历史，而Transformer-XL则仅缓存了最后一个片段的KV状态

Infini-attention通过结合局部压缩记忆和线性注意力机制来处理无限长的上下文。Infini-attention在每个Transformer块中实现以下功能：

多头缩放点积注意力

多头缩放点积注意力是LLMs中的关键组件，它通过并行计算多个注意力头来捕获序列的不同方面。在每个头中，输入序列首先被转换为查询（Q）、键（K）和值（V）状态，然后通过点积操作和softmax函数计算注意力权重，最后得到加权的值状态，形成注意力上下文。

压缩记忆

Infini-attention不丢弃历史片段的KV状态，而是将它们存储在压缩记忆中。这样，当处理后续序列时，可以通过当前的查询状态从压缩记忆中检索值状态。这种方法不仅减少了内存占用，还使得模型能够利用长期依赖信息。

压缩记忆的实现依赖于关联矩阵，这是一种参数化的内存系统，用于存储键值对的绑定。通过特定的更新规则和检索机制，Infini-attention能够高效地进行记忆更新和检索，同时保持训练过程的稳定性。

Infini-Transformer通过压缩记忆实现了一个无界上下文窗口，同时保持了有界的记忆占用。这一点与传统的Transformer模型不同，后者的上下文窗口大小受限于内存容量。Infini-Transformer的记忆占用是固定的，这意味着它可以在不增加额外内存成本的情况下处理更长的输入序列。

在Infini-Transformer中，每个注意力层都维护着全局压缩记忆和局部细粒度状态。这种设计允许模型在保持对当前上下文敏感的同时，也能够访问和整合长期的历史信息。通过这种方式，Infini-Transformer显著扩展了模型的上下文窗口，同时避免了传统模型在处理长序列时遇到的内存和计算瓶颈。

不同Transformer模型的内存占用、上下文长度、内存更新和检索方式，包括Transformer-XL、Compressive Transformer、Memorizing Transformers、RMT、AutoCompressors和Infini-Transformers

Infini-attention技术通过引入压缩记忆和改进的注意力机制，为LLMs提供了一种高效处理长文本数据的新方法。这种方法不仅提高了模型的性能，也为未来在更大规模数据集上的应用提供了可能性。

为了验证Infini-Transformer模型在处理长文本数据方面的有效性，研究者在PG19和Arxiv-math这两个广泛认可的长文本基准数据集上进行了实验。他们构建的Infini-Transformer模型在结构上与Memorizing Transformers相似，拥有12个层和8个注意力头，每个头的维度设置为128，而Feed Forward Networks（FFNs）的隐藏层则为4096。这样的配置旨在确保模型在保持足够复杂度的同时，能够有效地捕获长距离依赖关系。

比较了不同模型在PG19和Arxiv-math基准测试上的平均token级困惑度，包括Transformer-XL、Memorizing Transformers、RMT和Infini-Transformers（线性和线性+Delta模型）

实验结果表明，Infini-Transformer在这些基准测试中表现出色，相比于传统的Transformer-XL和其他变体，它在长上下文语言建模任务上取得了显著的性能提升。这证明了Infini-attention机制在处理长序列时的有效性，尤其是在内存使用效率和计算速度方面。

为了进一步探索Infini-attention在实际应用中的潜力，研究者对现有的大型语言模型进行了持续预训练。预训练的数据集包括PG19、Arxiv-math以及C4文本，这些文本的长度均超过了4K令牌。通过这种轻量级的持续预训练，旨在使模型能够适应不断增长的上下文长度，同时保持对新信息的快速适应能力。

预训练过程中，研究者特别关注模型对长上下文的学习能力，以及在不断累积知识的同时，如何有效地整合新信息。实验结果表明，经过持续预训练的模型在处理长文本时展现出了更强的泛化能力和适应性。

在1M passkey检索任务中，研究者将一个1B参数的大型语言模型中的多头注意力（MHA）替换为Infini-attention，并在4K长度的输入上继续进行预训练。这一任务要求模型在大量干扰信息中准确地检索出隐藏的密码。在经过30K步的预训练后，模型在passkey检索任务上进行了微调，以进一步提高其检索精度。

在不同长度输入（从32K到1M）的测试子集上，Infini-Transformers在解决passkey任务时的token级检索准确率，包括零样本准确率和微调准确率

实验结果显示，Infini-Transformer不仅能够在极长的文本中准确地找到passkey，而且在经过微调后，其检索能力得到了显著提升，证明了Infini-attention在处理具有挑战性的长文本检索任务中的潜力。

为了测试Infini-Transformer在更复杂任务上的表现，研究者在BookSum数据集上进行了书籍摘要生成的实验。他们首先对一个8B参数的LLM模型进行了持续预训练，输入长度为8K，并进行了30K步的训练。随后，在BookSum任务上进行了微调，目标是生成整本书内容的摘要。

在这个任务中，模型需要处理和理解长达500K令牌的文本，并生成准确、连贯且内容丰富的摘要。实验结果表明，Infini-Transformer在生成高质量书籍摘要方面表现出色，其性能超越了现有的最佳方法，展示了其在处理长文本摘要任务上的巨大潜力。

比较了Infini-Transformers与其他为摘要任务设计的编码器-解码器模型在BookSum任务上的表现，包括BART、PRIMERA及其基于检索的长上下文扩展

通过这些实验验证了Infini-attention技术的有效性，展示了其在多种长文本处理任务中的应用前景，为构建更加智能、高效的语言处理系统奠定了基础。随着技术的不断进步和应用的深入，我们有理由相信，Infini-attention将在人工智能领域发挥越来越重要的作用。

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