自监督训练提升大模型长时记忆能力,无需标注数据。
MemTrain: Self-Supervised Context Memory Training

- 用未标注维基文本设计双重自监督任务,强化记忆保持与过程还原。
- 在长文本问答和搜索问答任务上,性能最高提升17.67分。
- 适合需要长程推理的智能体开发,尤其资源受限场景。
长时序大模型智能体依赖记忆能力来保存并利用跨长时间交互积累的信息。现有方法通常采用强化学习在下游任务上端到端训练,但收集高质量的记忆密集型标注数据成本高昂,且训练数据多样性不足,难以覆盖通用记忆行为。本文提出MemTrain,一种自监督训练框架,用于普遍增强大模型智能体的上下文记忆能力,以提升下游后训练效果。MemTrain在未标注的维基百科语料上引入两个耦合的代理任务:(1) 端到端掩码重构目标,要求模型在多次记忆更新后恢复被掩码的实体,从最终结果角度鼓励记忆维持;(2) 中间记忆召回目标,要求模型使用中间记忆状态重构被掩码的历史信息,鼓励整个交互过程中记忆的忠实压缩与完整性。两个目标通过GRPO联合优化。在长文本问答和基于搜索的问答基准上的大量实验表明,MemTrain在不同模型上均显著提升下游记忆密集型推理性能,相较直接任务特定后训练最高提升17.67点。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。