🤖 machine learning

Looped Latent Attention: Cross-Loop KV Compression for Looped Transformers

Este artículo presenta la Atención Latente en Bucle (LLA, por sus siglas en inglés), un códec de post-entrenamiento que aprovecha la estructura de bajo rango de los cachés KV indexados por bucle en Transformers con pesos vinculados para lograr una compresión extrema (hasta 32x) e incrementar significativamente la capacidad de lote manteniendo un rendimiento casi sin pérdidas mediante la inicialización por SVD y la destilación.

Autores originales: James O' Neill, Fergal Reid

Publicado 2026-07-20

📖 6 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: James O' Neill, Fergal Reid

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

El problema de la memoria en los cerebros de IA

Imagina que estás intentando enseñarle a un robot a escribir una historia. El robot tiene un cerebro hecho de millones de diminutos interruptores (parámetros) que saben cómo encajan las palabras. Pero aquí está el truco: para escribir una historia larga, el robot no puede limitarse a mirar su cerebro; necesita un bloc de notas. Cada vez que escribe una palabra nueva, tiene que recordar todo lo que acaba de escribir para no repetirse o perder el hilo de la trama. En el mundo de la Inteligencia Artificial, este bloc de notas se llama KV cache (caché de clave-valor). Es un cuaderno privado para cada una de las conversaciones, y crece más con cada palabra que el robot pronuncia.

Ahora, imagina un tipo especial de robot diseñado para ser súper eficiente. En lugar de tener un cerebro enorme con muchas habitaciones diferentes (capas) para distintas tareas, este robot tiene una única habitación pequeña y astuta a la que visita una y otra vez. Reutiliza el mismo conjunto de instrucciones, recorriéndolas en bucle para pensar más profundamente. Esto se llama un Transformer con Bucle (Looped Transformer). Ahorra espacio en el propio cerebro, lo cual es genial. Pero hay un problema escurridizo: aunque el robot reutiliza la misma habitación, sigue escribiendo una página nueva en su bloc de notas cada vez que hace el bucle. Si el robot hace cuatro bucles para resolver un problema matemático, termina con cuatro páginas separadas de notas para el mismo momento de la historia. Esto significa que el robot "con bucle" todavía necesita una cantidad masiva de memoria para funcionar, invalidando el propósito de ser eficiente. La gran pregunta para los científicos es: ¿Podemos reducir ese bloc de notas sin hacer que el robot olvide cómo pensar?

El descubrimiento del artículo: Un atajo secreto en las notas

Este artículo presenta un truco ingenioso llamado Atención Latente con Bucle (Looped Latent Attention o LLA) para resolver ese problema de memoria. Los investigadores descubrieron que las notas del robot no son tan desordenadas como parecen. Cuando el robot recorre su proceso de pensamiento en bucle, las notas que escribe para la misma palabra no cambian drásticamente; siguen un camino suave y predecible, como una canica rodando por una colina suave. En lugar de guardar cada una de las páginas del bloc de notas del robot, el nuevo método guarda solo un pequeño "resumen" de ese camino y un conjunto de instrucciones sobre cómo reconstruir las páginas específicas solo cuando el robot realmente necesita leerlas.

Piénsalo de esta manera: Imagina que estás viendo una película donde el personaje principal camina por un pasillo. En una computadora normal, guardarías una foto de alta definición completa del personaje en cada paso que da. Eso ocupa una enorme cantidad de almacenamiento. Pero con LLA, te das cuenta de que el personaje simplemente está caminando en línea recta. Así que, en lugar de guardar miles de fotos, guardas una foto del personaje y una nota sencilla que dice: "Mover 1 pulgada hacia adelante para el paso 2, 2 pulgadas para el paso 3". Cuando la computadora necesita ver al personaje en el paso 3, dibuja rápidamente esa imagen usando la nota. Esto es exactamente lo que hace el artículo por la IA: comprime la parte del "bucle" de la memoria, convirtiendo una pila de cuatro páginas en un único y diminuto resumen que puede expandirse de nuevo en las páginas completas instantáneamente.

Los investigadores probaron esto en varios modelos de IA, incluyendo uno llamado Ouro-1.4B. Descubrieron que este método funciona increíblemente bien. De hecho, en un chip de computadora potente llamado H200, pudieron meter 768 conversaciones diferentes en la memoria a la vez, mientras que antes solo podían meter 32. Ese es un salto masivo en cuántas personas pueden usar la IA al mismo tiempo. También demostraron que este truco funciona incluso cuando la IA resuelve problemas matemáticos muy largos y difíciles, aunque necesita un poco de entrenamiento adicional para asegurar que no se confunda cuando esté escribiendo su propia historia en lugar de solo copiar a un profesor.

A lo que este artículo dice "No"

Es importante saber qué es lo que este método no es. Los investigadores probaron primero una idea muy simple: "¿Qué pasa si simplemente tiramos las tres primeras páginas de notas y solo nos quedamos con la última?". Pensaron que tal vez el robot se estabiliza hacia una respuesta final después de unos pocos bucles, por lo que las notas anteriores no importarían. Intentaron esto, y fue un desastre. La IA olvidó por completo cómo hacer matemáticas, obteniendo una puntuación de cero en las pruebas. Esto demuestra que el "camino" que siguen las notas es importante; no puedes simplemente tomar el punto final. Necesitas capturar el viaje.

También compararon su método con otras formas de reducir la memoria, como comprimir las notas (cuantización) o compartir notas entre diferentes partes del cerebro (compresión de eje de cabezales). Aunque esos métodos ayudan un poco, el artículo muestra que comprimir el eje del "bucle" es el ingrediente secreto. Si intentas comprimir las otras partes en lugar de los bucles, la IA pierde su capacidad de razonar. El artículo es muy claro: el bucle es la parte más redundante de la memoria, y ese es el único lugar donde se debe presionar.

¿Qué tan seguros están?

Los autores están bastante seguros de estos resultados porque los midieron directamente. No solo adivinaron; ejecutaron la IA en tareas reales como resolver problemas matemáticos (GSM8K) y escribir código (HumanEval). Demostraron que con una compresión de 4x (haciendo la memoria 4 veces más pequeña), la IA todavía rinde casi tan bien como la versión original no comprimida. Cuando llevaron la compresión al extremo (21.3x), la IA todavía podía manejar muchas tareas, aunque empezó a tener dificultades con los pasos de razonamiento más largos y complejos.

También probaron esto en un tipo diferente de modelo de IA llamado Huginn-3.5B, que tiene una estructura cerebral distinta y realiza hasta 32 bucles. Incluso allí, el método funcionó, demostrando que esto no es solo una casualidad de un robot específico. El artículo sugiere que este "trayecto de bajo rango" (el camino suave de las notas) es una propiedad fundamental de cómo piensan estos robots con bucles. Aunque han demostrado que funciona para ahorrar memoria e incrementar la capacidad, señalan que para obtener el mejor rendimiento en tareas muy largas, la IA necesita una segunda ronda de entrenamiento donde practique con sus propias notas comprimidas, en lugar de solo copiar a un profesor. Este entrenamiento "on-policy" corrige los pequeños errores que ocurren cuando el robot comienza a escribir su propia historia.

En resumen, este artículo muestra que podemos hacer que los modelos de IA con bucles sean mucho más eficientes al darnos cuenta de que su memoria es repetitiva de una manera específica. Al almacenar un pequeño resumen de esa repetición en lugar de las notas completas, podemos meter cientos de veces más conversaciones en un solo chip de computadora, haciendo que la IA poderosa esté disponible para más personas sin necesidad de hardware supercostoso.

Resumen Técnico: Atención Latente en Bucle (Looped Latent Attention - LLA)

Planteamiento del Problema

En los Transformers con bucles y pesos compartidos (weight-tied), un único bloque se aplica repetidamente (recurrencia) para reducir el recuento de parámetros. Si bien esto comprime eficazmente los pesos del modelo, el cuello de botella de memoria en la inferencia sigue siendo el caché de Clave/Valor (KV). En la inferencia con bucles estándar, el modelo vuelve a computar la atención en cada paso de recurrencia, almacenando un caché KV separado para cada paso, capa y cabeza. En consecuencia, un modelo con $T$ bucles sobre $D$ capas conlleva un caché KV equivalente al de un decodificador de $T \times D$ capas, lo que anula los beneficios de memoria de la compartición de pesos.

El artículo postula que este caché indexado por el bucle está altamente estructurado. Para un token, capa y cabeza fijos, los vectores K/V a través de los pasos de recurrenza trazan una trayectoria corta y de bajo rango en lugar de explorar un espacio de alta dimensión. La redundancia es específica del eje del bucle; los espectros a lo largo de los ejes de cabeza y capa permanecen significativamente más planos (rango más alto). El desafío es comprimir este eje de recurrencia sin colapsar la trayectoria en un único punto fijo, lo que degradaría el rendimiento.

Metodología: Atención Latente en Bucle (LLA)

Los autores introducen la Atención Latente en Bucle (LLA), un codec de caché post-entrenamiento que reemplaza los $T$ vectores K/V por por cada bucle con una representación latente compacta y mapas de reconstrucción específicos para cada bucle.

Mecanismo Central

Almacenamiento Latente: En lugar de almacenar $T$ $T$ vectores ( $x_1, \dots, x_T$ $x_{1}, \dots, x_{T}$ ) para un token, capa y cabeza específicos, LLA almacena un único vector latente centrado de bajo rango $c \in \mathbb{R}^r$ $c \in R^{r}$ .
- $c = (x - \mu)W_{\text{down}}$
- Aquí, $\mu$ es la media de un conjunto de calibración, y $W_{\text{down}}$ es una matriz de proyección descendente.
Reconstrucción Específica del Bucle: El vector K/V específico para el paso de recurrencia $t$ $t$ se reconstruye utilizando una proyección ascendente específica del bucle:
- $\hat{x}_t = c W_{\text{up},t}^\top + \mu_t$
Rangos Asimétricos: El método utiliza rangos separados para las Claves ( $r_k$ ) y los Valores ( $r_v$ ). El análisis empírico muestra que los vectores de Valor se estabilizan más lentamente a través de los bucles que los de Clave, lo que motiva que $r_v > r_k$ .
Variante LLA-2D: Para una compresión extrema, LLA-2D pliega el eje de la cabeza en el latente, almacenando un único latente por token y capa para reconstruir todas las $H$ cabezas. Esto reduce el tamaño del caché por un factor de $H$ , pero acopla las cabezas.

Entrenamiento e Inicialización

Inicialización mediante SVD: El codec se inicializa mediante la Descomposición en Valores Singulares (SVD) de las activaciones del profesor en un conjunto de calibración. La proyección descendente se establece con los $r$ vectores singulares derechos superiores. Esta inicialización es crítica; una inicialización aleatoria resulta en una divergencia KL significativamente mayor.
Destilación: El codec se refina congelando el modelo profesor y minimizando la divergencia KL forward entre la distribución del profesor y la distribución del modelo con el codec aplicado, aumentada con un término de coincidencia de la salida de atención.
Refinamiento On-Policy: Para abordar el sesgo de exposición en rollouts autorregresivos largos (donde el modelo condiciona sobre sus propios prefijos generados en lugar de los del profesor), se emplea una segunda etapa de destilación on-policy. El codec se actualiza utilizando prefijos muestreados del propio modelo estudiante, estabilizando el rendimiento en generaciones largas.

Contribuciones Clave

Definición de LLA: Un codec de K/V latente trans-bucle para Transformers con pesos compartidos que logra una tasa de compresión $\rho = \frac{2Td_{\text{head}}}{r_k + r_v}$ . Almacena $(r_k + r_v)H$ escalares por token por capa en lugar de $2Td_{\text{head}}$ .
Compresión Específica de Ejes: El artículo demuestra que el eje de recurrencia (bucle) es la dimensión más compresible en modelos con bucles (Ouro-1.4B, Ouro-2.6B-Thinking y Huginn-3.5B). Comprimir el eje del bucle supera en rendimiento a comprimir el eje de la cabeza (MLA), el eje de la capa (compartición entre capas) o aplicar cuantización con presupuestos de caché equivalentes.
Rechazo del Colapso de Punto Fijo: Los autores muestran que la trayectoria del bucle no puede ser reemplazada por un estado único (el último bucle). Un control de cero parámetros que reutiliza solo el K/V del último bucle colapsa el rendimiento de la generación (la precisión en GSM8K cae a 0.000), demostando que la "trayectoria corta de bajo rango" es necesaria.
Impacto en el Despliegue: La reducción exacta de escalares se traduce directamente en un aumento de la capacidad de servicio. En una sola GPU H200, LLA aumenta el tamaño de lote máximo para Ouro-1.4B con contexto de 4k de 32 a 768 secuencias (con una compresión de 21.3×).

Resultados Experimentales

Compresión vs. Precisión

Rendimiento en GSM8K: En Ouro-1.4B, la compresión ligera (1.33× a 2×) iguala la precisión del profesor (aprox. 0.79–0.80). Con una compresión de 4×, LLA mantiene una precisión de 0.59, superando significativamente a la compresión del eje de la cabeza (0.52) y al reúso del último bucle (0.00).
Generalización de Escala: En Ouro-2.6B-Thinking, LLA permanece casi sin pérdidas (0.81–0.83) a través de 4× a 10.67× de compresión, mientras que las líneas base del eje de la cabeza pierden precisión sustancialmente.
Transferencia a Huginn: El método se transfiere a Huginn-3.5B (32 pasos de recurrencia). Un codec SVD sin entrenamiento logra una fidelidad casi sin pérdidas hasta 32× en métricas independientes del decodificador.

Estabilidad de Generación Larga

MATH-500: Los rollouts matemáticos largos son sensibles a la deriva del prefijo. El entrenamiento off-policy (usando prefijos del profesor) conduce a la degradación de la precisión y a altas tasas de "no respuesta" a medida que aumenta la longitud de la generación.
Ganancia On-Policy: Aplicar el refinamiento on-policy eleva la precisión de MATH-500 de 0.43 a 0.66 con 4× de compresión y reduce significativamente las generaciones sin respuesta.

Capacidad de Servicio

Reducción de Memoria: La reducción del caché es exacta. Con una compresión de 21.3×, el caché KV para Ouro-1.4B con 1M de contexto cae de 768 GiB a 36 GiB.
Throughput: La "ruta de capacidad de almacenamiento latente" (almacenar solo latentes) está limitada por cómputo pero permite lotes y contextos previamente imposibles debido a OOM (Out of Memory). La "ruta rápida" (reconstruir y almacenar el K/V completo) preserva el throughput pero no realiza el ahorro de memoria.

Significación y Reivindicaciones

El artículo afirma que la recurrencia es el eje de caché más compresible en Transformers con bucles y pesos compartidos. La significación radica en:

Visión Mecanicista: Identifica que el eje del bucle es espectralmente distinto (rango más bajo) de los ejes de cabeza y capa, permitiendo una compresión eficiente sin colapsar el comportamiento dinámico del modelo en un punto fijo estático.
Despliegue Práctico: Proporciona una solución post-entrenamiento que aumenta directamente el tamaño de lote y las capacidades de longitud de contexto en el hardware existente sin reentrenar la arquitectura base del modelo.
Dirección Futura: Los resultados sugieren que las arquitecturas recurrentes nativas podrían beneficiarse de cachés latentes desacoplados, donde la recurrencia computa contenido compartido más deltas económicos por bucle, en lugar de re-materializar estados K/V completos.

Los autores enfatizan que su método no depende de implementaciones de atención absorbida para sus afirmaciones de precisión; el resultado central es la validez de la suposición de la trayectoria de bajo rango y la eficacia del codec para preservar el comportamiento del profesor bajo restricciones estrictas de memoria.