Resumen Técnico: Composición de Expertos de Memoria para Modelos Mundiales Difusivos

Enunciado del Problema

Los modelos mundiales (WM) son esenciales para el aprendizaje por refuerzo (RL) y la planificación, ya que predicen estados futuros plausibles consistentes con las observaciones pasadas. Sin embargo, las arquitecturas existentes enfrentan un compromiso fundamental entre la fidelidad de la memoria y la eficiencia computacional. Los modelos basados en transformadores preservan el detalle local pero sufren costos de atención cuadráticos, limitando la longitud del contexto. Por el contrario, las redes recurrentes y los modelos de espacio de estados (SSM) escalan eficientemente pero comprimen la historia en estados ocultos, descartando inevitablemente detalles finos con el tiempo. Simplemente aumentar la longitud del contexto en una sola arquitectura conduce a un entrenamiento inestable y costos de inferencia prohibitivos. Además, la composición ingenua de múltiples modelos a menudo resulta en "colapso de modos", donde la distribución combinada afila en exceso las regiones comunes y pierde diversidad.

Metodología

Los autores proponen Composición de Expertos de Memoria (CoME), un marco basado en difusión que desacopla la consistencia futuro-pasado de una arquitectura monolítica única. En su lugar, CoME aprovecha un conjunto de expertos especializados, cada uno manejando un rol de memoria específico, y los fusiona probabilísticamente en el momento de la inferencia.

1. Producto de Expertos Contrastivos (PoCE)

Para abordar las limitaciones del Producto Estándar de Expertos (PoE), que puede suprimir la diversidad y amplificar modos espurios, los autores introducen una formulación de Producto de Expertos Contrastivos.

• Mecanismo: Para un conjunto de expertos heterogéneos $p_i(x)$, la composición se define como $p(x) \propto \prod \tilde{p}_i(x)$, donde $\tilde{p}_i(x) \propto p_i(x)^{\alpha_i} \bar{p}_i(x)^{1-\alpha_i}$.
• Línea Base Contrastiva: $\bar{p}_i(x)$ representa la distribución de línea base incondicional del experto. El coeficiente $\alpha_i > 1$ actúa como un factor de contraste.
• Justificación Teórica: A diferencia del temple ingenuo (que reduce la varianza del kernel y distorsiona la geometría local), esta mezcla contrastiva repondera los componentes de la mezcla mientras preserva la forma de los kernels. Suprime efectivamente los modos espurios (donde un experto asigna alta probabilidad sin evidencia sólida) mientras mantiene modos secundarios válidos, siempre que el experto condicional asigne una probabilidad más alta a los modos válidos que la línea base.

2. Expertos de Memoria Especializados

El marco instancia tres expertos complementarios:

• Memoria a Corto Plazo (STM): Un modelo de difusión condicionado a una ventana de contexto reciente ($10\text{--}100$ cuadros). Captura dinámicas locales finas y dependencias temporales inmediatas utilizando una arquitectura con longitud de contexto extendida (por ejemplo, atención con ventana deslizante).
• Memoria a Largo Plazo (LTM): Un modelo de difusión separado que almacena la historia episódica ($100\text{--}1000$ cuadros) directamente en sus pesos. Esto se logra mediante ajuste fino en tiempo de prueba de adaptadores ligeros LoRA en el contexto a largo plazo. Este enfoque amortiza el costo de la recuperación, permitiendo la reutilización en tiempo constante de la experiencia pasada sin alimentar toda la historia a través del modelo en cada paso. Evita el olvido catastrófico congelando la base y actualizando solo los adaptadores.
• Memoria Espacial a Largo Plazo (SLTM): Una extensión de LTM que incorpora priores espaciales (por ejemplo, poses de cámara, mapas de puntos) extraídos de la historia. Este experto impone coherencia geométrica y espacial, desambiguando ubicaciones en entornos repetitivos o visualmente similares.

3. Generación Unificada

La distribución final es un producto del prior preentrenado y los tres expertos contrastivos (STM, LTM, SLTM). Durante el muestreo, las funciones de puntuación de estos expertos se combinan aditivamente (en espacio logarítmico) para guiar el proceso de eliminación de ruido, asegurando que el futuro generado respete las dinámicas recientes, la historia a largo plazo y la consistencia espacial.

Contribuciones Clave

1. Arquitectura de Memoria Desacoplada: El artículo propone resolver el compromiso de la memoria distribuyéndola entre expertos especializados en lugar de una sola base, formulado como un problema de Producto de Expertos.
2. Producto de Expertos Contrastivos (PoCE): Una estrategia novedosa para fusionar expertos heterogéneos. Los autores proporcionan evidencia teórica y empírica de que este método previene el colapso de modos y la sobreconfianza al factorizar modos redundantes, a diferencia del PoE estándar.
3. Memoria Difusiva Externa: La introducción de un módulo LTM que utiliza ajuste fino en tiempo de prueba de adaptadores LoRA para almacenar conocimiento episódico en los pesos. Esto se extiende a SLTM para incorporar estructura espacial, mejorando la precisión en tareas de navegación.

Resultados Experimentales

El método fue evaluado en benchmarks sintéticos (Memory Maze, DMLab, Minecraft) y del mundo real (RealEstate10K, RECON).

• Calidad de Reconstrucción: En el conjunto de datos Memory Maze, CoME logró un rendimiento de vanguardia en LPIPS (0.097), SSIM (0.892) y PSNR (23.07), superando significativamente a las líneas base, incluidos transformadores de atención completa, SSM y modelos con ventana deslizante.
• Escalabilidad: A diferencia de los modelos de atención completa que requieren $60\times$ más cómputo que la línea base STM, CoME logra una consistencia superior con una fracción del costo de entrenamiento e inferencia.
• Longitud de Contexto: El modelo no mostró saturación en la consistencia del rendimiento hasta 480 cuadros de contexto, demostrando la eficacia del componente LTM.
• Planificación y Navegación: En tareas de navegación condicionadas a objetivos (benchmark RECON), CoME redujo el Error de Trayectoria Absoluta (ATE) a 0.96, superando a los modelos mundiales de navegación especializados (NWM) y otras líneas base.
• Capacidad de Recuperación: En un experimento de trayectoria inversa en RealEstate10K, CoME recuperó con éxito los cuadros iniciales después de rodajes hacia adelante, mientras que las líneas base sin memoria generaron escenas plausibles pero inconsistentes.
• Estudios de Ablación: Eliminar la formulación contrastiva causó el colapso del rendimiento, confirmando la necesidad del mecanismo PoCE. Aumentar la rango de LoRA y la longitud del contexto para LTM mejoró consistentemente los resultados sin sobreajuste.

Significado y Afirmaciones

El artículo afirma establecer un paradigma novedoso para construir modelos mundiales difusivos aumentados con memoria. Al distribuir la memoria entre expertos especializados y utilizar una estrategia de composición contrastiva, CoME equilibra la eficiencia con una consistencia temporal de alta fidelidad. Los autores argumentan que este enfoque permite que los modelos mundiales escalen a contextos largos sin incurrir en costos cuadráticos o sufrir la pérdida de detalles inherente a la compresión recurrente. El método se presenta como un marco general aplicable a diversos modelos mundiales basados en difusión, mejorando su fiabilidad para la planificación y la toma de decisiones en entornos de aprendizaje por refuerzo. El trabajo enfatiza que la memoria debe ser una propiedad distribuida del sistema en lugar de una restricción de una sola arquitectura.