MemTrain: Zelf-gesuperviseerde Contextgeheugentraining

Samenvatting

Geheugen is een onmisbare capaciteit voor LLM-agenten met een lange horizon, omdat het hen in staat stelt informatie te bewaren en te gebruiken die is opgebouwd gedurende uitgebreide interacties. Bestaande geheugen-agentbenaderingen worden doorgaans end-to-end getraind met behulp van reinforcement learning op stroomafwaartse taken. Het verzamelen van hoogwaardige geannoteerde problemen voor geheugenintensieve scenario's is echter kostbaar, en de resulterende trainingsdata missen vaak voldoende diversiteit om algemeen geheugengedrag te dekken. In dit werk stellen we MemTrain voor, een zelfgecontroleerd trainingsraamwerk voor het algemeen verbeteren van de contextgeheugencapaciteit van LLM-agenten, ten behoeve van effectievere stroomafwaartse post-training. MemTrain introduceert twee gekoppelde proxtaken over ongelabelde Wikipedia-corpora: (1) een end-to-end gemaskeerde reconstructiedoelstelling, die het model vereist om gemaskeerde entiteiten te herstellen na meerdere rondes van geheugenupdates, waardoor geheugenbehoud vanuit het perspectief van het uiteindelijke resultaat wordt gestimuleerd; en (2) een tussentijdse geheugenherinneringsdoelstelling, die het model vereist om gemaskeerde historische informatie te reconstrueren met behulp van tussentijdse geheugentoestanden, wat getrouwe compressie en volledigheid van het geheugen gedurende het interactieproces aanmoedigt. De twee doelstellingen worden gezamenlijk geoptimaliseerd met behulp van GRPO. Uitgebreide experimenten op long-text QA- en op zoeken gebaseerde QA-benchmarks tonen aan dat MemTrain consequent de stroomafwaartse geheugenintensieve redeneerprestaties verbetert bij verschillende modellen, met winsten tot 17,67 punten ten opzichte van directe taakspecifieke post-training.

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.