MemTrain : Entraînement auto-supervisé de la mémoire de contexte

Résumé

La mémoire est une capacité indispensable pour les agents LLM à long horizon, leur permettant de conserver et d'exploiter les informations accumulées lors d’interactions prolongées. Les approches existantes d’agents à mémoire sont généralement entraînées de bout en bout par apprentissage par renforcement sur des tâches aval. Cependant, la collecte de problèmes annotés de haute qualité pour des scénarios à forte sollicitation mémoire est coûteuse, et les données d’entraînement qui en résultent manquent souvent de diversité suffisante pour couvrir les comportements mémoire généraux. Dans ce travail, nous proposons MemTrain, un cadre d’entraînement auto-supervisé pour améliorer de manière générale la capacité de mémoire contextuelle des agents LLM, en vue d’un post-entraînement aval plus efficace. MemTrain introduit deux tâches proxies couplées sur des corpus Wikipédia non étiquetés : (1) un objectif de reconstruction masquée de bout en bout, qui exige du modèle qu’il reconstitue des entités masquées après plusieurs tours de mise à jour mémoire, encourageant ainsi le maintien de la mémoire du point de vue du résultat final ; et (2) un objectif de rappel mémoire intermédiaire, qui exige du modèle qu’il reconstitue des informations historiques masquées à l’aide des états mémoire intermédiaires, encourageant une compression fidèle et une complétude mémoire tout au long du processus d’interaction. Les deux objectifs sont optimisés conjointement à l’aide de GRPO. Des expériences approfondies sur des références de QA sur textes longs et de QA par recherche démontrent que MemTrain améliore systématiquement les performances de raisonnement mémoire intensif aval sur différents modèles, avec des gains pouvant atteindre 17,67 points par rapport à un post-entraînement direct spécifique à la tâche.

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.