Résumé technique : PowerStep, optimisation adaptative économe en mémoire via la descente de la plus forte pente de norme ℓp

Énoncé du problème

Les optimiseurs adaptatifs, en particulier Adam et AdamW, sont devenus la norme pour l'entraînement de réseaux de neurones à grande échelle (par exemple, les Transformers) en préconditionnant les mises à jour de gradient avec des statistiques de premier et de second moment accumulées. Bien qu'efficaces, ces méthodes entraînent une surcharge mémoire substantielle car elles maintiennent un accumulateur de second moment ($v_t$) ayant la même dimensionnalité que les paramètres du modèle, nécessitant généralement une précision flottante 32 bits complète (fp32) pour la stabilité. Ce « mur de la mémoire » double l'empreinte de l'état de l'optimiseur par rapport à la descente de gradient stochastique (SGD) avec momentum, limitant ainsi l'évolutivité pour les modèles comportant des centaines de milliards de paramètres.

Les tentatives existantes pour atténuer ce problème incluent la factorisation de la matrice de second moment, le partage par blocs, ou l'application d'une quantification de basse précision. Cependant, la quantification conduit souvent à une instabilité de l'entraînement ou à une divergence dans les méthodes adaptatives en raison de la haute sensibilité de l'estimateur de second moment (spécifiquement l'opération de racine carrée réciproque) aux erreurs de quantification. Alternativement, certaines méthodes éliminent entièrement le second moment mais échouent souvent à égaler la vitesse de convergence d'Adam sur des tâches à grande échelle.

Méthodologie

Les auteurs introduisent PowerStep, un optimiseur économe en mémoire qui réalise une adaptativité au niveau des coordonnées sans stocker de statistiques de second moment. La méthode est dérivée des principes de la descente de la plus forte pente de norme $\ell_p$.

Fondement théorique

Dans les espaces normés généraux, la direction de descente de la plus forte pente $v^*$ qui diminue au maximum un objectif $f(\theta)$ sous une contrainte de norme $\ell_p$ est donnée par :
$$v^*_i \propto -\text{sign}(g_i) \cdot |g_i|^{\frac{1}{p}-1}$$
où $g_i$ est la composante du gradient. Cela conduit à une « transformation de puissance signée » $\Phi_\beta(x) = \text{sign}(x) \odot |x|^\beta$, où $\beta = 1/(p-1) \in [0, 1]$.

Conception de l'algorithme

Contrairement à la méthode Powerball, qui applique cette transformation aux gradients bruts, PowerStep applique la transformation à un tampon de momentum à bille lourde (heavy-ball momentum buffer) ($m_t$). La règle de mise à jour est :

1. Mise à jour du momentum : $m_t = \gamma m_{t-1} + g_t$
2. Mise à jour transformée : $\theta_t = \theta_{t-1} - \eta_t \Phi_\beta(m_t)$

Cette conception produit un préconditionneur diagonal $D_t = \text{diag}(|m_t|^{\beta-1})$. Pour $\beta \in (0, 1)$, l'exposant $\beta-1$ est négatif, ce qui signifie que la taille de l'étape est inversement proportionnelle à la magnitude du momentum accumulé. Les petites magnitudes de momentum (directions plates) sont amplifiées, tandis que les grandes magnitudes (directions raides) sont atténuées. Cela fournit une adaptativité au niveau des coordonnées similaire à celle d'Adam, mais calculée instantanément à partir du tampon de momentum, éliminant ainsi le besoin d'un estimateur de second moment.

Implémentation en basse précision

La conception à tampon unique de PowerStep se prête naturellement à une quantification agressive. Les auteurs proposent de stocker uniquement le tampon de momentum $m_t$ en précision int8, tout en effectuant la transformation de puissance signée et les mises à jour de paramètres en précision complète. Cela contraste avec Adam, où la quantification du tampon de second moment $v_t$ provoque une accumulation catastrophique d'erreurs en raison de la non-linéarité des opérations de racine carrée et de division.

Contributions clés

1. Conception de l'algorithme : PowerStep est dérivé des principes de la descente de la plus forte pente de norme $\ell_p$, démontrant que l'application d'une transformation de puissance signée à un tampon de momentum à bille lourde produit une adaptativité au niveau des coordonnées sans statistiques de second moment.
2. Analyse théorique : L'article démontre que PowerStep atteint le taux de convergence optimal $O(1/\sqrt{T})$ pour l'optimisation stochastique non convexe sous des hypothèses de régularité standard (lissité L, gradients bornés, etc.).
3. Validation empirique : Des expériences étendues sur des modèles Transformer allant de 124 M à 235 milliards de paramètres (incluant des architectures denses et à mélange d'experts) montrent que PowerStep égale la vitesse de convergence d'AdamW tout en n'utilisant que la moitié de la mémoire d'état de l'optimiseur.
4. Robustesse à la quantification : Combinée à une quantification agressive en int8, PowerStep reste numériquement stable et réduit la mémoire de l'optimiseur d'environ 8 fois par rapport à AdamW en pleine précision. En revanche, AdamW diverge dans les mêmes conditions de quantification.

Résultats expérimentaux

• Vitesse de convergence : Sur des modèles allant de 124 M à 235 milliards de paramètres, PowerStep correspond aux trajectoires de perte d'entraînement d'AdamW. Les méthodes économes en mémoire apparentées (pbSGDM, SignSGD, AdamS) ont montré une convergence plus lente ou une instabilité sur les modèles plus grands.
• Efficacité mémoire : PowerStep réduit l'empreinte mémoire de l'état de l'optimiseur de 50 % par rapport à AdamW (fp32) car il ne stocke qu'un seul tampon au lieu de deux.
• Performance de quantification : Sous quantification int8 de l'état de l'optimiseur :
  • AdamW : L'entraînement s'effondre immédiatement en raison de l'accumulation d'erreurs de quantification dans l'estimateur de second moment.
  • PowerStep : Reste stable et correspond à la convergence de sa contrepartie en pleine précision.
  • Réduction de mémoire : PowerStep (int8) réalise une réduction d'environ 8 fois de la mémoire de l'optimiseur par rapport à AdamW (fp32). Par exemple, sur un modèle de 235 milliards de paramètres, la mémoire par NPU passe de 6768 Mo (AdamW) à 872,4 Mo (PowerStep int8).
• Sensibilité aux hyperparamètres : La méthode est robuste au coefficient de momentum $\gamma$ dans la plage $[0,85, 0,95]$. Un exposant de puissance $\beta = 0,1$ offre le meilleur compromis entre progrès rapide au début et stabilité à long terme ; $\beta=0$ (équivalent à SignSGD) conduit à un effondrement, tandis que $\beta=0,2$ converge trop lentement.

Importance et affirmations

L'article affirme que PowerStep fournit une alternative fondée sur des principes, évolutif et économe en ressources pour l'entraînement à grande échelle. En établissant un lien direct entre l'optimisation adaptative et la géométrie de la norme $\ell_p$, PowerStep atteint l'optimalité théorique des méthodes adaptatives sans le coût mémoire de l'estimation du second moment.

Les auteurs soulignent que l'élimination du tampon de second moment non seulement divise par deux l'utilisation de la mémoire, mais améliore fondamentalement la robustesse à la quantification de basse précision. Cela permet une réduction supplémentaire d'un ordre de grandeur de la mémoire (environ 8 fois au total par rapport à AdamW) sans sacrifier la vitesse de convergence ni la qualité finale du modèle. Ce travail suggère que le momentum à bille lourde combiné à des transformations non linéaires peut remplacer la machinerie complexe du second moment d'Adam, offrant une voie plus simple et plus évolutive pour l'entraînement des futurs modèles à grande échelle.

Les axes de travail futurs identifiés par les auteurs incluent l'intégration de PowerStep avec des mises à jour structurées à valeurs matricielles (par exemple, l'optimiseur Muon) et l'exploration de quantifications encore plus agressives (par exemple, int4).