Dynamic Long Context Reasoning over Compressed Memory via End-to-End Reinforcement Learning

📄 arXiv: 2602.08382v1 📥 PDF

作者: Zhuoen Chen, Dongfang Li, Meishan Zhang, Baotian Hu, Min Zhang

分类: cs.CL, cs.AI

发布日期: 2026-02-09

备注: 26 pages, 7 figures. Code and models will be released

💡 一句话要点

提出基于强化学习的动态长文本推理框架，解决长文本处理中的效率和信息遗忘问题。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本推理 强化学习 记忆压缩 选择性召回 大型语言模型

📋 核心要点

  1. 现有长文本处理方法面临计算成本高、信息遗忘和上下文碎片化等挑战，限制了大型语言模型在长文本任务中的应用。
  2. 论文提出一种基于分块压缩和选择性记忆召回的框架，通过学习压缩器和动态门控机制，实现高效的长文本推理。
  3. 实验结果表明，该方法在长文本推理任务上具有竞争力，并显著降低了GPU内存使用量和推理时间。

📝 摘要（中文）

大型语言模型(LLMs)在长文本处理中面临着显著的挑战，包括二次方级别的计算成本、信息遗忘以及检索增强生成(RAG)中固有的上下文碎片化问题。我们提出了一种受认知启发的框架，用于高效的长文本推理，该框架基于分块压缩和选择性记忆召回，而不是处理所有原始token。该框架将长输入分割成块，并使用学习到的压缩器将每个块编码成压缩的记忆表示。一个门控模块动态地选择相关的记忆块，然后由一个推理模块迭代地处理这些记忆块，该推理模块具有不断演化的工作记忆，以解决下游任务。压缩器和推理器通过端到端强化学习联合优化，而门控模块则作为分类器单独训练。实验结果表明，该方法在多跳推理基准测试（如RULER-HQA）上实现了具有竞争力的准确率，将上下文长度从7K外推到1.75M token，并且与强大的长文本基线相比，提供了良好的准确率-效率权衡。特别是，与MemAgent相比，它实现了高达2倍的峰值GPU内存使用量减少和6倍的推理速度提升。

🔬 方法详解

问题定义：现有的大型语言模型在处理长文本时，面临着计算复杂度高（通常是二次方级别）、信息容易遗忘，以及在使用检索增强生成（RAG）时出现的上下文碎片化问题。这些问题限制了LLM在需要处理大量上下文信息的任务中的应用，例如多跳推理、文档总结等。

核心思路：论文的核心思路是模仿人类的认知过程，将长文本分割成块，并对每个块进行压缩，形成记忆表示。然后，通过一个门控机制动态地选择与当前任务相关的记忆块，并使用一个推理模块迭代地处理这些记忆块，从而实现高效的长文本推理。这种方法避免了处理所有原始token，降低了计算成本，并减少了信息遗忘的风险。

技术框架：该框架主要包含三个模块：压缩器、门控模块和推理器。首先，压缩器将长输入分割成块，并将每个块编码成压缩的记忆表示。然后，门控模块根据当前任务和工作记忆，动态地选择相关的记忆块。最后，推理器迭代地处理这些记忆块，并更新工作记忆，最终完成下游任务。整个框架采用端到端的方式进行训练，其中压缩器和推理器通过强化学习联合优化，而门控模块则作为分类器单独训练。

关键创新：该论文的关键创新在于将分块压缩、选择性记忆召回和强化学习相结合，构建了一个高效的长文本推理框架。与传统的处理整个长文本的方法相比，该方法显著降低了计算成本和内存占用。与静态的记忆召回方法相比，该方法通过动态的门控机制，能够更好地选择与当前任务相关的记忆块。

关键设计：压缩器可以使用各种编码器，例如Transformer或LSTM。门控模块可以使用分类器，例如MLP或Transformer。推理器可以使用循环神经网络或Transformer。压缩器和推理器的联合优化通过强化学习实现，奖励函数可以根据下游任务的性能进行设计。具体来说，论文中使用了Actor-Critic算法进行训练，Actor负责选择动作（例如选择哪些记忆块），Critic负责评估动作的价值。

📊 实验亮点

实验结果表明，该方法在RULER-HQA等多跳推理基准测试上取得了具有竞争力的准确率，并且能够将上下文长度从7K外推到1.75M token。与MemAgent相比，该方法实现了高达2倍的峰值GPU内存使用量减少和6倍的推理速度提升。这些结果表明，该方法在长文本推理任务上具有显著的优势。

🎯 应用场景

该研究成果可应用于需要处理长文本的各种场景，例如：长文档问答、多跳推理、医学报告分析、法律文件处理等。通过降低计算成本和内存占用，该方法使得大型语言模型能够在资源受限的环境中处理更长的文本，从而扩展了LLM的应用范围。未来，该方法还可以与其他技术相结合，例如知识图谱和检索增强生成，以进一步提高长文本处理的性能。

📄 摘要（原文）

Large Language Models (LLMs) face significant challenges in long-context processing, including quadratic computational costs, information forgetting, and the context fragmentation inherent in retrieval-augmented generation (RAG). We propose a cognitively inspired framework for efficient long-context inference based on chunk-wise compression and selective memory recall, rather than processing all raw tokens. The framework segments long inputs into chunks and encodes each chunk into compressed memory representations using a learned compressor. A gating module dynamically selects relevant memory blocks, which are then iteratively processed by a reasoning module with an evolving working memory to solve downstream tasks. The compressor and reasoner are jointly optimized via end-to-end reinforcement learning, while the gating module is trained separately as a classifier. Experimental results show that the proposed method achieves competitive accuracy on multi-hop reasoning benchmarks such as RULER-HQA, extrapolates context length from 7K to 1.75M tokens, and offers a favorable accuracy-efficiency trade-off compared to strong long-context baselines. In particular, it achieves up to a 2 times reduction in peak GPU memory usage and a 6 times inference speedup over MemAgent.