Resumen Técnico: Inferencia Estadística y Medidas de Calidad de Cuantizaciones de Cache KV Inspiradas en TurboQuant

Enunciado del Problema

El cache KV (Key-Value) es el principal cuello de botella de memoria en la inferencia de transformadores, creciendo linealmente con la longitud del contexto. La compresión mediante cuantización es esencial, sin embargo, la cuantización escalar estándar (MSE) a menudo falla en preservar las propiedades geométricas requeridas para los mecanismos de atención. Aunque TurboQuant propone un pipeline de dos etapas (rotación aleatoria + cuantización escalar + boceto QJL de 1 bit para corrección de sesgo), la asignación óptima de bits entre los caches de Clave (K) y Valor (V) permanece poco clara. Específicamente, se desconoce si aplicar la corrección residual QJL a las Claves (además de los Valores) mejora o degrada la calidad de la atención, particularmente bajo presupuestos de bits variables y distribuciones de datos (por ejemplo, de bajo rango, con colas pesadas).

Las métricas de evaluación existentes, como el Error Cuadrático Medio (MSE) o incluso la divergencia KL, a menudo confunden errores de escala y dirección o fallan en capturar fallos catastróficos de "desplazamiento del argmax" donde una sola clave dominante mal cuantizada colapsa la distribución de atención.

Metodología

Marco Teórico

Los autores fundamentan su análisis en la geometría de esferas de alta dimensión. Establecen que para un vector distribuido uniformemente en una hipersfera $S^{d-1}$, la distribución marginal de las coordenadas sigue una distribución Beta $\text{Beta}(\frac{d-1}{2}, \frac{d-1}{2})$. A medida que la dimensión $d \to \infty$, esto converge a una Gaussiana.

• Desigualdad de Jensen y Softmax: El artículo destaca una no linealidad crítica: la función softmax amplifica la varianza en los productos internos de las Claves. Incluso si los errores de cuantización tienen media cero, la convexidad de la función exponencial ($e^x$) introduce un sesgo (sesgo de Jensen) que infla los pesos de atención para las claves con alta varianza de cuantización.
• Asimetría K-V: El tensor de Claves entra en el mecanismo de atención mediante productos internos (no lineales), mientras que el tensor de Valores entra mediante una suma ponderada lineal. En consecuencia, la inflación de varianza en las Claves es catastrófica, whereas la corrección de sesgo sin sesgo en los Valores es beneficiosa.

Esquemas de Cuantización

Tres esquemas se comparan bajo una restricción estricta y de presupuesto de bits igual ($n$ bits por vector):

1. KV: Cuantización escalar MSE (Lloyd-Max) tanto en K como en V. Sin rotación, sin boceto.
2. KQV: Aplica una Transformada de Walsh-Hadamard (WHT) a ambos. K se cuantiza mediante MSE escalar ($n$ bits). V se cuantiza mediante MSE escalar ($n-1$ bits) + boceto QJL de 1 bit.
3. QKQV: Aplica WHT a ambos. Tanto K como V utilizan ($n-1$ bits) + boceto QJL de 1 bit.

Métricas de Evaluación

Para ir más allá de las métricas escalares, el artículo introduce un Marco de Error 6D que descompone el error en Escala (SNR) y Dirección ($\phi$) para tres sitios: cache de Claves, cache de Valores y Salida ($T$).

• Plano Shannon-$\phi$: Grafica el error de escala frente al error de dirección.
• Instrumentos Estadísticos: U de Mann-Whitney (correlación biserial de rangos), Kolmogorov-Smirnov (separación de CDF) y Distancia de Energía (métrica 6D).
• Proxies Mecanicistas: Divergencia KL entre las distribuciones de softmax de referencia y cuantizadas, y recuperación de enrutamiento top-5.

Contribuciones Clave

1. Prueba Teórica del Daño del QJL en Claves

El artículo demuestra que reemplazar la cuantización escalar en las Claves con un boceto QJL de 1 bit infla la varianza del producto interno por un factor de al menos $\pi/2$ (aprox. 1.57) en relación con el residual escalar, y por un factor de $2\pi$ al comparar un esquema escalar completo de $n$ bits con un esquema escalar de $(n-1)$ bits + QJL de 1 bit.

• Mecanismo: Aunque QJL proporciona un estimador sin sesgo del residual, descarta la información de magnitud, reemplazándola con bits de signo. La varianza de este estimador es $\frac{\pi}{2} \frac{\|r\|^2}{d}$.
• Consecuencia: Esta varianza inflada se amplifica no linealmente por la exponencial de softmax (a través de la desigualdad de Jensen), lo que lleva a una distorsión sistemática de los pesos de atención. Por lo tanto, aplicar QJL a las Claves es estructuralmente inferior a la cuantización escalar bajo presupuestos justos.

2. La Asimetría K-V y el Esquema Óptimo (KQV)

Los resultados empíricos confirman que KQV (QJL solo en V) es superior a QKQV (QJL en ambos) en todas las distribuciones y presupuestos probados.

• Divergencia KL: En el presupuesto de 4 bits ($n=4$) que domina en la práctica, KQV logra una divergencia KL 2.25$\times$ menor que QKQV ($p < 10^{-23}$).
• Enrutamiento: Una menor KL en KQV se correlaciona directamente con una mayor recuperación de enrutamiento top-5 y un menor error de dirección de salida.
• Conclusión: El boceto QJL de 1 bit debe reservarse para el cache de Valores (donde corrige el sesgo sin penalidad de varianza) y nunca aplicarse al cache de Claves (donde introduce una varianza fatal).

3. La Patología de Bajo Rango

El artículo identifica un modo de fallo crítico para las distribuciones de Claves de bajo rango que las métricas escalares pasan por alto.

• Marginal vs. Conjunta: Mientras que WHT hace que las coordenadas marginales sigan una distribución Beta (coincidiendo con el código), no destruye la estructura conjunta de bajo rango.
• KL Catastrófico: En altas dimensiones ($d=1024$), las claves de bajo rango producen una divergencia KL órdenes de magnitud peor que las claves de colas pesadas, a pesar de tener histogramas marginales idénticos. El softmax amplifica pequeños errores en las claves dominantes dentro del subespacio, causando fallos de "desplazamiento del argmax".
• Implicación: La cuantización basada en rotación estándar es insuficiente para matrices de atención de bajo rango; la estructura conjunta debe abordarse.

4. Cruce Dependiente del Presupuesto

Un hallazgo contraintuitivo es un cruce en la calidad de reconstrucción geométrica de K:

• En presupuestos bajos ($n \in \{2, 3\}$), QKQV (QJL en K) a veces logra una mejor reconstrucción geométrica de K (menor error de dirección) que KQV.
• Sin embargo, esta mejora geométrica no se traduce en una mejor calidad de atención (KL o enrutamiento) porque la inflación de varianza aún daña el softmax.
• En $n=4$, KQV gana en todas las medidas. Este cruce es invariante al rango y al peso de la cola, sugiriendo un compromiso fundamental tasa-distorsión entre la precisión escalar y la corrección residual JL.

Resultados

• Dominio de KQV: En $n=4$, KQV supera a QKQV y KV en divergencia KL, error geométrico de K y distancia 6D en todas las distribuciones (Gaussiana, cola pesada, bajo rango, enfocada).
• QJL en Claves es Dañino: Reemplazar bits escalares con QJL en las Claves aumenta consistentemente la divergencia KL debido a la inflación de varianza, excepto en escenarios específicos de reconstrucción geométrica de bajo presupuesto que no benefician la salida final.
• Robustez a Colas Pesadas: WHT proporciona protección incondicional contra la concentración marginal (colas pesadas). KQV mantiene un error de dirección estable ($\epsilon_{dir}^K \approx 0.027$) incluso a medida que aumenta el peso de la cola ($\nu \to 1.05$), mientras que KV falla catastróficamente.
• Fallo de Bajo Rango: Las distribuciones de bajo rango causan la degradación más severa en la calidad de la atención, superando con creces los problemas de colas pesadas, debido a la preservación de la estructura del subespacio conjunto a pesar de la normalización marginal.

Significado y Afirmaciones

El artículo afirma proporcionar una explicación estadística y geométrica rigurosa de por qué la aplicación original de TurboQuant de QJL a las Claves (si se extendiera simétricamente) sería perjudicial.

1. Refinamiento de Métricas: Argumenta que la KL media es una métrica optimista; la KL del peor caso y la recuperación de enrutamiento son las verdaderas medidas de fiabilidad para la cuantización de cache KV.
2. Principio de Diseño: Establece que los caches K y V requieren estrategias de cuantización fundamentalmente diferentes debido a la no linealidad de softmax. QJL es una herramienta de corrección de sesgo adecuada para salidas lineales (V) pero una herramienta de inyección de varianza que daña las salidas no lineales (K).
3. Limitación de las Garantías Marginales: Demuestra que asegurar que las distribuciones marginales coincidan con el código (mediante WHT) es necesario pero insuficiente; la estructura conjunta (específicamente el bajo rango) dicta el modo de fallo último.
4. Problema Abierto: El artículo identifica el cruce dependiente del presupuesto en la reconstrucción geométrica como un problema abierto de tasa-distorsión, señalando que, aunque QKQV gana geométricamente en bits muy bajos ($n=2,3$), pierde funcionalmente debido a la amplificación de varianza por Jensen.

El trabajo concluye que una cuantización robusta de cache KV requiere el manejo explícito de estructuras de bajo rango (por ejemplo, mediante SVD o códigos conscientes del subespacio) y que el esquema KQV (WHT + Escalar en K, WHT + Escalar + QJL en V) es la configuración óptima para los presupuestos de bits actuales.