通过端到端强化学习在压缩记忆上实现动态长上下文推理

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