Resumen Técnico: Ablación, Inferencia Estadística y Validación para la Compresión del KV-Cache

Declaración del Problema

La inferencia de Transformers a escala está limitada por el ancho de banda de memoria, siendo el acceso al caché de Clave-Valor (KV) lo que domina el tráfico total de memoria durante la generación de contextos largos. Aunque la cuantización de claves y valores a 2–4 bits por elemento es un enfoque estándar para reducir el tamaño del caché, evaluar estos esquemas de compresión presenta un desafío significativo. Las diferencias de calidad observadas suelen estar confundidas por las propiedades distribucionales del tráfico real de los Grandes Modelos de Lenguaje (LLM), efectos de hardware y suposiciones algorítmicas, lo que dificulta atribuir las ganancias o fallos de rendimiento a decisiones de diseño específicas.

Metodología

Los autores introducen un marco de evaluación controlado que consta de seis regímenes estadísticos sintéticos diseñados para aislar suposiciones estructurales específicas del pipeline de compresión (por ejemplo, la forma de la distribución, la alineación del subespacio K/V y la caída de los autovalores). Este marco se combina con una metodología de validación estadística para distinguir las diferencias algorítmicas sistemáticas de la varianza inducida por la implementación.

Configuración Experimental

• Esquemas Evaluados: Se comparan dos familias primarias de compresión:
  1. TurboQuant (TQ): Un enfoque ajeno a los datos que utiliza una Transformada de Walsh-Hadamard (WHT) aleatorizada para dispersar la energía de los tokens de manera uniforme, seguida de un código de libro de códigos (codebook) Beta Lloyd-Max ajeno a los datos. Opcionalmente, utiliza un boceto residual de 1 bit de Johnson-Lindenstrauss (QJL) en las rutas de la clave (K) o del valor (V).
  2. SpectralQuant (SQ): Un enfoque adaptativo a los datos que calibra una base propia (eigenbasis) por cabeza a partir de tokens representativos, asigna bits mediante el método de llenado de agua (water-filling) a las dimensiones de alta varianza, y utiliza un boceto QJL selectivo en las dimensiones efectivas superiores.
• Métricas: El estudio utiliza una métrica de error 2D ($d^2$) que combina el error cuadrático medio relativo (relMSE) y el error de dirección (similitud de coseno). También emplea métricas de aislamiento (relMSE_K, relMSE_V) y una geometría de error 6D para descomponer los errores en el caché K, el caché V y la salida final de atención.
• Validación: Para separar el rendimiento algorítmico del ruido de implementación, los autores comparan implementaciones en Python (oráculo) y HIP/GPU (producción) utilizando diferentes generadores de números aleatorios (PCG64 vs. mt19937). Se aplica la prueba de Kolmogorov-Smirnov (KS) para identificar diferencias sistemáticas frente a la varianza dependiente de la semilla.

Contribuciones Clave

1. Ablación Sistemática de Esquemas de Compresión

A través de una ablación completa de ocho variantes de TQ y múltiples configuraciones de SQ, el estudio identifica tres esquemas no dominados:

• Cuantización Escalar (Plain-KV): Sin rotación, efectivo para datos isotrópicos.
• TQ-KV: Rotación WHT con un codebook Beta Lloyd-Max.
• TQ-QKV: TQ-KV aumentado con un boceto residual QJL de 1 bit en la ruta de la clave (K).
  El estudio elimina las variantes de QJL en la ruta V, demostrando que aplicar QJL a los valores (que se suman sobre los tokens) es un heurístico que degrada consistentemente el rendimiento en comparación con las correcciones en la ruta K.

2. Validación Estadística e Desigualdad de Jensen

El documento introduce una metodología para separar las diferencias del codec del ruido de implementación. Un hallazgo crítico es la identificación de la desigualdad de Jensen como un mecanismo de amplificación del error.

• La no linealidad de la función softmax es convexa; por lo tanto, la varianza de un error de corrección de media cero ($\epsilon$) añadido a las puntuaciones de atención se amplifica exponencialmente ($E[\exp(s+\epsilon)] \approx \exp(s) \cdot \exp(\sigma^2/2)$).
• La corrección QJL en la ruta K introduce dicha varianza. El estudio muestra que la varianza de la brecha Python-HIP para TQ-QKV es significativamente mayor que para TQ-KV, confirmando que la ruta K es una firma directa de la desigualdad de Jensen amplificando la varianza de la puntuación. Por el contrario, las correcciones en la ruta V entran linealmente en la salida y no activan esta amplificación.

3. Rendimiento Específico por Régimen y Sensibilidad de Calibración

• Datos de Cola Pesada (Heavy-Tailed): Para datos con colas pesadas (por ejemplo, distribuciones Student-t), los métodos basados en la base propia (SQ) fallan catastróficamente. Las covarianzas de la muestra se desestabilizan por los valores atípicos (outliers), lo que conduce a una base calibrada sistemáticamente desalineada. Ningún aumento en el presupuesto de bits recupera esta pérdida. TQ-KV domina en estos regímenes porque la WHT dispersa los valores atípicos de cola pesada, haciendo que las marginales proyectadas sean aproximadamente gaussianas.
• Datos Estructurados: En regímenes de bajo rango y estructurados, SQ supera a TQ cuando K y V se calibran en conjuntos representativos separados. Sin embargo, el estudio encuentra que la asignación de bits mediante llenado de agua (water-filling) proporciona un beneficio insignificante sobre la asignación uniforme en todos los regímenes probados, ya que las razones de autovalores dentro del subespacio efectivo rara vez superan el umbral requerido para justificar una asignación no uniforme.
• Dimensión Semántica Efectiva ($d_{eff}$): El documento caracteriza $d_{eff}$ como algo que se autocalibra al presupuesto de calibración disponible en lugar de recuperar el rango real de los datos. Un único autovalor espurio puede causar un fallo catastrófico (salida anticorrelacionada), lo que requiere un criterio de brecha espectral para podar la dimensión estimada.

Resultados y Recomendaciones

Rendimiento Comparativo

• Presupuesto $b=2$: Ningún esquema único domina. Las variantes de SQ ganan en regímenes estructurados (bajo rango), TQ-KV gana en datos de cola pesada, y Plain-KV es competitivo en datos isotrópicos.
• Presupuesto $b \ge 4$: TQ-KV establece su dominio en la mayoría de los regímenes. Las variantes de SQ solo mantienen una ventaja en regímenes altamente estructurados donde la base propia calibrada puede concentrar los bits de manera eficiente, pero fallan en recuperarse de la desalineación de la base propia en datos de cola pesada.
• Almacenamiento vs. Calidad: Aunque SQ puede lograr costos de almacenamiento más bajos con presupuestos más altos debido a su capacidad para ignorar las dimensiones de la cola, el "suelo de la cola" (bits fijos para dimensiones no modeladas) se convierte en una restricción vinculante. Para secuencias largas, los errores acumulados de la cola erosionan la ventaja semántica de SQ, convirtiendo a TQ-KV en la opción más segura para tareas de generación.

Recomendaciones Prácticas

Los autores concluyen con directrices específicas para el despliegue:

1. Usar TQ-KV cuando la distribución de los datos sea desconocida, de cola pesada o esté sujeta a cambios de distribución, y para cualquier presupuesto superior a 2 bits.
2. Usar SQU-KV (SpectralQuant con asignación uniforme) solo cuando se sepa que los datos son estructurados (subespacio KV de bajo rango), el presupuesto sea bajo (2–3 bits), la longitud de la secuencia sea corta (cargas de trabajo de prefill) y se disponga de conjuntos de calibración separados para K y V.
3. Descartar el Llenado de Agua (Water-Filling): En datos sintéticos, la asignación uniforme tiene un rendimiento equivalente al llenado de agua.
4. Evitar QJL en K para Colas Pesadas: La corrección de la puntuación QJL en las claves amplifica la varianza a través del softmax (efecto Jensen) en datos de cola pesada y debe evitarse en esos regímenes.

Significancia

La importancia del documento reside en su metodología rigurosa y controlada para evaluar la compresión del caché KV. Al aislar las suposiciones estructurales mediante regímenes sintéticos y aplicar la inferencia estadística para separar la señal algorítmica del ruido de implementación, los autores proporcionan un marco de diagnóstico claro. Este trabajo va más allá del simple reporte de precisión para explicar por qué los métodos tienen éxito o fallan, revelando las limitaciones fundamentales de los métodos basados en la base propia en datos de cola pesada y el comportamiento no monotónico de la dimensionalidad efectiva bajo presupuestos de calibración finitos. Los hallazgos ofrecen una guía definitiva para seleccionar esquemas de compresión basados en las características de los datos y las restricciones operativas.