Convergence Rate Analysis of the AdamW-Style Shampoo: Unifying One-sided and Two-Sided Preconditioning

1. Énoncé du problème

Le papier aborde l'analyse théorique de la convergence de l'optimiseur Shampoo de style AdamW, un algorithme de pointe qui a remporté le concours de formation de réseaux neuronaux AlgoPerf. Bien que Shampoo (et ses variantes telles que SOAP et Muon) ait démontré des performances empiriques supérieures aux méthodes de préconditionnement diagonal (comme Adam) en utilisant des préconditionneurs à matrice complète, ses garanties théoriques de convergence pour l'optimisation non convexe faisaient défaut ou étaient limitées.

Plus précisément, la littérature existante présentait trois lacunes principales :

1. Non-convexité : La plupart des analyses antérieures étaient restreintes à des cadres convexes.
2. Portée du préconditionnement : Les analyses non convexes précédentes (par exemple, Xie et al., 2026) étaient limitées au préconditionnement unilatéral, ignorant le préconditionnement bilatéral plus complexe et plus efficace utilisé dans les implémentations pratiques.
3. Composantes pratiques : Les théories existantes omettaient souvent des caractéristiques pratiques critiques comme l'impulsion (momentum) et la décroissance des poids découplée (style AdamW), qui sont essentielles pour la formation moderne des grands modèles de langage (LLM).

Le papier vise à établir un taux de convergence rigoureux pour l'algorithme général Shampoo de style AdamW, unifiant le préconditionnement unilatéral et bilatéral dans un cadre non convexe.

2. Méthodologie

Les auteurs analysent l'Algorithme 1, qui représente l'optimiseur Shampoo de style AdamW. La règle de mise à jour implique :

• Moyennes mobiles exponentielles (EMA) : Suivi des premiers moments ($M_k$) et des seconds moments ($L_k, R_k$) des gradients.
• Préconditionnement bilatéral : L'étape de mise à jour est $X_{k+1} = (1-\lambda\eta)X_k - \eta L_{k,\epsilon}^{-\frac{1}{2p}} M_k R_{k,\epsilon}^{-\frac{1}{2q}}$, où $L_k$ et $R_k$ sont des produits externes accumulés de gradients (approchant $E[GG^T]$ et $E[G^TG]$).
• Exposants unifiés : Les paramètres $p$ et $q$ (où $1/p + 1/q = 1$) permettent à l'algorithme de basculer entre le préconditionnement bilatéral ($p, q < \infty$) et unilatéral ($p=1, q=\infty$ ou vice versa).

Innovations techniques clés :
Pour gérer la complexité du préconditionnement matriciel bilatéral dans des cadres non convexes, les auteurs ont développé de nouvelles techniques de preuve :

1. Inégalité de Hölder de Schatten-p : Au lieu des inégalités de Hölder vectorielles standard utilisées dans l'analyse d'Adam, les auteurs utilisent l'inégalité de Hölder de Schatten-p pour borner la norme nucléaire ($\|\cdot\|_*$) du gradient. Ceci est crucial car le préconditionneur agit sur les valeurs singulières des paramètres matriciels.
2. Inégalité de Cauchy-Schwarz matricielle : Pour borner la norme spectrale de l'étape de mise à jour ($L^{-1/2p} M R^{-1/2q}$), ils emploient une inégalité de Cauchy-Schwarz spécifique aux matrices pour gérer l'interaction entre les préconditionneurs gauche et droit.
3. Analyse du plancher spectral : Ils introduisent une condition (Condition 4) concernant la borne inférieure des préconditionneurs ($\hat{\epsilon}$). Ils soutiennent que dans la pratique, le plancher spectral effectif des préconditionneurs est considérablement plus grand que la régularisation numérique $\epsilon$ (par exemple, $10^{-12}$), alignant les bornes théoriques avec les observations empiriques.

3. Contributions clés

• Taux de convergence unifié : Le papier établit le premier taux de convergence pour Shampoo de style AdamW en optimisation non convexe, couvrant à la fois le préconditionnement unilatéral et bilatéral.
• Dérivation d'un taux optimal : Ils prouvent que la norme nucléaire moyenne attendue du gradient converge à un taux de :
  $$\frac{1}{K} \sum_{k=1}^K \mathbb{E}[\|\nabla f(X_k)\|_*] \leq O\left( \frac{\sqrt{m+n} \cdot C}{K^{1/4}} \right)$$
  où $m, n$ sont les dimensions de la matrice, $K$ est le nombre d'itérations, et $C$ correspond à la constante du taux optimal de SGD.
• Justification théorique pour des paramètres pratiques : L'analyse explique pourquoi le réglage du régularisateur numérique $\epsilon$ à des valeurs extrêmement faibles (par exemple, $10^{-12}$) ne nuit pas à la convergence. La théorie montre que le plancher effectif du préconditionneur est déterminé par l'échelle du bruit du gradient, et non par le $\epsilon$ numérique.
• Relation de norme : Les auteurs démontrent que le taux de norme nucléaire dérivé est analogue au taux optimal de norme de Frobenius de SGD ($O(K^{-1/4})$) dans la condition idéale où $\|\nabla f(X)\|_* \approx \Theta(\sqrt{\min(m,n)}) \|\nabla f(X)\|_F$.

4. Résultats

• Bornes théoriques :
  • La borne du pire cas dépend de $\sqrt{m+n}$, ce qui est significativement plus serré que la dépendance $\sqrt{mn}$ (dimension totale $d$) trouvée dans certaines méthodes diagonales ou bornes inférieures à matrice complète.
  • Le taux de convergence correspond à la borne inférieure connue pour l'optimisation stochastique non convexe ($O(K^{-1/4})$) jusqu'à un facteur de $\sqrt{m+n}$.
• Validation empirique (Préformation GPT-2) :
  • Ratio de norme de gradient : Les expériences sur la préformation de GPT-2 ont montré que le ratio $\|\nabla f(X)\|_* / \|\nabla f(X)\|_F$ reste constamment proche de la borne supérieure théorique $\sqrt{\min(m,n)}$. Cela valide l'hypothèse que le taux de norme nucléaire est effectivement équivalent au taux de norme de Frobenius dans la pratique.
  • Plancher spectral : Les valeurs propres minimales des préconditionneurs ($L_k, R_k$) ont été observées comme étant plusieurs ordres de grandeur plus grandes que le $\epsilon$ numérique ($10^{-12}$) et comparables à l'échelle dépendante du bruit $\hat{\sigma}^2/(m+n)$. Cela confirme que l'hypothèse "non pratique" d'un plancher spectral plus grand tient dans la formation réelle de LLM.

5. Signification

• Pont entre théorie et pratique : Ce travail fournit la première fondation théorique rigoureuse pour la variante spécifique de Shampoo qui a remporté le concours AlgoPerf, expliquant pourquoi elle fonctionne si bien en pratique.
• Avance sur les méthodes diagonales : Il valide théoriquement la supériorité du préconditionnement non diagonal (à matrice complète) par rapport aux méthodes diagonales (comme Adam) pour les paramètres structurés en matrice, montrant que les corrélations inter-paramètres capturées par Shampoo conduisent à de meilleures directions de convergence.
• Guide pour les hyperparamètres : L'analyse offre une justification théorique pour le choix des hyperparamètres (par exemple, la relation entre l'impulsion $\theta$ et la décroissance du second moment $\beta$, et le rôle de $\epsilon$), guidant les praticiens dans le réglage de ces optimiseurs pour les modèles à grande échelle.
• Fondement pour les futurs optimiseurs : En unifiant le préconditionnement unilatéral et bilatéral et en incorporant l'impulsion et la décroissance des poids, ce cadre établit une nouvelle norme pour l'analyse des optimiseurs avancés de type second ordre en apprentissage profond.

En résumé, le papier unifie avec succès l'analyse théorique de Shampoo, prouvant qu'il atteint un taux de convergence optimal pour les problèmes non convexes tout en capturant les structures matricielles complexes inhérentes aux modèles d'apprentissage profond modernes.