Resumen Técnico: HeadQ: Distorsión Visible para el Modelo y Corrección en el Espacio de Puntuación para la Cuantización de KV-Cache

1. Declaración del Problema

Los métodos actuales de cuantización de KV-cache optimizan predominantemente para la reconstrucción del espacio de almacenamiento, minimizando típicamente el Error Cuadrático Medio (MSE) entre las claves y los valores originales y cuantizados. Sin embargo, los autores argumentan que este objetivo no se alinea con cómo los mecanismos de atención consumen realmente los datos de la caché.

• Las Claves se leen a través de logits (puntuaciones consulta-clave), lo que significa que solo proyecciones lineales específicas de las claves afectan la distribución de atención.
• Los Valores se leen a través de una lectura ponderada por atención ($\sum p_j v_j$).

El artículo postula que el error persistente de la caché debe medirse en coordenadas visibles para el modelo en lugar de coordenadas de almacenamiento euclidianas crudas. Una perturbación puede tener un gran MSE de almacenamiento pero cero impacto en la atención (si se encuentra en una dirección "nula para softmax"), o un pequeño MSE de almacenamiento pero un daño catastrófico en la atención (si afecta la dirección "visible para la puntuación"). Esta discrepancia conduce a estrategias de cuantización subóptimas que desperdician bits en componentes irrelevantes mientras fallan en proteger aquellos críticos para el comportamiento.

2. Fundamento Teórico: Geometría Visible para el Modelo

El artículo establece un marco geométrico que distingue entre el error de almacenamiento y el error de comportamiento:

2.1 Espacio de Puntuación de Claves (Geometría de Fisher)

Para una consulta fija $q$, una perturbación de clave $\delta k_j$ induce una perturbación de puntuación $\delta s_j = q^\top \delta k_j / \sqrt{d}$.

• Direcciones Nulas para Softmax: Una traducción compartida por token (añadir un vector constante $\mu$ a todas las claves) añade una constante a todas las puntuaciones. Dado que softmax es invariante a desplazamientos constantes, esta perturbación es nula para softmax y causa cero daño en la atención, a pesar de un potencial MSE de almacenamiento grande.
• Espacio Visible para la Puntuación: La atención depende únicamente de los coeficientes activos centrados de la perturbación de la clave con respecto a la consulta.
• Información de Fisher: La geometría local de segundo orden de la distribución softmax se captura mediante la Matriz de Información de Fisher. El artículo deriva que la divergencia KL entre las distribuciones de atención original y perturbada es localmente proporcional a una varianza ponderada por Fisher de los errores de puntuación (Ecuación 6), y no al MSE crudo de la clave.

2.2 Geometría de Lectura de Valores

Los valores se consumen mediante una suma ponderada. Para pesos de atención fijos $p$, el error de lectura es $\Delta o = \sum p_j \Delta v_j$.

• El artículo propone una distorsión de token ponderada por $A^2$ como el sustituto apropiado para el error de valor. A diferencia del MSE de valor crudo, esta métrica pondera la distorsión de cada token por el cuadrado de su probabilidad de atención ($p_j^2$), reflejando el hecho de que los errores en los tokens altamente atendidos contribuyen más al error de salida.

3. Metodología: HeadQ

HeadQ es un método de cuantización del lado de las claves diseñado para preservar el residuo "visible para la puntuación" en lugar del vector de clave crudo.

Descripción General del Algoritmo

1. Calibración: Para cada capa y cabeza, se aprende una base de bajo rango $U_h$ (rango $r$) mediante PCA a partir de consultas de calibración (post-RoPE).
2. Escritura en Caché:
   • Un cuantizador estándar de bajo bit $B_b$ genera una clave base $\hat{k}^{base}_j$.
   • Se calcula el residuo $e_j = k_j - \hat{k}^{base}_j$.
   • El residuo se proyecta sobre la base aprendida: $z_j = Q_q(U_h^\top e_j)$, donde $Q_q$ es un cuantizador lateral de bajo bit.
   • La caché almacena la clave base y el código lateral $z_j$.
3. Lectura de Caché (Inferencia):
   • La consulta actual $q$ se utiliza para descuantizar el código lateral.
   • Se aplica una corrección aditiva de logits:
     $$\hat{s}_j = \frac{q^\top \hat{k}^{base}_j}{\sqrt{d}} + \frac{q^\top U_h \hat{z}_j}{\sqrt{d}}$$
   • Esta corrección se dirige directamente al error de puntuación, evitando la necesidad de reconstruir el vector de clave completo.

Distinción Clave: HeadQ no almacena la consulta actual; almacena un código residual en una base aprendida mediante calibración que la consulta de tiempo de ejecución puede decodificar en una corrección de logits.

4. Contribuciones Clave

1. Formulación de la Distorsión Visible para el Modelo: El artículo define la distorsión de claves como "error visible para la puntuación módulo desplazamientos constantes" y la distorsión de valores como "error visible para la lectura con un sustituto diagonal $A^2$". Esto explica por qué el MSE crudo es un mal proxy para el rendimiento de comportamiento.
2. Algoritmo HeadQ: Un algoritmo práctico de escritura/lectura de caché que almacena un código lateral residual de bajo rango en una base de consultas y lo aplica como una corrección aditiva de logits.
3. Cadena de Evidencia Falsable:
   • Superioridad de la Métrica: El error en el espacio de Fisher/puntuación predice la divergencia KL de atención significativamente mejor (ganancia de Spearman $\rho \approx 0.62\text{--}0.69$) que el MSE crudo de claves en seis modelos.
   • Contraejemplos: Experimentos emparejados muestran que un cuantizador con menor MSE de almacenamiento puede tener mayor KL de atención y error de salida, demostrando que el MSE de almacenamiento no es un proxy suficiente.
   • Intervenciones en el Espacio Nulo: Restaurar un componente de media "nulo para softmax" (que la atención densa ignora) aumenta drásticamente el daño de Fisher para cuantizadores de cuadrícula fija no equivariantes, demostrando que el desperdicio de coordenadas de almacenamiento rompe el rendimiento de comportamiento.
   • Implementabilidad: Las bases de PCA de consultas aprendidas a partir de datos de calibración aproximan con éxito el residuo activo, mientras que bases aleatorias de igual presupuesto fallan.
   • Direccionalidad: Invertir el signo de la corrección HeadQ (control de "signo incorrecto") causa un fallo catastrófico, confirmando que el método funciona corrigiendo el objeto de puntuación específico en la dirección correcta.

5. Resultados Experimentales

Los experimentos se realizaron en WikiText-103 utilizando seis modelos (GPT-2, Pythia, Qwen2, Qwen2.5, TinyLlama, Mistral).

• Experimentos Solo-K (Valores Densos):
  • En las filas de 2 bits más dañadas, HeadQ (rango 4) elimina aproximadamente 84–94% de la perplejidad (PPL) excesiva en comparación con una línea base tipo KIVI para modelos como Qwen2, Pythia y GPT-2.
  • TinyLlama muestra una mejora material; Mistral ya tiene un rendimiento cercano al denso en este protocolo.
• Composición KV Completa (2 bits):
  • Combinar HeadQ (lado de claves) con una política de valores ponderada por $A^2$ mejora el rendimiento en los seis modelos a una precisión de 2 bits.
• Análisis de Anomalías:
  • Pythia-160M mostró una correlación más débil entre Fisher y KL en regímenes específicos de baja entropía de "cambio de ruta" (donde pequeñas perturbaciones invierten el token top-1). Esta anomalía se localizó en un límite específico de la familia y se debilitó en puntos de control Pythia más grandes (410M, 1B).

6. Significado y Afirmaciones

El artículo afirma que el ajuste de operadores es la lección de diseño crítica para la cuantización de KV-cache:

• Las claves son objetos de puntuación; la compresión debe preservar los residuos visibles para la consulta y evitar desperdiciar resolución en desplazamientos nulos para softmax.
• Los valores son objetos de lectura; la compresión debe priorizar el error ponderado por $A^2$.

Los autores declaran explícitamente que este trabajo es mecanicista y diagnóstico, no una afirmación de un sistema de servicio listo para producción. No afirman haber optimizado kernels empaquetados, latencia o rendimiento. El significado radica en:

1. Corregir el objetivo teórico de la cuantización de KV-cache de la reconstrucción de almacenamiento a la distorsión visible para el modelo.
2. Proporcionar un método implementable (HeadQ) que logra una recuperación sustancial de PPL en regímenes de bajo bit al dirigirse a la geometría de error correcta.
3. Demostrar que las métricas de almacenamiento estándar (MSE) pueden ser engañosas desde el punto de vista del comportamiento, haciendo necesaria una transición hacia proxies del espacio de puntuación.

El artículo concluye que el siguiente paso es extender esta visión ajustada a operadores hacia un códec KV completo totalmente optimizado, en lugar de tratar el MSE de almacenamiento como el objetivo predeterminado.