Mem-α: Construcción de Memoria mediante Aprendizaje por Refuerzo

Mem-α: Learning Memory Construction via Reinforcement Learning

September 30, 2025

Autores: Yu Wang, Ryuichi Takanobu, Zhiqi Liang, Yuzhen Mao, Yuanzhe Hu, Julian McAuley, Xiaojian Wu

cs.AI

Resumen

Los agentes de modelos de lenguaje de gran escala (LLM, por sus siglas en inglés) están limitados por ventanas de contexto reducidas, lo que hace necesario el uso de sistemas de memoria externa para la comprensión de información a largo plazo. Los agentes actuales potenciados con memoria suelen depender de instrucciones y herramientas predefinidas para actualizar la memoria. Sin embargo, los modelos de lenguaje pueden carecer de la capacidad para determinar qué información almacenar, cómo estructurarla y cuándo actualizarla, especialmente a medida que los sistemas de memoria se vuelven más complejos. Esto resulta en una construcción subóptima de la memoria y pérdida de información. Para abordar este problema, proponemos Mem-alpha, un marco de aprendizaje por refuerzo que entrena a los agentes para gestionar eficazmente sistemas de memoria complejos mediante interacción y retroalimentación. Además, construimos un conjunto de datos de entrenamiento especializado que abarca diversos patrones de interacción multiturno, acompañados de preguntas de evaluación exhaustivas diseñadas para enseñar una gestión eficaz de la memoria. Durante el entrenamiento, los agentes procesan fragmentos secuenciales de información, aprenden a extraer y almacenar contenido relevante, y luego actualizan el sistema de memoria. La señal de recompensa se deriva de la precisión en la respuesta a preguntas basadas en el historial completo de interacción, optimizando directamente la construcción de la memoria. Para ilustrar la efectividad de nuestro marco de entrenamiento, diseñamos una arquitectura de memoria que incluye componentes centrales, episódicos y semánticos, equipados con múltiples herramientas para operaciones de memoria. La evaluación empírica demuestra que Mem-alpha logra mejoras significativas en comparación con los agentes potenciados con memoria existentes. A pesar de haber sido entrenados exclusivamente en instancias con una longitud máxima de 30k tokens, nuestros agentes exhiben una notable generalización a secuencias que superan los 400k tokens, más de 13 veces la longitud de entrenamiento, lo que resalta la robustez de Mem-alpha.

English

Large language model (LLM) agents are constrained by limited context windows, necessitating external memory systems for long-term information understanding. Current memory-augmented agents typically depend on pre-defined instructions and tools for memory updates. However, language models may lack the ability to determine which information to store, how to structure it, and when to update it, especially as memory systems become more complex. This results in suboptimal memory construction and information loss. To this end, we propose Mem-alpha, a reinforcement learning framework that trains agents to effectively manage complex memory systems through interaction and feedback. We also construct a specialized training dataset spanning diverse multi-turn interaction patterns paired with comprehensive evaluation questions designed to teach effective memory management. During training, agents process sequential information chunks, learn to extract and store relevant content, then update the memory system. The reward signal derives from downstream question-answering accuracy over the full interaction history, directly optimizing for memory construction. To illustrate the effectiveness of our training framework, we design a memory architecture comprising core, episodic, and semantic components, equipped with multiple tools for memory operations. Empirical evaluation demonstrates that Mem-alpha achieves significant improvements over existing memory-augmented agent baselines. Despite being trained exclusively on instances with a maximum length of 30k tokens, our agents exhibit remarkable generalization to sequences exceeding 400k tokens, over 13x the training length, highlighting the robustness of Mem-alpha.