MemTrain: Treinamento Auto-Supervisionado de Memória de Contexto

Resumo

A memória é uma capacidade indispensável para agentes LLM de horizonte longo, permitindo-lhes preservar e utilizar informações acumuladas ao longo de interações estendidas. As abordagens existentes de agentes com memória são tipicamente treinadas de ponta a ponta com aprendizado por reforço em tarefas downstream. No entanto, coletar problemas anotados de alta qualidade para cenários intensivos em memória é caro, e os dados de treinamento resultantes frequentemente carecem de diversidade suficiente para cobrir comportamentos gerais de memória. Neste trabalho, propomos o MemTrain, uma estrutura de treinamento auto-supervisionado para melhorar genericamente a capacidade de memória de contexto de agentes LLM, visando um pós-treinamento downstream mais eficaz. O MemTrain introduz duas tarefas proxy acopladas sobre corpora da Wikipedia não rotulados: (1) um objetivo de reconstrução mascarada de ponta a ponta, que exige que o modelo recupere entidades mascaradas após múltiplas rodadas de atualizações de memória, incentivando assim a manutenção da memória sob a perspectiva do resultado final; e (2) um objetivo de recordação de memória intermediária, que exige que o modelo reconstrua informações históricas mascaradas usando estados intermediários de memória, incentivando a compressão fiel e a completude da memória ao longo do processo de interação. Os dois objetivos são otimizados conjuntamente usando GRPO. Experimentos extensivos em benchmarks de QA de texto longo e QA baseado em busca demonstram que o MemTrain melhora consistentemente o desempenho downstream de raciocínio intensivo em memória em diferentes modelos, alcançando ganhos de até 17,67 pontos em relação ao pós-treinamento direto específico para a tarefa.

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.