Resumen Técnico: RATEQUANT

Enunciado del Problema

Los Modelos de Lenguaje Grande (LLM) requieren almacenar en caché todos los pares clave-valor (KV) previamente calculados durante la generación, creando un cuello de botella en la memoria que escala linealmente con la longitud de la secuencia, el tamaño del lote y la profundidad del modelo. Aunque la cuantización de la caché KV reduce esta huella, los métodos existentes asignan uniformemente el mismo ancho de bits a cada cabeza de atención. Este enfoque ignora la heterogeneidad bien documentada en la importancia de las cabezas.

Los intentos recientes de asignación de precisión mixta a nivel de capa o canal dependen de reglas heurísticas vinculadas a cuantizadores específicos, careciendo de garantías teóricas y transferibilidad. Los autores identifican un obstáculo más profundo, previamente no reconocido: la desajuste del modelo de distorsión. Diferentes cuantizadores siguen curvas de distorsión-tasa distintas, modeladas como $D(b) = \alpha \beta^{-b}$. La tasa de decaimiento $\beta$ varía significativamente entre diseños (por ejemplo, de 3.6 en TURBOQUANT a 5.3 en QuaRot). Aplicar ingenuamente una estrategia de asignación de precisión mixta derivada del modelo de distorsión de un cuantizador a otro invierte el orden de la ganancia marginal de las cabezas de atención, dando a menudo como resultado un rendimiento peor que la cuantización uniforme.

Metodología

RATEQUANT aborda estos desafíos formalizando la asignación de bits de la caché KV por cabeza como un problema de optimización de distorsión-tasa. El marco consta de cuatro fases:

1. Formulación del Problema y Asignación Óptima

Los autores modelan el Error Cuadrático Medio (MSE) de la cuantización por cabeza como una función exponencial del ancho de bits: $D(b) = \alpha \beta^{-b}$. El objetivo es minimizar la suma ponderada de distorsiones a través de $N$ grupos de cuantización (cabezas) sujeto a un presupuesto total de bits $B$:
$$\min_{b} \sum_{i=1}^{N} w_i \cdot D(b_i) \quad \text{sujeto a} \quad \sum b_i = B$$
donde $w_i$ representa la sensibilidad de la cabeza $i$.

Utilizando la estacionariedad lagrangiana, los autores derivan una solución de forma cerrada mediante relleno inverso de agua:
$$b^*_i = \bar{b} + \frac{\ln w_i - \ln \bar{w}}{\ln \beta}$$
donde $\bar{b}$ es el ancho de bits promedio y $\bar{w}$ es la media geométrica de las sensibilidades. Esta solución asigna más bits a las cabezas con mayor sensibilidad. Teóricamente, la reducción de distorsión alcanzable está acotada por la relación entre la media aritmética y la media geométrica (AM/GM) de las sensibilidades de las cabezas.

2. Estimación de Sensibilidad

Para determinar los pesos $w_i$, RATEQUANT emplea sensibilidad basada en gradientes. Calcula las normas cuadradas de los gradientes de las salidas de proyección KV con respecto a la pérdida de modelado de lenguaje causal:
$$w_{l,h} = \mathbb{E}_{x \sim \mathcal{D}} \left[ \frac{1}{T} \sum_{t=1}^{T} \left\| \frac{\partial \mathcal{L}}{\partial K_{l,h,t}} \right\|^2 \right]$$
Los autores demuestran que los proxies basados en gradientes son superiores a las métricas basadas en activaciones (por ejemplo, magnitud de canal), que miden la amplificación del error en lugar del impacto directo en la pérdida. La sensibilidad basada en gradientes revela un patrón en forma de U de importancia (alta sensibilidad en capas tempranas y tardías), mientras que las métricas basadas en activaciones muestran un aumento monótono con la profundidad.

3. Asignación Entera y Extensiones

Dado que los anchos de bits deben ser enteros, la solución continua se refina utilizando un algoritmo de ganancia marginal voraz (Algoritmo 1). Este algoritmo asigna iterativamente bits a la cabeza que ofrece la mayor reducción en la distorsión ponderada hasta agotar el presupuesto.

Dos extensiones críticas hacen que RATEQUANT sea agnóstico al cuantizador:

• Calibración de Distorsión: En lugar de asumir un $\beta$ universal, RATEQUANT ajusta un modelo de distorsión por cuantizador ($\alpha, \beta$) utilizando un pequeño conjunto de calibración (16 secuencias). Esto previene el modo de fallo de "desajuste del modelo de distorsión".
• Separación K/V: Reconociendo que las claves y los valores a menudo utilizan esquemas de cuantización diferentes (por ejemplo, simétrica por canal para claves frente a asimétrica por token para valores en KIVI), el marco los trata como $2N$ componentes independientes con curvas de distorsión distintas. Esto permite que el algoritmo descubra presupuestos de bits asimétricos (por ejemplo, asignar más bits a las claves que a los valores).

4. Pipeline

El proceso implica costos fuera de línea únicos:

1. Calibración de Sensibilidad: Pasos hacia adelante y hacia atrás en un conjunto de datos pequeño (~1.6s para un modelo de 8B).
2. Modelado de Distorsión: Ajuste de $\alpha$ y $\beta$ en 5 anchos de bits (<0.1s).
3. Asignación: Ejecución del algoritmo voraz (<0.01s).
4. Inferencia: Aplicación de los anchos de bits estáticos con cero sobrecarga en tiempo de ejecución.

Resultados Clave

Se realizaron experimentos en modelos Qwen3 (4B, 8B, 32B) utilizando la perplejidad (PPL) de WikiText-2 y tareas posteriores (ARC-Challenge, HellaSwag, etc.).

• Resolución del Desajuste: Aplicar un modelo de distorsión desajustado (por ejemplo, usar el $\beta$ de TURBOQUANT para KIVI) degradó el PPL de KIVI de 49.3 a 87.0 a 2.5 bits promedio. RATEQUANT calibrado resolvió esto, reduciendo el PPL a 14.9 (una reducción del 70%).
• Ganancias de Rendimiento:
  • KIVI (2.5 bits): RATEQUANT mejoró el PPL de 49.3 a 14.9.
  • QuaRot (2.5 bits): Mejoró el PPL en 6.6 puntos.
  • TURBOQUANT (4.0 bits): Recuperó el 66% de la degradación inducida por la cuantización para Qwen3-8B.
• Separación K/V: Para KIVI a 2.5 bits, la asignación óptima descubierta por RATEQUANT fue de 2.85 bits para claves y 2.15 bits para valores, superando significativamente a la asignación conjunta.
• Eficiencia: El método no añade sobrecarga en tiempo de ejecución. La calibración toma menos de 2 segundos en una sola GPU.
• Precisión Posterior: A 4.0 bits, RATEQUANT logró una precisión promedio del 64.2% en tareas posteriores, igualando casi la línea base FP16 (64.4%) y recuperando el 89.8% de la brecha de precisión.

Significado y Afirmaciones

El artículo afirma que RATEQUANT es el primer marco que aplica la teoría de distorsión-tasa a la cuantización de la caché KV a granularidad por cabeza, abordando simultáneamente las diferencias estructurales entre claves y valores.

El hallazgo central es que el desajuste del modelo de distorsión es un modo de fallo crítico en la asignación de precisión mixta. Sin calibrar la curva de distorsión ($\beta$) al cuantizador base específico, las estrategias de precisión mixta pueden ser perjudiciales. Al combinar:

1. Sensibilidad basada en gradientes (superior a los proxies basados en activaciones),
2. Calibración empírica de distorsión (para igualar el cuantizador), y
3. Separación K/V (para manejar características de error asimétricas),

RATEQUANT transforma la cuantización KV de precisión mixta en una capa de asignación general y agnóstica al cuantizador. Los autores postulan que la relación AM/GM de las sensibilidades de las cabezas sirve como un predictor práctico de cuándo la asignación de precisión mixta producirá ganancias significativas. El trabajo sugiere que la perspectiva de distorsión-tasa puede generalizarse a la asignación de recursos en otros componentes heterogéneos de redes neuronales más allá de las cachés KV.