Delving into Muon and Beyond: Deep Analysis and Extensions

Résumé technique : Approfondir Muon et au-delà : Analyse approfondie et extensions

Énoncé du problème
L'optimiseur Muon a récemment attiré l'attention pour ses performances empiriques dans l'entraînement de grands modèles de langage, en utilisant des mises à jour orthogonalisées au niveau matriciel. Cependant, les mécanismes sous-jacents de Muon et sa relation avec les optimiseurs adaptatifs établis comme Adam restent insuffisamment compris. La littérature existante confond souvent les effets de Muon avec des techniques auxiliaires (par exemple, QK-Norm, QK-Clip) ou manque d'une comparaison systématique par rapport aux bases de référence standards. De plus, il n'est pas clair si l'orthogonalisation spectrale de Muon représente une stratégie d'optimisation universellement supérieure ou une forme spécifique de stabilisation, et si elle peut être améliorée ou généralisée.

Méthodologie
Les auteurs proposent un cadre spectral unifié pour analyser Muon et ses variantes. Ils définissent Muon comme le point final $p=0$ d'une famille de transformations spectrales de la forme $U\Sigma^p V^\top$, où $U\Sigma V^\top$ est la décomposition en valeurs singulières (SVD) de la matrice de mise à jour. L'étude examine les exposants intermédiaires $p \in \{1, 1/2, 1/4, 0\}$ :

• $p=1$ : Identité (mises à jour standards).
• $p=1/2, 1/4$ : Compression spectrale partielle.
• $p=0$ : Orthogonalisation complète (Muon).

Ces transformations sont appliquées à deux types d'entrées distincts :

1. Moment du premier ordre ($O_{mom}$) : Équivalent aux entrées de momentum SGD.
2. Mises à jour normalisées par RMS ($O_{rms}$) : Équivalent aux entrées d'Adam, où les gradients sont normalisés par le second moment.

Pour permettre un calcul efficace des puissances fractionnaires ($p=1/2, 1/4$) sans SVD explicite, les auteurs développent une itération de Newton-Schulz couplée. Cet algorithme calcule simultanément la racine carrée de la matrice et son inverse en utilisant uniquement des multiplications matricielles, garantissant la stabilité numérique et l'extensibilité.

L'évaluation expérimentale est menée sur un modèle GPT-2 de 124 millions de paramètres entraîné sur OpenWebText. L'étude utilise un protocole hautement contrôlé :

• Taux d'apprentissage découplés : Des taux d'apprentissage distincts sont utilisés pour les paramètres matriciels et vectoriels afin d'isoler l'effet de l'optimiseur.
• Aucune technique auxiliaire : Le weight decay, QK-Norm et QK-Clip sont désactivés pour assurer des comparaisons équitables.
• Réglage systématique : Une recherche de taux d'apprentissage en deux étapes (balayage logarithmique grossier suivi d'un raffinement précis) est effectuée pour les huit variantes de l'optimiseur (4 exposants spectraux $\times$ 2 types d'entrées).

Principales contributions

1. Cadre spectral unifié : Le papier introduit un cadre formel $U\Sigma^p V^\top$ qui positionne Muon comme un cas spécifique ($p=0$) et propose des variantes intermédiaires ($p=1/2, 1/4$) pour interpoler entre les mises à jour standards et l'orthogonalisation complète.
2. Algorithme efficace : Une itération de Newton-Schulz couplée est développée pour calculer efficacement les mises à jour spectrales fractionnaires ($p=1/2, 1/4$), évitant le coût de calcul d'une SVD explicite.
3. Comparaison empirique rigoureuse : L'étude fournit une comparaison contrôlée entre Muon, Adam et leurs variantes spectrales, isolant les effets spectraux de facteurs de confusion tels que le weight decay ou les techniques de normalisation auxiliaires.

Résultats clés
Les expériences produisent trois conclusions principales :

1. Stabilisation des mises à jour du premier moment : Sur les entrées de premier moment (mSGD), Muon ($p=0$) apporte une stabilisation significative, permettant des taux d'apprentissage beaucoup plus élevés et une plus grande robustesse par rapport au mSGD standard. Cependant, la compression partielle ($p=1/4$) atteint souvent de meilleures performances finales une fois la stabilité établie.
2. Gains limités sur les entrées normalisées par RMS : Lorsqu'elles sont appliquées à des entrées normalisées par RMS (style Adam), les compressions spectrales offrent des rendements décroissants. L'orthogonalisation complète (AdamZ/Muon appliquée à $O_{rms}$) ne surpasse pas Adam standard et peut être inférieure à la compression partielle. Cela suggère que la normalisation RMS élément par élément atténue déjà une grande partie de l'anisotropie que la compression spectrale cible.
3. Supériorité des méthodes du second moment : À travers tous les exposants spectraux, les optimiseurs utilisant la normalisation par le second moment (basée sur RMS) surpassent systématiquement leurs homologues de premier moment (uniquement momentum) en termes de stabilité et de perte finale.

Signification et affirmations
Le papier conclut que Muon est mieux compris comme une forme efficace de normalisation spectrale plutôt que comme une méthode d'optimisation universellement supérieure.

• Stabilisation vs Performance : La contribution primaire de Muon est de stabiliser les mises à jour du premier moment en écartant les magnitudes des valeurs singulières, ce qui empêche les directions avec de grandes valeurs singulières de dominer. Cependant, cet "aplatissement" est un mécanisme rudimentaire ; il supprime le filtrage basé sur la magnitude, ce qui peut amplifier les directions de gradient bruitées (petites valeurs singulières) et sous-estimer les modes dominants informatifs.
• Utilité contextuelle : Bien que la compression spectrale offre de réels bénéfices de stabilisation, surtout lorsque l'entrée de mise à jour n'est pas déjà normalisée, l'orthogonalisation complète ($p=0$) n'est pas un remplacement pour les optimiseurs modernes à second moment normalisés par RMS comme Adam. Les auteurs soutiennent que les avantages rapportés de Muon dans les travaux précédents peuvent être attribués à la combinaison spécifique de techniques utilisées (par exemple, QK-Norm) plutôt qu'à l'orthogonalisation elle-même.

Le travail souligne que, bien que les transformations spectrales soient des outils puissants pour contrôler l'anisotropie, elles ne surpassent pas intrinsèquement les bénéfices de la mise à l'échelle adaptative du second moment dans les régimes d'entraînement à grande échelle.