Résumé Technique : CAGE (Curvature-Aware Gradient Estimation)

1. Énoncé du Problème

La quantification consciente de l'entraînement (Quantization-Aware Training - QAT) vise à entraîner des réseaux de neurones pour qu'ils s'adaptent aux contraintes de quantification lors du processus d'optimisation, ce qui offre souvent une meilleure précision que la quantification post-entraînement (PTQ). Cependant, un écart de précision significatif subsiste entre les modèles QAT à faible nombre de bits et leurs homologues en pleine précision.

Le principal goulot d'étranglement du QAT est la nature non dérivable de l'opérateur de quantification $Q(\cdot)$, qui possède des gradients nuls presque partout. La norme industrielle, l'estimateur de passage direct (Straight-Through Estimator - STE), approxime le gradient de la fonction de quantification par la matrice identité lors de la rétropropagation. Bien que polyvalent, le STE est connu pour souffrir d'une convergence lente, d'instabilité et de solutions sous-optimales car il ignore la courbure du paysage de perte et la dynamique spécifique de l'erreur de quantification. Les raffinements existants (par exemple, LSQ, ProxQuant) reposent souvent sur des heuristiques ou des paramètres apprenables sans offrir de garanties de convergence formelles dans des contextes non convexes.

2. Méthodologie : CAGE

Les auteurs proposent CAGE (Curvature-Aware Gradient Estimation), une méthode qui augmente le gradient STE standard avec un terme de correction dérivé d'une perspective d'optimisation multi-objectif.

2.1 Formulation Multi-Objectif

L'article reformule le QAT comme un problème multi-objectif équilibrant deux objectifs concurrents :

1. Minimiser la perte de tâche $f(x)$.
2. Satisfaire la contrainte de quantification $x \in \mathcal{Q}$ (minimiser la distance entre les paramètres $x$ et leur version quantifiée $Q(x)$).

Puisque ces objectifs sont souvent en conflit (un point minimisant la perte peut ne pas être quantifiable, et vice versa), les auteurs définissent l'objectif comme la recherche d'une solution Pareto-optimale. Un point $x^*$ est défini comme $\lambda$-Pareto optimal si :
$$\nabla f(x^*) + \lambda(x^* - Q(x^*)) = 0$$
où $\lambda > 0$ équilibre les deux objectifs. Le terme $(x^* - Q(x^*))$ représente l'erreur de quantification.

2.2 Estimation du Gradient

Sur la base de la condition de Pareto, CAGE introduit un terme de correction au gradient de mise à jour. Au lieu de se fier uniquement au gradient STE $\nabla f(Q(x))$, la mise à jour incorpore l'erreur de quantification $e_t = x_t - Q(x_t)$ :
$$g_t = \nabla f(x_t) + \lambda_t e_t$$
Cette correction contrecarre explicitement l'augmentation de la perte induite par l'étape de quantification en exploitant l'information de second ordre locale (courbure).

2.3 Variantes d'Implémentation

Le papier propose deux manières d'intégrer cette correction, particulièrement pour les optimiseurs à état (stateful) comme Adam :

• Correction Couplée (Coupled Correction) : Le terme d'erreur est ajouté au gradient avant d'être traité par les statistiques de momentum et de variance de l'optimiseur. Cela préconditionne efficacement la correction en utilisant les estimations du second moment de l'optimiseur (approximant la courbure).
• Correction Découplée (Decoupled Correction) : Le terme d'erreur est ajouté après l'étape de l'optimiseur. Cela applique la correction à une échelle unitaire, évitant la distorsion par les estimations de variance par coordonnée qui peuvent fluctuer dans les régimes à faible nombre de bits.

Choix Pratique : Les auteurs implémentent la variante Découplée sur AdamW (Algorithme 1). Ils soutiennent que cela préserve le comportement de l'optimiseur, empêche la direction de la correction d'être déformée et évite l'instabilité numérique.

2.4 Calendrier d'Entraînement

Pour éviter l'instabilité durant les premières étapes de l'entraînement (où les résidus de quantification se comportent comme du bruit aléatoire), CAGE emploie une période de silence. Le coefficient de correction $\lambda_t$ est fixé à 0 pour les $s \cdot T$ premières étapes (où $s \in [0,8, 0,95]$) puis augmente linéairement jusqu'au $\lambda$ cible. Cela permet au modèle de se stabiliser dans le paysage de perte avant d'imposer la contrainte de stationnarité de Pareto.

3. Contributions Clés

1. Cadre Théorique : L'article établit un cadre formel pour le QAT en tant que problème d'optimisation multi-objectif et définit les solutions Pareto-optimales pour l'optimisation quantifiée.
2. Garanties de Convergence : Sous des hypothèses de lissage et de variance bornée, les auteurs prouvent que CAGE atteint une convergence ergodique forte vers un point Pareto-optimal dans le cadre lisse non convexe. Cela répond à une limitation primaire des méthodes basées sur le STE, qui manquent de garanties formelles.
3. Correction Sensible à la Courbure : La méthode introduit un terme de correction principalisé dérivé directement de la condition de Pareto, qui exploite la courbure locale (approximée via les statistiques d'Adam dans la variante couplée ou appliquée directement dans la variante découplée) pour contrecarre l'augmentation de la perte induite par la quantification.
4. Agnoticisme de l'Optimiseur : Bien qu'implémentée efficacement avec AdamW, la méthode est théoriquement agnostique à l'optimiseur et a été validée à travers AdamW, Muon, Shampoo et SOAP.

4. Résultats Expérimentaux

Les auteurs valident CAGE via des expériences synthétiques, du fine-tuning et du pré-entraînement de modèles de type Llama (jusqu'à 800M de paramètres).

• Expériences Synthétiques : Sur des objectifs quadratiques avec des Hessiennes non isotropes, CAGE réduit systématiquement l'erreur stationnaire et converge plus près de la solution optimale par rapport à STE-SGD, STE-Adam et LOTION (un travail concomitant).
• Fine-Tuning (MXFP4) : Sur Llama-3.2-3B fine-tuné avec des poids et activations en 4 bits (MXFP4), CAGE réduit de moitié l'erreur de quantification par rapport à la méthode de l'état de l'art précédent (QuEST) sur les benchmarks GSM8K, HellaSwag et WinoGrande.
• Pré-entraînement :
  • Précision : Les modèles entraînés avec CAGE avec des poids et activations en 3 bits (W3A3) atteignent une perte de validation inférieure aux modèles en 4 bits (W4A4) entraînés avec la meilleure méthode précédente (QuEST).
  • Lois d'Échelle (Scaling Laws) : L'analyse de l'efficacité des paramètres ($eff(P)$) montre que CAGE améliore la capacité effective du modèle de plus de 10 % à 4 bits et de 20 % à 2 bits par rapport à QuEST.
  • Robustesse de l'Optimiseur : CAGE produit des améliorations constantes à travers AdamW, Muon, Shampoo et SOAP sans ajustement supplémentaire des hyperparamètres.
• Efficacité : CAGE introduit un surcoût computationnel négligeable. Elle ne nécessite qu'un appel de quantification supplémentaire pour calculer le résidu et des opérations élément par élément. Les temps d'itération (wall-clock) sont statistiquement indiscernables de la ligne de base.

5. Signification et Revendications

L'article affirme que CAGE représente une étape significative vers le comblement de l'écart de performance entre les modèles quantifiés à faible nombre de bits et leurs homologues en pleine précision.

• Fondement Théorique : Contrairement aux approches heuristiques précédentes, CAGE fournit une dérivation principalisée de l'optimisation multi-objectif avec des garanties de convergence formelles.
• Efficacité Pratique : La méthode atteint une précision de l'état de l'art avec des coûts computationnels similaires, permettant une quantification plus agressive (par exemple, 3 bits) sans sacrifier la performance.
• Généralité : L'approche est agnostique au quantificateur et à l'optimiseur, ce qui en fait une amélioration directe pour les pipelines QAT existants.

Les auteurs concluent qu'incorporer l'information de courbure locale et la régularisation via CAGE est une direction prometteuse pour l'entraînement précis à faible nombre de bits, bien qu'ils notent que des travaux futurs pourraient explorer les régimes de très bas bits et l'intégration avec d'autres techniques de compression comme la sparsification.