KV Cache Compression Through the Lens of Transform Coding

Resumen Técnico: Compresión del Caché KV a través de la Lente de la Codificación por Transformada

Planteamiento del Problema
En la inferencia de contexto largo para Modelos de Lenguaje de Gran Escala (LLMs), el caché de Clave-Valor (KV), que almacena las representaciones de los tokens pasados para acelerar la generación, se convierte en el cuello de botella dominante de la memoria. A medida que la longitud del contexto crece linealmente, el tamaño del caché se vuelve prohibitivo para el despliegue bajo restricciones realistas de memoria y latencia. Si bien los métodos de cuantización existentes (por ejemplo, KIVI, KVQuant) reducen la memoria representando el caché KV con tipos de datos de menor precisión, normalmente minimizan el error de reconstrucción bruto del propio caché. Estos métodos a menudo no tienen en cuenta cómo los errores de cuantización se propagan a través del mecanismo de atención para afectar el resultado final. En consecuencia, pueden asignar bits de manera ineficiente, preservando información que tiene poco impacto en la predicción del modelo mientras descartan información crítica.

Metodología: Codificación por Transformada Consciente de la Atención (AATC)
Los autores proponen la Codificación por Transformada Consciente de la Atención (AATC), un método que tiende un puente entre la teoría del procesamiento de señales y la inferencia de LLM. El enfoque se fundamenta en tres componentes principales:

1. Fundamento Teórico (Descomposición de la Distorsión):
   El artículo deriva una nueva medida de distorsión consciente de la atención. Bajo un modelo de cuantización de ruido blanco, los autores demuestran que la distorsión del error cuadrático esperado de la salida de la atención se descompone en contribuciones aditivas de clave y valor. Crucialmente, esta distorsión se factoriza a través de tokens y canales.

   • Distorsión de Clave ($D_K$): Ponderada por los pesos de atención al cuadrado ($a_i^2$), el residuo de la proyección de valor ($\|v_i W_O - o\|^2$) y los componentes de consulta al cuadrado ($q_c^2$).
   • Distorsión de Valor ($D_V$): Ponderada por los pesos de atención al cuadrado ($a_i^2$) y las normas de la proyección de salida ($\|W_{O,c}\|^2$).
     Esta descomposición revela que no todos los canales o tokens contribuyen de igual manera al error de salida; algunos son "conscientes de la atención" y requieren mayor precisión, mientras que otros pueden comprimirse agresivamente.

2. Codificación por Transformada (Blanqueamiento):
   Inspirado en la codificación por transformada clásica (por ejemplo, JPEG), el método aplica primero una transformada lineal para decorrelacionar las características de clave y valor. Utilizando un conjunto de calibración, los autores computan una transformada de blanqueamiento (mediante la descomposición de Cholesky y la SDE de las matrices de peso) que transforma el caché KV en una representación donde las dimensiones no están correlacionadas y están ordenadas por fuerza (varianza). Este paso concentra la energía de la señal en menos coeficientes, haciendo que la subsiguiente asignación de bits sea más efectiva.

3. Asignación de Bits (Agua Inversa/Reverse Waterfilling):
   Una vez que los datos han sido blanqueados, el método emplea el llenado de agua inverso (reverse waterfilling) para asignar bits a través de los canales transformados. Este problema de optimización minimiza la distorsión consciente de la atención sujeta a un presupuesto global de bits.

   • Los canales con un alto "peso consciente de la atención" (derivado de la fórmula de distorsión) reciben más bits.
   • Los canales con bajos pesos reciben menos bits o cero bits (realizando efectivamente una reducción de rango).
     La asignación se realiza por separado para claves y valores, reconociendo sus roles asimétricos en la atención.

Contribuciones Clave

• Medida de Distorsión: La derivación de una medida de distorsión consciente de la atención que separa los errores de cuantización en factores dependientes de tokens y canales, unificando puntos de vista previamente desconectados sobre la expulsión de tokens, la cuantización y la proyección de bajo rango.
• Perspectiva Unificada: Un marco teórico que relaciona los métodos existentes de compresión de caché KV (como KIVI, PALU y KVTC) con factores específicos dentro de la derivada descomposición de la distorsión.
• Algoritmo AATC: El desarrollo de un algoritmo práctico que combina transformadas de blanqueamiento con el llenado de agua inverso para lograr una asignación de bits por canal óptima bajo la métrica de distorsión consciente de la atención.
• Validación Empírica: Evaluación exhaustiva que demuestra que AATC logra una precisión casi sin pérdidas en altas tasas de compresión donde otros métodos se degradan.

Resultados Experimentales
El método fue evaluado en Llama-3.1-8B-Instruct y Qwen-2.5-7B-Instruct a través de benchmarks estándar incluyendo LongBench, RULER, GSM8K, MMLU-Pro y MATH-500.

• Eficiencia de Compresión: AATC logra aproximadamente una compresión de 5.8× (reduciendo la memoria de ~1.07 GB a ~184 MB para Llama-3.1-8B) mientras mantiene un rendimiento estadísticamente indistinguible de la línea base de precisión completa (FP16) en todos los benchmarks.
• Comparación: A diferencia de las líneas base (KIVI, KVQuant, PALU), que se degradan en al menos algunas configuraciones (particularmente en la recuperación de contexto largo en Qwen), AATC mantiene una precisión casi sin pérdidas.
• Beneficio de la Conciencia de la Atención: Un estudio de ablación ("AATC var-only") mostró que, si bien el método completo y la variante de solo varianza funcionan de manera similar en Llama con compresión moderada, el método completo consciente de la atención supera significativamente a la versión de solo varianza en Qwen bajo compresión agresiva (presupuesto de 2 bits), particularmente en escenarios de contexto largo (RULER-32k).

Significancia y Reivindicaciones
El artículo afirma que, al cambiar el objetivo de optimización de la reconstrucción bruta del caché a la distorsión consciente de la atención, se pueden derivar sistemáticamente estrategias de asignación de bits que preserven el rendimiento del modelo incluso en altas tasas de compresión. Los autores posicionan a AATC como un enfoque unificador que aprovecha herramientas clásicas de procesamiento de señales (codificación por transformada y llenado de agua inverso) para resolver el problema específico de los cuellos de botella de la inferencia de LLM. Afirman que su método proporciona una alternativa "casi sin pérdidas" a las técnicas de compresión existentes, permitiendo el despliegue de LLMs basados en transformadores en entornos con restricciones de memoria sin sacrificar la calidad del razonamiento o la recuperación de contexto largo.

Limitaciones
Los autores señalan que su análisis teórico asume la independencia de los tokens, lo cual puede no cumplirse estrictamente en el lenguaje natural, aunque observan que esto no impacta negativamente la calidad de la codificación en la práctica. Además, la implementación actual carece de optimizaciones a nivel de sistema (por ejemplo, kernels de CUDA dedicados), y la evaluación está limitada a dos familias de modelos (Llama y Qwen).