Resumen Técnico: KVpop — Compresión de Caché de Clave-Valor con Poda Predictiva en Línea

1. Planteamiento del Problema

La decodificación autorregresiva en los modelos de lenguaje basados en Transformers (LLMs) depende de un caché de Clave-Valor (KV) para evitar la recomputación de las representaciones de los tokens pasados. Sin embargo, el caché KV crece linealmente con la longitud del contexto, lo que crea un cuello de botella significativo para la memoria y el ancho de banda durante la inferencia de contexto largo.

Las soluciones existentes para reducir este costo generalmente se dividen en dos categorías:

• Evicción Estática/Heurística: Métodos como la Ventana Deslizante o la preservación de Tokens de Sumidero (por ejemplo, StreamingLLM) dependen de reglas fijas, mientras que los métodos basados en puntuación utilizan proxies en línea (por ejemplo, atención acumulada). Estos suelen fallar al rastrear la utilidad de los tokens futuros, lo que conduce a una evicción frágil cuando la relevancia cambia.
• Evicción Aprendida: Enfoques como la Compresión de Memoria Dinámica (DMC) o la Esparcidad de Memoria Dinámica (DMS) entrenan políticas para fusionar o descartar tokens. Sin embargo, muchos de estos métodos toman decisiones de evicción en el momento de la inserción del token o dependen de relajaciones diferenciables que no supervisan directamente la decisión específica de "mantener o descartar" basada en la utilidad futura. Además, los métodos con ventanas protegidas retrasan la eliminación de tokens, pero a menudo deciden su destino inmediatamente al entrar, ignorando la evidencia que se acumula mientras el token permanece en la ventana protegida.

El desafío central es predecir qué tokens serán útiles para consultas futuras, ya que los tokens localmente salientes pueden volverse irrelevantes, mientras que los tokens con baja atención inmediata pueden volverse críticos más tarde (por ejemplo, en trazas de razonamiento).

2. Metodología: KVpop

KVpop (Compresión de KV con Poda Predictiva en Línea) es un retrofit de atención dispersa diseñado para imponer un presupuesto fijo por cabeza de KV sin alterar la arquitectura del modelo base. Aborda las limitaciones del trabajo previo mediante dos mecanismos principales: Supervisión de Atención Futura y Puntuación de Estado Retrasada.

2.1 Arquitectura y Presupuesto

KVpop mantiene un caché acotado por cabeza de KV definido por el presupuesto $B = s + w + k$:

• $s$ Tokens de Sumidero (Sink Tokens): Los primeros tokens se mantienen siempre.
• $w$ Ventana Protegida: Los $w$ tokens más recientes están protegidos de la evicción.
• $k$ Top-K de Largo Alcance: Los tokens más antiguos compiten por un presupuesto fijo de $k$ ranuras basado en puntuaciones de importancia aprendidas.

2.2 Supervisión de Objetivo de Atención Futura

A diferencia de los métodos que puntúan tokens basándose en la atención pasada o estadísticas locales, KVpop supervisa al evaluador utilizando la masa de atención futura que un token recibe después de salir de la ventana protegida.

• Definición del Objetivo: Para un token $t$, el objetivo es la masa de log-atención agregada que recibe de las consultas futuras (posiciones $d \ge t + w$).
• Computación Eficiente: Para evitar el costo $O(S^2)$ de materializar mapas de atención densos durante el entrenamiento, KVpop utiliza una pasada de atención transpuesta. Al intercambiar los roles de las consultas (queries) y las claves (keys) y reutilizar los normalizadores log-sum-exp (LSE) devueltos por los kernels de atención dispersa, el sistema computa el objetivo de atención futura de manera eficiente sin la materialización de matrices densas.
• Función de Pérdida: El objetivo de entrenamiento es una pérdida logística de par de frontera sensible. Compara la puntuación de un token recién elegible ($t_{new}$) contra la puntuación del token de corte del profesor ($t_{bnd}$). La pérdida empuja al estudiante a clasificar correctamente $t_{new}$ en relación con la frontera, asegurando que la selección top-$k$ coincida con el ranking de utilidad futura del profesor.

2.3 Puntuación de Estado con Decisión Retrasada

KVpop introduce un evaluador de estado basado en memoria (implementado mediante un mLSTM) que retrasa la decisión de puntuación hasta que el token alcanza la frontera de evicción.

• Mecanismo: En lugar de puntuar un token inmediatamente tras su inserción, el evaluador actualiza su estado interno con la representación del token y el contexto subsecuente. Cuando el token sale de la ventana protegida ($q - w$), el evaluador lee la memoria acumulada para generar la puntuación de importancia.
• Ventaja: Esto permite que el evaluador incorpore el contexto de "futuro cercano" (evidencia que se acumula mientras el token está en la ventana protegida) antes de tomar la decisión de mantener o descartar.

2.4 Eficiencia de Inferencia

Durante la inferencia, KVpop impone el presupuesto fijo $B$. La selección de los $k$ tokens de largo alcance se optimiza utilizando un umbral de ejecución sobre prioridades estáticas de tokens y un árbol de Fenwick, lo que permite generar la máscara dispersa sobre la marcha dentro del kernel de atención (por ejemplo, FlexAttention) sin almacenar máscaras densas $S \times S$.

3. Contribuciones Clave

1. Política de Evicción Supervisada: Los autores formulan la evicción de KV de presupuesto fijo como un problema de predicción supervisada en la frontera de evicción, utilizando un objetivo de atención futura para entrenar las decisiones top-$k$ de largo alcance.
2. Computación de Objetivo Eficiente: Introducen una pasada de atención transpuesta exclusiva para el entrenamiento para computar los objetivos de atención futura sin materializar mapas de atención densos, garantando que no haya sobrecarga en la inferencia.
3. Puntuación Retrasada: Demuestran que un evaluador con estado puede explotar el contexto de futuro cercano al retrasar la decisión de puntuación hasta que el token sea evictible, superando a los métodos de puntuación inmediata.
4. Validación Empírica: El método se evalúa en razonamiento matemático (AIME, HMMT) y tareas fuera de dominio (GPQA-D, LiveCodeBench), mostrando que las políticas de evicción aprendidas pueden mantener un alto rendimiento bajo una compresión agresiva.

4. Resultados Experimentales

Los autores evaluaron KVpop en modelos Qwen3-4B y Qwen3-8B, entrenados con el conjunto de datos Nemotron-Math v2 con una longitud de secuencia de 16,384.

4.1 Razonamiento Matemático (AIME y HMMT)

• Relación de Compresión 75% (Presupuesto 2048):
  • Qwen3-4B: KVpop retuvo el 98% del rendimiento del profesor de atención completa.
  • Qwen3-8B: KVpop retuvo el 100% (rendimiento cercano al del profesor completo).
  • Comparación: KVpop superó significativamente a las líneas base sin entrenamiento (StreamingLLM, TOVA) y a la línea base aprendida DMS. Por ejemplo, con un 75% de compresión, StreamingLLM cayó a un ~76% de rendimiento relativo en Qwen3-4B, mientras que KVpop mantuvo el 98%.
• Relación de Compresión 88% (Presupuesto 4096):
  • Qwen3-4B: KVpop retuvo el 97% del rendimiento del profesor.
  • Qwen3-8B: KVpop retuvo el 100% del rendimiento del profesor.
  • Comparación: Los métodos sin entrenamiento colapsaron significativamente (por ejemplo, StreamingLLM cayó a ~47% de rendimiento relativo), mientras que KVpop se mantuvo robusto.

4.2 Generalización Fuera de Dominio

Evaluado en GPQA-Diamond (razonamiento científico) y LiveCodeBench (generación de código):

• Los modelos KVpop entrenados solo con datos de razonamiento matemático generalizaron bien a estos dominios.
• Con una relación de compresión del 75%, KVpop en Qwen3-4B logró una precisión comparable a la del profesor denso en GPQA-D (0.59 vs 0.59) y resultados competitivos en LCB.

4.3 Eficiencia de Inferencia

• Memoria: A diferencia de la atención densa, donde el uso de VRAM escala linealmente con la longitud de la generación (alcanzando 36GB a 131k tokens), KVpop y DMS mantuvieron una huella de memoria casi constante (19GB) independientemente de la longitud de la secuencia.
• Latencia: KVpop demostró una menor latencia de decodificación de extremo a extremo en comparación con DMS y la atención densa en longitudes de generación largas (hasta 131k tokens). Los autores atribuyen esto al presupuesto fijo por cabeza de KVpop, que evita las estructuras de caché "irregulares" de DMS que complican la compilación y ejecución.

4.4 Estudios de Ablación

• Puntuación Retrasada: Eliminar el retraso (puntuar inmediatamente tras la inserción) resultó en una caída de 0.2 puntos en la precisión de los tokens, confirmando el beneficio de integrar el contexto de futuro cercano.
• Estadual vs. Sin Estado: La variante con estado (KVpop) generalmente igualó o superó ligeramente a la variante sin estado (KVpopmlp), particularmente en altas relaciones de compresión.

5. Significancia y Reivindicaciones

El artículo afirma que supervisar la evicción con señales de atención futura es una estrategia superior para la compresión del caché KV en comparación con las heurísticas estáticas o las puntuaciones proxy. Al entrenar directamente una política para predecir qué tokens importarán para las consultas futuras, KVpop logra:

1. Alta Fidelidad: Mantiene un rendimiento cercano al del profesor incluso con altas relaciones de compresión (hasta el 88%), un régimen donde otros métodos fallan.
2. Robustez: La política aprendida se transfiere eficazmente a tareas fuera de dominio, lo que sugiere que el modelo aprende principios generales de utilidad de los tokens en lugar de sobreajustarse a distribuciones de entrenamiento específicas.
3. Eficiencia: Impone un límite de memoria estricto, permitiendo una decodificación de contexto largo estable con uso de memoria constante y un mejor rendimiento (throughput) en comparación con los métodos de gating dinámico como DMS.

Los autores concluyen que KVpop demuestra que las políticas de evicción aprendidas, cuando se supervisan adecuadamente mediante la utilidad futura y se les permite utilizar el contexto de futuro cercano, pueden resolver eficazmente el cuello de botella de la memoria en la generación autorregresiva sin sacrificar la calidad del modelo.