Résumé Technique : RATEQUANT

Énoncé du Problème

Les modèles de langage de grande taille (LLM) nécessitent la mise en cache de toutes les paires clé-valeur (KV) précédemment calculées durant la génération, créant un goulot d'étranglement mémoire qui évolue linéairement avec la longueur de la séquence, la taille du lot et la profondeur du modèle. Bien que la quantification du cache KV réduise cette empreinte, les méthodes existantes attribuent uniformément la même largeur de bits à chaque tête d'attention. Cette approche ignore l'hétérogénéité bien documentée de l'importance des têtes.

Les tentatives récentes d'allocation de précision mixte au niveau de la couche ou du canal reposent sur des règles heuristiques liées à des quantificateurs spécifiques, manquant de garanties théoriques et de transférabilité. Les auteurs identifient un obstacle plus profond, précédemment non reconnu : le décalage du modèle de distorsion. Différents quantificateurs suivent des courbes distorsion-débit distinctes, modélisées comme $D(b) = \alpha \beta^{-b}$. Le taux de décroissance $\beta$ varie considérablement selon les conceptions (par exemple, de 3,6 dans TURBOQUANT à 5,3 dans QuaRot). Appliquer naïvement une stratégie d'allocation de précision mixte dérivée du modèle de distorsion d'un quantificateur à un autre inverse l'ordre du gain marginal des têtes d'attention, conduisant souvent à des performances inférieures à celles de la quantification uniforme.

Méthodologie

RATEQUANT répond à ces défis en formalisant l'allocation de bits du cache KV par tête comme un problème d'optimisation distorsion-débit. Le cadre se compose de quatre phases :

1. Formulation du Problème et Allocation Optimale

Les auteurs modélisent l'erreur quadratique moyenne (MSE) de quantification par tête comme une fonction exponentielle de la largeur de bits : $D(b) = \alpha \beta^{-b}$. L'objectif est de minimiser la somme pondérée des distorsions sur $N$ groupes de quantification (têtes) sous contrainte d'un budget total de bits $B$ :
$$\min_{b} \sum_{i=1}^{N} w_i \cdot D(b_i) \quad \text{s.t.} \quad \sum b_i = B$$
où $w_i$ représente la sensibilité de la tête $i$.

En utilisant la stationnarité lagrangienne, les auteurs dérivent une solution sous forme close via le remplissage inverse (reverse waterfilling) :
$$b^*_i = \bar{b} + \frac{\ln w_i - \ln \bar{w}}{\ln \beta}$$
où $\bar{b}$ est la largeur de bits moyenne et $\bar{w}$ est la moyenne géométrique des sensibilités. Cette solution attribue plus de bits aux têtes ayant une sensibilité plus élevée. Théoriquement, la réduction de distorsion réalisable est bornée par le rapport entre la moyenne arithmétique et la moyenne géométrique (AM/GM) des sensibilités des têtes.

2. Estimation de la Sensibilité

Pour déterminer les poids $w_i$, RATEQUANT emploie une sensibilité basée sur le gradient. Il calcule les normes carrées des gradients des sorties de projection KV par rapport à la perte de modélisation du langage 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]$$
Les auteurs démontrent que les proxies basés sur le gradient sont supérieurs aux métriques basées sur l'activation (par exemple, l'amplitude du canal), qui mesurent l'amplification de l'erreur plutôt que l'impact direct sur la perte. La sensibilité basée sur le gradient révèle un motif en forme de U d'importance (forte sensibilité dans les couches précoces et tardives), tandis que les métriques basées sur l'activation montrent une augmentation monotone avec la profondeur.

3. Allocation Entière et Extensions

Puisque les largeurs de bits doivent être des entiers, la solution continue est affinée en utilisant un algorithme glouton de gain marginal (Algorithme 1). Cet algorithme attribue itérativement des bits à la tête offrant la plus grande réduction de distorsion pondérée jusqu'à épuisement du budget.

Deux extensions critiques rendent RATEQUANT agnostique au quantificateur :

• Calibration de la Distorsion : Au lieu de supposer un $\beta$ universel, RATEQUANT ajuste un modèle de distorsion par quantificateur ($\alpha, \beta$) en utilisant un petit ensemble de calibration (16 séquences). Cela prévient le mode de défaillance du « décalage du modèle de distorsion ».
• Séparation K/V : Reconnaissant que les clés et les valeurs utilisent souvent des schémas de quantification différents (par exemple, symétrique par canal pour les clés contre asymétrique par token pour les valeurs dans KIVI), le cadre les traite comme $2N$ composants indépendants avec des courbes de distorsion distinctes. Cela permet à l'algorithme de découvrir des budgets de bits asymétriques (par exemple, allouer plus de bits aux clés qu'aux valeurs).

4. Pipeline

Le processus implique des coûts hors ligne uniques :

1. Calibration de la Sensibilité : Passées avant et arrière sur un petit jeu de données (~1,6 s pour un modèle 8B).
2. Modélisation de la Distorsion : Ajustement de $\alpha$ et $\beta$ à 5 largeurs de bits (<0,1 s).
3. Allocation : Exécution de l'algorithme glouton (<0,01 s).
4. Inférence : Application des largeurs de bits statiques avec zéro surcharge d'exécution.

Résultats Clés

Des expériences ont été menées sur les modèles Qwen3 (4B, 8B, 32B) en utilisant la perplexité WikiText-2 (PPL) et des tâches en aval (ARC-Challenge, HellaSwag, etc.).

• Résolution du Décalage : L'application d'un modèle de distorsion inadapté (par exemple, utiliser le $\beta$ de TURBOQUANT pour KIVI) a dégradé la PPL de KIVI de 49,3 à 87,0 à 2,5 bits en moyenne. RATEQUANT calibré a résolu ce problème, réduisant la PPL à 14,9 (une réduction de 70 %).
• Gains de Performance :
  • KIVI (2,5 bits) : RATEQUANT a amélioré la PPL de 49,3 à 14,9.
  • QuaRot (2,5 bits) : Amélioration de la PPL de 6,6 points.
  • TURBOQUANT (4,0 bits) : Récupération de 66 % de la dégradation induite par la quantification pour Qwen3-8B.
• Séparation K/V : Pour KIVI à 2,5 bits, l'allocation optimale découverte par RATEQUANT était de 2,85 bits pour les clés et 2,15 bits pour les valeurs, surpassant significativement l'allocation conjointe.
• Efficacité : La méthode n'ajoute aucune surcharge d'exécution. La calibration prend moins de 2 secondes sur un seul GPU.
• Précision en Aval : À 4,0 bits, RATEQUANT a atteint une précision moyenne de 64,2 % sur les tâches en aval, correspondant presque à la référence FP16 (64,4 %) et récupérant 89,8 % de l'écart de précision.

Importance et Revendications

L'article revendique que RATEQUANT est le premier cadre à appliquer la théorie distorsion-débit à la quantification du cache KV au niveau de chaque tête tout en traitant les différences structurelles entre les clés et les valeurs.

La découverte centrale est que le décalage du modèle de distorsion est un mode de défaillance critique dans l'allocation de précision mixte. Sans calibrer la courbe de distorsion ($\beta$) au quantificateur de base spécifique, les stratégies de précision mixte peuvent être préjudiciables. En combinant :

1. La sensibilité basée sur le gradient (supérieure aux proxies basés sur l'activation),
2. La calibration empirique de la distorsion (pour correspondre au quantificateur), et
3. La séparation K/V (pour gérer les caractéristiques d'erreur asymétriques),

RATEQUANT transforme la quantification KV de précision mixte en une couche d'allocation générale et agnostique au quantificateur. Les auteurs postulent que le rapport AM/GM des sensibilités des têtes sert de prédicteur pratique pour déterminer quand l'allocation de précision mixte produira des gains significatifs. Le travail suggère que la perspective distorsion-débit peut se généraliser à l'allocation de ressources dans d'autres composants de réseaux de neurones hétérogènes au-delà des caches KV.