Ortogonalizar la memoria mejora el recuerdo asociativo en modelos recurrentes

Fuentes: Matrix Orthogonalization Improves Memory in Recurrent Models

Las redes neuronales recurrentes (RNN) tienen dificultades para competir con los transformers en tareas de recuerdo asociativo, una capacidad que la atención facilita al dar a cada token acceso directo a los anteriores. Sin embargo, en dominios como el aprendizaje por refuerzo de horizonte largo —al estilo de Dreamer— el coste cuadrático de la atención resulta prohibitivo, por lo que conviene potenciar el rendimiento de las RNN sin renunciar a esa capacidad.

La arquitectura recurrente más conocida para recuerdo asociativo es mLSTM, una variante de LSTM con memoria matricial. Funciona bien en MQAR, pero los autores proponen evaluar los modelos con recuerdo asociativo ruidoso (NAR), un escenario más exigente en el que aparecen tokens distractores entre claves y valores. Tomando como base los benchmarks ruidosos de MAD, el equipo de Paradigm plantea una intervención inspirada en el optimizador Muon —que ortogonaliza los momentos para igualar las direcciones representadas—: aplicar ortogonalización a la matriz de memoria de mLSTM durante las lecturas, pero sin reescribir la memoria ortogonalizada.

Los experimentos usan AdamW (con dos mil pasos y tamaño de lote 64) sobre tareas NAR con frac_noise=0,8, variando tamaños de vocabulario y longitudes de secuencia. La ortogonalización se aplica con cinco iteraciones de Newton-Schulz y normalización por la norma de Frobenius.

Los resultados muestran mejoras generalizadas en tasa de éxito y precisión media, con una brecha que se amplía en regímenes difíciles (vocabulario 96, secuencias de 768 y 1.024). En esos casos, mLSTM ortogonalizado pasa de apenas 4 de 24 semillas resueltas a 14-16. El código es reproducible.

Los autores piden cautela: los resultados se circunscriben a modelos pequeños y a tareas sintéticas, y queda por verificar si las ganancias se trasladan a benchmarks reales con modelos mayores.