MemTrain:自监督上下文记忆训练
阅读原文· arxiv.orgMemTrain 是一个专为增强大语言模型智能体上下文记忆能力而设计的自监督训练框架。它基于未标注的 Wikipedia 语料,引入两个耦合代理任务:端到端掩码重建(要求模型在多轮记忆更新后恢复被掩码实体)与中间记忆召回(利用中间记忆状态重建被掩码历史信息),并通过 GRPO 联合优化。在长文本 QA 和搜索型 QA 基准上,MemTrain 一致提升不同模型的记忆密集型推理性能,最高达 17.67 个百分点的增益。
Memory is an indispensable capability for long-horizon LLM agents, enabling them to preserve and utilize information accumulated across extended interactions. Existing memory-agent approaches are typically trained end-to-end with reinforcement learning on downstream tasks. However, collecting high-quality annotated problems for memory-intensive scenarios is costly, and the resulting training data often lack sufficient diversity to cover general memory behaviors. In this work, we propose MemTrain, a self-supervised training framework for generally enhancing the context-memory capability of LLM agents for more effective downstream post-training. MemTrain introduces two coupled proxy tasks over unlabeled Wikipedia corpora: (1) an end-to-end masked reconstruction objective, which requires the model to recover masked entities after multiple rounds of memory updates, thereby encouraging memory maintenance from the final outcome perspective; and (2) an intermediate memory recall objective, which requires the model to reconstruct masked historical information using intermediate memory states, encouraging faithful compression and memory completeness throughout the interaction process. The two objectives are jointly optimized using GRPO. Extensive experiments on long-text QA and search-based QA benchmarks demonstrate that MemTrain consistently improves downstream memory-intensive reasoning performance across different models, achieving gains of up to 17.67 points over direct task-specific post-training.