De los transformadores a la memoria asociativa, cómo los titanes y MIRAS repensan el modelado de contexto largo

¿Qué viene después de Transformers? Google Research propone una nueva forma de proporcionar modelos de secuencia con memoria a largo plazo utilizable con Titans y MIRAS, manteniendo al mismo tiempo el entrenamiento paralelo y la inferencia casi lineal.

Titans es una arquitectura concreta que agrega una memoria neuronal profunda a una columna vertebral estilo Transformer. MIRAS es un marco general que ve la mayoría de los modelos de secuencia modernos como instancias de optimización en línea sobre una memoria asociativa.

¿Por qué Titanes y MIRAS?

Los transformadores estándar utilizan la atención sobre un caché de valores clave. Esto proporciona un sólido aprendizaje en contexto, pero el costo crece cuadráticamente con la longitud del contexto, por lo que el contexto práctico es limitado incluso con FlashAttention y otros trucos del kernel.

Las redes neuronales recurrentes lineales eficientes y los modelos de espacio de estados como Mamba-2 comprimen la historia en un estado de tamaño fijo, por lo que el costo es lineal en la longitud de la secuencia. Sin embargo, esta compresión pierde información en secuencias muy largas, lo que perjudica tareas como el modelado genómico y la recuperación de contextos extremadamente largos.

Titans y MIRAS combinan estas ideas. La atención actúa como una memoria precisa a corto plazo en la ventana actual. Un módulo neuronal separado proporciona memoria a largo plazo, aprende en el momento de la prueba y está entrenado para que su dinámica sea paralelizable en los aceleradores.

https://research.google/blog/titans-miras-helping-ai-have-long-term-memory/

Titanes, una memoria neuronal a largo plazo que aprende en el momento de la prueba

El artículo de investigación de Titans presenta un módulo neuronal de memoria a largo plazo que es en sí mismo un perceptrón multicapa profundo en lugar de un estado vectorial o matricial. La atención se interpreta como memoria a corto plazo, ya que sólo ve una ventana limitada, mientras que la memoria neuronal actúa como memoria persistente a largo plazo.

Para cada token, Titans define una pérdida de memoria asociativa.

ℓ(Mₜ₋₁; kₜ, vₜ) = ‖Mₜ₋₁(kₜ) − vₜ‖²

donde Mₜ₋₁ es la memoria actual, kₜ es la clave y vₜ es el valor. El gradiente de esta pérdida con respecto a los parámetros de la memoria es la “métrica sorpresa”. Los gradientes grandes corresponden a tokens sorprendentes que deben almacenarse, los gradientes pequeños corresponden a tokens esperados que pueden ignorarse en su mayoría.

Los parámetros de la memoria se actualizan en el momento de la prueba mediante el descenso del gradiente con el impulso y la caída del peso, que en conjunto actúan como una puerta de retención y un mecanismo de olvido. Para mantener eficiente esta optimización en línea, el artículo de investigación muestra cómo calcular estas actualizaciones con multiplicaciones de matrices por lotes sobre fragmentos de secuencia, lo que preserva el entrenamiento paralelo en secuencias largas.

Arquitectónicamente, Titans utiliza tres ramas de memoria en la red troncal, a menudo instanciadas en la variante MAC de Titans:

una rama central que realiza el aprendizaje estándar en contexto con atención una rama de la memoria contextual que aprende de la secuencia reciente una rama de la memoria persistente con pesos fijos que codifica el conocimiento previo al entrenamiento

La memoria a largo plazo comprime los tokens pasados ​​en un resumen, que luego se pasa a la atención como contexto adicional. La atención puede elegir cuándo leer ese resumen.

Resultados experimentales para titanes

En referencias de modelado de lenguaje y razonamiento de sentido común como C4, WikiText y HellaSwag, las arquitecturas de Titans superan las líneas de base lineales recurrentes de última generación Mamba-2 y los modelos Gated DeltaNet y Transformer++ de tamaño comparable. La investigación de Google atribuye esto al mayor poder expresivo de la memoria profunda y su capacidad para mantener el rendimiento a medida que crece la longitud del contexto. Los recuerdos neuronales profundos con el mismo presupuesto de parámetros pero mayor profundidad producen consistentemente una menor perplejidad.

Para recordar contextos extremadamente largos, el equipo de investigación utiliza el punto de referencia BABILong, donde los hechos se distribuyen en documentos muy extensos. Titans supera todas las líneas de base, incluidos modelos muy grandes como GPT-4, mientras utiliza muchos menos parámetros y escala a ventanas de contexto más allá de 2.000.000 de tokens.

El equipo de investigación informa que Titans mantiene un entrenamiento paralelo eficiente y una inferencia lineal rápida. La memoria neuronal por sí sola es ligeramente más lenta que los modelos lineales recurrentes más rápidos, pero las capas híbridas Titans con Sliding Window Attention siguen siendo competitivas en rendimiento y al mismo tiempo mejoran la precisión.

https://arxiv.org/pdf/2504.13173

MIRAS, un marco unificado para modelos de secuencia como memoria asociativa

El artículo de investigación de MIRAS, “Todo está conectado: un viaje a través de la memorización del tiempo de prueba, el sesgo de atención, la retención y la optimización en línea”, generaliza esta visión. Observa que los modelos de secuencia modernos pueden verse como recuerdos asociativos que asignan claves a valores mientras equilibran el aprendizaje y el olvido.

MIRAS define cualquier modelo de secuencia a través de cuatro opciones de diseño:

Estructura de la memoria, por ejemplo, un vector, un mapa lineal o un MLP. Sesgo de atención, la pérdida interna que define las similitudes que le interesan a la memoria. Puerta de retención, el regularizador que mantiene la memoria cerca de su estado anterior. Algoritmo de memoria, la regla de optimización en línea, a menudo descenso de gradiente con impulso.

Mediante esta lente, MIRAS recupera varias familias:

Modelos recurrentes lineales de estilo hebbiano y RetNet como memorias asociativas basadas en productos escalares. Modelos de reglas delta como DeltaNet y Gated DeltaNet como memorias basadas en MSE con reemplazo de valor y puertas de retención específicas. Titans LMM como memoria no lineal basada en MSE con retención local y global optimizada mediante descenso de gradiente con impulso.

Fundamentalmente, MIRAS va más allá de los objetivos habituales de MSE o de productos punto. El equipo de investigación construye nuevos sesgos de atención basados ​​en normas Lₚ, pérdida robusta de Huber y optimización robusta, y nuevas puertas de retención basadas en divergencias sobre probabilidades simples, regularización neta elástica y divergencia de Bregman.

Desde este espacio de diseño, el equipo de investigación crea instancias de tres modelos sin atención:

Moneta usa una memoria MLP de 2 capas con sesgo atencional Lₚ y una puerta de retención híbrida basada en normas generalizadas. Yaad usa la misma memoria MLP con sesgo atencional de pérdida de Huber y una puerta de olvido relacionada con Titans Memora usa pérdida de regresión como sesgo atencional y una puerta de retención basada en divergencia KL sobre una memoria de estilo de probabilidad simplex.

Estas variantes de MIRAS reemplazan los bloques de atención en una columna vertebral estilo Llama, utilizan convoluciones separables en profundidad en la capa Miras y se pueden combinar con Sliding Window Attention en modelos híbridos. El entrenamiento permanece en paralelo fragmentando secuencias y calculando gradientes con respecto al estado de la memoria del fragmento anterior.

En experimentos de investigación, Moneta, Yaad y Memora igualan o superan los fuertes modelos recurrentes lineales y Transformer++ en modelado de lenguaje, razonamiento de sentido común y recuperación de tareas intensivas, manteniendo al mismo tiempo la inferencia de tiempo lineal.

Conclusiones clave

Titans introduce una memoria neuronal profunda a largo plazo que aprende en el momento de la prueba, utilizando un descenso de gradiente en una pérdida de memoria asociativa L2, de modo que el modelo almacena selectivamente solo tokens sorprendentes mientras mantiene las actualizaciones paralelizables en los aceleradores. Titans combina la atención con la memoria neuronal para un contexto prolongado, utilizando ramas como el núcleo, la memoria contextual y la memoria persistente para que la atención maneje la precisión de corto alcance y el módulo neuronal mantenga información en secuencias de más de 2.000.000 de tokens. Titans supera las sólidas líneas de base lineales de RNN y Transformer++, incluidos Mamba-2 y Gated DeltaNet, en modelado de lenguaje y puntos de referencia de razonamiento de sentido común en escalas de parámetros comparables, sin dejar de ser competitivo en rendimiento. En puntos de referencia de recuperación de contexto extremadamente largos, como BABILong, Titans logra una mayor precisión que todas las líneas de base, incluidos modelos de atención más grandes como GPT 4, al tiempo que utiliza menos parámetros y aún permite un entrenamiento e inferencia eficientes. MIRAS proporciona un marco unificador para modelos de secuencia como memorias asociativas, definiéndolos por estructura de memoria, sesgo de atención, puerta de retención y regla de optimización, y produce nuevas arquitecturas sin atención como Moneta, Yaad y Memora que igualan o superan a los RNN lineales y Transformer++ en tareas de razonamiento y contexto prolongado.

Consulta los detalles técnicos. No dude en consultar nuestra página de GitHub para tutoriales, códigos y cuadernos. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 100.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.

Michal Sutter es un profesional de la ciencia de datos con una Maestría en Ciencias de Datos de la Universidad de Padua. Con una base sólida en análisis estadístico, aprendizaje automático e ingeniería de datos, Michal se destaca en transformar conjuntos de datos complejos en conocimientos prácticos.

🙌 Siga MARKTECHPOST: agréguenos como fuente preferida en Google.