← 返回 2026-05-20

用于高效长上下文生成的上下文记忆 Context Memorization for Efficient Long Context Generation

Yasuyuki Okoshi, Hao Mark Chen, Guanxi Lu, Hongxiang Fan, Masato Motomura, Daichi Fujiki 📅 2026-05-18 👍 6 2026-07-13 08:36

KV缓存推理优化注意力分解注意力机制长上下文

arXiv HuggingFace PDF

提出注意力状态记忆，无训练地将长前缀外部化为查找记忆

前置知识

前缀增强推理

前缀增强推理是指在推理时为模型提供长上下文前缀（如示例、规则、文档）来引导模型生成。常见场景包括上下文学习（ICL）在提示中包含多个示例，以及检索增强生成（RAG）插入检索到的文档。这种方式的核心思想是将知识通过上下文传递给模型，而不需要修改模型参数。然而，前缀越长，推理时的注意力计算开销就越大，且随着生成长度增加，前缀的影响力会逐渐衰减

理解前缀增强推理是理解本文动机的基础，因为本文针对的核心问题就是长前缀带来的两个结构缺陷：前缀衰减和计算开销线性增长

Online-Softmax恒等式

Online-Softmax恒等式是一个数学性质，它允许将注意力分解为多个独立块的组合。具体来说，对于查询向量$q$和分块的键值对$K=[K_1,...,K_B]$、$V=[V_1,...,V_B]$，注意力可以表示为$ ext{attn}(q,K,V) = \sum_{b} \alpha_b \text{attn}(q,K_b,V_b)$，其中$\alpha_b$是第b块的归一化权重。这个恒等式意味着我们可以分别计算每个块的注意力，然后在推理时合并，无需重新计算完整注意力

这是本文方法的理论基础，它证明了存储注意力状态$(a(q), Z(q))$足以重构完整注意力，使前缀可以被外部化到查找记忆中

KV缓存

KV缓存是自回归生成中常用的优化技术，它存储每个token的键和值向量，避免在生成新token时重新计算历史token的表示。在标准注意力机制中，每个解码步骤都需要加载所有历史token的KV缓存来计算注意力，这导致推理延迟和内存消耗与上下文长度线性相关。分组查询注意力（GQA）通过让多个查询头共享KV头进一步减少缓存大小

理解KV缓存对于理解本文的效率分析和内存评估至关重要，本文直接与KV缓存压缩方法对比，并计算了在不同内存预算下的性能

上下文学习（ICL）

上下文学习是一种无需训练的提示方法，通过在输入中包含任务示例让模型学习新任务。例如，要教模型做分类，可以在提示中提供多个（输入，标签）对，然后让模型预测新输入的标签。研究表明，ICL随着示例数量增加可以显著提升性能，但这种提升在推理时需要为每个示例计算注意力，开销巨大。本文在ManyICLBench上评估方法，这是一个专门测试多示例学习的基准

ICL是本文的主要应用场景之一，理解ICL的工作原理和局限性有助于理解为什么需要将长前缀外部化到查找记忆中

研究动机

现代大语言模型应用越来越依赖长上下文前缀来控制推理时的模型行为，包括上下文学习（ICL）、外部知识源和智能体指令。然而，这种前缀增强方法引入了两个结构性代价。第一个是前缀衰减：随着生成继续，模型的注意力被分散到各个token上，导致前缀对模型行为的影响力衰减，这在长上下文场景中尤为明显。第二个是推理低效：前缀长度增加时，对前缀的注意力计算在prefill和每个解码步骤都引入与长度线性相关的延迟和内存开销。虽然前缀缓存可以摊销prefill成本，但仍消耗大量内存，这对部署的智能体系统（如Claude Code）是重大瓶颈。现有方法要么在推理时仍需关注压缩后的前缀，要么需要基于梯度的训练将前缀内化到参数中，都不理想

本文的目标是本文的目标是提出一种训练免费的方法，完全消除推理时对前缀的注意力计算，同时避免梯度训练的高昂代价。具体而言，作者希望将长前缀外部化为轻量级的基于查找的记忆，使得推理时只需要$O(\log K)$的检索开销，而不是$O(L)$的注意力计算（$K$是记忆条目数，$L$是前缀长度）。同时，记忆应该可以通过纯前向计算构建和更新，无需反向传播，从而支持前缀的动态更新。在性能方面，作者希望该方法在上下文学习和检索增强生成任务上达到与完整注意力相当甚至更好的准确率

与已有工作不同的是，本文的独特切入角度是从注意力分解的数学性质出发，将前缀注意力问题从"如何压缩前缀"或"如何将前缀编码到参数中"转变为"如何预计算并检索注意力状态"。与保留前缀在注意力中的压缩方法不同，本文完全消除了推理时对前缀的注意力；与将前缀内化到参数中的方法不同，本文无需梯度训练，而是通过纯前向计算构建记忆。这种视角转变使本文首次同时实现了前缀长度无关的解码延迟、训练免费的前缀构建和无需辅助模型这三个特性，填补了现有方法在权衡空间中的空白

核心方法

注意力状态记忆（Attention-State Memory，ASM）是一种每层字典，存储预计算的注意力状态，并通过查询向量索引和聚类实现跨查询共享。整体思路分为两个阶段：离线校准阶段和在线推理阶段。在离线阶段，模型对前缀和响应轨迹对进行前向传播，收集响应中每个查询对前缀的注意力状态$(a, Z)$，然后通过K-means聚类查询向量将它们压缩为固定数量的条目。在推理阶段，对于每个用户查询token，在记忆中检索最相似的质心，获取预计算的注意力状态，然后将其与查询的自注意力合并，无需真正计算对前缀的注意力。这种方法的核心洞察是online-softmax恒等式证明了注意力状态的充分性和可组合性，使得前缀可以被外部化到查找记忆中

核心创新点在于利用online-softmax恒等式将前缀注意力外部化到查询索引的字典中，而不是保持前缀在注意力中或内化到模型参数中。与现有方法的本质区别在于：KVzip等压缩方法仍然需要在推理时计算对压缩前缀的注意力，虽然减少了前缀长度但仍有线性开销；上下文蒸馏等内化方法需要梯度训练，不适合前缀更新；而ASM通过预计算注意力状态并存储在查找表中，推理时只需要$O(\log K)$的检索成本，完全消除了对前缀的注意力计算。另一个关键区别是ASM通过合并算子支持分块前缀的并行编码，可以降低校准时的峰值内存，同时保持性能

方法步骤详情

方法分为四个主要步骤。第一步是离线校准的收集阶段：对于每个前缀-轨迹对$(p, t)$，对拼接序列$[p, t]$进行前向传播，获取每层的KV缓存。对于响应轨迹中的每个查询$q$，记录其对前缀的注意力状态$(a_p(q), Z_p(q))$，其中$a_p(q) = \text{attn}(q, K_p, V_p)$是注意力输出，$Z_p(q) = \sum_{k \in K_p} \exp(q \cdot k / \sqrt{d_h})$是归一化因子。第二步是离线校准的聚类阶段：对每层收集的查询向量集合$\{q : (q, \cdot, \cdot) \in P^{(i)}\}$应用K-means聚类，将相似的查询分组到K个簇中。对每个簇，使用注意力感知聚合计算质心：查询质心为簇内查询的平均，注意力状态质心通过合并算子$\text{merge}((a_A, Z_A), (a_B, Z_B)) = \frac{Z_A a_A + Z_B a_B}{Z_A + Z_B}, Z_A + Z_B$得到。第三步是在线推理的检索阶段：对于每个用户查询token，使用查询向量作为查找键，通过余弦相似度找到最近的簇质心$c^* = \arg\max_k \cos(q, \bar{q}_k)$，检索对应的压缩注意力状态$(\bar{a}_{c^*}, \bar{Z}_{c^*})$。第四步是在线推理的合并阶段：将检索到的注意力状态与用户查询的自注意力$(a(q), Z(q))$合并，得到最终输出$a_{merge}(q) = \frac{Z(q)}{Z(q) + \bar{Z}_{c^*}}a(q) + \frac{\bar{Z}_{c^*}}{Z(q) + \bar{Z}_{c^*}}\bar{a}_{c^*}$。这个合并过程根据online-softmax恒等式是无损的，完全恢复了对前缀和生成的自注意力的组合效果

技术新颖性

技术新颖性体现在三个方面：首先，将online-softmax恒等式从高效注意力计算扩展到跨查询的前缀重用场景，这是该恒等式的一个新应用。其次，提出了注意力感知聚合，在计算簇质心时利用合并算子保持注意力的数学结构，而不是简单的平均。第三，设计了分层查找索引，将检索复杂度从$O(K)$降低到$O(\log K)$，使得记忆规模可以独立于前缀长度增长而不增加推理延迟。此外，ASM还支持分组查询注意力（GQA）的扩展，通过聚合组内查询并按KV头存储质心，进一步减少了记忆条目数量

**Figure 2**: Overview of attention-state memory

实验结果

实验在ManyICLBench和RuleArena上使用LLaMA 3.1-8B验证了ASM在上下文学习（ICL）和检索增强生成（RAG）场景中的有效性。在ICL的五个基准任务上，ASM在1K到8K记忆条目范围内平均准确率超过了ICL基线，并在4K条目内扩展效率更高。这源于每个条目存储了对整个前缀的注意力输出，因此可以保持ICL在上下文大小受限时会丢失的信息。在推理任务（math和science）上，ASM与ICL基线保持相近的准确率，表明记忆有效保留了前缀信息。在RAG的NBA基准上，ASM在最佳设置（K=4K）达到了25.5%的准确率，超过了使用完整20K规则书的ICL基线（24.1%），而只使用了20%的内存占用。效率分析显示，ASM在8K条目时将注意力延迟降低了1.36倍，在16K条目时达到了1.8倍加速。分层查找验证显示，当扩展top-16第一级质心时，分层检索可以匹配线性查找的top-1准确率。分块前缀构建实验表明，用四个4K前缀块构建16K前缀可以达到与单次16K前缀相近的准确率（78.5% vs 79.0%），而峰值内存仅为1/4

**Table 1**: Accuracy of ICL and ASM (ours) on reasoning benchmarks (math_counting and gpqa_cot) across different memory sizes

**Table 2**: Accuracy on NBA benchmark in RuleArena

**Figure 3**: Accuracy of five benchmarks from ManyICL Bench

**Figure 4**: Heatmap for hierarchical lookup

查看结构化数据

任务	指标	本文	基线	提升
ManyICLBench（5个任务平均）	准确率	ASM (4K条目)	ICL (4K KV条目)	超过ICL，具体数值因任务而异
NBA benchmark (RuleArena)	精确匹配准确率	ASM (4K条目)	ICL (20K规则书)	25.5% vs 24.1%，同时减少80%内存
注意力延迟（生成512 token平均）	相对延迟	ASM (8K条目，分层查找)	完整注意力	1.36倍加速
banking77准确率（分块构建）	准确率	ASM (4个4K块构建16K)	单次16K前缀	78.5% vs 79.0%，峰值内存减少75%

局限与改进

作者承认的主要局限性是当记忆条目数量较大时（如16K），某些任务（如bbh_geometric_shapes只有150个训练样本）的校准数据不足以填充代码本，阻碍了ASM的有效扩展。此外，准确率并非单调递增，而是在中间代码本大小达到峰值，表明最优条目数是任务特定的超参数，需要调优而非设置得越大越好。从方法本身来看，ASM假设相似查询的注意力输出会相似，这在语义相似的token上成立，但对于需要精确匹配的场景（如规则应用中的特定条件）可能不够精确。另一个限制是ASM需要离线校准阶段，虽然只是前向传播，但仍需要额外时间和计算资源来构建记忆，且当频率更新前缀时可能需要频繁重建记忆。最后，ASM目前主要针对固定前缀场景，对于频繁变化的动态前缀，其优势可能减弱

独立分析的弱点

第一个弱点是聚类质量依赖于校准数据的代表性和数量，对于数据稀缺的任务，高容量记忆可能无法得到有效填充，导致性能下降。改进方向可以包括自适应聚类算法，根据数据密度自动调整簇的数量，或者使用半监督方法增强小数据集的校准质量。第二个弱点是当前使用余弦相似度作为检索度量，这可能不够精确，特别是在GQA场景中聚合查询时可能丢失重要信息。改进方向可以包括学习感知的检索度量，或者根据查询和质心的交互历史动态调整检索权重。第三个弱点是分层查找的参数（如第一级质心数、扩展的top-m）需要启发式调优，不同任务可能需要不同配置。改进方向可以包括自适应分层索引，根据记忆大小和查询模式自动优化索引结构。第四个弱点是当前方法主要针对固定前缀，对于需要频繁更新的动态场景（如实时知识库更新），重建记忆的成本可能成为瓶颈。改进方向可以包括增量更新机制，只更新受新前缀影响的部分记忆条目

未来方向

作者提出的未来工作方向包括将ASM扩展到更多注意力变体（如滑动窗口注意力、局部敏感哈希注意力）以及探索跨层记忆共享以进一步减少内存占用。基于论文成果，可以延伸的研究方向包括：1）自适应记忆构建，根据任务难度和查询分布动态调整记忆条目数和聚类策略；2）跨任务记忆迁移，将一个任务构建的记忆迁移到相关任务，减少校准开销；3）多模态前缀记忆，将ASM扩展到视觉-语言场景中的图像-文本前缀；4）理论分析，深入研究注意力状态的分布特性，为记忆设计提供更多理论指导；5）硬件加速，为分层查找设计专用硬件加速器，进一步降低检索延迟；6）与模型并行结合，探索在分布式推理场景中如何高效共享和同步注意力状态记忆

复现评估

论文提供了开源代码（https://github.com/yasu0001/AttentionMemory），这大大降低了复现难度。实验使用公开基准ManyICLBench和RuleArena，数据集可获取。模型使用LLaMA 3.1-8B的指令调优版本，可通过HuggingFace获取。校准数据使用各任务的训练集，NBA基准使用合成数据并过滤与测试集重叠的序列。算力要求方面，离线校准需要前向传播过前缀和轨迹，峰值内存随前缀长度线性增长，但通过分块构建可以降低。实验没有报告具体的GPU型号和时间，但考虑到只是前向传播而非训练，算力需求相对较低。代码应该包含所有超参数设置（如K值、RoPE处理方式、白化选择等），论文也报告了各任务的最佳配置。总体而言，复现难度中等，主要挑战可能是校准数据的准备和超参数调优，但有开源代码和明确设置应该可以成功复现

论文图表

**Figure 1**: Comparison of three approaches for handling a long fixed prefix that is reused across many user requests
这张图比较了三种处理跨请求复用的长固定前缀的方法：前缀压缩（Prefix Compression，如Prefix Cache）、前缀内部化（Prefix Internalization，如Fine-tuning）和注意力状态记忆（Attention-State Memory，本文方法）。顶部展示了一个长固定前缀（如上下文示例）被多个用户请求重用的场景。中间部分展示了三种方法的架构：前缀压缩仍然在每次请求中计算对前缀的注意力，复杂度$O(K')$；前缀内部化通过梯度训练将前缀吸收到参数中，无需推理时注意力但需要参数更新；注意力状态记忆构建K条目的查找表，推理时只需$O(\log K)$的检索开销。底部总结了每种方法的优缺点：前缀压缩有延迟开销和前缀衰减；前缀内部化无延迟但需要参数更新；注意力状态记忆有对数级延迟开销且训练免费

这张图对理解论文至关重要，因为它清晰展示了本文方法与现有方法的根本区别，直观说明了为什么ASM可以同时实现低延迟和训练自由

Figure 5: Attention latency comparison between existing in-context learning and attention-state memory with linear lookup and hierarchical lookup

这张图展示了不同方法的注意力延迟对比。X轴是记忆条目数，Y轴是归一化延迟。比较了三种配置：完整注意力（蓝色）、ASM线性查找（绿色）、ASM分层查找（红色）。结果显示完整注意力的延迟随着条目数线性增长，因为每个解码步骤都要关注整个前缀。ASM线性查找也有一定的增长，但较缓慢。ASM分层查找在小记忆预算时有较高开销，但随着预算增加，其开销增长远慢于完整注意力。综合线性和分层查找的优化配置，ASM在约4K条目时变得比完整注意力更快，在16K条目时达到1.8倍加速。这证明了ASM可以抑制推理时间随内存预算的增长，与标准自注意力不同

这张图直接证明了ASM的效率优势，展示了如何将对数检索复杂度转化为实际的推理加速，是论文的核心贡献