Resumen Técnico: Evicción de la Caché KV Consciente de la Epifanía sin la Matriz de Atención

Declaración del Problema

A medida que los modelos de razonamiento (por ejemplo, DeepSeek-R1) generan cadenas de pensamiento que abarcan decenas de miles de tokens, la caché de Clave-Valor (KV) crece linealmente, convirtiéndose en el principal cuello de botella de memoria para el despliegue. Las estrategias actuales de evicción de la caché KV suelen clasificar los tokens basándose en los pesos de atención. Este enfoque presenta dos limitaciones críticas:

1. Proxy de Importancia Ruidoso: Los pesos de atención son indicadores poco fiables de la importancia de un token en las trazas de razonamiento largas. Los "sumideros de atención" (attention sinks) absorben peso independientemente del contenido, y los tokens de relleno atraen peso durante la generación, pero rara vez son referenciados posteriormente.
2. Incompatibilidad Arquitectónica: Leer los pesos de atención requiere la materialización de la matriz de atención $n \times n$. Esto obliga al uso de kernels "eager", desactivando los kernels fusionados como FlashAttention. En consecuencia, el uso de memoria escala cuadráticamente ($O(n^2)$), causando errores de fuera de memoria (OOM) en longitudes de contexto (por ejemplo, 8192 tokens) que son mucho más cortas que las trazas de razonamiento típicas.

Los métodos existentes que intentan mitigar esto (por ejemplo, ThinKV, RaaS) siguen dependiendo de señales derivadas de la atención, mientras que otros (como LagKV) son genéricos pero carecen de optimización específica para el razonamiento.

Metodología: EPIKV

Los autores proponen EPIKV (KV Cache con Conciencia de Epifanía), una política de evicción libre de entrenamiento que califica los tokens basándose en los cambios en la representación interna del modelo en lugar de los pesos de atención.

1. La Puntuación de Epifanía

En lugar de la atención, EPIKV mide el cambo en el flujo residual (estado oculto) entre pasos de decodificación consecutivos. Para un token $t$ en la capa $l$, la señal se define como la diferencia L2:
$$g_l(t) = \|h_l(t) - h_l(t-1)\|_2$$
Los cambios grandes indican "tokens de epifanía": puntos de transición donde el modelo cambia de un razonamiento exploratorio a uno convergente, concluye un paso o se compromete con una idea.

2. Anatomía de Capas de Dos Bandas

A través de experimentos de oclusión contrafáctica, los autores identifican una estructura de capas específica en modelos de razonamiento de 32 capas:

• Banda A (Capas 7–13): Un alto cambio en el estado oculto aquí correlaciona positivamente con la importancia del token. Estas capas se alinean con sitios conocidos de recuperación de hechos y enrutamiento de características.
• Banda B (Capas 18–25): Un alto cambio en el estado oculto aquí correlaciona negativamente con la importancia. Estas capas medias-superiores están activas incluso para continuaciones fluidas y de baja sorpresa; por lo tanto, los cambios grandes aquí señalan tokens predecibles y prescindibles.

La puntuación final combina estas bandas:
$$s(t) = \bar{g}_{10}(t) - \bar{g}_{21}(t)$$
donde $\bar{g}$ representa una media móvil sobre una ventana causal.

3. Corrección de Tendencia Temporal

La señal bruta exhibe una tendencia posicional monotónica (rastreando la posición tanto como el contenido). Para corregir esto, los autores aplican una puntuación z de media móvil causal:
$$z_l(t) = \frac{g_l(t) - \mu_l(t)}{\sigma_l(t) + \epsilon}$$
La puntuación de evicción final es $z_{10}(t) - z_{21}(t)$. Esta normalización elimina el sesgo posicional, asegurando que la puntuación refleje la desviación del contenido local en lugar de la magnitud absoluta.

4. Compatibilidad con FlashAttention

Crucialmente, EPIKV calcula las puntuaciones utilizando solo:

1. Claves y valores en caché (ya presentes en la memoria de alto ancho de banda).
2. Estados ocultos por capa expuestos por los pasos de propagación hacia adelante (forward passes) estándar.
   Nunca solicita output_attentions=True y nunca materializa la matriz de atención. Esto permite que EPIKV se ejecute dentro de los stacks de inferencia de FlashAttention estándar sin forzar un retroceso a kernels eager.

Contribuciones Clave

1. Identificación de la Anatomía de Dos Bandas: Descubrimiento de que los cambios en el estado oculto en las capas medias (7–13) correlacionan positivamente con la importancia, mientras que los cambios en las capas medias-superiores (18–25) correlacionan negativamente, superando las señales basadas en la atención.
2. Puntuación Independiente de la Posición: Demostración de que una puntuación z de media móvil causal elimina las tendencias posicionales inherentes a la varianza del estado oculto bruto, recuperando un alto rendimiento de evicción.
3. Paridad de Rendimiento sin Atención: EPIKV iguala o supera a los baselines más fuertes basados en atención (ThinKV, H2O) en MATH-500 y AIME-2024 con presupuestos de caché desplegables, sin haber materializado nunca la matriz de atención.
4. Eficiencia de Ingeniería: El método es hasta 2.8 veces más rápido que la evicción basada en atención con presupuestos iguales y evita el muro de memoria que limita la longitud del contexto en la puntuación basada en atención.

Resultados

Los experimentos se realizaron en DeepSeek-R1-Distill-LLaMA-8B utilizando los conjuntos de datos MATH-500 y AIME-2024.

• Precisión:
  • Con una caché de 4096 tokens en MATH-500, EPIKV alcanza un 72% de precisión, igualando a ThinKV (71%) y superando a H2O (67%).
  • Con una caché de 8192 tokens en AIME-2024, la variante de EPIKV con normalización de retardo alcanza un 37%, superando al mejor baseline basado en atención (33%).
• Velocidad:
  • En AIME-2024 (8192 tokens), EPIKV procesa problemas en 440 segundos, comparado con 721 segundos para el baseline de atención más rápido (ThinKV) y 1239 segundos para RaaS. Esto representa una aceleración de 1.6× a 2.8×.
• Escalabilidad de Memoria:
  • Mientras que la puntuación basada en atención agota una GPU A100 de 80 GB a ~8192 tokens debido a la matriz de atención $O(n^2)$, EPIKV (vía FlashAttention) escala a 65,536 tokens (un aumento de 16× en la longitud de contexto factible) utilizando solo 48 GB de memoria.
• Modos de Fallo de los Baselines: Los métodos basados en atención como H2O colapsan en las trazas de razonamiento, produciendo salidas vacías o texto sin estructura con presupuestos ajustados (por ejemplo, 1024 tokens), en lugar de degradarse de forma gradual.

Significado y Reivindicaciones

El artículo afirma que EPIKV hace que la evicción de la caché KV sea desplegable en stacks de servicio FlashAttention estándar. Al desacoplar la puntuación de importancia de la matriz de atención, el método:

• Elimina el costo arquitectónico de materializar los pesos de atención.
• Proporciona una señal más robusta para las trazas de razonamiento donde los pesos de atención son ruidosos.
• Permite ventanas de contexto significativamente más largas y mayores tamaños de lote (batch sizes) en el hardware existente.

Los autores señalan explícitamente que la anatomía de dos bandas es específica para la matemática de competición y cambia con la dificultad de la tarea (por ejemplo, moviéndose hacia capas anteriores en GSM8K), lo que sugiere que la señal rastrea la consolidación del razonamiento en lugar de índices de capa fijos. Liberan etiquetas de importancia contrafáctica como un recurso de validación para trabajos futuros.