SPECTRA: Pushing the KV Cache Beyond the 2-Bit Cliff via Spectral Transform Coding

Resumen Técnico: SPECTRA

Planteamiento del Problema

En la era de los Modelos de Lenguaje de Gran Escala (LLM) agénticos, la inferencia está dominada cada vez más por entradas de contexto largo, que van desde bases de código completas hasta diálogos de múltiples turnos. A medida que la longitud del contexto crece, el caché de Clave-Valor (KV)—el almacén de ejecución de las claves y valores de atención—se convierte en el cuello de botella principal, superando a menudo el tamaño de los pesos del modelo. Este caché debe residir en una memoria rápida y ser releído en cada paso de generación, lo que limita tanto la longitud máxima del contexto como el tamaño del lote (batch size) que un GPU puede soportar.

Los métodos de compresión existentes generalmente se dividen en tres categorías: la eliminación de tokens (descartar tokens poco importantes), la proyección de bajo rango (reducir la dimensionalidad) y la cuantización (reducir la precisión). Si bien la cuantización es ampliamente utilizada, esta golpea un "acantilado de 2 bits" donde el rendimiento se degrada drásticamente. A 2 bits por valor (cuatro niveles), el rango dinámico limitado no puede acomodar la distribución con alta presencia de valores atípicos (outliers) del caché; las entradas grandes consumen los niveles disponibles, colapsando los valores más pequeños y críticos en ruido. Además, la cuantización estándar aplica un ancho de bits uniforme en todos los canales, ignorando el hecho de que los canales del caché KV están fuertemente correlacionados, lo que dificulta identificar qué canales específicos transportan la mayor cantidad de información.

Metodología: SPECTRA

Los autores proponen SPECTRA, un códec de sustitución directa (drop-in) que no requiere entrenamiento, que replantea la compresión del caché KV como un problema de tasa-distorsión. En lugar de fijar un presupuesto de antemano (por ejemplo, un rango o ancho de bits específico), SPECTRA asigna dinámicamente los bits a las partes del caché que más importan. El método se basa en tres observaciones clave derivadas de sondear LLMs preentrenados:

1. Concentración en la Base Correcta: En la base de canales crudos, los canales del caché KV están altamente correlacionados. Sin embargo, cuando se rotan hacia una base computada a partir de las propias estadísticas del caché (específicamente, la matriz de Gram de las activaciones), los canales se vuelven no correlacionados y la energía se concentra bruscamente: una pequeña fracción de los canales transporta la gran mayoría de la señal.
2. Asignación Ponderada por Importancia: La asignación de bits basada en la importancia del canal (energía) supera significativamente a la precisión uniforme, particularmente en el régimen inferior a 2 bits. Los canales con cero bits son efectivamente descartados, unificando la cuantización y la reducción de rango en un único mecanismo de control.
3. Importancia Ponderada por Activación: La importancia de un canal se determina por la fuerza con la que las activaciones de entrada lo impulsan, no meramente por la magnitud de sus pesos de proyección. Por lo tanto, la transformación debe estar ponderada por las estadísticas de activación ($G = E[h^\top h]$) en lugar de solo por las matrices de pesos.

El Pipeline de SPECTRA

SPECTRA implementa estas observaciones a través de un proceso de tres etapas:

1. Transformación Ponderada por G (Caché Latente):
   El método factoriza las proyecciones de clave y valor ($W_K, W_V$) utilizando una Descomposición en Valores Singulares (SVD) ponderada por la matriz de Gram de activación $G$. Al blanquear los pesos con $G^{1/2}$, el método encuentra una factorización de bajo rango que minimiza el error en la métrica de las activaciones reales producidas por el modelo, en lugar de en el espacio de pesos crudos. Esto genera una representación latente compacta y ordenada por energía ($C$) donde las coordenadas son ortogonales y están ordenadas por varianza.

   • Refinamiento: Se aplica una corrección de forma cerrada y sin entrenamiento a la proyección de salida ($W_O$) para compensar el residuo estructurado introducido por la truncación de bajo rango.

2. Cuantización de Tasa-Distorsión (Water-Filling):
   Una vez que el caché se transforma en este espacio latente ordenado por energía, SPECTRA aplica un algoritmo de "llenado de agua inverso" (reverse water-filling) para asignar bits. Basándose en las varianzas marginales medidas de las coordenadas latentes, el algoritmo asigna más bits a los canales de alta varianza (alta energía) y menos (hasta cero) a los canales de baja varianza.

   • Agrupación: Para mantener una estructura de caché de solo adición (append-only) adecuada para la inferencia de flujo (streaming), las coordenadas adyacentes se agrupan para compartir un único ancho de bits.

3. Integración del Sistema:
   Las matrices de transformación ($W_{down}, W_{up}$) y la corrección de $W_O$ se integran de nuevo en los pesos del modelo, convirtiendo a SPECTRA en un reemplazo directo que no requiere ajuste fino (fine-tuning). Durante la inferencia, los tokens son proyectados hacia abajo, cuantizados por token y añadidos al caché latente. Las claves y valores completos se reconstruyen sobre la marcha para la atención estándar.

Contribuciones Clave

• Visión Unificada de la Compresión: El artículo redefine la compresión del caché KV como un problema único de gasto de un presupuesto fijo de bits a través de los canales, unificando la proyección de bajo rango (eliminar canales) y la cuantización (reducir bits) como dos extremos de un mismo espectro de control.
• Transformación Ponderada por Activación: Identifica que minimizar el error de reconstrucción en la métrica de las activaciones del modelo (mediante una transformación ponderada por $G$) es superior a la factorización en el espacio de pesos, produciendo un latente que es tanto ortogonal de segundo momento como ordenado por energía.
• Códec SPECTRA: Una implementación práctica y sin entrenamiento que concentra el presupuesto de bits en los canales que portan la señal, logrando superar el "acantilado de 2 bits" que afecta a la cuantización uniforme.

Resultados Experimentales

Los autores evaluaron SPECTRA en Llama-3.1-8B, Mistral-7B y Qwen2.5-7B utilizando los benchmarks LongBench y RULER (Needle-in-a-Haystack).

• Rendimiento vs. Compresión:
  • Con una compresión de 4x, SPECTRA es casi sin pérdida, igualando o superando el rendimiento del modelo denso (fp16) y superando a baselines como KIVI, OTT y PolarQuant.
  • Con una compresión de 8x, donde los métodos de cuantización uniforme colapsan (cayendo significativamente por debajo de la línea base densa), SPECTRA sigue siendo competitivo. Para Llama-3.1-8B, SPECTRA logra un puntaje promedio de LongBench de 51.89 con 8.84x de compresión, mientras que los baselines como TurboQuant y RotateKV caen a ~47 y ~46 respectivamente.
  • SPECTRA sigue siendo utilizable hasta una compresión de 12x (por ejemplo, un puntaje de 48.81 con 11.12x para Llama), un régimen donde otros métodos fallan por completo.
• Capacidad de Memoria:
  • En un GPU H200 con 24GB de memoria, un modelo denso Llama-3.1-8B soporta un contexto de ~65k tokens. SPECTRA con 12x de compresión extiende esto a 793k tokens, multiplicando efectivamente la capacidad de contexto por más de 12 veces sin cambiar el hardware.
• Recuperación (Retrieval):
  • En tareas de Needle-in-a-Haystack, SPECTRA mantiene una precisión de recuperación casi perfecta hasta 12x de compresión, mientras que los baselines de 2 bits se degradan bruscamente en contextos largos.

Significancia y Reivindicaciones

El artículo sostiene que la palanca principal para la compresión del caché KV no es simplemente qué tan finamente cuantizar, sino qué almacenar. Al rotar el caché hacia un sistema de coordenadas derivado de las propias estadísticas del modelo, SPECTRA revela que la información está altamente concentrada. Esto permite al método gastar bits donde más importan, superando efectivamente el "acantilado de 2 bits" que plaga a la cuantización uniforme.

Los autores enfatizan que SPECTRA es una solución de sustitución directa y sin entrenamiento que no requiere preentrenamiento, ajuste fino ni cambios arquitectónicos. Aunque la implementación actual se centra en el ahorro de memoria (reconstruyendo claves/valores sobre la marcha), el artículo señala que la estructura del espacio latente abre teóricamente un camino para el ahorro de cómputo mediante la atención directa en el espacio latente, lo cual queda como trabajo futuro. La contribución central es demostrar que es posible empujar la compresión más allá del acantilado de 2 bits tratando el caché como una señal para ser codificada mediante principios de codificación de transformada, en lugar de un tensor que se recorta uniformemente.