1. Introducción
Durante dos años, fuimos testigos de una carrera por la longitud de las secuencias en los modelos de lenguaje de IA. Evolucionamos gradualmente desde una longitud de contexto de 4k a 32k, luego a 128k, hasta la ventana masiva de 1 millón de tokens prometida por primera vez por modelos como Gemini 1.5 pro. La promesa era atractiva: volcar bases de código o novelas completas en el modelo y dejar que razone todo.
Pero hay un costo oculto en esta longitud de contexto prácticamente “infinita”, que rara vez se menciona: la memoria.
En una arquitectura Transformer estándar, memorizar y razonar en todo el mensaje no es gratis. A medida que crece la secuencia de entrada, el modelo debe almacenar los estados de clave y valor (KV) de cada token para calcular las puntuaciones de atención. Para una secuencia de 1 millón de tokens, este KV Cache puede crecer rápidamente hasta alcanzar cientos de gigabytes, lo que a su vez requiere grandes grupos de GPU en múltiples centros de datos, todo para mantener la conversación en la memoria.
2. La motivación
En un mecanismo de atención estándar (Vaswani et al., 2017)6, cada nuevo token que genera el modelo debe "mirar hacia atrás" a cada token anterior en el mensaje para comprender completamente el contexto. Para que esto sea eficiente a lo largo de varias generaciones, el modelo almacena en caché los vectores clave (K) y valor (V) de tokens anteriores en la VRAM de la GPU. Esto se conoce como caché KV.
La trampa del crecimiento lineal
Si bien el almacenamiento en caché de los vectores de clave y valor (caché KV) puede ahorrar tiempo (ya que no tenemos que volver a calcular el pasado para cada token nuevo), ocupa una enorme cantidad de memoria, que crece linealmente con la longitud de la secuencia de entrada.
Para poner esto en perspectiva: almacenar la caché KV para un modelo de parámetros estándar de 500B para un contexto de solo 20.000 tokens requiere alrededor de 126 GB de memoria. Si escalamos eso al recuento de parámetros de los parámetros 1T+ del LLM moderno y atendemos a millones de usuarios en un momento dado, la huella de memoria total se convierte en una cifra astronómicamente grande.
Históricamente, hemos tenido dos formas de manejar datos secuenciales, ninguna de las cuales es perfecta:
RNN: las redes neuronales recurrentes procesan el mensaje de entrada token por token, actualizando un estado oculto único y fijo. Si bien esto puede reducir en gran medida los requisitos de memoria, tienen dificultades para retener información y detalles en solicitudes prolongadas. Esto hace que los modelos finalmente olviden el comienzo de la secuencia de entrada cuando llegan al final. Transformers: Los Transformers, a diferencia de los RNN, no sufren este problema ya que recuerdan todo perfectamente al mantener todo el historial de la conversación en KV Cache. Tienen una memoria perfecta, pero debido al gran caché KV, consumen mucha memoria.
Ésta es la compensación que Infini-attention pretende cubrir.
3. La solución: atención infinita
Para resolver la paradoja de la memoria, investigadores de Google formularon la atención infinita (Munkhdalai et al., 2024)1. El principio básico de este enfoque es que en lugar de almacenar la conversación completa, podemos almacenar un resumen de la misma.
Infini-attention divide la producción de atención en dos mecanismos distintos, que funcionan simultáneamente:
Atención local: Igual que un transformador estándar. Ve el contexto inmediato y calcula una matriz de atención para cada token para capturar detalles en alta resolución. Atención lineal global: una memoria comprimida que almacena un resumen de todo el historial pasado en una matriz de tamaño fijo, para que el modelo haga referencia a ella.
Repasemos el proceso de cómo esto procesa una entrada larga.
Visualización de cómo funciona la atención infinita (Recuperación)
Paso 1: segmentación
En primer lugar, toda la secuencia de entrada se divide en segmentos más pequeños (digamos, N = 2048 tokens). Dentro de cada segmento, el modelo utiliza la Atención de producto punto estándar para comprender el contexto. Esto garantiza que, para tareas inmediatas, la resolución siga siendo perfecta.
Paso 2: La compresión (actualización de memoria)
Para pasar al siguiente segmento, el modelo almacena los estados comprimidos de la clave (K) y el valor (V) del segmento actual en una matriz de memoria (M) de tamaño fijo. Esto permite que el modelo consulte la Matriz de memoria (en lugar del caché KV más grande) para obtener información sobre los segmentos anteriores.
Sin embargo, agregar nuevos datos a ciegas a Memory Matrix puede corromper rápidamente la información anterior que contenía. Para evitarlo, los autores utilizan la Regla Delta (Schlag et al., 2021)7. La intuición detrás de esto es: antes de agregar cualquier información nueva, verifique si la memoria ya la almacena o no. Esto evita actualizaciones redundantes. Todo el proceso de actualización se explica a continuación:
A. El “Peek” (Calculando V recuperado)
En primer lugar, el modelo recupera valores de la memoria existente utilizando las claves actuales (K) como si fueran consultas. El modelo hace esto para evaluar qué tipo de información (valores) la memoria ya asocia con las claves actuales.
K: Claves generadas para el segmento actual
Moho: el estado actual de la memoria global
σ: Función de activación no lineal (ELU+1)
z: factor de normalización
VRecuperado: Matriz de valores de la memoria global.
B. El paso de actualización
Luego, el modelo compara los nuevos valores reales (V) con los valores recuperados (Vretrived). Calcula la diferencia (el residual) y solo la agrega a la memoria. Esto evita actualizar la memoria con lo que ya sabe.
Mnew: Memoria global actualizada
KT: Matriz de clave transpuesta del segmento actual
V: Matriz de valores del segmento actual
Vretrived: Vector de matriz recuperado de la memoria global
Esto implica que si la memoria ya contiene perfectamente la información del segmento actual, la actualización es cero. Esto mantiene la memoria estable y "limpia" durante numerosas actualizaciones.
Paso 3: Recuperación global (atención lineal)
Para generar el siguiente token, el modelo necesita la información contextual de todo el mensaje, es decir, de todos los segmentos. Para obtener la información relevante, el modelo consulta la Matriz de Memoria realizando una multiplicación de matrices.
Amem: Salida de atención de la memoria global
P: Matriz de consulta del segmento actual
M: matriz de memoria global
z: factor de normalización
La matriz Amem resultante contiene la información relevante de todos los segmentos anteriores para generar el siguiente token.
Paso 4: La Agregación (El “Mezclador”)
Finalmente, el modelo tiene dos salidas:
Adot: El contexto local detallado del segmento actual. Amem: el historial global comprimido de todos los segmentos anteriores de la matriz de memoria.
Para combinar los dos, utiliza un escalar de activación aprendido, β (beta):
Sigmoide: activación no lineal para unir β entre 0 y 1
Amem y Adot: salidas de atención de la memoria global y del producto escalar, respectivamente
β: parámetro de activación aprendido para controlar la influencia de Amem y Adot en el resultado final
El parámetro β actúa como un coeficiente de mezcla que determina el equilibrio entre los flujos de información a largo plazo (Amem) y a corto plazo (Adot):
Cuando β es bajo: la función sigmoidea se acerca a 0. Esto hace que el factor de ponderación complementario (1−sigmoide(β)) se vuelva dominante, lo que hace que el modelo priorice la atención local del producto escalar (Adot) más que la memoria de compresión global. Cuando β es alto: la función sigmoidea se acerca a 1. El modelo prioriza el contenido de la memoria recuperada (Amem), permitiendo que el contexto global anule la información local del segmento actual.
4. Los resultados: por qué es importante la atención infinita
Los autores pusieron a prueba Infini-attention con modelos de contexto largo existentes, como Transformer-XL (Dai et al., 2019)2 y Memorizing Transformers (Wu et al., 2022)3. Los siguientes son los resultados:
1. La compresión de memoria “114x”
El logro más impactante de este artículo es la reducción masiva de los recursos de memoria utilizados. Como Infini-Attention almacena todo el contexto histórico en una matriz de memoria de tamaño fijo en lugar de una caché KV que crece linealmente, puede almacenar 114 veces menos parámetros en la VRAM de la GPU en comparación con Memorizing Transformers. Como se muestra en la siguiente tabla, para una longitud de contexto de 65 000 tokens, Infini-Attention logra puntuaciones de perplejidad SOTA en puntos de referencia como PG19 y Arxiv-math mientras necesita almacenar solo 1,6 millones de parámetros (tamaño de la matriz de memoria), a diferencia de las arquitecturas de la competencia.
La atención infinita reduce notablemente el uso de memoria y al mismo tiempo logra la perplejidad SOTA en los puntos de referencia PG19 y Arxiv-math
2. La prueba de la “clave de acceso” de 1 millón de tokens
Para una arquitectura de contexto largo, el desafío de la aguja en un pajar es convencional. Los autores probaron esto ocultando una clave aleatoria en un corpus masivo de texto y pidiendo al modelo que la recuperara. Como se muestra en la siguiente tabla, en una configuración de disparo cero, el modelo lucha por encontrar la clave y logra en su mayoría una precisión <20 %.
Luego, los autores ajustaron el modelo para 400 pasos con secuencias que tenían una longitud de solo 5000 tokens. Sorprendentemente, el modelo pudo generalizar el ajuste para trabajar con secuencias de hasta 1 millón de tokens de longitud, con una precisión de recuperación drásticamente mejorada en todos los ámbitos.
Las tres puntuaciones por entrada denotan la precisión de la recuperación en relación con la posición de la clave oculta en el corpus (inicio/medio/final).
3. Resumen de libros de última generación (contexto de 500k)
Además de las pruebas sintéticas, los autores también probaron el modelo en el punto de referencia BookSum (Kryściński et al.)5, donde se requiere que el modelo genere un resumen de una novela larga. El modelo Infini-Attention con parámetros 8B estableció un nuevo rendimiento de vanguardia en el punto de referencia, al generar resúmenes exitosos de libros de hasta 500.000 tokens de longitud.
Los resultados también muestran una clara tendencia a que las capacidades de resumen del modelo mejoren a medida que se le introducen contextos más largos. El gráfico que se muestra a continuación valida esta hipótesis, que en lugar de olvidar información previa (un modo de falla común conocido como "perdido en el medio"), el modelo puede usar efectivamente Memory Matrix para generar resúmenes precisos.
Colorete vs longitud de entrada. Rouge mide qué tan cerca está un resumen generado por IA de un resumen escrito por humanos basándose en la similitud léxica.
4. Visualización del escalar de puerta
Como estudio de ablación adicional, los autores visualizaron el escalar de activación aprendido (β) para ver cómo el modelo estaba usando su nueva memoria. A continuación se muestra el mapa de calor de la visualización resultante. Las cabezas de atención se dividen en dos funciones distintas:
Jefes especializados: jefes que tienen una puntuación cercana a 1 o 0, lo que indica que eligen centrarse en el contexto local (dentro del segmento) o en la historia global (segmentos anteriores). Cabezales mezcladores: Cabezas que tienen puntuaciones cercanas a 0,5, lo que indica que su función principal es fusionar información de ambas vías de manera eficiente.
Esto sugiere que el modelo puede aprender a cambiar entre recuerdos a corto y largo plazo y mezclar información a lo largo de toda la secuencia.
La visualización de β revela que las cabezas de atención tienden a especializarse para la atención global o local bajo la arquitectura de atención infinita.
5. Conclusión
Si bien es posible que no reemplace completamente las bases de datos vectoriales externas y los sistemas RAG para razonar sobre el conocimiento estático, sí cambia la forma en que los modelos procesan las consultas de los usuarios estándar. La integración de dichas arquitecturas podría ser el siguiente paso para dar rienda suelta a la creatividad en la investigación, que anteriormente tuvo que verse obstaculizada por los avances del hardware, acelerando en última instancia el progreso en el campo del modelado de lenguajes.
👉Si te gustó este artículo, comparto reseñas más breves y actualizadas en Substack.
👉Y si desea apoyar la redacción de investigaciones independientes, BuyMeACoffee le ayuda a seguir adelante.
6. Referencias
Atención infinita (artículo principal): Munkhdalai, T., Faruqui, M. y Gopal, S. (2024). No deje ningún contexto atrás: transformadores de contexto infinitos eficientes con atención infinita. arXiv preimpresión arXiv:2404.07143. Transformer-XL: Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, QV y Salakhutdinov, R. (2019). Transformer-XL: modelos de lenguaje atento más allá de un contexto de longitud fija. Preimpresión de arXiv arXiv:1901.02860. Memorización de transformadores: Wu, Y., Rabe, MN, Hutchins, D. y Szegedy, C. (2022). Memorizando transformadores. Preimpresión de arXiv arXiv:2203.08913. Atención lineal (La base de las matemáticas): Katharopoulos, A., Vyas, A., Pappas, N. y Fleuret, F. (2020). Los transformadores son RNN: transformadores autorregresivos rápidos con atención lineal. Congreso Internacional sobre Aprendizaje Automático. Punto de referencia BookSum: Kryściński, W., Rajani, N., Agarwal, D., Xiong, C. y Radev, D. (2021). BookSum: una colección de conjuntos de datos para resúmenes narrativos extensos. Preimpresión de arXiv arXiv:2105.08209. Atención estándar: Vaswani, Ashish, et al. "Atención es todo lo que necesitas". Avances en sistemas de procesamiento de información neuronal 30 (2017). Regla delta: Schlag, Imanol, Kazuki Irie y Jürgen Schmidhuber. "Los transformadores lineales son programadores de peso secretamente rápidos". Congreso internacional sobre aprendizaje automático. PMLR, 2021.