Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering

Résumé technique : « Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering »

Énoncé du problème
L'optimiseur Muon a récemment démontré une forte performance empirique lors de l'entraînement de grands modèles de langage (LLM) en orthogonalisant approximativement une matrice de momentum via une itération de Newton–Schulz. Cependant, le rôle théorique du momentum au sein de Muon reste obscur. Les analyses existantes soit suppriment le momentum pour étudier les mises à jour spectrales de manière isolée, soit le conservent sans expliquer ses bénéfices empiriques. De plus, la littérature théorique sur Muon suggère souvent qu'un momentum important peut dégrader les taux de convergence ou n'offrir aucun avantage, ce qui contredit les observations pratiques où le momentum est essentiel. Par ailleurs, l'ordre spécifique des opérations — appliquer le momentum avant l'orthogonalisation (Pre-polar) versus orthogonaliser avant le momentum (Post-polar) — manque d'une explication rigoureuse, bien que le mode Pre-polar surpasse nettement le mode Post-polar en pratique.

Méthodologie
Les auteurs proposent un modèle de gradient signal-plus-perturbation structuré pour analyser Muon. Ils supposent que le gradient $G_t$ à l'itération $t$ se décompose en une composante de signal cohérente et de faible rang $G^{\text{sig}}_t$ et une perturbation BVMZOS (bounded-variance, mean-zero orthogonal) $\Xi_t$ à variance bornée et de moyenne nulle.

La contribution théorique centrale est l'analyse du tampon de momentum $M_t = \beta M_{t-1} + (1-\beta)G_t$ en tant que filtre spectral. Les auteurs étudient trois variantes de pipeline pour la mise à jour des poids $\Delta W_t$ :

1. Pre-polar (Muon standard) : $\Delta W_t \propto O(M_t)$, où le momentum est appliqué en premier, suivi du facteur polaire $O(\cdot)$.
2. Post-polar (Orthogonal-SGDM) : $\Delta W_t \propto \tilde{M}_t$, où le facteur polaire est appliqué d'abord au gradient brut, puis le momentum est accumulé.
3. Polar-only : $\Delta W_t \propto O(G_t)$, une ligne de base sans momentum.

L'analyse s'appuie sur la taille de fenêtre effective $T = 1/(1-\beta)$ et utilise les outils de la théorie des matrices aléatoires, notamment le théorème $\sin \Theta$ de Wedin pour la perturbation des sous-espaces singuliers, afin de quantifier comment le momentum affecte les valeurs singulières et les sous-espaces de la matrice de gradient.

Contributions clés

1. Le momentum comme filtre spectral (Théorème 1 et Corollaire 1) :
   Les auteurs prouvent que, sous leur modèle de gradient, le momentum agit comme un filtre passe-bas qui supprime la perturbation (le bruit) tout en préservant le signal cohérent. Spécifiquement, à mesure que la taille de fenêtre effective $T$ augmente, les valeurs singulières de la queue (dominées par le bruit) décroissent à un taux de $O(T^{-1/4})$, tandis que les valeurs singulières de tête (le signal) restent éloignées de zéro. Cela crée un écart spectral élargi entre les composantes de signal et de bruit. Par conséquent, les sous-espaces singuliers de rang supérieur $r$ du tampon de momentum $M_t$ s'alignent plus de manière fiable avec les sous-espaces de signal réels que les sous-espaces du gradient brut.

2. Non-commutativité et ordre des pipelines (Théorème 2) :
   Le papier établit que le momentum et l'orthogonalisation ne commutent pas.

   • Dominance du Pre-polar : Lorsque le momentum est appliqué en premier, il débruite le flux de gradients, permettant au facteur polaire suivant d'opérer sur une matrice possédant une structure de sous-espace singulier stabilisée. Les auteurs prouvent que le Pre-polar atteint un alignement de signal attendu plus élevé que les lignes de base Post-polar et Polar-only.
   • Échec des alternatives : Pour le Post-polar, l'orthogonalisation du gradient bruité en premier « efface » l'écart d'amplitude entre les directions de signal et de bruit. Le facteur polaire remplace toutes les valeurs singulières par des uns, rendant impossible la distinction entre le signal et le bruit lors des étapes de momentum ultérieures. Théoriquement, les pipelines Post-polar et Polar-only échouent à récupérer efficacement la direction du signal, particulièrement dans les régimes de faible rapport signal sur bruit (SNR).

3. Séparation quantitative en faible SNR (Théorème 3) :
   Dans un modèle gaussien à pointe (spiked) de rang 1 avec un faible SNR, les auteurs montrent que l'alignement du signal du pipeline Polar-only s'annule lorsque la dimension de la matrice $m \to \infty$. En revanche, le pipeline Pre-polar continue de récupérer le signal cible pourvu que la fenêtre de momentum $T$ soit suffisamment grande.

Résultats
Les conclusions théoriques sont validées par des expériences approfondies :

• Visualisation de l'écart spectral : En utilisant des sondes stationnaires sur NanoGPT, les auteurs démontrent qu'augmenter le coefficient de momentum $\beta$ élargit l'écart entre les valeurs singulières de tête et de queue du tampon de momentum, ce qui est cohérent avec le Théorème 1.
• Fiabilité des sous-espaces : Les mesures empiriques des erreurs d'alignement des sous-espaces montrent que la déviation entre les sous-espaces singuliers du tampon de momentum et les véritables sous-espaces de signal diminue à mesure que $T$ augmente, validant le Corollaire 1.
• Comparaison des pipelines : À travers des modèles synthétiques, CIFAR-10 et l'entraînement de LLaMA 350M/NanoGPT, le pipeline Pre-polar surpasse systématiquement le Post-polar et le Polar-only en termes d'alignement de signal et de perte de validation. Notamment, la performance du Post-polar reste plate quel que soit $\beta$, confirmant la prédiction théorique selon laquelle l'orthogonalisation avant la moyenne ne peut pas récupérer les directions de signal perdues à cause du bruit.
• Entraînement de bout en bout : Dans les tâches de pré-entraînement de LLM (NanoGPT et LLaMA 350M), le pipeline Pre-polar (Muon standard) atteint une perte de validation plus faible par rapport aux autres variantes, soutenant l'affirmation que le momentum est un composant essentiel pour la récupération du signal, et non un simple tour de force d'accélération.

Signification
Ce travail fournit la première justification théorique de la raison pour laquelle Muon nécessite du momentum et pourquoi l'ordre des opérations (débruiter avant d'orthogonaliser) est critique. Il recadre le momentum dans les optimiseurs basés sur les matrices non pas seulement comme un mécanisme d'accélération, mais comme un filtre spectral qui stabilise les sous-espaces singuliers contre le bruit. Le principe directeur « débruiter d'abord, puis orthogonaliser » offre une règle de conception pour les futurs optimiseurs conscients des matrices utilisant des étapes spectrales non linéaires, suggérant que la séparation du signal et du bruit doit précéder toute normalisation ou orthogonalisation spectrale. Les auteurs notent que si leur théorie explique les bénéfices du momentum dans Muon, elle souligne également les limites des bornes théoriques actuelles qui ne parviennent pas à prédire ces bénéfices, pointant ainsi un écart entre la théorie de l'optimisation standard et le comportement pratique des optimiseurs spectraux.