如何用Infini-attention实现无限上下文Transformer

Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention¶

1. 论文信息¶

https://arxiv.org/pdf/2404.07143.pdf

2. 引言¶

在这篇论文中，作者探讨了大型语言模型（LLMs）在处理长序列数据时面临的存储和计算挑战。传统的Transformer模型使用的注意力机制虽然强大，但是随着输入序列的增长，其内存和计算复杂度呈二次方增长，这限制了模型处理长序列的能力。

论文解决的主要问题是如何让Transformer模型有效处理无限长的输入序列，同时保持有限的内存和计算资源消耗。这是一个关键问题，因为目前的大型语言模型（如GPT和BERT系列）通常只能处理有限长度的序列，这限制了它们在某些应用场景中的效用，特别是那些需要长期记忆和大量上下文信息的任务。现有的解决方案通常涉及更复杂的模型架构或是牺牲一部分性能以换取更长的输入处理能力。例如，通过引入稀疏注意力机制或是分层注意力机制来减少计算和内存需求，但这些方法往往会引入新的问题，如信息损失或是实现复杂度的增加。

论文提出了一种新的注意力机制——Infini-attention，该机制将压缩内存集成到标准的注意力机制中。通过这种设计，模型能够存储旧的注意力键值对（KV状态）而不是像传统注意力机制那样丢弃它们。这允许模型在处理后续序列时回顾以前的上下文信息，从而支持无限长的输入处理。此外，Infini-attention结合了局部掩蔽注意力和长期线性注意力机制，使其能够高效处理长范围和短范围的上下文依赖Infini-attention通过重用标准注意力计算的所有键、值和查询状态，实现了对长期记忆的巩固和检索，而不需要增加额外的内存或计算开销。这种方法允许模型以流式方式处理极长的输入，从而实现在有限的资源下扩展到无限长的上下文。所以论文的贡献可以概括为：

1. 提出了一种新的注意力机制——Infini-attention，该机制有效整合了长期压缩记忆和局部因果注意力，能够高效地建模长距离和短距离的上下文依赖。
2. Infini-attention对标准的点积注意力机制进行了微小的改动，支持即插即用的持续预训练和长上下文适应。
3. 使Transformer LLMs能够通过以流式方式处理极长的输入，在有限的内存和计算资源下扩展到无限长的上下文。

这些贡献显示了Infini-attention在处理长序列语言建模任务中的潜力，如在持续的预训练和任务微调后，一个8B模型在500K长度的书籍摘要任务上达到了新的最佳状态。

3. 方法¶

3.1 Infini-attention¶

Infini-attention是一种创新的注意力机制，结合了标准的多头点积注意力（Scaled Dot-product Attention）和压缩内存（Compressive Memory）技术，以优化长期和局部信息的处理效率。

Scaled Dot-product Attention¶

多头点积注意力（MHA）是大型语言模型中的核心构建块，它利用自注意力机制（a variant of self-attention）来模拟上下文依赖的动态计算。在MHA中，每个头首先计算注意力查询（$Q$）、键（$K$）和值（$V$）状态： $$ K = XW_K, \quad V = XW_V \quad \text{和} \quad Q = XW_Q. $$ 其中，$X \in \mathbb{R}^{N \times d_{model}}$ 是输入序列，$W_K, W_V, W_Q$ 是可训练的投影矩阵。然后，通过下面的公式计算注意力上下文： $$ A_{dot} = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_{model}}}\right) V. $$ 在多头注意力机制中，为每个序列元素并行计算$H$个注意力上下文向量，然后将这些向量沿第二维度连接，并最终投影回模型空间，以获得注意力输出。

Compressive Memory¶

与传统方法不同，Infini-attention在计算点积注意力后，不是创建新的内存条目，而是重用$Q, K, V$状态。这种状态共享和重用不仅提高了插拔式长上下文适应的效率，也加速了训练和推理过程。

压缩内存中，通过以下公式从内存 $M_{s-1}$ 中检索新内容 $A_{mem}$： $$ A_{mem} = \frac{\sigma({Q}) M_{s-1}}{{\sigma(Q)} z_{s-1}}. $$ 其中，$\sigma$ 是非线性激活函数，$z_{s-1}$ 是归一化项，以保持训练的稳定性。

内存的更新则通过以下公式实现： $$ M_{s} \leftarrow M_{s-1} + \sigma(K)^T V \quad \text{和} \quad z_{s} \leftarrow z_{s-1} + \sum_{t=1}^N \sigma(K_t). $$ 通过这种方式，每个注意力层都建立了一个递归，将当前段的状态传递到下一段。

Long-term context injection¶

通过学习的门控标量 $\beta$，Infini-attention聚合局部注意力状态 $A_{dot}$ 和内存检索内容 $A_{mem}$： $$ A = \textit{sigmoid} (\beta) \odot A_{mem} + (1 - \textit{sigmoid}(\beta)) \odot A_{dot}. $$ 这种方式在模型中引入了长期和局部信息流之间的可学习折中，每个头只增加一个标量值作为训练参数。

如同标准的多头注意力，对于多头Infini-attention，论文并行计算$H$个上下文状态，然后将它们连接并投影以获得最终的注意力输出 $O \in \mathbb{R}^{N \times d_{model}}$： $$ O = [A^1; \dots; A^H] W_O $$ 其中 $W_O \in \mathbb{R}^{H \times d_{value} \times d_{model}}$ 是可训练的权重。

通过这种设计，Infini-attention有效地增强了Transformer模型处理长期依赖关系的能力，同时保持对局部上下文的

敏感度。

3.2 Memory and Effective Context Window¶

在这一部分中，文章讨论了不同Transformer模型在段级别内存管理上的策略，并展示了各模型的内存占用（Memory (cache) footprint）、有效上下文长度（Context length）、内存更新（Memory update）和内存检索（Memory retrieval）机制。这些模型被设计用来处理序列数据的长期依赖问题，每种模型都有其独特的方法和挑战。

Transformer-XL¶

Transformer-XL 使用缓存的机制存储上一个段的键值（KV）状态，从而在每个层级扩展上下文窗口。这种方法将上下文窗口从 $N$ 扩展到 $N \times l$ 个tokens，增加了额外的内存占用 $(d_{key} + d_{value}) \times H \times N \times l$。

Compressive Transformer¶

Compressive Transformer 在Transformer-XL的基础上添加了一个额外的缓存来存储过去段激活的压缩表示。这种方法通过增加 $c \times r \times l$ 扩展了上下文窗口，但内存复杂度依旧较高。

Memorizing Transformers¶

Memorizing Transformers 选择存储整个序列的KV状态作为上下文。这种方法虽然能够覆盖整个序列历史，但在单一层面上的存储变得非常昂贵，因此仅限于单层计算，并使用快速的kNN检索器来构建上下文窗口。

RMT 和 AutoCompressors¶

这两种模型允许理论上无限的上下文长度，它们将输入压缩成摘要向量，然后将这些向量作为后续段的额外软提示（soft-prompt）输入。然而，这种技术的成功在实践中高度依赖于软提示向量的大小。随着软提示向量数量的增加，内存和计算复杂度迅速增长，从而降低了效率。

Infini-Transformers¶

与其他模型不同，Infini-Transformers 设计了一种能够以固定的内存复杂度 $(d_{key} \times (d_{value} + 1) \times H \times l$) 处理无界上下文窗口的机制。这是通过在每个头部的单个层上存储压缩的上下文 $M_s$ 和 $z_s$ 实现的。Infini-Transformer 的内存更新机制为增量式，使用线性注意力机制进行内存检索。这种设计显著提高了模型在长上下文建模任务上的表现，同时实现了超过100倍的压缩比，并在困惑度评分上取得了进一步改进。

通过这些不同的内存和上下文窗口管理策略，各模型试图在内存效率和长期依赖建模能力之间找到平衡。Infini-Transformer特别在这方面展示了显著的优势，通过其创新的内存管理技术有效地扩展了模型的应用范围。

4. 实验¶

在最近的实验中，不同的Transformer模型针对长文本建模任务的性能进行了比较，结果显示Infini-Transformer在效率和有效性方面均表现出显著的优势。该研究的目的是评估不同模型在处理具有长上下文依赖的文本序列时的表现，主要通过平均tokens级困惑度来衡量，该指标反映了模型预测文本序列的能力。困惑度越低，模型的预测能力越强。在对比中，传统的Transformer-XL使用较大的缓存来扩展其上下文处理能力，但在PG19和Arxiv-math数据集上的困惑度相对较高，表明其在处理长文本序列上存在局限。Memorizing Transformers尽管在困惑度上略有改进，但其内存使用量大，缺乏有效的内存压缩，这限制了其在资源受限环境下的应用。RMT模型通过极端压缩内存至2.5M，并实现了高达73倍的压缩比，但这种极度的内存压缩牺牲了其在两个数据集上的预测性能，困惑度较高。这表明，在追求内存效率的同时，也可能会对模型的预测能力产生负面影响。相比之下，Infini-Transformer以其极低的内存使用（仅1.6M）和高达114倍的压缩比，在所有比较的模型中表现最佳。其在PG19数据集上的困惑度为9.65和9.67，以及在Arxiv-math数据集上的2.24和2.23，展示了其出色的预测能力和极高的内存效率。这一显著的性能得益于其创新的内存管理策略，特别是在Infini-Transformer模型中引入的压缩内存和增量更新机制，这些设计使模型能够有效地处理长序列，同时维持低资源消耗。Infini-Transformer的出色表现不仅突显了其在长文本建模任务中的优势，还展示了在保持资源效率的同时，如何通过技术创新提升模型的预测能力。这些发现对于设计未来的高效和实用的语言处理模型具有重要意义。

这项实验研究了Infini-Transformer模型在处理长达1M上下文长度的序列时，对于隐藏在不同部分（起始、中间、末尾）的通行密钥的检索准确性。实验结果在两种不同的训练设置下呈现：零样本（Zero-shot）和经过400步微调（FT, Fine-Tuned）的情况。实验中报告了在各种长度（从32K到1M）的输入中，不同部位通行密钥的tokens级检索准确性。在零样本设置中，Infini-Transformer（Linear）和Infini-Transformer（Linear + Delta）两种变体在处理较短序列（如32K和128K）时表现出较高的准确性，但随着序列长度的增加到256K及以上，准确性有所下降。具体来说，Linear变体在32K序列长度时在开始/中间/末尾的准确性为14/13/98，而到了1M长度时，准确性降至8/6/98。Linear + Delta变体则显示出稍微均衡的性能，尤其是在较长的序列中，例如在1M长度时达到7/6/97的准确性。

在经过400步微调后，两个变体的表现显著提高，几乎在所有长度和部位达到了100%的准确性。Linear变体在最长的1M序列长度下，开始/中间/末尾的准确性分别为96/94/100，而Linear + Delta变体在相同条件下实现了完美的100/100/100的准确性。这表明微调过程显著提升了模型对于复杂序列中特定信息点的检索能力，特别是在处理极长的输入序列时。这些结果强调了Infini-Transformer模型在处理极长序列信息和复杂上下文时的强大能力，尤其是在经过针对性微调后，能够有效提升模型对难以捕捉信息的检索准确性，显示出其在实际应用中的巨大潜力。

5. 讨论¶

在讨论Infini-Transformer及其Infini-attention机制时，我们可以从多个维度探讨其优点和潜在的缺点。该模型通过引入压缩内存和创新的注意力机制来处理长序列，表现出明显的优势，尤其在处理大规模数据集时的内存效率和处理速度上。

首先，Infini-Transformer的显著优点之一是其能够有效处理极长的输入序列，而不受传统Transformer模型因注意力机制的二次方复杂度增长而限制的问题。通过压缩内存技术，Infini-Transformer可以持续积累并利用历史信息，这使得模型在涉及大上下文依赖的任务，如文本摘要和长文本生成中表现尤为突出。此外，该模型采用线性注意力机制，相比于标准的Transformer，大大减少了计算资源的消耗，提高了处理速度。

然而，尽管Infini-Transformer在理论和实验上都展示了其优越性，但该方法也存在一些潜在的缺点。首先，虽然压缩内存允许模型存储更长时间的信息，但这种机制可能在某些情况下导致信息损失，特别是在极端压缩比例下。信息的损失可能会影响模型对上下文的理解能力，从而影响任务的最终性能。此外，压缩和解压过程的设计需要精细的平衡，以确保既不过度压缩以损失重要信息，也不过少压缩而浪费存储和计算资源。另一个考虑是，虽然Infini-attention提供了对长期依赖关系的支持，但这种机制的实现复杂度高于传统的注意力机制，可能需要更多的调优和维护。在实际应用中，这可能会增加模型部署的难度，尤其是在资源受限的环境中。

也就是说Infini-Transformer提供了一种有效的解决方案，用于处理传统Transformer在长序列处理上的限制。它通过创新的内存管理和注意力机制，提高了模型在多个应用场景中的表现和效率。然而，信息损失和实现复杂度的潜在问题也是未来优化和应用过程中需要考虑的重要因素。

6. 结论¶

在本研究中，论文提出了Infini-Transformer，一种新颖的Transformer架构，它通过引入Infini-attention机制解决了传统Transformer模型在处理长序列数据时遇到的挑战。通过结合压缩内存和线性注意力技术，Infini-Transformer能够有效地处理无限长的输入序列，同时保持有限的内存和计算资源消耗。实验结果表明，该模型在长文本建模任务上具有出色的表现，尤其是在困惑度和上下文窗口扩展方面，显著优于现有的Transformer变体。

Infini-Transformer的设计优化了信息存储和检索过程，使得模型能够在维持高效内存使用的同时，提供对长期依赖关系的支持。尽管该模型在信息压缩方面可能面临信息损失的风险，但其整体性能和灵活性表明它是解决长序列处理问题的一个有力候选。未来的工作将集中于进一步优化压缩内存机制，减少潜在的信息损失，并探索其在更广泛应用场景中的有效性和可扩展性。