用压缩记忆动态推理,让大模型高效处理超长文本。
Dynamic Long Context Reasoning over Compressed Memory via End-to-End Reinforcement Learning
- 将长文本分块压缩成记忆,只召回相关部分进行推理。
- 在7K到175万词元间外推上下文长度,准确率媲美强基线。
- 适合需要长文本理解的场景,如法律、医疗问答。
大语言模型在长上下文处理中面临计算开销大、信息遗忘和检索增强生成中上下文碎片化等问题。本文提出一种受认知启发的高效长上下文推理框架,通过分块压缩和选择性记忆召回,避免处理全部原始词元。该框架将长输入分块,使用可学习压缩器将每块编码为压缩记忆表示;门控模块动态选择相关记忆块,由具备动态工作内存的推理模块迭代处理以解决下游任务。压缩器与推理器通过端到端强化学习联合优化,门控模块单独作为分类器训练。实验表明,该方法在多跳推理基准RULER-HQA上表现优异,可将上下文长度外推至1.75M词元,相比强基线模型实现最高2倍的峰值显存降低和6倍的推理加速,显著提升准确性与效率的平衡。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。