Résumé Technique : Quantification Stochastique Adaptative à Contrainte d'Entropie (ECASQ)

1. Formulation du Problème

L'article traite des limites des méthodes de quantification existantes dans les charges de travail de l'apprentissage automatique moderne, spécifiquement concernant le compromis entre impartialité (unbiasedness), adaptativité et efficacité de l'entropie.

• Contexte : La quantification stochastique adaptative (ASQ) optimise l'erreur quadratique moyenne (MSE) pour un intrant donné tout en préservant l'impartialité ($E[\hat{X}] = X$). Cependant, l'ASQ standard sélectionne les valeurs de quantification sans tenir compte de l'étape ultérieure de codage d'entropie sans perte, qui assigne des chaînes de bits plus courtes aux valeurs fréquentes.
• La Lacune : Les méthodes sensibles à l'entropie existantes (par exemple, la quantification vectorielle sous contrainte d'entropie) n'imposent pas l'impartialité, tandis que les méthodes impartiales (par exemple, QSGD, ASQ) sont souvent non adaptatives ou ignorent les contraintes d'entropie.
• Le Problème ECASQ : Les auteurs formulent le problème de la Quantification Stochastique Adaptative à Contrainte d'Entropie (ECASQ). Étant donné un vecteur d'entrée $X \in \mathbb{R}^d$, un alphabet admissible $P$, une taille maximale de codebook $s$, et un budget d'entropie $b$, l'objectif est de sélectionner un codebook $Q \subseteq P$ ($|Q| \le s$) pour minimiser la MSE sous réserve de :
  1. Impartialité : $E[\hat{X}] = X$.
  2. Contrainte d'Entropie : L'entropie moyenne par entrée $H(\hat{X}) \le b$.

2. Méthodologie

2.1 Relaxation Lagrangienne

Les auteurs abordent le problème à contrainte stricte en introduisant un multiplicateur de Lagrange $\lambda \ge 0$, transformant l'objectif en la minimisation de :
$$\text{MSE}(Q, X) + \lambda \cdot H(\hat{X})$$
sous réserve de $|Q| \le s$ et de l'impartialité. En faisant varier $\lambda$, on peut tracer la courbe de compromis Pareto-optimale entre la MSE et l'entropie. L'article note que, bien qu'un seul $\lambda$ puisse ne pas atteindre exactement un budget d'entropie $b$ (en raison de la nature discrète du problème), le partage de temps bidirectionnel (exécuter deux quantificateurs sur une partition aléatoire de l'intrant) permet de récupérer la solution optimale pour tout $b$.

2.2 Programmation Dynamique (DP) Optimale

La DP standard de l'ASQ ne peut pas être étendue à l'ECASQ car l'entropie d'une valeur de quantification dépend à la fois de ses voisins gauche et droit, ce qui rend la fonction de coût non séparable.

• Définition de l'état : Les auteurs définissent un nouvel état de DP $dp_{OPT}(i, \ell, j)$ représentant le coût minimal pour quantifier les $\ell$ premières entrées en utilisant $i-1$ valeurs de quantification (incluant $x_\ell$), plus la MSE pour l'intervalle $[x_\ell, x_j]$ en supposant que $x_\ell$ et $x_j$ sont des valeurs de quantification consécutives.
• Complexité : Une implémentation directe nécessite un temps de $O(d^3 s)$ et un espace de $O(d^2 s)$.
• Accélérations :
  • Algorithme SMAWK : Utilisé pour trouver les minima de lignes en $O(d)$ en exploitant l'inégalité du quadrangle (propriété de Monge) de la matrice de coût, réduisant le temps à $O(d^2 s)$.
  • Algorithme de Hirschberg : Utilisé pour réduire la complexité spatiale à $O(d^2)$ en employant une stratégie de type diviser pour régner qui ne stocke que deux couches de la table DP à la fois.
• Généralisation : L'algorithme s'étend à des ensembles admissibles arbitraires $P$ (taille $p$) avec une complexité de $O(p^2 s + d)$ en temps et $O(p^2 + d)$ en espace.

2.3 Algorithme d'Approximation à Espace Efficient

Pour gérer de grands intrants et des ensembles admissibles où $O(d^2)$ d'espace est prohibitif, les auteurs proposent une approximation qui optimise une métrique d'entropie de substitution $H(I, B)$ (entropie jointe de l'indice d'intervalle et du bit d'arrondi) plutôt que l'entropie exacte de la sortie $H(\hat{X})$.

• Garantie : La solution satisfait $H(\hat{X}) \le H(I, B) \le H(\hat{X}) + 1$. Par conséquent, l'approximation avec un budget d'entropie $b$ produit une MSE non plus grande que la solution optimale avec un budget $b-1$.
• Complexité : Cette approche réduit l'espace à $O(d)$ (ou $O(p+d)$ pour un $P$ général) tout en maintenant un temps de $O(d^2 s)$. Elle est conçue pour être compatible avec les GPU.

2.4 Raffinement Itératif

Pour combler l'écart de précision entre l'approximation et la solution optimale, les auteurs introduisent une procédure de raffinement. Partant de la solution approximative, elle optimise de manière itérative des sous-ensembles de valeurs de quantification ($Q_0, Q_1, Q_2$) en parallèle via une recherche binaire dans leurs intervalles locaux. Ce processus maintient l'impartialité et converge généralement vers des résultats quasi-optimaux en quelques itérations.

3. Contributions Clés

1. Formulation du Problème : Définition de l'ECASQ, le premier cadre optimisant conjointement la quantification stochastique adaptative pour la MSE, l'impartialité et les contraintes d'entropie.
2. Algorithme Optimal : Développement d'un algorithme de programmation dynamique de $O(d^2 s)$ en temps et $O(d^2)$ en espace pour la solution exacte, utilisant les techniques SMAWK et Hirschberg.
3. Approximation avec Garanties : Proposition d'un algorithme d'approximation efficace en espace ($O(d)$) avec une borne prouvable : la MSE est bornée par la solution optimale avec un bit d'entropie de moins.
4. Stratégie de Raffinement : Introduction d'un raffinement itératif rapide et parallélisable qui atteint une précision quasi-optimale avec un coût computationnel nettement inférieur à celui du solveur optimal.
5. Complétude Théorique : Démonstration que le partage de temps bidirectionnel permet à la relaxation lagrangienne de récupérer la solution optimale pour tout budget d'entropie, même lorsqu'un quantificateur unique ne peut atteindre la contrainte exacte.

4. Résultats Expérimentaux

Les auteurs ont évalué leurs méthodes sur des vecteurs synthétiques (Normale, Lognormale, Student-t, Outlier parcimonieux) et des tenseurs de modèles réels (poids, activations, KV-caches de modèles tels que Qwen2.5, Gemma, Llama, Phi).

• Précision :
  • L'Approx ECASQ raffinée réduit la vNMSE moyenne excédentaire (par rapport à l'Optimal ECASQ) de ~4,97 % (non raffinée) à 1,37 % sur les données synthétiques.
  • Sur les tenseurs de modèles réels, l'approximation raffinée surpasse systématiquement les bases (QSGD, QUIVER, Uniform + Entropie) à travers divers taux de bits (1,5 à 4 bits/valeur).
• Temps d'exécution :
  • Pour une taille de vecteur $d = 262\,144$, l'Optimal ECASQ prend environ 6,39 secondes.
  • L'Approx ECASQ prend environ 0,26 seconde (accélération de 24x).
  • L'Approx raffinée prend environ 0,33 seconde (accélération de 19x).
• Partage de Temps (Time-Sharing) : Les expériences indiquent que le gain potentiel du partage de temps par rapport à un quantificateur optimal unique est faible (médiane < 2 %, max ~21 % dans les intervalles de faible taux étroits), suggérant qu'un quantificateur unique bien réglé est souvent suffisant en pratique.

5. Signification et Revendications

L'article affirme que l'ECASQ fournit une solution fondée sur des principes au problème de la "précision sur la table" où les quantificateurs impartiaux existants ignorent les coûts de codage d'entropie.

• Praticité : La combinaison de l'approximation efficace en espace et de la procédure de raffinement rapide rend l'approche viable pour les charges de travail de ML à grande échelle (par exemple, compression de gradient, réduction de KV-cache) où la DP optimale est trop lente ou trop gourmande en mémoire.
• Généralité : Le cadre supporte des alphabets admissibles arbitraires (par exemple, des valeurs représentables en BF16 ou FP32), permettant des contraintes matérielles-conscientes.
• Modestie : Les auteurs reconnaissent que, bien que le partage de temps garantisse théoriquement l'optimalité pour tout budget d'entropie, les gains empiriques sont souvent marginaux, et l'approche par quantificateur unique (via l'optimisation lagrangienne) est généralement suffisante. Ils notent également que la garantie "un-bit" de l'approximation est une borne dans le pire des cas, les résultats empiriques montrant des performances bien plus serrées.