Investigadores del MIT, NVIDIA y la Universidad de Zhejiang proponen TriAttention: un método de compresión de caché KV que iguala la atención total con un rendimiento 2,5 veces mayor

El razonamiento de cadena larga es una de las tareas con mayor uso intensivo de cómputo en los modelos de lenguajes grandes modernos. Cuando un modelo como DeepSeek-R1 o Qwen3 resuelve un problema matemático complejo, puede generar decenas de miles de tokens antes de llegar a una respuesta. Cada uno de esos tokens debe almacenarse en lo que se llama caché KV, una estructura de memoria que contiene los vectores de clave y valor a los que el modelo debe prestar atención durante la generación. Cuanto más larga es la cadena de razonamiento, mayor crece la caché KV y, en muchos escenarios de implementación, especialmente en hardware de consumo, este crecimiento acaba agotando por completo la memoria de la GPU.

Un equipo de investigadores del MIT, NVIDIA y la Universidad de Zhejiang propuso un método llamado TriAttention que aborda directamente este problema. En el punto de referencia de razonamiento matemático AIME25 con generación de tokens de 32K, TriAttention iguala la precisión de Atención total al tiempo que logra un rendimiento 2,5 veces mayor o una reducción de memoria de 10,7 veces KV. Las líneas de base líderes logran sólo aproximadamente la mitad de la precisión con el mismo nivel de eficiencia.

https://arxiv.org/pdf/2604.04921

El problema con la compresión de caché KV existente

Para comprender por qué TriAttention es importante, es útil comprender el enfoque estándar para la compresión de caché KV. La mayoría de los métodos existentes, incluidos SnapKV, H2O y R-KV, funcionan estimando qué tokens en la caché de KV son importantes y desalojando el resto. La importancia normalmente se estima observando las puntuaciones de atención: si una clave recibe mucha atención en consultas recientes, se considera importante y se conserva.

El problema es que estos métodos operan en lo que el equipo de investigación llama espacio post-RoPE. RoPE, o Rotary Position Embedding, es el esquema de codificación posicional utilizado por la mayoría de los LLM modernos, incluidos Llama, Qwen y Mistral. RoPE codifica la posición rotando los vectores de consulta y clave de forma dependiente de la frecuencia. Como resultado, un vector de consulta en la posición 10.000 se ve muy diferente de la misma consulta semántica en la posición 100, porque su dirección ha sido rotada por la codificación de posición.

Esta rotación significa que sólo las consultas generadas más recientemente tienen orientaciones que están “actualizadas” para estimar qué claves son importantes en este momento. Trabajos anteriores han confirmado esto empíricamente: aumentar la ventana de observación para la estimación de importancia no ayuda: el rendimiento alcanza su punto máximo alrededor de 25 consultas y disminuye después de eso. Con una ventana tan pequeña, algunas llaves que luego serán importantes quedan desalojadas permanentemente.

Este problema es especialmente grave para lo que el equipo de investigación llama cabezas de recuperación: cabezas de atención cuya función es recuperar tokens fácticos específicos de contextos prolongados. Los tokens relevantes para una cabeza de recuperación pueden permanecer inactivos durante miles de tokens antes de volverse repentinamente esenciales para la cadena de razonamiento. Los métodos posteriores a RoPE, que operan en una ventana de observación estrecha, ven poca atención a esos tokens durante el período de inactividad y los desalojan permanentemente. Cuando más tarde el modelo necesita recordar esa información, ésta ya ha desaparecido y la cadena de pensamiento se rompe.

La observación previa al RoPE: concentración de Q/K

La información clave en TriAttention proviene de observar los vectores de consulta y clave antes de que se aplique la rotación RoPE: el espacio anterior a RoPE. Cuando el equipo de investigación visualizó los vectores Q y K en este espacio, encontraron algo consistente y sorprendente: en la gran mayoría de las cabezas de atención y en múltiples arquitecturas de modelos, tanto los vectores Q como K se agrupan estrechamente alrededor de puntos centrales fijos distintos de cero. El equipo de investigación denomina esta propiedad concentración Q/K y la mide utilizando la longitud media resultante R, una medida estadística direccional estándar donde R → 1 significa agrupamiento estrecho y R → 0 significa dispersión en todas las direcciones.

En Qwen3-8B, aproximadamente el 90% de las cabezas de atención exhiben R > 0,95, lo que significa que sus vectores Q/K anteriores a RoPE están casi perfectamente concentrados alrededor de sus respectivos centros. Fundamentalmente, estos centros son estables en diferentes posiciones simbólicas y en diferentes secuencias de entrada: son una propiedad intrínseca de los pesos aprendidos del modelo, no una propiedad de ninguna entrada en particular. El equipo de investigación confirma además que la concentración de Q/K es independiente del dominio: medir la longitud resultante media en los dominios de matemáticas, codificación y chat en Qwen3-8B arroja valores casi idénticos de 0,977 a 0,980.

Esta estabilidad es lo que los métodos post-RoPE no pueden explotar. La rotación de RoPE dispersa estos vectores concentrados en patrones de arco que varían con la posición. Pero en el espacio anterior al RoPE, los centros permanecen fijos.

De la concentración a una serie trigonométrica

Luego, el equipo de investigación muestra matemáticamente que cuando los vectores Q y K se concentran alrededor de sus centros, el logit de atención (la puntuación bruta antes de softmax que determina en qué medida una consulta atiende a una clave) se simplifica drásticamente. Al sustituir los centros Q/K en la fórmula de atención de RoPE, el logit se reduce a una función que depende únicamente de la distancia QK (la brecha posicional relativa entre la consulta y la clave), expresada como una serie trigonométrica:

logit(Δ)≈∑f‖q‾f‖‖k‾f‖⏟amplitudecos⁡(ωfΔ+ϕ‾f⏟fase)=∑f[afcos⁡(ωfΔ)+bfsin⁡(ωfΔ)] \text{logit}(\Delta) \approx \sum_{f} \underbrace{\|\bar{q}_f\| \|\bar{k}_f\|}_{\text{amplitud}} \cos(\omega_f \Delta + \underbrace{\bar{\phi}_f}_{\text{fase}}) = \sum_{f} [a_f \cos(\omega_f \Delta) + b_f \sin(\omega_f \Delta)]

Aquí, Δ es la distancia posicional, ωf son las frecuencias de rotación de RoPE para cada banda de frecuencia f, y los coeficientes af y bf están determinados por los centros Q/K. Esta serie produce una curva característica de atención versus distancia para cada cabeza. Algunas cabezas prefieren claves cercanas (atención local), otras prefieren claves muy distantes (atención disipadora). Los centros, calculados fuera de línea a partir de los datos de calibración, determinan completamente qué distancias se prefieren.

El equipo de investigación validó esto experimentalmente en 1152 cabezas de atención en Qwen3-8B y en arquitecturas Qwen2.5 y Llama3. La correlación de Pearson entre la curva trigonométrica predicha y los logits de atención reales tiene una media superior a 0,5 en todas las cabezas, y muchas cabezas logran correlaciones de 0,6 a 0,9. El equipo de investigación valida esto aún más en GLM-4.7-Flash, que utiliza atención latente de múltiples cabezas (MLA) en lugar de atención de consultas agrupadas estándar, una arquitectura de atención significativamente diferente. En MLA, el 96,6% de las cabezas exhiben R > 0,95, en comparación con el 84,7% para GQA, lo que confirma que la concentración de Q/K no es específica de un diseño de atención, sino que es una propiedad general de los LLM modernos.

Cómo TriAttention utiliza esto

TriAttention es un método de compresión de caché KV que utiliza estos hallazgos para calificar claves sin necesidad de observaciones de consultas en vivo. La función de puntuación tiene dos componentes:

La puntuación de la serie trigonométrica (Strig) utiliza el centro Q calculado fuera de línea y la representación de la clave real almacenada en caché para estimar cuánta atención recibirá la clave, en función de su distancia posicional respecto de consultas futuras. Debido a que una clave puede ser atendida por consultas en muchas posiciones futuras, TriAttention promedia esta puntuación sobre un conjunto de compensaciones futuras usando espaciado geométrico.

Strig(k,Δ)=∑f‖𝔼[qf]‖⋅‖kf‖⋅cos⁡(ωfΔ+ϕf)S_{\text{trig}}(k, \Delta) = \sum_{f} \|\mathbb{E}[q_f]\| \cdot \|k_f\| \cdot \cos(\omega_f \Delta + \phi_f)

La puntuación basada en normas (Snorm) maneja la minoría de cabezas de atención donde la concentración Q/K es menor. Pondera cada banda de frecuencia según la contribución esperada de la norma de consulta, proporcionando información complementaria sobre la importancia del token más allá de la simple preferencia de distancia.

Norma(0)(k)=∑f𝔼[‖qf‖]⋅‖kf‖S_{\text{norma}}^{(0)}(k) = \sum_{f} \mathbb{E}[\|q_f\|] \cdot \|k_f\|

Las dos puntuaciones se combinan utilizando la longitud media resultante R como ponderación adaptativa: cuando la concentración es alta, Strig domina; cuando la concentración es menor, Snorm contribuye más. Cada 128 tokens generados, TriAttention puntúa todas las claves en el caché y retiene solo la B superior, expulsando al resto.

Resultados sobre razonamiento matemático

En AIME24 con Qwen3-8B, TriAttention logra una precisión del 42,1% frente al 57,1% de Full Attention, mientras que R-KV alcanza solo el 25,4% con el mismo presupuesto de KV de 2048 tokens. En AIME25, TriAttention logra un 32,9% frente al 17,5% de R-KV, una diferencia de 15,4 puntos porcentuales. En MATH 500, con solo 1.024 tokens en la caché KV de 32.768 posibles, TriAttention logra una precisión del 68,4 % frente al 69,6 % de Full Attention.

https://arxiv.org/pdf/2604.04921

El equipo de investigación también presenta un punto de referencia de consulta de estado recursivo basado en simulación recursiva mediante búsqueda en profundidad. Las tareas recursivas hacen hincapié en la retención de la memoria porque el modelo debe mantener estados intermedios a lo largo de cadenas largas y retroceder hasta ellos más tarde; si se elimina cualquier estado intermedio, el error se propaga a través de todos los valores de retorno posteriores, corrompiendo el resultado final. Bajo una presión de memoria moderada hasta la profundidad 16, TriAttention funciona de manera comparable a la Atención total, mientras que R-KV muestra una degradación catastrófica de la precisión: cae de aproximadamente el 61 % en la profundidad 14 al 31 % en la profundidad 16. Esto indica que R-KV desaloja incorrectamente estados críticos de razonamiento intermedio.

En cuanto al rendimiento, TriAttention logra 1.405 tokens por segundo en MATH 500 frente a los 223 tokens por segundo de Full Attention, una aceleración de 6,3 veces. En AIME25, alcanza 563,5 tokens por segundo frente a 222,8, una aceleración de 2,5 veces con una precisión igualada.

https://arxiv.org/pdf/2604.04921

Generalización más allá del razonamiento matemático

Los resultados van mucho más allá de los puntos de referencia matemáticos. En LongBench, un punto de referencia de 16 subtareas que cubre respuestas a preguntas, resúmenes, clasificación de pocas tomas, recuperación, conteo y tareas de codificación, TriAttention logra la puntuación promedio más alta de 48,1 entre todos los métodos de compresión con un presupuesto de KV del 50 % en Qwen3-8B, ganando 11 de 16 subtareas y superando la siguiente mejor línea de base, Ada-KV+SnapKV, por 2,5 puntos. En el punto de referencia de recuperación de RULER en una longitud de contexto de 4K, TriAttention logra 66,1, una brecha de 10,5 puntos sobre SnapKV. Estos resultados confirman que el método no está adaptado únicamente al razonamiento matemático: el fenómeno subyacente de concentración de Q/K se transfiere a tareas de lenguaje general.

Conclusiones clave

Los métodos de compresión de caché KV existentes tienen un punto ciego fundamental: métodos como SnapKV y R-KV estiman la importancia del token utilizando consultas recientes posteriores a RoPE, pero debido a que RoPE rota los vectores de consulta con la posición, solo se puede utilizar una pequeña ventana de consultas. Esto hace que los tokens importantes, especialmente aquellos que necesitan los jefes de recuperación, sean desalojados permanentemente antes de que se vuelvan críticos. Los vectores clave y de consulta anteriores a RoPE se agrupan alrededor de centros estables y fijos en casi todos los cabezales de atención: esta propiedad, llamada concentración Q/K, se mantiene independientemente del contenido de entrada, la posición del token o el dominio, y es consistente en Qwen3, Qwen2.5, Llama3 e incluso arquitecturas de atención latente de cabezales múltiples como GLM-4.7-Flash. Estos centros estables hacen que los patrones de atención sean matemáticamente predecibles sin observar ninguna consulta activa: cuando los vectores Q/K se concentran, la puntuación de atención entre cualquier consulta y clave se reduce a una función que depende únicamente de su distancia posicional, codificada como una serie trigonométrica. TriAttention usa esto para calificar cada clave almacenada en caché fuera de línea utilizando únicamente datos de calibración. TriAttention iguala la precisión del razonamiento de Atención total a una fracción del costo de memoria y computación: en AIME25 con generación de tokens de 32 000, logra un rendimiento 2,5 veces mayor o una reducción de memoria KV de 10,7 veces mientras iguala la precisión de Atención total, casi duplicando la precisión de R-KV con el mismo presupuesto de memoria en AIME24 y AIME25. El método se generaliza más allá de las matemáticas y funciona en hardware de consumo. TriAttention supera todas las líneas de base en LongBench en 16 subtareas generales de PNL y en el punto de referencia de recuperación RULER, y permite que un modelo de razonamiento de 32 B se ejecute en un solo RTX 4090 de 24 GB a través de OpenClaw, una tarea que causa errores de falta de memoria bajo Atención total.

Consulte la página de artículos, repositorios y proyectos. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 120.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.

¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros