基于端到端强化学习的压缩记忆动态长上下文推理 Dynamic Long Context Reasoning over Compressed Memory via End-to-End Reinforcement Learning
通过压缩KV缓存和选择性记忆召回实现高效长上下文多跳推理
前置知识
KV缓存（Key-Value Cache）
在Transformer推理过程中，每一层的注意力机制需要计算查询（Query）与所有历史键（Key）的注意力分数，再用这些分数对值（Value）加权求和。为了避免重复计算，系统会将已计算过的Key和Value张量缓存起来，这就是KV缓存。它直接决定了长上下文推理时的显存占用——序列长度为N时，KV缓存大小与N成正比增长，这成为处理百万级token文档的核心瓶颈之一。本文的压缩方法本质上是将一段文本的KV缓存压缩到远少于原文token数的「记忆token」隐状态中。
KV缓存是长上下文LLM推理的显存瓶颈。理解其结构才能理解本文为何选择在KV缓存层面做压缩，以及压缩比alpha如何影响内存占用和推理速度。
LoRA（Low-Rank Adaptation）
LoRA是一种参数高效的微调方法，其核心思想是在预训练模型的权重矩阵旁插入一个低秩分解的旁路矩阵。具体来说，对于原始权重矩阵W，LoRA添加一个旁路 $\Delta W = BA$，其中 B 和 A 是远小于W的低秩矩阵，训练时只更新B和A而冻结W。本文中三个组件（Compressor、Gate、Reasoner）各自使用一个独立的LoRA模块叠加在同一个基础模型上，共享基础模型参数但拥有独立的适配器权重。
本文的关键架构设计之一就是用同一个基础模型通过不同的LoRA模块扮演三个角色。理解LoRA才能理解这种「共享主干+独立适配器」的设计如何实现参数效率和功能解耦。
端到端强化学习（End-to-End RL）
端到端强化学习指将整个系统的多个组件作为一个整体策略进行联合优化，让最终任务的奖励信号能够反向传播更新所有可训练参数。本文中，压缩阶段生成的记忆表示theta会影响后续推理阶段的准确率，端到端RL使得推理阶段的奖励梯度可以穿过推理器、穿过压缩后的记忆、一直传到压缩器的LoRA参数，从而让压缩器学会生成「对推理最有用」的表示，而不仅仅是「信息保留最完整」的表示。
这是本文最核心的技术贡献。传统方法分别优化压缩和推理，而本文通过GSPO算法实现联合策略优化，使压缩表示直接面向下游任务优化，这是与先前工作的本质区别。
多跳推理（Multi-hop Reasoning）
多跳推理是指回答一个问题需要串联多个证据片段、经过多步逻辑推导才能得到答案的推理任务。例如：「张三的导师的博士论文研究了什么？」需要先找到张三的导师是谁（第一跳），再找到该导师的博士论文主题（第二跳）。在长上下文场景下，这些证据可能分布在文档的不同位置，且后续跳的证据只有在前一跳推理完成后才能确定其相关性，这对检索和推理系统提出了很高的要求。
本文的核心评测任务就是多跳推理。多跳推理要求系统不仅能找到相关文本，还要能基于推理状态动态调整检索策略——这正是Gate模块（条件于工作记忆状态）设计的动机。
GSPO（Group Sequence Policy Optimization）
GSPO是本文采用的强化学习优化算法，属于策略梯度方法的一种变体。与标准PPO（Proximal Policy Optimization）在token级别计算重要性采样权重不同，GSPO在序列级别定义重要性采样权重，并对采样轨迹进行分组（group size G），通过裁剪（clip）机制约束策略更新幅度，同时加入KL散度惩罚项防止策略偏离参考策略太远。这种设计更适合长序列生成任务中的策略优化。
GSPO是本文端到端RL训练的具体实现算法。理解其序列级优化特性才能理解为什么压缩器和推理器可以被联合训练而不出现训练不稳定。
研究动机
大语言模型在处理长上下文时面临三个核心困难。首先是计算复杂度问题：标准自注意力机制的计算和显存开销与序列长度呈二次方关系，当文档长度达到数十万甚至百万token时，直接处理全部原始token在工程上几乎不可行。以Qwen2.5-7B为例，在64K长度的上下文上推理就已经显著慢于基于记忆的方法，在128K时甚至直接OOM。其次是信息遗忘问题：即使模型理论上支持长上下文，随着输入长度增加，模型对关键信息的注意力会被大量无关内容稀释，导致性能急剧下降——DeepSeek-R1-Distill-Qwen-7B在7K上下文时准确率为30.47%，到56K时已经跌至0%。最后是RAG的上下文碎片化问题：传统检索增强生成将文档切分为独立块进行检索，破坏了跨块的逻辑依赖关系，难以支持需要串联多个证据的多跳推理，且无法利用推理过程中积累的中间状态来指导后续检索。
本文的目标是本文的目标是设计一个受人类记忆机制启发的框架LycheeMemory，通过分块压缩和选择性记忆召回来实现高效的长上下文推理。具体而言，该框架需要满足三个目标：第一，将上下文长度从训练时的7K成功外推到1.75M tokens；第二，在多跳推理基准上保持具有竞争力的准确率（目标不低于80% sub-EM）；第三，相比现有的流式处理方法（如MemAgent），在推理效率上实现数量级的提升，包括降低峰值GPU显存占用和加速推理延迟。
与已有工作不同的是，现有方法存在三个关键空白。第一，显式记忆方法（如标准RAG和MemAgent）直接操作原始token，计算开销大且无法利用压缩后的高效表示，MemAgent虽然通过流式处理实现了线性复杂度，但对每个文本块都需要进行完整的token生成操作。第二，隐式记忆方法（如KV缓存压缩、LoRA适配器编码）虽然减少了内存占用，但激进压缩往往导致长尾推理能力退化，且缺乏结构化的检索机制。第三，最重要的是，现有方法中压缩和推理是分离优化的——压缩器只学习「保留信息」，推理器只学习「利用记忆」，二者之间缺乏端到端的梯度连接。LycheeMemory的独特切入角度是：通过端到端强化学习将压缩和推理联合优化，让压缩器学会生成「对推理最有用」的表示，而非仅仅是「信息保留最完整」的表示，同时引入条件于推理状态的动态门控机制实现选择性记忆召回。
核心方法
LycheeMemory的核心直觉来自人类记忆系统的双重结构：长期记忆存储压缩后的信息，工作记忆在推理过程中动态更新。技术路线上，整个框架由三个组件构成，它们共享同一个基础模型但各自配备独立的LoRA适配器：压缩器（Compressor）将输入文档分块后编码为压缩的KV缓存表示，形成静态的压缩记忆库；门控模块（Gate）在推理时对每个压缩记忆块进行相关性评分，决定是否值得激活；推理器（Reasoner）在工作记忆中迭代处理被选中的记忆块，逐步积累推理所需的信息。推理过程从空的工作记忆开始，顺序遍历压缩记忆库中的每个块，Gate判断该块是否与当前查询和工作记忆状态相关，如果相关则推理器更新工作记忆，否则跳过。最终基于更新后的工作记忆和查询生成答案。这种设计将长上下文处理从「一次性处理所有原始token」转变为「在压缩记忆上迭代推理」，核心优势在于推理阶段的计算量不再与文档长度成正比，而只与被选中的相关块数量成正比。
本文的核心创新有两个层面。第一个创新是压缩记忆的KV缓存风格表示：不同于直接截断或丢弃token，本文在每个文本块中插入可训练的记忆token（memory tokens），与原始token交错排列后通过压缩器进行前向传播，让记忆token的隐状态吸收其前面alpha个原始token的语义信息，形成紧凑的压缩表示。压缩比alpha=4时，每4个token压缩为1个记忆token的隐状态，在保持接近无损精度的同时将显存占用减半。第二个更关键的创新是端到端联合策略优化：将压缩器的参数和推理器的参数定义为统一策略，使得推理阶段的奖励梯度可以反向传播穿过推理器、穿过压缩后的记忆表示、一直到达压缩器的LoRA参数。这意味着压缩器不再独立地学习「保留信息」，而是学习「保留对推理最有用的信息」。与先前工作（如分别训练压缩器和推理器的两阶段方法）相比，这种联合优化在实验中带来了+8.37 sub-EM的性能提升（从60.38到68.75）。
方法步骤详情
LycheeMemory的处理流程分为三个阶段。第一阶段是压缩器预训练（Stage 1）：将输入文档D分割为K个大小为sz=4096的顺序块，对每个文本块确定压缩比alpha后定义可训练记忆token，将它们与原始token交错排列形成序列，通过压缩器进行单次前向传播，提取记忆token对应的隐状态作为压缩表示。预训练时使用三个辅助任务（文本重建、QA生成、创意生成）来优化压缩LoRA，同时冻结基础模型。文本重建损失要求模型仅用压缩表示重建原始文本；QA生成损失要求模型基于压缩表示和预生成的合成问题生成答案；创意生成损失要求模型基于压缩表示执行摘要等高级语义任务。第二阶段是端到端RL联合优化（Stage 2）：定义联合策略，使用GSPO算法通过最大化期望奖励来联合优化压缩器和推理器的参数。GSPO在序列级别定义重要性采样权重，通过裁剪机制约束策略更新幅度，同时加入KL散度惩罚项。第三阶段是Gate训练（Stage 3）：将文本块与下游任务对齐构建训练数据，标注记忆块是否包含回答查询所需的证据，使用二元交叉熵损失训练Gate模块作为二分类器。推理时，Gate输出相关性概率，只有概率超过阈值时才激活推理器更新工作记忆。
技术新颖性
LycheeMemory的技术新颖性体现在四个方面。第一，它是首个将KV缓存风格压缩与端到端RL联合优化相结合的长上下文框架，实现了从感知（压缩）到推理的全链路梯度连接，这与MemAgent分别优化压缩和推理的做法形成本质区别。第二，记忆token的设计巧妙地将压缩问题转化为序列到序列的映射，避免了引入外部编码器，压缩器直接复用基础模型加LoRA的方式既保持了架构简洁性又确保了压缩表示与推理器的兼容性。第三，Gate模块的条件检索设计——同时以查询和演化中的工作记忆为条件——实现了状态依赖的动态检索，解决了多跳推理中「后续证据只有在前一跳推理完成后才能确定相关性」的核心挑战。实验表明这种设计在112K长度下实现了86.3%的召回率，显著优于静态嵌入检索的82.1%。第四，三个组件共享基础模型但各自使用独立LoRA的设计，通过参数高效的方式实现了功能解耦，避免了为每个角色部署独立模型带来的计算和存储开销。
实验结果
本文在多个维度上进行了系统实验，核心发现如下。首先，在主实验（Table 1）中，LycheeMemory-7B w/o Gate变体在从7K到1.75M的所有上下文长度上均保持了优异性能，在112K长度下达到80.47%的sub-EM准确率，显著优于Qwen2.5-Instruct-7B-1M在同一长度下的51.56%。带Gate的完整模型在112K长度下达到75.78%，以约6%的准确率代价换取了3.5倍的推理加速。其次，在推理效率分析（Figure 3）中，LycheeMemory展现出接近常数的推理延迟曲线：当上下文从8K增长到128K时，推理时间仅有微幅上升，而Qwen2.5-7B基线呈二次爆炸式增长，MemAgent呈线性增长。在128K长度下，LycheeMemory相比MemAgent实现了6倍加速，峰值GPU显存占用降低2倍。第三，在零样本泛化实验（Table 2）中，LycheeMemory在未见过的2WikiMultihopQA数据集上达到了75.0%（14K）、70.3%（28K）、73.4%（56K）的sub-EM，在StreamingQA上达到80.8%的F1和73.4%的sub-EM，展示了良好的迁移能力。第四，消融实验（Table 4）系统验证了各训练阶段的贡献：朴素分块（Stage-1）仅36.72%平均sub-EM，记忆压缩（Stage-2）提升至38.28%，SFT（Stage-3）大幅提升至59.38%，RL优化达到66.80%，端到端RL最终达到68.75%，证明联合优化的关键价值。第五，压缩比实验（Figure 4）表明alpha=4是最优选择，与alpha=2相比几乎没有精度损失但内存减半，而alpha=16在128K长度下性能从71.5%骤降至42.0%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| RULER-HQA（多跳问答，112K上下文） | Normalized Sub-EM (%) | LycheeMemory-7B: 75.78; w/o Gate: 80.47 | MemAgent-7B: 79.69; Qwen2.5-Instruct-7B-1M: 51.56 | w/o Gate比MemAgent高0.78；带Gate以6%精度换取3.5倍加速 |
| RULER-HQA（多跳问答，1.75M上下文） | Normalized Sub-EM (%) | LycheeMemory-7B: 71.09; w/o Gate: 78.12 | MemAgent-7B: 75.78; 所有全上下文模型: OOM | 唯一能在1.75M长度下运行的方案（全上下文模型均OOM） |
| 2WikiMultihopQA（零样本，56K上下文） | Sub-EM (%) | 73.4 | MemAgent: 71.1; RAG: 59.4; Qwen2.5-Instruct-7B: 37.5 | 比MemAgent高2.3%，比RAG高14.0% |
| StreamingQA（零样本） | F1 / Sub-EM (%) | 80.8 / 73.4 | RAG: 84.3 / 67.2; MemAgent: 77.9 / 60.2 | Sub-EM比MemAgent高13.2%，比RAG高6.2% |
| 推理效率（128K上下文，128样本） | 推理时间 (秒) | 约2000秒 | MemAgent: 约12000秒 | 6倍加速，峰值GPU显存降低2倍 |
局限与改进
本文存在以下局限性。首先，压缩过程本身是一次性的预处理开销，虽然可以并行化，但对于实时流式输入场景（如在线对话），需要在推理前等待所有块完成压缩，引入额外延迟。其次，压缩比的选择是一个全局超参数（实验中固定为alpha=4），无法自适应地根据每个文本块的信息密度调整压缩程度——对于包含关键细节的段落和冗余重复的段落使用相同的压缩比，可能导致关键信息丢失。第三，Gate模块作为二分类器单独训练而非与压缩器和推理器联合优化，这意味着Gate的决策标准与推理器的最优策略可能存在偏差。第四，实验仅在Qwen2.5系列模型上进行验证，基座模型规模局限于3B和7B，在更大规模模型（如70B、405B）上的效果尚不清楚。第五，评测任务主要集中在多跳问答和信息检索类任务，在需要复杂推理链（如数学证明、代码调试）的任务上未做验证。此外，作者在失败模式分析（Appendix C）中承认了三类主要错误：单向依赖不匹配（35%）、过早推理锚定（21%）、压缩诱导幻觉（17%），这些问题反映了压缩记忆机制的固有缺陷。
独立分析的弱点
本文存在几个值得深入讨论的弱点。第一，固定压缩比alpha=4对所有文本块一视同仁，这在实践中可能导致信息密度不均：对于事实密集的关键段落（如包含具体数字、日期、名称的段落），4倍压缩可能丢失关键细节；对于叙事性、冗余较多的段落，4倍压缩可能过于保守。改进方向是设计自适应压缩比机制，让压缩器根据每个块的信息熵或与查询的相关性动态调整压缩率。第二，Gate单独训练为二分类器而非参与端到端优化，导致其决策标准与推理器的最优策略可能存在偏差。作者承认这是因为「离散召回决策导致的梯度不连续」，但可以通过Gumbel-Softmax、REINFORCE等方法实现梯度近似，将Gate纳入联合优化。第三，工作记忆m使用固定长度的token窗口，限制了推理过程中能积累的信息量。当需要跨大量记忆块进行复杂推理时，固定容量的工作记忆可能成为瓶颈。第四，压缩器预训练阶段使用的三个辅助任务（文本重建、QA生成、创意生成）需要大量标注数据和计算资源，且预训练质量直接影响后续RL阶段的效果。
未来方向
基于本文的框架，有几个有前景的研究方向。第一，自适应压缩机制：可以设计一个信息密度评估模块，根据每个文本块的内容复杂度、与查询的相关性预测、以及在文档中的结构位置（如标题、摘要、结论等关键段落）动态调整压缩比，实现「按需压缩」。第二，Gate的端到端集成：利用连续松弛技术（如Gumbel-Softmax）或策略梯度方法将Gate纳入联合RL优化，消除当前分阶段训练带来的优化不一致。第三，多轮对话和交互式推理：当前框架假设单次查询-回答模式，可以扩展为支持多轮交互的工作记忆管理，让系统在对话过程中持续积累和更新记忆。第四，多模态扩展：将压缩记忆机制从纯文本扩展到多模态文档（包含图像、表格、公式），这对科学论文理解和复杂文档分析具有重要意义。第五，更大规模模型验证：在70B甚至405B规模的基座模型上验证LycheeMemory的效果，探索压缩记忆在更大模型上的缩放规律。第六，将工作记忆从固定token窗口扩展为可增长的结构化记忆，支持更复杂的推理链路。
复现评估
从复现角度来看，本文提供了相对充分的实现细节。附录A详细描述了三个训练阶段的超参数设置：压缩器预训练使用学习率1e-4、rank 16的LoRA；RL阶段使用GSPO算法、group size G=4、clip系数epsilon=0.2、KL惩罚系数beta=0.01；Gate训练使用二元交叉熵损失、rank 8的LoRA。数据集构建方面，作者遵循MemAgent的方法从RULER-HQA合成训练数据，平均约20K token的长度。然而，复现面临几个挑战：第一，端到端RL训练需要大量GPU资源，实验在2xA100 (80GB)上进行，推理效率测试使用128个样本；第二，训练数据的合成过程涉及多个步骤（混合相关文章与干扰项、生成QA对等），需要专门的数据处理流水线；第三，代码和模型权重是否开源未在论文中明确说明。总体而言，具备足够算力的团队可以基于论文中的超参数描述进行复现，但从零实现完整的训练流水线需要较大的工程投入。
论文图表
