Arrondir dans l'espace du préconditionneur : repenser la quantification 4 bits de l'état d'AdamW

AdamW est l'optimiseur par défaut de l'entraînement des grands modèles, et son coût est facile à sous-estimer. Outre les poids et les gradients, chaque paramètre porte deux états supplémentaires, un premier moment et un second moment. Stockés en 32 bits, ces états occupent en mémoire résidente autant que les poids, souvent davantage. Les comprimer à 4 bits est l'un des moyens les plus directs de réduire la mémoire d'entraînement. TorchAO propose déjà un AdamW 4 bits, mais un écart mesurable de perte de validation avec AdamW 32 bits subsiste. L'article de Hanyang Li et de quatre coauteurs, publié sur arXiv cs.LG le 8 octobre 2026, ne règle ni la forme du livre de codes ni la taille des blocs. Il pose une question plus élémentaire : lorsqu'un quantificateur choisit entre deux niveaux de reconstruction voisins, dans quel repère doit-il faire ce choix ?

Les auteurs appellent ce repère l'espace d'arrondi. Il compte parce que l'erreur de quantification ne reste pas en place. Elle se propage dans les récurrences des moments et perturbe chaque mise à jour adaptative suivante. Pour le premier moment, la mise à jour dépend à peu près linéairement de l'état, et arrondir dans l'espace des états pose peu de problèmes. Le second moment est différent. Adam divise par la racine carrée du second moment, et cet inverse est le préconditionneur. L'article analyse la cellule de quantification adjacente à zéro et montre qu'une faible erreur moyenne sur l'état n'implique pas une faible erreur moyenne sur le préconditionneur du pas suivant. Un état qui s'écarte un peu de zéro passe par une fonction inverse très raide, et l'erreur est amplifiée du côté du préconditionneur. Une construction quadratique à une dimension montre ensuite qu'arrondir dans l'espace des états ou dans celui du préconditionneur produit des dynamiques d'optimisation qualitativement différentes, et pas seulement des chiffres différents.

Ces observations mènent à la première méthode, ZIP-SR, pour arrondi stochastique dans l'espace du préconditionneur avec inclusion de zéro. Elle garde zéro dans le livre de codes du second moment et calcule les probabilités d'arrondi stochastique dans l'espace du préconditionneur plutôt que dans celui des états. L'arrondi stochastique est non biaisé en espérance et transforme une erreur systématique déterministe en bruit de moyenne nulle. Bien choisir l'espace place cette absence de biais sur la grandeur qui fixe réellement la taille du pas. La seconde méthode, ZE-EDEN, suit une voie complémentaire. Son livre de codes du second moment exclut zéro, de sorte que les valeurs quantifiées ont un plancher positif et que l'inverse ne diverge pas près de zéro. Ce plancher déforme lui-même le préconditionneur, et la méthode remet donc à l'échelle chaque bloc quantifié du second moment par un étalonnage de type EDEN pour compenser la distorsion. Les deux configurations utilisent NormalFloat 4 bits (NF4) pour le premier moment, et appliquent un arrondi stochastique ciblé au premier moment de la tête du modèle de langue pendant les 10 % finaux de l'entraînement.

Les expériences couvrent des préentraînements de style GPT et Llama, de 130M à 2,7B de paramètres. À chaque taille évaluée, les deux méthodes réduisent l'écart moyen de perte de validation entre AdamW 4 bits de TorchAO et AdamW 32 bits, et la plus forte réduction rapportée atteint 70 %. En réglage fin supervisé de tous les paramètres, les deux recettes obtiennent une perte de validation inférieure à celle de TorchAO tout en restant proches d'AdamW 32 bits sur les tâches aval. La constance compte plus qu'une seule grande victoire. Un gain présent à chaque taille suggère un mécanisme à l'œuvre, et non la chance d'une exécution.

Notre appréciation est que l'apport principal est un redécoupage de l'espace de conception. Les travaux antérieurs portaient sur la conception du livre de codes et le découpage en blocs. Cet article ajoute un troisième axe : le repère dans lequel l'arrondi a lieu. Ce point de vue pourrait bien s'appliquer à d'autres optimiseurs qui appliquent une transformation non linéaire à leur état stocké. De la retenue reste de mise. Le résumé ne donne aucun chiffre sur la mémoire économisée ni sur le débit, le plus grand modèle compte 2,7B de paramètres, et la persistance des gains à plus grande échelle et sur des entraînements plus longs relève d'une réplication indépendante. Pour une équipe à court de mémoire d'entraînement, le pas raisonnable consiste à traiter ZIP-SR et ZE-EDEN comme des candidats pour un petit essai face à TorchAO à sa propre échelle, et non comme un remplacement direct d'une recette de production.