Resumen Técnico: Memoria Delta Dispersa (SDM)

Planteamiento del Problema

Las Redes Neuronales Recurrentes (RNN) lineales, como Gated DeltaNet (GDN) y Mamba2, ofrecen la ventaja de una memoria y un cómputo fijos por token, lo que las hace teóricamente adecuadas para contextos arbitrariamente largos. Sin embargo, su rendimiento práctico en la recuperación de contexto largo y el aprendizaje en contexto está limitado por sus pequeños tamaños de estado oculto. Aumentar el tamaño del estado en las RNN lineales estándar conduce a un incremento lineal de los FLOPs debido a las actualizaciones densas de estado (multiplicaciones de matriz-vector), creando un cuello de botella computacional prohibitivo. Por el contrario, las arquitecturas Transformer estándar con atención softmax mantienen una alta capacidad de recuperación, pero sufren costos crecientes de memoria y cómputo del caché de Clave-Valor (KV) a medida que aumenta la longitud de la secuencia.

El desafío central abordado es cómo escalar la capacidad del estado oculto de las RNN lineales en órdenes de magnitud para mejorar el rendimiento en contextos largos sin aumentar el costo computacional (FLOPs) por token.

Metodología: Memoria Delta Dispersa (SDM)

Los autores introducen la Memoria Delta Dispersa (SDM), una arquitectura que extiende a Gated DeltaNet reemplazando el producto exterior denso de clave-valor por un esquema de direccionamiento disperso basado en Product-Key Memory (PKM).

Mecanismo Central

1. Memoria Dispersa Explícita: En lugar de mantener una matriz densa $M_t \in \mathbb{R}^{d_{qk} \times d_v}$, SDM mantiene una tabla de memoria explícita $M_t \in \mathbb{R}^{N \times d_v}$ con $N$ ranuras, donde $N$ puede ser órdenes de magnitud mayor que la dimensión oculta (por ejemplo, $N \approx 10^6$).
2. Selección de Claves Dispersa: En cada paso de tiempo, las claves proyectadas ($k'$) y las consultas ($q'$) se dividen en dos mitades. Su suma exterior genera $N$ puntuaciones para las ranuras de memoria. Utilizando una estrategia de selección top-$W$ (escritura) y top-$R$ (lectura), el modelo identifica ranuras de memoria específicas para actualizar o leer. Esta selección se realiza de manera eficiente sin materializar la matriz de puntuación completa de tamaño $N$, aprovechando la propiedad de que $\text{topk}(s_1 \oplus s_2) = \text{topk}(\text{topk}(s_1) \oplus \text{topk}(s_2))$.
3. Actualizaciones Delta con Compuerta (Gated): Para las $W$ ranuras de escritura seleccionadas, el modelo aplica la regla de actualización de Gated DeltaNet:
   • Compuerta de Olvido (Forget Gate): $\tilde{M}_t[i] = \alpha_t \cdot M_{t-1}[i]$
   • Actualización Delta: $M_t[i] = \tilde{M}_t[i] + \beta_t \cdot k^{(i)}_t \cdot (v_t - \tilde{M}_t[i])$
   • Las ranuras no seleccionadas permanecen sin cambios.
4. Lectura Dispersa: La salida es una suma ponderada de las $R$ ranuras de lectura seleccionadas.

Diseño IsoFLOP

Para asegurar una comparación justa, los autores diseñan SDM para que coincida con el recuento de parámetros y FLOPs de una línea base de GDN densa:

• FLOPs: El costo por token es $O((W+R) \times d_v)$, independiente del tamaño de la memoria $N$. Al establecer $W$ y $R$ iguales a la dimensión de consulta de la línea base de GDN, el costo computacional permanece idéntico.
• Parámetros: Las proyecciones lineales tienen tamaños idénticos a la línea base. El tamaño de la memoria $N$ se controla mediante el número de cabezales $H$ y las dimensiones de proyección, permitiendo que el tamaño del estado escale como $O(1/H^2)$ sin afectar los FLOPs.

Estado Inicial Aprendido

A diferencia de GDN, que típicamente inicializa su pequeño estado a cero, SDM trata el estado de memoria inicial $M_0$ como un parámetro aprendible. Esto permite al modelo almacenar conocimiento de preentrenamiento directamente en la estructura de la memoria, actuando como una memoria paramétrica que persiste a través de la inferencia.

Contribuciones Clave

1. Arquitectura: La propuesta de SDM, que escala la capacidad del estado de las RNN lineales por tres órdenes de magnitud (por ejemplo, de kilobytes a gigabytes) manteniendo FLOPs constantes por token.
2. Validación IsoFLOP: Demostración de que aumentar la capacidad del estado por sí solo (sin aumentar el cómputo) mejora significamente el rendimiento en tareas de contexto largo.
3. Memoria Paramétrica: La introducción de un estado inicial $M_0$ aprendido para las RNN lineales, mostrando que este mecanismo almacena efectivamente el conocimiento común y mejora las tareas de razonamiento.
4. Leyes de Escalamiento: Validación empírica de que SDM sigue leyes de escalamiento predecibles, superando a la densa GDN y, a una escala mayor (8B parámetros), superando a los modelos de Atención Completa (Full Attention).

Resultados Experimentales

Los autores entrenaron modelos que van desde 1.4B hasta 8B de parámetros sobre más de 1 billón de tokens.

• Leyes de Escalamiento: En una escalera de escalamiento, SDM logró consistentemente una pérdida de entrenamiento menor que GDN en todos los niveles de cómputo. A la escala de 8B, SDM alcanzó una pérdida menor que tanto GDN como un modelo de Atención Completa.
• Recuperación de Contexto Largo (Benchmark RULER): SDM superó significativamente a GDN en tareas de recuperación de contexto largo. A escalas de 1.4B y 8B, SDM logró puntuaciones de 31.2 y 50.2 respectivamente, comparado con el 20.0 y 34.2 de GDN. Notablemente, SDM igualó o superó el rendimiento de la Atención Completa en varias subtareas de RULER a pesar de utilizar una representación de memoria fija.
• Contexto Corto y Razonamiento: SDM logró una menor perplejidad en datos de código y una mayor precisión promedio en benchmarks de razonamiento (por ejemplo, MMLU, GSM8K, HumanEval) en comparación con GDN. A la escala de 8B, SDM superó a la Atención Completa en precisión promedio.
• Estudios de Ablación:
  • Tamaño del Estado: El rendimiento degradó de forma monótona a medida que se reducía el tamaño de la memoria, confirmando que el gran estado es el principal motor de las mejoras.
  • Inicialización Aprendida: Aprender $M_0$ mejoró el rendimiento para SDM pero proporcionó ganancias insignificantes para GDN, sugiriendo que la gran capacidad es necesaria para utilizar eficazmente la inicialización aprendida.
  • Eficiencia: Aunque el rendimiento de entrenamiento de SDM fue 1.49x más lento que el de GDN a la escala de 8B (debido al acceso a la memoria HBM frente a SRAM), fue 6x más rápido que la Atención Completa durante la inferencia.

Significado y Reivindicaciones

El artículo afirma que SDM logra cerrar la brecha entre la eficiencia de las RNN lineales y las capacidades de contexto largo de los Transformers. Al desacoplar la capacidad de memoria del costo computacional mediante la dispersión, SDM permite:

• Espacio/Cómputo Constante: Procesar secuencias arbitrariamente largas (se demostró hasta 1 millón de tokens) sin el crecimiento ilimitado de los cachés KV.
• Rendimiento Superior en Contexto Largo: Lograr resultados de vanguardia entre las RNN lineales y resultados competitivos contra la Atención Completa en recuperación de contexto largo y aprendizaje en contexto.
• Almacenamiento de Conocimiento Paramétrico: La capacidad de aprender y almacenar conocimiento de preentrenamiento dentro del estado de la memoria misma, mejorando las tareas de razonamiento general y de conocimiento común.

Los autores reconocen limitaciones, específicamente que la gran huella de memoria (comparable a los parámetros del modelo) requiere recursos significativos de HBM, y que las implementaciones actuales de kernels son menos eficientes que los kernels optimizados de GDN denso. Sin embargo, postulan que SDM abre nuevas posibilidades para desarrollar agentes con mejor memoria a largo plazo y para procesar videos largos o trazas de razonamiento extendidas.