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

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