Memoir: Should a Model Write to Its Memory While It Thinks?

Resumen Técnico: Memoir – ¿Debería un modelo escribir en su memoria mientras piensa?

Planteamiento del Problema
Los modelos de lenguaje actuales suelen mantener los parámetros entrenados fijos durante una pasada hacia adelante (forward pass), dependiendo de la recuperación externa o de contextos largos para expandir su memoria efectiva sin alterar su estado de pesos internos. Si bien el entrenamiento en tiempo de prueba (test-time training) y la recurrencia adaptativa permiten a los modelos actualizar los estados ocultos o dedicar más computación a entradas difíciles, persiste un riesgo técnico específico no abordado: la estabilidad de un proceso latente recurrente que lee y escribe simultáneamente en su propia memoria rápida.

La cuestión central investigada es si se debería permitir a un modelo reescribir su nivel de memoria rápida durante cada iteración de ponderación (recurrente). Los autores hipotetizan un modo de fallo específico: si la memoria se reescribe mientras el latente está siendo refinado, un cabezal de energía podría parecer mejorar simplemente porque la memoria cambió para coincidir con el latente actual, en lugar de porque la representación latente misma se haya vuelto más coherente. Esta escritura "acoplada" crea un operador no estacionario durante el bucle de computación, lo que potencialmente corrompe la señal de aprendizaje.

Metodología
El artículo presenta Memoir, un marco que integra cuatro niveles de memoria, recurrencia latente de profundidad variable y un objetivo de energía de latente futuro.

1. Interfaz de Memoria de Cuatro Niveles:

   • $W_0$ (Rápida): Matriz por muestra, escribible en cada pasada hacia adelante.
   • $W_1$ (Episódica): Matriz por muestra, actualizada mediante una compuerta aprendida por segmento.
   • $W_2$ (Lenta): Parámetros entrenables compartidos, actualizados raramente (consolidación).
   • $W_3$ (Congelada): Parámetros compartidos, nunca actualizados.
     La interfaz define explícitamente la propiedad y los límites de escritura, distinguiendo entre estados plásticos por muestra y estados compartidos lentos/congelados.

2. Ponderación de Profundidad Variable:
   El modelo emplea un proceso latente recurrente donde un bloque compartido lee los latentes actuales y los niveles disponibles para proponer un nuevo estado ($h^{(\ell+1)}$). Un cabezal de parada determina cuándo dejar de iterar.

   • Brazo Acoplado (Brazo A): Durante las iteraciones activas de ponderación, el modelo lee el nivel rápido ($W_0$) e inmediatamente escribe en él utilizando una actualización basada en la regla delta según el latente actual.
   • Brazo de Solo Lectura (Brazo B): El modelo lee el nivel rápido durante la ponderación pero no escribe en él hasta alcanzar el límite de la pasada hacia adelante. El operador permanece estacionario durante el bucle.

3. Objetivos:

   • Predicción de Latente Futuro: El modelo predice un latente objetivo $z_{t+\Delta}$ a partir de la representación actual, entrenado contra un profesor de Media Móvil Exponencial (EMA) con detención de gradiente (stop-gradient).
   • Objetivo de Energía: Un cabezal de energía califica pares de representaciones. El margen se define como $m = E[E(\text{corrupto})] - E[E(\text{coherente})]$. Un margen positivo indica que los pares coherentes reciben menor energía.

4. Diseño Experimental:
   El estudio utiliza una tarea de recuperación asociativa procedimental con interferencia de clave (64 valores, 16 distractores, interferencia 0.85).

   • Control: Ambos brazos tienen recuentos de parámetros idénticos (81,738 en total, 76,362 entrenables), recuentos de multiplicaciones-acumulaciones (multiply-accumulate) de paso hacia adelante iguales, datos, optimizador y semillas coincidentes.
   • Presupuestos: Una prueba primaria a 240 pasos (12 semillas) y un control más largo a 960 pasos (8 semillas).
   • Implementación: El sistema incluye una ablación arquitectónica emparejada donde la única diferencia es la presencia de escrituras dentro del bucle.

Resultados Clave

1. Penalización de Velocidad de Aprendizaje: Al presupuesto fijo de 240 pasos, el brazo de solo lectura (Brazo B) superó significativamente al brazo acoplado (Brazo A).

   • Recuperación de Solo Lectura: 0.6557 (IC 95%: [0.5953, 0.7160])
   • Recuperación Acoplada: 0.5203 (IC 95%: [0.4522, 0.5883])
   • Diferencia: 0.1354. Una prueba t pareada en 12 semillas arrojó $t=3.23$ (df=11), con un intervalo de confianza del 95% para la diferencia de [0.0431, 0.2277], excluyendo el cero. El brazo de solo lectura ganó en 10 de 12 semillas.

2. Comportamiento de Convergencia: A 960 pasos, ambos brazos alcanzaron el techo de la tarea de 1.0000, y la brecha de rendimiento desapareció. Esto indica que la arquitectura acoplada no tiene una capacidad alcanzable inferior en esta tarea específica, sino que sufre una penalización de velocidad de aprendizaje bajo el presupuesto medido.

3. Estabilidad del Margen de Energía: El modo de fallo predicho —donde la reescritura de memoria corrompe la señal de energía al hacer que la memoria coincida con el latente actual— no ocurrió. El margen de energía creció de casi cero a 0.0231 durante el entrenamiento, demostrando que las escrituras acopladas no desestabilizaron el objetivo de energía en este entorno.

4. Mejoras de Ingeniería: La reestructuración del kernel redujo el tiempo de ejecución de la regla delta de 0.907 ms a 0.351 ms (una aceleración de 2.59×) en una GPU NVIDIA GeForce RTX 5090 Laptop, con un ancho de banda efectivo que aumentó de 18.6 GB/s a 48.0 GB/s.

Significancia y Reivindicaciones
El artículo realiza cinco contribuciones específicas y verificables:

1. Especificación de Interfaz: Una interfaz de memoria de cuatro niveles con propiedad, tiempo de vida y límites de escritura explícitos.
2. Definición de Proceso: Un proceso de ponderación de profundidad variable vertical donde la política de escritura del nivel rápido es intercambiable sin alterar el resto del modelo.
3. Ablación Rigurosa: Una ablación emparejada de 12 semillas a 240 pasos y un control de 8 semillas a 960 pasos, aislando el efecto de las escrituras dentro del bucle.
4. Validación de Resultado Negativo: Confirmación de que el colapso predicho del margen de energía no ocurrió, refutando la hipótesis específica de que la ponderación escribible corrompe inherentemente la señal de energía en esta configuración.
5. Transparencia de Implementación: Código completo disponible (5,920 líneas en memoir/, 2,061 líneas de pruebas) con mediciones de tiempo y memoria a nivel de dispositivo.

Limitaciones y Reivindicaciones Modestas
Los autores declaran explícitamente que los resultados no prueban que la arquitectura acoplada tenga una capacidad última inferior, ya que ambos brazos alcanzaron el techo en la tarea saturada. La brecha observada es una penalización de velocidad de aprendizaje, no una penalización de capacidad. Además, el experimento no separa el tiempo de las escrituras de su volumen; el brazo acoplado realiza aproximadamente siete veces más escrituras por episodio que el brazo de solo lectura. El estudio no prueba lenguaje natural, aprendizaje continuo de largo horizonte, resistencia al olvido catastrófico o reglas de actualización aprendidas. El margen de energía se reporta como un diagnóstico específico de la tarea, no como un verificador general de alucinaciones.

En conclusión, Memoir proporciona una interfaz concreta para probar la memoria escribible durante la inferencia. El hallazgo principal es que, bajo las restricciones específicas de esta tarea de recuperación procedimental y el presupuesto de entrenamiento, permitir que un modelo reescriba su memoria rápida mientras pondera ralentiza el aprendizaje en comparación con un enfoque de solo lectura, aunque no impide la convergencia eventual.