KVSculpt: KV Cache Compression as Distillation

Resumen Técnico: KVSCULPT

1. El Problema

En la inferencia de Modelos de Lenguaje Grande (LLM) autoregresivos con contextos largos, el caché de claves y valores (KV cache) se convierte en el principal cuello de botella de memoria, consumiendo decenas de gigabytes para una sola secuencia. Para mitigar esto, existen dos dimensiones de compresión:

1. Reducción de tamaño por par: Cuantización y descomposición de bajo rango.
2. Reducción de la longitud de la secuencia: Disminuir el número de pares KV almacenados.

Los métodos actuales para reducir la longitud se basan en dos paradigmas restrictivos:

• Evicción (Selección): Seleccionar un subconjunto discreto de los pares KV originales y descartar el resto (ej. H2O, StreamingLLM).
• Fusión (Merging): Combinar pares similares en promedios ponderados, pero permaneciendo anclados a las entradas originales.

Limitación clave: Ambos enfoques asumen que el caché comprimido debe ser un subconjunto o una combinación lineal de las entradas originales. Sin embargo, tras la codificación RoPE (Rotary Positional Embedding), los pares KV son vectores en un espacio de incrustación continuo sin orden inherente. No hay razón teórica para que el caché comprimido deba limitarse a las posiciones originales; cualquier conjunto de vectores que replique el comportamiento de atención debería ser válido.

2. Metodología: KVSCULPT

KVSCULPT reformula la compresión del caché KV como un problema de destilación en un espacio continuo, en lugar de una selección discreta.

• Formulación del Problema:
  Dado un caché completo de $N$ pares, se divide en una zona de retención (los tokens más recientes, $m$) y una zona de compresión (los tokens más antiguos, $N-m$). El objetivo es optimizar un conjunto de $k$ pares KV libres y no restringidos ($K_c, V_c$) que, al combinarse con la zona de retención, reproduzcan la salida de atención del caché completo para cualquier consulta futura.

• Función de Pérdida:
  Se optimiza para igualar la salida de atención completa. La pérdida incluye dos términos:

  1. Error Cuadrático Medio (MSE) de la salida: $\| \hat{Y} - Y \|_F^2$.
  2. Coincidencia Log-Sum-Exp (LSE): $\| \hat{\ell} - \ell \|_F^2$. Esto asegura que la "masa de atención" asignada al bloque de contexto sea correcta, lo cual es crítico para la combinación con futuros tokens de decodificación.

• Algoritmo de Optimización Alternada:

  • Optimización de Claves ($K_c$): Se utiliza L-BFGS (un método cuasi-Newton). Esto es crucial porque el paisaje de pérdida de la función softmax es no convexo y tiene valles estrechos; los optimizadores de primer orden como Adam tienden a quedar atrapados en mínimos locales pobres.
  • Solución de Valores ($V_c$): Dadas las claves fijas, la solución óptima para los valores se obtiene analíticamente mediante regresión de mínimos cuadrados con penalización Ridge (solución en forma cerrada).
  • Se alternan pasos: actualizaciones de $K$ con L-BFGS y resolución de $V$ cada pocos pasos.

• Construcción de Consultas de Entrenamiento:
  Dado que las consultas futuras no están disponibles durante la compresión, se utilizan:

  1. Consultas de retención (reales).
  2. Consultas sintéticas futuras: Se extraen vectores de contenido del contexto, se les aplica la inversión de RoPE y luego se les reasignan posiciones futuras ($N, N+1, \dots$). El muestreo uniforme de estos vectores ha demostrado ser la estrategia más robusta.

• Asignación Adaptativa de Presupuesto:
  Se introduce un "vuelo de prueba" (pilot run) rápido con asignación uniforme para medir el MSE por capa y por cabeza de KV.

  • Estrategia: Se redistribuye el presupuesto total de pares comprimidos ($k$) proporcionalmente a la raíz cuadrada del MSE del piloto ($MSE^{0.5}$).
  • Objetivo: Asignar más recursos a capas y cabezas difíciles de comprimir y menos a las fáciles. El uso de la raíz cuadrada evita que las capas con errores extremos (outliers) consuman todo el presupuesto.

3. Contribuciones Clave

1. Reformulación de Distilación: Cambia el paradigma de selección discreta a optimización continua en $\mathbb{R}^d$, eliminando la búsqueda combinatoria sobre posiciones.
2. Optimizador Híbrido: Propone un optimizador alternado (L-BFGS para claves + Mínimos Cuadrados para valores) que logra una reducción significativa en la divergencia KL.
3. Asignación Adaptativa Granular: Introduce una asignación de presupuesto basada en un piloto a nivel de capa y cabeza, mejorando la calidad sin costo adicional en la inferencia.
4. Análisis de Dificultad: Demuestra que la dificultad de compresión es altamente no uniforme (varía hasta 100x entre capas y 467x entre cabezas dentro de una misma capa), justificando la necesidad de asignación dinámica.

4. Resultados Experimentales

Los experimentos se realizaron en Qwen2.5-1.5B-Instruct con contextos de 2048 tokens.

• Rendimiento General:

  • KVSCULPT reduce la divergencia KL entre 3.5x y 4.1x en comparación con la mejor línea base de evicción (Select+Fit) en ratios de compresión $r \in \{0.3, 0.5, 0.7\}$.
  • En secuencias "fáciles", logra compresión casi sin pérdida (KL < 0.01), superando a las líneas base en más de 40x.

• Importancia de las Claves Continuas:

  • La optimización conjunta de valores manteniendo las claves en posiciones originales (Joint Optimization) apenas mejora sobre la evicción simple.
  • La ventaja decisiva proviene de permitir que las claves se muevan libremente en el espacio de incrustación para "resumir" el efecto combinado de múltiples tokens originales.

• Optimizador:

  • Reemplazar L-BFGS por Adam degrada el rendimiento entre 17x y 95x en MSE por capa, confirmando la necesidad de métodos de segundo orden para este paisaje de pérdida.

• Asignación Adaptativa:

  • La asignación basada en el piloto reduce el KL medio en un 25% (factor de 1.3x) sin costo de inferencia.
  • La asignación por cabeza (dentro de la capa) aporta mejoras adicionales, especialmente en secuencias difíciles donde la asimetría entre cabezas es extrema.

5. Significado y Limitaciones

• Significado: KVSCULPT demuestra que la compresión del caché KV no necesita ser una operación de selección o fusión de datos existentes, sino un problema de optimización de representaciones. Esto permite una preservación mucho más fiel del comportamiento de atención, especialmente en compresiones agresivas.
• Escenario de Despliegue: Está diseñado para entornos offline o amortizados (ej. compresión de un documento largo que se consultará muchas veces). El costo de compresión (~170s en una A100 para 2048 tokens) no es viable para decodificación en línea de un solo paso, pero se amortiza en múltiples pasos de generación.
• Limitaciones:
  • Los errores se acumulan a través de las capas del transformador (propagación de error), lo que limita la calidad en secuencias muy difíciles incluso con optimización por capa casi óptima.
  • El enfoque actual es por capa; el siguiente paso natural sería una optimización "consciente de la cascada" que considere los errores de capas anteriores.
  • Validación limitada a un modelo (Qwen2.5-1.5B) y longitud de contexto fija (2048).

En conclusión, KVSCULPT establece un nuevo estado del arte al tratar la compresión del caché como un problema de aprendizaje continuo, superando significativamente a los métodos basados en selección y ofreciendo una vía para la inferencia de LLMs de contexto largo más eficiente.