The Spectral Dynamics and Noise Geometry of Muon

Résumé Technique : La Dynamique Spectrale et la Géométrie du Bruit de Muon

Énoncé du Problème
L'article étudie le biais de régularisation implicite induit par l'optimiseur Muon, qui met à jour les poids en utilisant le facteur polaire de la matrice de gradient plutôt que le gradient lui-même. Étant donné un gradient $G = U\Sigma V^\top$, Muon effectue une mise à jour via $UV^\top$, écartant ainsi les magnitudes relatives des valeurs singulières tout en préservant les directions singulières actives. Bien que Muon ait démontré un succès empirique dans le pré-entraînement de transformers, son mécanisme de sélection théorique reste obscur. Deux hypothèses concurrentes sont abordées :

1. Muon est-il simplement une forme de normalisation du gradient matriciel ?
2. La géométrie de la norme spectrale implique-t-elle un biais vers des solutions de faible rang (minimisation de la norme nucléaire) ?

Les auteurs soutiennent que les deux sont incomplètes. Au lieu de cela, ils proposent que Muon induit un biais de spectre de mise à jour plat (flat-update-spectrum bias), allouant une amplitude égale à toutes les directions singulières actives. L'article cherche à caractériser ce biais de manière théorique et empirique, spécifiquement dans le contexte de la régression sous-déterminée où l'optimiseur doit sélectionner un interpolant fini parmi une variété affine de solutions de perte nulle.

Méthodologie
L'étude emploie une combinaison d'analyse de systèmes dynamiques en temps continu, de calcul variationnel et d'expériences empiriques contrôlées.

1. Cadre Théorique :

   • Flux Polaire Projeté : Les auteurs analysent le flux en temps continu $\dot{W} = -\eta P_\perp P(W)$, où $P(W)$ est le facteur polaire de $W$ et $P_\perp$ est le projecteur orthogonal sur l'espace tangent de la variété d'interpolation $\mathcal{M} = \{W : XW = Y\}$.
   • Hypothèse de Cadre Partagé (SF - Shared-Frame) : Une condition structurelle où les matrices de gradient et de poids partagent des cadres singuliers ($U_G = U_W, V_G = V_W$) est introduite pour simplifier l'analyse. Des théorèmes sont prouvés sans cette hypothèse pour le flux auto-polaire, tandis que l'hypothèse SF comble l'écart avec le flux réel piloté par le gradient de Muon.
   • Fonctionnelles : L'analyse utilise une fonctionnelle spectrale par paire $R_{pw}(W) = -\sum_{i \neq j} \log(\sigma_i + \sigma_j)$ et sa version invariante d'échelle $\tilde{R}_{pw}$ pour suivre la forme spectrale.
   • Géométrie du Bruit : Une dérivation de la taille de lot critique est réalisée à l'aide d'un croisement signal-sur-bruit linéarisé en une étape, en analysant la dérivée de Fréchet de la carte polaire.

2. Validation Empirique :

   • Détection de Matrices (Matrix Sensing) : Des expériences contrôlées sur des tâches de régression linéaire sous-déterminée comparent Muon aux minimiseurs de norme nucléaire (résolus via CVXPY) et à la descente de gradient.
   • Pré-entraînement de Transformers : Pré-entraînement à petite échelle de NanoGPT (124M) et d'un contrôle ViT apparié sur CIFAR-10 pour observer la dynamique spectrale (rang stable, entropie spectrale) dans les régimes de deep learning.
   • Ablations : Tests sur le couplage de la décroissance des poids (couplé vs découplé) et le désalignement du cadre partagé pour vérifier la robustesse.

Contributions Clés

1. Théorème de Biais Spectral en une Étape :
   Sous des hypothèses d'alignement explicites, les auteurs prouvent que la mise à jour polaire est le choix unique maximisant l'entropie en une étape parmi les mises à jour bornées qui suivent les directions singulières du gradient mais ignorent le spectre actuel du poids. Ils établissent un critère d'aplatissement spectral : le spectre normalisé s'aplatit si et seulement si la géométrie de la mesure préserve préférentiellement les directions correspondant aux plus grandes valeurs singulières (quantifié par un critère de signe impliquant $\alpha_i q_i$).

2. Dynamique Exacte des Valeurs Singulières :
   L'article dérive les dynamiques exactes pour le flux polaire projeté, montrant que les valeurs singulières décroissent à des taux $\dot{\sigma}_i = -\eta \alpha_i$, où $\alpha_i$ dépend de l'alignement entre les vecteurs singuliers gauches du poids et l'espace nul de la matrice de mesure. Cela conduit à une caractérisation variationnelle où, à une norme de Frobenius fixée, l'unique minimiseur de la fonctionnelle spectrale est un spectre plat (toutes les valeurs singulières non nulles sont égales).

3. Séparation des Explications Existantes :

   • Normalisation : Les expériences montrent que la descente de gradient avec norme appariée ne reproduit pas le comportement de Muon.
   • Norme Nucléaire : Les résultats théoriques et empiriques démontrent que les points limites de Muon ne correspondent pas aux minimiseurs de norme nucléaire. Alors que la minimisation de la norme nucléaire favorise la concentration spectrale (faible rang), Muon favorise l'aplatissement spectral (rang stable élevé).

4. Utilité Dépendante du Régime :
   L'article établit que le biais de spectre plat de Muon n'est pas universellement supérieur. Dans le petit pré-entraînement de NanoGPT, Muon préserve le rang stable et améliore la perte de validation par rapport à AdamW. Cependant, dans un contrôle ViT apparié, le classement s'inverse, suggérant que le biais est bénéfique uniquement lorsque de nombreuses directions spectrales doivent rester actives.

Résultats

• Théoriques : Le Théorème 1 prouve que le flux auto-polaire projeté conduit le système vers un spectre plat si la géométrie de la mesure satisfait un critère de signe spécifique. Le Théorème 2 limite l'écart lorsque l'hypothèse de Cadre Partagé n'est que partiellement respectée. Le Théorème 4 fournit une formule pour la taille de lot critique $B_{crit}$ dans le SGD de la carte polaire, liant cela à la sensibilité au bruit de la carte polaire $S(\mu) = \sum_{i \neq j} (\sigma_i + \sigma_j)^{-2}$.
• Empiriques (Détection de Matrices) : Muon converge vers des solutions ayant des normes nucléaires significativement plus élevées (1,29x–2,02x) et une entropie spectrale plus élevée que l'interpolant de norme nucléaire minimale, falsifiant l'hypothèse selon laquelle Muon minimise la norme nucléaire.
• Empiriques (Transformers) : Dans NanoGPT, Muon maintient une entropie spectrale et un rang stable par couche plus élevés que AdamW, qui tend à faire s'effondrer le rang dans les tâches de langage. Inversement, dans le contrôle ViT, AdamW surpasse Muon, soutenant la conclusion de la "dépendance au régime".
• Sensibilité au Bruit : La formule dérivée de la taille de lot critique montre que $B_{crit}$ est gouvernée par la sensibilité de la carte polaire au bruit, qui est gonflée par des valeurs singulières petites ou presque dégénérées. Les spectres plats à rayon de Frobenius fixe minimisent cette sensibilité.

Signification et Revendications
L'article prétend fournir la première caractérisation structurelle du biais implicite de Muon comme un mécanisme de spectre de mise à jour plat plutôt qu'un effet de faible rang ou de simple normalisation.

• Portée Modeste : Les auteurs déclarent explicitement que les résultats sont dépendants du régime. Muon n'est pas un optimiseur universel ; son utilité provient des régimes spectraux où maintenir l'activité de nombreuses directions est bénéfique.
• Limites Théoriques : Les théorèmes sont locaux et structurels, reposant sur des hypothèses telles que des écarts de valeurs singulières simples et des géométries de mesure spécifiques. L'article ne prétend pas expliquer le succès de Muon dans tous les scénarios d'entraînement de LLM à grande échelle, notant que la réduction des couches de transformer au modèle de régression sous-déterminée étudié reste une "observation conditionnelle" nécessitant des hypothèses non vérifiées sur les rangs d'activation.
• Aperçu Pratique : Le travail offre un outil de diagnostic aux praticiens : si une application nécessite de préserver le rang stable et une activité spectrale large, le biais de spectre plat de Muon est avantageux ; si le problème est intrinsèquement de faible rang ou bien conditionné, le biais peut être inutile ou préjudiciable.

L'article conclut en identifiant des problèmes ouverts, notamment la caractérisation des géométries de mesure spécifiques qui garantissent globalement le critère d'aplatissement et l'extension de l'analyse de sensibilité au bruit aux cas rectangulaires et de rang déficient.