Resumen Técnico: Compactación Rápida de KV mediante Correspondencia de Atención

1. Declaración del Problema

Escalar los modelos de lenguaje grandes (LLM) hacia contextos largos está actualmente limitado por los requisitos de memoria de la memoria caché de Clave-Valor (KV). En las arquitecturas Transformer autoregresivas, la caché KV debe retener claves y valores para todos los tokens anteriores, alcanzando a menudo gigabytes por solicitud en configuraciones de horizonte largo (por ejemplo, diálogos multi-sesión, codificación con agentes).

Los enfoques existentes para reducir el tamaño de la caché KV generalmente caen en dos categorías, ambas con limitaciones significativas:

• Métodos en el espacio de tokens: Técnicas como la expulsión de tokens (H2O, SnapKV), la fusión de tokens (KVMerger) y la resumen degradan rápidamente a altas tasas de compresión. El resumen, en particular, es altamente con pérdida y perjudica el rendimiento aguas abajo en tareas que requieren extracción precisa de información.
• Métodos en el espacio latente: Trabajos recientes como Cartridges (Eyuboglu et al., 2025) demostraron que las cachés KV compactas pueden entrenarse en el espacio latente para igualar el rendimiento de contexto completo. Sin embargo, este enfoque depende de una optimización basada en gradientes de extremo a extremo (prefix-tuning) para cada contexto, lo cual es prohibitivamente lento y costoso (requiriendo varias horas de GPU por contexto).

El problema abierto es lograr una compactación de contexto rápida y de alta calidad en el espacio latente que preserve el comportamiento del modelo aguas abajo sin el costo computacional de la optimización de gradientes por contexto.

2. Metodología: Correspondencia de Atención

Los autores proponen una familia de métodos basada en la Correspondencia de Atención, que construye una caché KV compacta $(C_k, C_v)$ para reproducir las salidas de atención y la masa de atención de la caché original $(K, V)$ para un conjunto de consultas de referencia.

Objetivo Central

El objetivo es encontrar una caché compacta $(C_k, C_v)$ de tamaño $t \ll T$ tal que, para cualquier consulta $q$ y cualquier token fijo/futuro añadido $(K_{fixed}, V_{fixed})$, la salida de atención permanezca aproximadamente preservada. Los autores aprovechan la descomposición de la atención sobre bloques concatenados, mostrando que basta con igualar dos cantidades sobre un conjunto de consultas de referencia $Q_{ref}$:

1. Salida de Atención Local: La suma ponderada de valores.
2. Masa de Atención: La suma de logits de atención no normalizados (el denominador en el softmax).

Para manejar la discrepancia donde la caché compactada tiene menos tokens ($t < T$) y, por lo tanto, no puede igualar naturalmente la masa total de atención (por ejemplo, para $q=0$, la masa es $T$ frente a $t$), el método introduce un sesgo escalar por token $\beta$. Los objetivos de optimización son:
$$\frac{\exp(qK^\top)V}{\sum \exp(qK^\top)} \approx \frac{\exp(qC_k^\top + \beta)C_v}{\sum \exp(qC_k^\top + \beta)}$$
$$\sum \exp(qK^\top) \approx \sum \exp(qC_k^\top + \beta)$$

Pipeline Algorítmico

El marco descompone el problema en subproblemas eficientes y de forma cerrada, evitando el descenso de gradiente en el momento de la compactación:

1. Generación de Consultas de Referencia ($Q_{ref}$):

   • Prefill repetido: Ejecutar un pase de prefill sobre el contexto mismo (o una versión repetida) para extraer activaciones de consultas.
   • Autoestudio: Generar interacciones sintéticas condicionadas al contexto para ampliar la distribución de consultas.
   • En política: Compactar capas secuencialmente y extraer consultas del modelo parcialmente compactado para reducir el desplazamiento de distribución.

2. Selección de Claves ($C_k$):

   • Claves de Mayor Atención: Seleccionar claves originales que reciben las puntuaciones de atención más altas a través de las consultas de referencia (agregadas mediante RMS).
   • Búsqueda de Ajuste Ortogonal (OMP): Seleccionar claves codiciosamente para satisfacer mejor el objetivo de igualación de masa de atención, seguido de Mínimos Cuadrados No Negativos (NNLS) para ajustar los pesos.

3. Ajuste de Sesgo y Valores ($\beta, C_v$):

   • Ajuste de $\beta$: Resuelto mediante Mínimos Cuadrados No Negativos (NNLS) para igualar la masa de atención de las claves originales.
   • Ajuste de $C_v$: Resuelto mediante Mínimos Cuadrados Ordinarios (OLS) para igualar las salidas de atención dadas las claves seleccionadas y los sesgos ajustados.

Compactación No Uniforme

Los autores observan que diferentes cabezas de atención exhiben sensibilidades variables a la compactación. Proponen un programa de compactación no uniforme donde diferentes cabezas y capas se asignan a diferentes ratios de retención. Este programa se precalcula una vez por modelo utilizando un algoritmo de intercambio codicioso basado en curvas de sensibilidad por cabeza, permitiendo una aplicación robusta a través de diferentes contextos.

Compactación por Bloques

Para soportar contextos muy largos (por ejemplo, 60k+ tokens), el método aplica la compactación independientemente a bloques contiguos de la entrada. El enfoque de fragmentación basada en KV (prefill del contexto completo, corte de estados KV, compactación y fusión) se utiliza por defecto para preservar la semántica posicional global y las interacciones entre bloques.

3. Contribuciones Clave

• Formulación de Correspondencia de Atención: Un objetivo novedoso para la compactación en espacio latente que optimiza directamente la preservación del comportamiento de atención en lugar de las verosimilitudes de salida, descomponiéndose naturalmente en problemas de álgebra lineal eficientes.
• Soluciones de Forma Cerrada: El método elimina la necesidad de optimización basada en gradientes de extremo a extremo durante la compactación, reemplazándola con solucionadores de mínimos cuadrados y NNLS.
• Frontera de Pareto Velocidad-Calidad: El enfoque logra ratios de compactación de hasta 50× en segundos (en una sola GPU), órdenes de magnitud más rápido que Cartridges, mientras mantiene una precisión aguas abajo comparable o superior.
• Programación No Uniforme: Una estrategia de asignación de presupuesto precalculada específica del modelo que prioriza la capacidad para las cabezas de atención más sensibles a la compactación.
• Mecanismo de Sesgo: La introducción de sesgos escalares por token para corregir las discrepancias de masa de atención, permitiendo una igualación exacta de la masa con un conjunto reducido de claves.

4. Resultados

Los métodos fueron evaluados en QuALITY (QA de documentos largos), LongHealth (registros clínicos densos), QASPER (QA generativa) y LongBench v2.

• Rendimiento vs. Velocidad: En QuALITY con un ratio de compactación de 50×, los métodos de Correspondencia de Atención trazan la frontera de Pareto, superando a las líneas base de expulsión de tokens (H2O+, KVzip) y al resumen. Igualan o exceden la precisión de Cartridges mientras son dos órdenes de magnitud más rápidos (minutos frente a horas).
• Alta Compresión: En ratios extremos (hasta 100×), los métodos retienen una precisión aguas abajo significativa (por ejemplo, ~60% en LongHealth frente a ~30% para líneas base de resumen).
• Generalización: El enfoque se generaliza a tareas generativas (QASPER) y benchmarks intensivos en recuperación (RULER), donde la variante más rápida (AM-HighestAttnKeys-fast) superó a KVzip, el estado del arte, en RULER.
• Eficiencia: La compactación de un contexto de 60k tokens en Gemma-3-12B toma aproximadamente 100–600 segundos dependiendo de la estrategia de selección de claves (OMP frente a Mayor Atención), dominada por los pasos de generación de consultas y selección de claves.
• Compactación en Línea: Resultados preliminares en AIME muestran que la compactación repetida a mitad de trayectoria preserva el rendimiento de razonamiento, desacoplando la profundidad del razonamiento de los límites físicos de memoria.

5. Significado y Afirmaciones

El artículo afirma que la Correspondencia de Atención proporciona un primitivo práctico y eficiente para la compactación de caché KV que cierra la brecha entre los métodos rápidos y con pérdida en el espacio de tokens y los métodos lentos y de alta calidad en el espacio latente.

• Practicidad: Al evitar el descenso de gradiente por contexto, el método hace viable la compactación de alta fidelidad para la implementación en el mundo real donde la latencia y el costo computacional son restricciones.
• Integración de Sistemas: El enfoque es compatible con motores de inferencia existentes (por ejemplo, FlashAttention, empaquetado de KV de longitud variable) y puede realizarse fuera de la ruta crítica (compactación desagregada).
• Limitaciones: Los autores reconocen que, aunque su método es significativamente más rápido que Cartridges, sigue siendo más lento que las heurísticas simples. Además, en ratios de compresión extremos (por ejemplo, 100× en LongHealth), los métodos basados en gradientes como Cartridges pueden seguir superando debido a su capacidad para buscar un espacio de representación más amplio no restringido a las claves originales.
• Direcciones Futuras: El artículo sugiere trabajo futuro en la optimización de la generación de consultas, ir más allá de la selección de subconjuntos hacia la optimización directa de claves compactas, e integrar la compactación en motores de inferencia para configuraciones de agentes en línea y de horizonte largo.

En resumen, este trabajo establece que la compactación de contexto rápida y de alta calidad es alcanzable igualando el comportamiento de atención directamente, ofreciendo una solución escalable al cuello de botella de memoria en LLMs de contexto largo.