Compress, Gather, and Recompute: REFORMing Long-Context Processing in Transformers

作者: Woomin Song¹,†, Sai Muralidhar Jayanthi², Srikanth Ronanki², Kanthashree Mysore Sathyendra², Jinwoo Shin¹, Aram Galstyan², Shubham Katiyar², Sravan Babu Bodapati²
机构: ¹KAIST, ²Amazon AGI

A1 主要贡献

核心问题：大型语言模型（LLM）在处理超出其预训练上下文长度的极长上下文时面临严峻挑战。现有方法主要分为两类：循环压缩方法和随机访问方法。循环方法虽然节省计算资源，但因压缩或丢弃部分键值（KV）缓存而导致信息丢失（“遗忘”问题）；随机访问方法虽然保留了完整的上下文信息，但需要巨大的内存资源，导致高昂的内存开销和延迟，尤其是在需要CPU卸载的实际部署中。

研究目标：为了解决上述挑战，本文提出了一种名为 REFORM（REcurrent chunked Forwarding with On-demand cache RecoMputation）的新型推理框架，旨在结合循环方法的效率和随机访问方法卓越的召回能力。

创新点：REFORM 框架通过一个“压缩-收集-重计算”的流水线实现高效的长上下文处理，其核心创新在于：
1. 两阶段处理流程：
* 编码阶段（循环分块前向传播）：将输入分块并迭代处理。在每次迭代中，模型保留一个经压缩的KV缓存（保留最重要的token），并利用这个稀疏的KV缓存从多个中间层和注意力头中提取并存储轻量级的跨层上下文嵌入。此阶段还采用早退（early exit）策略，跳过用于嵌入提取的最高层之后的计算，进一步提升效率。
* 重计算阶段（按需缓存重计算）：当需要生成响应时，利用查询（输入的最后一部分）的嵌入，通过相似性匹配从存储的上下文嵌入中识别并收集历史上最相关的输入token。然后，仅对这些选定的token进行所有层的完整KV缓存重计算。

1. 用压缩构建检索嵌入而非直接生成：与现有循环方法不同，REFORM 不直接使用压缩后的KV缓存进行生成，而是用它来构建轻量级的、用于检索的token级嵌入。这使得模型能在保持低内存占用的同时，通过后续的重计算恢复高保真的上下文表示。
2. 高效率与高性能的平衡：通过仅对查询最相关的token进行重计算，REFORM 在避免为所有token维持完整表示的高昂成本的同时，确保了对长程依赖的精确处理，从而在计算效率和模型性能之间取得了更好的平衡。

主要成果：
* 在RULER和BABILong基准测试中，在1M上下文长度下，REFORM的性能分别比表现最好的基线高出52%和34%。
* 在∞-Bench、RepoEval和多模态MM-NIAH等多样化任务中也超越了基线。
* 在处理256k token输入时，与InfLLM和InfiniPot相比，REFORM将推理时间减少了30%，峰值内存使用量减少了5%，实现了效率和性能的双重提升。

A3 背景知识

循环上下文处理。为了应对长上下文处理的计算挑战，一些研究探索了使用循环来提高效率。一类工作【10, Transformer-xl: Attentive language models beyond a fixed-length context, 2019, arXiv; 11, Recurrent memory transformer, 2022, NeurIPS; 12, Leave no context behind: Efficient infinite context transformers with infini-attention, 2024, arXiv】对Transformer架构进行了修改，引入了块级循环操作，以便以更小、可管理的单元处理长上下文。然而，这些方法通常需要对模型进行大量训练，因此不能直接应用于现有的预训练大型语言模型。最近的努力则利用KV缓存驱逐来迭代地编码输入块并压缩KV缓存，从而避免了架构修改或改变模型参数。例如，StreamingLLM【3, Efficient streaming language models with attention sinks, 2023, arXiv】通过保留初始和最近的token同时压缩中间token来维持流畅的生成。后续的方法【4, H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2023, NeurIPS; 5, Transformers are multi-state rnns, 2024, arXiv; 6, Infinipot: Infinite context processing on memory-constrained llms, 2024, arXiv】从先前的上下文中识别重要的token，从而实现信息更丰富的缓存压缩。尽管这些方法效率很高，但压缩先前输入的过程常常导致关键信息的丢失，从而引发“遗忘”问题。因此，这些方法可能难以处理需要精确检索早期输入的任务。REFORM通过其收集和重计算阶段解决了这个问题，该阶段能够为所有与查询相关的输入token生成高保真的表示。

随机访问方法。另一个研究方向是实现对先前上下文的随机访问，类似于全注意力机制，但计算效率更高。这些方法通常将完整的KV缓存存储在内存中，并根据需要动态检索相关token。一些方法通过训练模型【13, Memorizing transformers, 2022, arXiv】或辅助的侧网络【14, Augmenting language models with long-term memory, 2023, arXiv】来有效利用检索到的token。最近，一些策略在不修改模型参数的情况下实现了相同的目标，即通过将完整的KV缓存存储在内存中并动态检索【7, Infllm: Unveiling the intrinsic capacity of llms for understanding extremely long sequences with training-free memory, 2024, arXiv; 8, Reattention: Training-free infinite context with finite attention scope, 2024, arXiv】。虽然这些方法允许随机访问输入序列的任何部分，但由于需要维护大型缓存，它们引入了显著的内存开销。在实践中，这通常需要CPU卸载，从而可能进一步增加延迟。此外，访问先前上下文的灵活性并不一定能带来高的检索性能。相比之下，REFORM使用KV缓存压缩，并仅利用表现优异的注意力头构建紧凑的token级嵌入，以减少内存开销，同时保持高检索性能。

A2 方法细节

3.1 使用循环分块前向传播和早退策略进行嵌入提取

使用循环分块前向传播进行嵌入提取。为了在有限的计算预算和上下文窗口下编码长上下文，本文采用了循环KV缓存压缩方法，这使得在有限资源和上下文窗口下能够处理无限长的上下文。具体而言，本文采用了迭代式KV缓存压缩方法【3, Efficient streaming language models with attention sinks, 2023, arXiv; 4, H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2023, NeurIPS; 5, Transformers are multi-state rnns, 2024, arXiv; 6, Infinipot: Infinite context processing on memory-constrained llms, 2024, arXiv】。首先将长输入分割成较大的块（实验中为32k tokens），并在处理完每个块后应用KV缓存压缩，以更好地利用并行计算。对于KV缓存压缩，本文采用了H2O【4, H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2023, NeurIPS】中基于注意力的token驱逐策略。压缩后，重新分配位置ID，使得压缩缓存中的token具有连续的位置ID。这种位置重分配使得模型能够处理超出其预训练上下文限制的更长序列。与现有旨在直接使用这些压缩KV缓存进行生成的方法不同，本文仅用它来构建上下文嵌入，这些嵌入稍后将用于识别生成响应所需的输入部分。

早退策略。使用压缩方法创建嵌入带来了额外的好处：可以通过采用早退策略进一步提高效率。正如3.2节所观察到的，高性能的嵌入通常在较低的Transformer层中就已可用。因此，在用于嵌入提取的最高层之后，再将输入通过剩余的层进行前向传播是不必要的。所提出的早退策略减少了计算和内存需求，因为我们不需要为上层保留KV缓存。

3.2 构建跨层上下文嵌入

寻找高性能嵌入的动机。先前的研究揭示，在Transformer的不同层中分布着专门的注意力头，它们能够从长上下文输入中准确检索相关信息【15, Retrieval head mechanistically explains long-context factuality, 2024】。因此，为了构建信息丰富的嵌入，本文分析了Transformer中各种头和嵌入的检索性能，以确定最适合我们方法的类型。具体来说，本文比较了注意力分数（不含位置编码，以适用于极长输入）、隐藏状态之间的余弦相似度，以及注意力QKV状态（即注意力层中QKV投影产生的嵌入）在不同Transformer层上的token级检索性能。

嵌入头识别。本文在一系列多跳问答数据集上进行实验，以确定用于相似性搜索的最佳嵌入；详情见C.1节。表1报告了表现最佳的层和头的MNR分数。有些出人意料的是，我们观察到使用注意力QKV状态进行余弦相似度搜索的性能，常常优于广泛使用的注意力分数。尽管其尺寸小得多，但其性能也超过了使用隐藏状态之间的余弦相似度。这一发现表明，通过仔细选择合适的头，直接使用注意力层的QKV状态进行余弦相似度搜索可以实现非常高的性能，同时只需要极少的内存。

Table 1: 比较相似性搜索方法。与不同相似性搜索方法（包括注意力以及使用隐藏状态（HS）或注意力QKV状态的余弦相似度搜索）对应的最佳3个MNR分数（越低越好）。分数是使用Mistral-Nemo-Instruct-2407测量，并在500个多跳问答样本上取平均值。

高性能头的分布与早退策略的动机。图2展示了表现最佳的值（value）头的分布。值得注意的是，表现最好的头不一定来自最后的Transformer层。这一观察意味着，为了构建有效的检索嵌入，可能不需要将输入通过上层进行前向传播，这成为我们在前一节中描述的早退策略的关键动机。

构建通用嵌入。为了创建一个对广泛任务都有用的通用嵌入，我们进行了额外的实验来识别能够有效表示复杂输入模式的头。具体来说，我们创建了一个合成的键值检索任务，其中涉及将多个格式为“The value corresponding to the id {key} is {value}.”的句子嵌入到WikiText【16, Pointer sentinel mixture models, 2016】语料库中，其中键和值是随机的10个字符的ASCII字符串。在评估了每个合成数据集的500个样本后，我们选择了表现最佳的嵌入。我们强调，尽管头的选择是基于相对较短的合成数据（8k tokens），但其好处可以外推到涉及数百万token的更长上下文中。

组合多个头。在确定表现最佳的头之后，我们将它们的嵌入组合起来，创建一个单一的、token级别的嵌入。在初步实验中，我们观察到使用从不同头获得的检索分数的平均值通常可以提高最终的检索性能。因此，我们将收集到的嵌入在归一化后进行拼接：

$$e_{\mathtt{comb}} = \mathtt{concat} \left( \left\{ \frac{e_{\mathtt{i}}}{\| e_{\mathtt{i}} \|}, \mathtt{i} \in \mathtt{selected\_heads} \right\} \right)$$

这种方法确保了使用结果嵌入进行余弦相似度搜索在数学上等同于为每个头独立计算余弦相似度分数然后将它们平均。

3.3 按需缓存重计算

核心流程。为了实现对先前输入的随机访问，我们利用跨层上下文嵌入来识别与输入最后部分相关的输入段。然后，我们收集相应的输入嵌入，并再次将它们通过模型前向传播，用最相关的输入重建KV缓存。详细过程如下。

重要输入的识别。在构建了与输入相对应的跨层上下文嵌入后，我们使用这些上下文嵌入，在查询（输入的最后一部分）和其余输入之间进行token级别的相似性搜索。然后，我们对查询token的相似性分数进行最大池化（max-pool），以确保每个token被赋予一个单一的分数。为了保持所识别输入的连续性，我们进一步将每个token的分数与相邻的128个token进行最大池化。在处理完重要性分数后，我们识别出得分最高的token。我们始终保留最初和最后的256个token以保持连贯性。

按需缓存重计算。一旦我们确定了相关段落，我们就收集相应的输入嵌入，并再次将它们通过模型前向传播，重新计算KV缓存。新的KV缓存随后用于后续的解码过程。通过引入按需缓存重计算方案，我们避免了存储完整KV缓存的需要，同时实现了对先前输入的随机访问，从而显著减少了内存需求。

A4 实验环境

通用设置与基线：
* 对比方法：主要与不改变模型参数的上下文外推方法进行比较，重点是基于循环和随机访问的方法。具体基线包括：StreamingLLM [3]、TOVA [5]、H2O [4]、InfiniPot [6] 和 InfLLM [7]，以及一个简单的截断基线。
* H2O实现细节：为提高效率，H2O的注意力分数计算限制在每个块的最后128个token。

模型配置：
* 文本实验模型：Mistral-NeMo-Instruct-2407 [2] 和 Qwen2.5-7B-Instruct [21]。
* 代码补全实验模型：Qwen2.5-Coder-1.5B 和 7B 模型。
* 多模态实验模型：Pixtral-12B-2409 [9]。

参数设置：
* 循环基线：KV缓存预算和块大小均为32k tokens。
* InfLLM：使用32k的活动KV缓存预算。
* 所有方法：为保持文本连贯性，所有基线和REFORM都保留初始和最近的256个token在缓存中。
* REFORM：重计算预算在Mistral-Nemo-Instruct-2407上为8k tokens，在其他所有模型上为16k tokens。

数据集：
* 精确检索：Needle-in-a-Haystack (NIAH) [22]。
* 复杂合成任务：RULER [17] 和 BABILong [18]，包含更具挑战性的NIAH、聚合、问答和多跳推理任务。RULER被扩展到1M token长度。
* 现实世界任务：∞-Bench [19]（来自长篇书籍和对话的任务）和 RepoEval [20]（仓库级代码补全）。
* 多模态任务：MM-NIAH [23]。

硬件与软件：
* 硬件：效率分析在单个H100 GPU上进行。
* 软件：所有效率测量均使用Flash Attention 2 [48]。

A4 实验结果

4.1 “大海捞针”评估（Needle-In-A-Haystack）

实验内容：使用“大海捞针”（NIAH）基准测试【22, gkamradt. Needle in a haystack - pressure testing llms. https://github.com/gkamradt/ LLMTest_NeedleInAHaystack, 2023】评估REFORM的精确检索性能。该任务将一个特定的“针”句子嵌入到由Paul Graham散文组成的无关上下文（“干草堆”）的不同深度中，模型需要回答一个关于“针”内容的问题。如果回答包含所有三个关键短语，则视为正确。

实验结果：如图3所示，REFORM在不同上下文长度和“针”深度下进行了测试。结果显示，在高达1M token的各种设置中，REFORM都表现出完美的性能。

分析结论：该结果凸显了REFORM在处理极长上下文时保持精确检索性能的鲁棒性。

4.2 在RULER和BABILong上的性能

实验内容：在扩展版的RULER【17, Ruler: What’s the real context size of your long-context language models?, 2024, arXiv】和BABILong【18, Babilong: Testing the limits of llms with long context reasoning-in-a-haystack, 2024, arXiv】基准上评估了不同方法的性能。RULER是一个包含多样化和挑战性“大海捞针”任务以及一些聚合和问答任务的合成长上下文基准。BABILong通过引入更困难的任务（如多跳推理）进一步挑战模型。本文将RULER数据集扩展至1M token以评估更长输入。

实验结果：如表2所示，在RULER和BABILong两个基准测试中，REFORM的性能均远超所有基线方法。

分析结论：这些结果表明REFORM在需要精确回忆上下文关键部分的任务中具有优越性。这得益于其能够从长输入中定位关键上下文的能力，以及通过重计算消除了基于循环或随机访问方法中常见的KV缓存分布偏移问题。

Table 2: 在RULER和BABILong上的评估。我们在扩展版的RULER [17] 和 BABILong [18] 基准上测量性能。我们报告了在不同上下文长度下所有任务的平均性能。最佳值以粗体突出显示。

4.3 在∞-bench、RepoEval和多模态评估上的性能

实验内容：为进一步评估REFORM在多样化长上下文任务上的性能，在三个更贴近现实场景的基准上进行了测试：
1. ∞-bench [19]: 一个现实的长上下文基准，任务源自长篇书籍和对话。
2. RepoEval [20]: 一个仓库级别的代码补全基准，评估模型在给定整个代码仓库作为上下文时的代码补全能力。
3. MM-NIAH [23]: 一个多模态“大海捞针”基准，用于展示REFORM的广泛适用性。

实验结果：如表3和表4所示，REFORM在所有三个基准测试中都持续优于所有基线方法。

分析结论：这些结果凸显了REFORM在现实任务中的卓越性能，以及其处理多样化输入（甚至跨不同模态）的灵活性。

Table 3: 在∞-Bench上的评估。我们评估了每种方法在∞-Bench [19] 的10个数据集上的表现。我们没有在C.Run和M.Calc数据集上进行评估，因为使用这些模型时没有方法能够获得非零分数。最佳值以粗体突出显示。

Table 4: 在RepoEval和MM-NIAH上的评估。对于RepoEval，我们报告了使用1.5B和7B模型在RepoEval API级补全任务和行级补全任务上的编辑相似度（ES）分数。对于MM-NIAH，我们报告了跨输入长度的归一化性能，以确保每个上下文长度范围的贡献相等。我们没有对InfLLM进行多模态评估，因为其实施仅支持基于文本的模型。最佳结果以粗体显示。

4.4 与检索增强生成（RAG）的比较

实验内容：将REFORM与流行的长输入处理方法——检索增强生成（RAG）【24, Retrieval augmented generation or long-context llms? a comprehensive study and hybrid approach, 2024, EMNLP; 25, In defense of rag in the era of long-context language models, 2024, arXiv】进行比较。RAG将输入分块独立编码，并使用外部检索模型识别相关段落。本文在RULER的四组“大海捞针”数据集（单针、多键、多值、多查询）上，在300k上下文长度下比较了REFORM与使用稀疏和密集检索器的RAG方法的性能。

RAG的局限性分析：
1. 上下文碎片化：RAG独立编码文本块，而REFORM的检索嵌入是基于整个输入进行条件化的，保证了全局上下文的连续性。
2. 领域和模态限制：RAG受限于检索模型的训练领域，需要针对不同领域和模态进行重新训练或高级适配。REFORM是内在灵活的，无需修改即可处理多模态应用。
3. 外部依赖：REFORM将检索功能直接集成到模型中，无需外部检索模型。

实验结果：如表5所示，REFORM在所有评估中都持续优于稀疏RAG和密集RAG。此外，一个混合方法（将REFORM的token级重要性分数与密集检索器的检索分数加权求和）表现更佳。

分析结论：REFORM在鲁棒性和效率上优于标准的RAG方法。混合方法的成功表明REFORM和RAG的优势是互补的。

Table 5: 与RAG的比较。我们使用Mistral-NeMoInstruct-2407模型，在300k上下文的RULER四组大海捞针数据集（single, multikey, multivalue, 和 multiquery）上比较了RAG方法和REFORM的性能。

4.5 消融研究

实验内容：进行消融研究以评估REFORM关键组件的有效性。在512k上下文的BABILong和300k上下文的RULER基准上，分析了以下三个方面的影响：
1. 循环压缩方法的选择：将默认的H2O替换为StreamingLLM和TOVA。
2. 用于检索的注意力头的选择：将精心挑选的头替换为（1）随机选择的四个头和（2）表现最差的四个头。
3. 收集阶段的最大池化核大小：将默认的129-token窗口大小减少到5或17。

实验结果（见表6a）：
1. 压缩方法：虽然H2O效果最好，但使用其他压缩方法（StreamingLLM, TOVA）也能达到可比较的性能。
2. 注意力头选择：精心挑选的头性能最佳；随机头性能尚可但较低；使用最差的头导致性能大幅下降。
3. 池化核大小：减小池化窗口大小显著降低了性能。

分析结论：
1. REFORM框架具有灵活性，不依赖于特定的压缩方法，并可能通过更先进的压缩技术获得更高性能。
2. 正确的注意力头选择对于构建有效的嵌入至关重要，验证了本文头选择机制的有效性。
3. 池化操作对于在重计算过程中维持上下文信息至关重要。

Table 6: 消融研究与效率分析。(a) 我们报告了使用Mistral-NeMo-Instruct-2407模型在RULER 300k和BABILong 512k数据集上的平均性能。(b) 我们比较了在256k输入条件下生成10个token所需的推理时间和峰值内存使用量。所有测量均在单个H100 GPU上使用Mistral-NeMo-Instruct-2407模型完成，并且是10个样本的平均值。最佳值以粗体突出显示。

(a) 消融研究。

(b) 效率分析。

4.6 效率分析

实验内容：测量处理长输入所需的峰值内存使用量和推理时间。

实验结果（见表6b）：
* InfLLM：由于频繁的CPU和GPU之间内存传输，推理时间很长，并且需要大量内存来存储缓存。
* 循环方法：由于使用固定大小的KV缓存，推理速度更快，内存成本更低。
* REFORM：与循环基线相比，REFORM显示出更低的延迟和内存需求。

分析结论：REFORM的效率优势归功于其“早退”策略，该策略通过移除对上层KV缓存的需求，节省了计算和内存。附录D.2的进一步分析表明，REFORM在首个token生成时间（TTFT）和每个输出token时间（TPOT）上都保持了最低水平。

A5 结论

本文介绍了REFORM，一个用于高效处理长上下文的新型推理框架。REFORM通过增量处理输入块并维护一个压缩的KV缓存，提取关键的QKV状态来构建跨层上下文嵌入。早退策略提高了效率，而基于相似性的选择机制则识别并收集必要的token以进行KV缓存重计算。实验证明，REFORM在长上下文基准测试中优于现有方法，同时减少了推理时间和内存使用。此外，其与模态无关的设计使其能够广泛应用于包括多模态在内的多种用例。

局限性与未来工作：
* 局限性：REFORM的有效性部分依赖于“收集”阶段识别信息丰富token的能力，在某些边缘情况下可能存在优化空间。此外，当前实现中存在冗余的注意力分数计算（用于驱逐分数），未来可将驱逐分数计算集成到Flash Attention内核中以提高效率。
* 未来工作：计划研究更复杂的、为构建上下文嵌入量身定制的压缩方法。另一个有前景的方向是将其应用于更多样化的数据模态，如音频和视频。

A6 附录

A REFORM算法

我们在算法1的伪代码中展示了REFORM的整个过程。

B 扩展相关工作

扩展LLM以处理极长输入。为有效扩展大型语言模型（LLM）的上下文窗口，已提出多种方法。大量工作集中于修改位置嵌入，包括缩放旋转位置嵌入（RoPE）【26, Roformer: Enhanced transformer with rotary position embedding, 2024, Neurocomputing】以超越模型的上下文限制【27, Extending context window of large language models via positional interpolation, 2023, arXiv; 28, kaiokendev. Things i’m learning while training superhot. https://kaiokendev.github. io/til#extending-context-to-8k./, 2023; 29, bloc97. Ntk-aware scaled rope allows llama models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation. https: //http://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_%20scaled_rope_ allows_llama_models_to_have/, 2023; 30, Yarn: Efficient context window extension of large language models, 2023, arXiv】，应用注意力掩码【31, Lm-infinite: Simple on-the-fly length generalization for large language models, 2023, arXiv】，或调整token间的相对距离使其落入预定义范围内【32, Jianlin Su. Rectified rotary position embeddings. https://github.com/bojone/rerope, 2023; 33, Llm maybe longlm: Self-extend llm context window without tuning, 2024】。另一类研究探索了微调技术以使模型适应更长的上下文【34, Pose: Efficient context window extension of llms via positional skip-wise training, 2023, arXiv; 35, Longlora: Efficient fine-tuning of long-context large language models, 2023, arXiv】。虽然这些方法使模型能够处理扩展的输入，但它们没有解决自注意力机制带来的显著计算和内存成本问题，限制了它们在极长上下文中的实际应用。因此，我们在实验中未将它们作为基线。

其他高效长上下文处理方法。与循环KV缓存压缩方法一起，大量近期工作专注于减小KV缓存的大小，以实现更高效的长上下文推理。例如，SnapKV【36, Snapkv: Llm knows what you are looking for before generation, 2024, arXiv】提出先将完整输入通过模型，然后根据注意力分数驱逐token来压缩缓存。虽然在解码时效率高，但它要求模型首先处理完整输入，因此不适用于超出模型预训练上下文窗口的极长输入。或者，HOMER【37, Hierarchical context merging: Better long context understanding for pre-trained llms, 2024, arXiv】提出使用分层分治方法来结合编码和驱逐过程。一些工作提出通过合并token而非驱逐它们来进一步增强KV缓存压缩【38, Model tells you where to merge: Adaptive kv cache merging for llms on long-context tasks, 2024, arXiv; 39, Get more with less: Synthesizing recurrence with kv cache compression for efficient llm inference, 2024, arXiv】，但它们的实验也只考虑了模型上下文限制内的输入，其外推能力仍然未知。一些近期工作提出了另一个方向，即仅为一些选定的被称为“检索头”的注意力头保留完整缓存【40, Razorattention: Efficient kv cache compression through retrieval heads, 2024, arXiv; 41, Duoattention: Efficient long-context llm inference with retrieval and streaming heads, 2024, arXiv】，减少了保留完整KV缓存的内存负担。其他工作研究了量化【42, Kvquant: Towards 10 million context length llm inference with kv cache quantization, 2024, arXiv; 43, Gear: An efficient kv cache compression recipefor near-lossless generative inference of llm, 2024, arXiv】和低秩缓存压缩【44, Loki: Low-rank keys for efficient sparse attention, 2024, arXiv】以进一步减少KV缓存的内存需求。然而，这些方法也无法外推到超出模型预训练上下文限制的更长序列。

C 实验细节

C.1 评估检索头和嵌入

数据集准备。为了评估嵌入，我们基于多跳问答构建了一个合成数据集。在这个设置中，我们将HotPotQA数据集【45, Hotpotqa: A dataset for diverse, explainable multi-hop question answering, 2018, arXiv】中的文档嵌入到从WikiText数据集【16, Pointer sentinel mixture models, 2016】派生的长文本语料库的随机位置。每个问题都附加在上下文的末尾，并创建了token级别的标签，其中来自黄金文档的token被标记为真实标签。所有样本都被设计为8k token长，这在Mistral-7B-Instruct-v0.2模型的上下文窗口内。

嵌入提取。为了模拟由于计算或内存限制而无法进行全注意力计算的长上下文场景，我们采用了基于H2O【4, H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2023, NeurIPS】的循环分块前向传播方法，详见3.1节。对于注意力，我们使用查询状态（Q）和键状态（K）之间的点积来计算检索分数，而不应用位置编码。对于所有其他嵌入，我们使用问题嵌入和上下文嵌入之间的余弦相似度计算重要性分数，然后对问题token进行最大池化。此外，每个上下文token的检索分数通过与20个相邻token进行均值池化来平滑。

性能测量。检索性能通过平均归一化排名（Mean Normalized Rank, MNR）来量化，该排名计算为黄金token的平均归一化排名。分数越低表示性能越高，因为黄金token的排名很高。

$$ \mathtt{MNR} = \frac{1}{\mathtt{len(gold\_doc)}} \sum_{t \in \mathtt{gold\_doc}} \frac{\mathtt{rank(t)}}{\mathtt{num\_tokens}} $$

C.2 详细的头选择过程

步骤1：数据准备。我们构建了两个用于注意力头选择的合成任务：一个模式匹配任务和一个多跳问答任务。对于模式匹配任务，我们将格式为“The value corresponding to the id key is value.”的句子插入到WikiText语料库中，其中key和value都是随机的10字符ASCII字符串。对于多跳问答任务，我们使用HotPotQA中包含回答给定问题所需信息的段落。对应这些关键信息的token被标记为“黄金token”。

步骤2：上下文编码。我们对数据集中的每个样本应用使用H2O的循环分块前向传播，从每个层和注意力头中提取token级别的QKV嵌入。

步骤3：重要性分数计算。对于每一层中的每个QKV头，我们计算一个token级的重要性分数。这是通过（1）计算上下文token和问题token之间的余弦相似度，（2）通过最大池化聚合问题token的相似度分数，以及（3）使用20个token窗口的均值池化来平滑分数来完成的。

步骤4：性能测量和头选择。我们使用平均归一化排名（MNR）分数来评估每个注意力头的有效性，如C.1节所述。基于这些分数，我们选择四个头：两个在模式匹配任务上表现最佳，两个在多跳问答任务上表现最佳。

C.3 多模态评估

基线细节。在多模态实验中，我们仅使用基于循环的方法评估模型性能，因为InfLLM的代码库仅支持基于文本的模型。对于InfiniPot，NuC（压缩下的新颖性）分数不能用于多模态模型的缓存压缩，因为视觉token不输出logit。因此，在多模态实验中，我们仅对InfiniPot基线应用CaP（催化剂提示）分数。

C.4 与RAG的比较

实验设置。我们遵循OP-RAG【25, In defense of rag in the era of long-context language models, 2024, arXiv】中的设置进行RAG实验，将输入分割成128个token的块，并保持块的顺序而不是根据检索分数重新排列它们。我们使用Mistral-NeMo-Instruct-2407作为基础LLM。我们使用BM25【46, Okapi at trec-3, 1995, Nist Special Publication Sp】作为稀疏检索器，使用bge-large-en-v1.5【47, C-pack: Packaged resources to advance general chinese embedding, 2023】作为密集检索器。对于每个样本，总共检索8k个token，与我们的方法中的KV大小相匹配，以确保公平比较。

C.5 效率测量

实验设置。我们测量在256k个token的条件下生成10个token的平均推理时间和峰值内存使用量。所有测量都在单个H100 GPU上进行，并且我们对所有测量应用了Flash Attention 2【48, FlashAttention-2: Faster attention with better parallelism and work partitioning, 2024, ICLR】。我们进一步详细说明了InfLLM的实验设置，因为其推理速度和内存消耗会根据配置有很大差异。我们使用了他们GitHub仓库中提供的默认配置，同时修改了检索块的数量以在缓存中保持32k个活动token。GPU中缓存的最大块数设置为检索块数的两倍，遵循他们官方配置文件中的惯例。

C.6 REFORM的嵌入构建和相似性搜索

嵌入头选择。我们通过组合四个QKV嵌入来构建上下文嵌入，其中两个头是使用模式匹配数据集识别的，另外两个是使用多跳问答数据集识别的。为了在性能和效率增益之间取得平衡，我们从深度低于70%的层中选择表现最佳的模式匹配头。更详细的讨论见D.1节。

对于Mistral-NeMo-Instruct-2407，使用了以下头：
1. 第15层的查询头9
2. 第19层的值头5
3. 第27层的值头0
4. 第27层的值头7

对于Qwen2.5-7B-Instruct，使用了以下头：
1. 第7层的值头3
2. 第14层的键头0
3. 第14层的值头3
4. 第19层的值头0

对于Qwen2.5-Coder-1.5B-Instruct，使用了以下头：
1. 第8层的查询头3
2. 第11层的值头1
3. 第14层的键头0
4. 第15层的值头0

对于Qwen2.5-Coder-7B-Instruct，使用了以下头：
1. 第13层的值头2
2. 第14层的键头0
3. 第14层的值头3
4. 第14层的查询头4

对于Pixtral-12B-2409，使用了以下头：
1. 第10层的值头3
2. 第19层的值头5
3. 第27层的值头0
4. 第27层的值头7

相似性搜索。REFORM在查询（输入的最后一部分）中的每个token与其余token之间执行余弦相似性搜索。为了更精确地识别相关输入，我们在计算相似性分数时移除了特殊token和生成前缀（例如，“the answer is”）。

D 附加结果

D.1 模式匹配任务的嵌入头识别

结果分析。在本节中，我们展示了使用我们的模式匹配数据集测量的MNR分数分布，类似于我们在表1和图2中展示的内容。模式匹配数据集的相应结果呈现在表7和图4中。与多跳问答数据集的情况类似，QKV头的检索性能通常优于隐藏状态。有趣的是，表现最佳的头的分布与多跳问答数据集相比呈现出不同的模式，较低层和中高层的头表现出最高的性能。这表明不同的头根据任务表现出不同的特性。这也激发了我们使用由不同任务识别的嵌入的方法，因为它能产生更通用的表示，并使基于相似性的检索更加准确。同样重要的是要注意，虽然上层有更多表现良好的头，但这些头也可以在较低层中被识别出来（例如，第16层，头1）。为了平衡性能与早退策略提供的效率增益，我们从深度低于70%的层中选择表现最佳的模式匹配头。这一策略确保我们既能利用高性能的头，又能享受到早退带来的计算节省。

Table 7: 比较不同LLM嵌入。对应于隐藏状态和注意力状态的最佳3个MNR分数（越低越好），由Mistral-Nemo-Instruct-2407测量。分数是在500个合成模式匹配样本上取平均值。

D.2 延迟分解

实验扩展。在这里，我们将延迟分析扩展到三种上下文长度（正文中仅有256k），并在表8a和表8b中分别展示预填充和解码延迟。

结果分析。在表8a中，由于动态缓存驱逐和检索带来的额外计算，InfiniPot和InfLLM的处理速度比StreamingLLM慢。另一方面，得益于早退优化，REFORM始终显示出比所有基线更快的预填充时间。Compress+Gather阶段和Recompute阶段的进一步延迟分解表明，重计算开销极小，占总预填充时间的不到1.52%。表8b显示了解码时间的类似趋势。由于额外的操作，InfiniPot和InfLLM的解码时间比StreamingLLM慢。另一方面，REFORM也显示出改进的解码速度，因为生成是基于用小重计算预算（8k）创建的KV缓存进行的，而基线则是在完整的压缩缓存（32k）上进行条件生成。（请注意，这些预算是我们正文中Mistral-Nemo实验使用的超参数。）

(a) 首个token生成时间（秒）。

Table 8: 延迟分解。(a) 首个token生成时间（秒）测量和 (b) 每个输出token时间（秒）测量（Mistral-Nemo-Instruct-2407，单个H200，20次运行平均，每次测量生成200个token）。

(b) 每个输出token时间（秒）

D.3 使用更大模型的评估

实验内容。本节展示了使用Qwen2.5-32B-Instruct模型在RULER和BABILong基准上的附加结果。

结果分析。如表9所示，REFORM在所有评估设置中都持续优于H2O和InfiniPot，证明了我们方法在更大模型上的可扩展性和鲁棒性。

Table 9: Qwen2.5-32B-Instruct的性能。我们报告了H2O、InfiniPot和REFORM在关键长上下文基准上的性能。最佳值以粗体突出显示。

E 进一步讨论

E.1 复杂性分析

时间和内存复杂性。假设块大小和查询长度固定（或有界），像StreamingLLM这样的循环基线的时间复杂度为O(L)，其中L是输入长度。每个循环步骤涉及恒定的计算量，总步数与输入长度线性相关。它们的内存复杂度为O(1)，因为始终只维护一个固定大小的KV缓存。相比之下，像InfLLM这样的随机访问基线的时间复杂度为O(L²)，主要是由于在整个输入长度上进行周期性的缓存查找。这些方法还需要O(L)的内存，因为它们将整个KV缓存存储在内存中。这种内存负担是巨大的，通常需要CPU卸载，从而进一步增加延迟。REFORM通过其循环分块处理保持了O(L)的时间复杂度。虽然由于存储token级嵌入而具有O(L)的内存成本，但这些嵌入在实践中非常小。此外，早退策略显著减少了内存和计算需求。因此，REFORM实现了比两种基线更快的速度和更低的内存使用（如表6b所示），突显了我们方法的实际效率。

离线计算量。虽然REFORM涉及用于头选择的离线计算，但头选择的计算开销在实践中是最小的。选择过程仅使用8k token的合成输入进行一次，并且选择的头在所有推理运行中重复使用，与下游任务或输入领域无关。通过使用短输入选择头并在长上下文推理中重复使用它们，REFORM将额外的计算量保持在最低水平。从数量上看，整个头选择过程所需的计算量（以FLOPs计）大致相当于用H2O处理两个1M-token输入的计算量。

E.2 头选择的理论分析

理论动机。我们的注意力头选择受到机械可解释性文献发现的启发，该文献表明不同的头专注于不同的功能【49, Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned, 2019, arXiv; 50, Identifying semantic induction heads to understand in-context learning, 2024, arXiv】。一些头在模式识别或推理等任务上比其他头更有用。我们通过仅选择少数高性能的头来利用这种多样性，既为了更好的识别性能，也为了保持token级嵌入的大小可控（包括所有头会占用大量内存）。因此，我们的方法旨在实现功能专业化和计算效率。

与现有工作的联系。虽然我们使用余弦相似度而不是注意力权重来评估头，但我们的方法与现有的可解释性工作共享关键见解。先前的研究表明，低层头通常检测句法或结构特征，中层头处理语义推理，高层头指导输出生成【51, Attention heads of large language models: A survey, 2024, arXiv】。这与我们的经验发现相符：用于模式匹配的最有用的头倾向于出现在低层和高层，而用于多跳问答（即语义理解）的头则集中在中层。这些趋势反映在我们的MNR可视化中（图2和图4）。此外，Skean等人【52, Layer by layer: Uncovering hidden representations in language models, 2025, arXiv】提出，对于32个MTEB（Massive Text Embedding Benchmark, 【53, Mteb: Massive text embedding benchmark, 2022, arXiv】）任务，中间层的表示比最终层的表示更有用，这与我们在多跳问答任务中观察到的类似。

任务选择的依据。此外，我们选择头识别任务的依据源于信息检索文献，该文献表明结合稀疏（结构化）和密集（语义）检索器的混合检索器【54, An analysis of fusion functions for hybrid retrieval, 2023, ACM Transactions on Information Systems; 55, Out-of-domain semantics to the rescue! zero-shot hybrid retrieval models, 2022, European Conference on Information Retrieval】的性能优于单独使用任何一种。受此启发，我们使用了两个不同的头评估任务：模式匹配和多跳问答。前者识别捕捉结构相似性的头，而后者则针对编码语义相关性的头。这种双任务策略帮助我们选择了一套互补的头，这些头在多样化的长上下文场景中具有鲁棒性。

F 更广泛的影响

我们相信，通过将模型能力扩展到高效且有效地处理非常长的上下文，其强大的能力和灵活性将有助于大型基础模型的日常使用。另一方面，REFORM的这种能力也可能被恶意方利用来分析大量数据，从而增强可用于操纵或传播虚假信息的自主系统的能力。

G 数据集和模型的许可证信息

在这里，我们提供了我们实验中使用的所有数据集和模型的许可证。Apache 2.0许可证适用于Babilong、RULER、Mistral-Nemo-Instruct-2407、Qwen2.5-Coder系列和Pixtral-12B-2409。BSD许可证也适用于Babilong数据集的某些部分。MIT许可证适用于Needle-in-a-Haystack、InfiniteBench、RepoEval、WikiText和bge-large-en-v1.5。CC BY-SA 4.0适用于HotPotQA。