Résumé technique : Muon en tant que connexion résiduelle

Énoncé du problème

Muon est apparu comme un optimiseur hautement efficace pour l'entraînement de grands réseaux de neurones, atteignant des performances de pointe à travers diverses architectures, allant des CNN aux LLM. Malgré son succès empirique, les explications existantes de l'efficacité de Muon reposent sur des cadres mathématiques divers et sophistiqués, incluant la descente de gradient de norme spectrale, l'optimisation de région de confiance non euclidienne, les méthodes de Frank–Wolfe sur des boules de norme spectrale, l'apprentissage de longue traîne (long-tail learning) et l'analyse de courbure.

Le présent article répond au besoin d'une explication plus intuitive et mécaniste qui ne nécessite pas une formation mathématique approfondie. Plus précisément, il cherche à expliquer pourquoi Muon fonctionne bien en recadrant la dynamique d'optimisation en termes d'utilisabilité de la représentation pour les couches en aval, plutôt qu'en se concentrant uniquement sur la fidélité immédiate du gradient ou la réduction de la perte locale.

Méthodologie

Les auteurs proposent une interprétation novatrice : Muon fonctionne comme une connexion résiduelle implicite.

Cadre théorique

1. L'orthogonalisation comme résidualité : Dans les réseaux résiduels standards, la mise à jour est explicite : $x_{l+1} = x_l + \text{Block}(x_l)$, ou dans l'espace des poids, $W \leftarrow W + I$ (ou une matrice orthogonale $R$). Muon remplace la mise à jour standard $W \leftarrow W + \Delta W$ par une mise à jour orthogonalisée $W \leftarrow W + \text{Orth}(\Delta W)$.
2. Le compromis : L'article pose un compromis entre la fidélité du gradient (rester proche du gradient brut pour la réduction immédiate de la perte) et la préservation de la représentation (maintenir une structure de représentation facile à utiliser pour les couches en aval). L'orthogonalisation de la mise à jour sacrifie une partie de la fidélité immédiate du gradient mais injecte implicitement une transformation de type résiduelle qui préserve l'utilisabilité de la représentation en amont.

Validation expérimentale

Pour tester cette hypothèse, les auteurs ont mené des expériences contrôlées utilisant des modèles linéaires, évitant ainsi les complexités des réseaux non linéaires profonds afin d'isoler le mécanisme.

1. Expérience linéaire en deux phases :

   • Phase 1 (Optimisation locale) : Un réseau linéaire à deux couches ($\hat{y} = W_2 W_1 x$) est entraîné pour correspondre à une transformation cible $A$. $W_2$ est gelé, et $W_1$ est optimisé vers une cible locale $W_1^*$.
     • Hypothèse : SGD devrait surpasser Muon ici car SGD suit la direction de la descente la plus raide pour l'objectif local, tandis que Muon orthogonalise la mise à jour.
   • Phase 2 (Adaptation en aval) : Le $W_1$ appris est gelé, et $W_2$ est réinitialisé et optimisé via SGD pour minimiser $\|W_2 W_1 - A\|_F$.
     • Hypothèse : Si l'effet de "résiduel implicite" de Muon se vérifie, le $W_1$ appris par Muon devrait être plus facile à adapter pour $W_2$, malgré l'éloignement de la cible locale en Phase 1.

2. Expérience de l'ordonnancement $\tau$ :

   • Pour simuler un entraînement de bout en bout, les auteurs ont introduit un ordonnancement où les couches sont mises à jour par segments alternés de longueur $\tau$ (par exemple, $\tau$ étapes sur $W_1$, puis $\tau$ étapes sur $W_2$).
   • Cette configuration permet d'observer comment le "coût local" de Muon (pendant les mises à jour de $W_1$) interagit avec le "bénéfice en aval" (pendant les mises à jour de $W_2$) dans une boucle d'entraînement continue.

Résultats clés

Expérience linéaire en deux phases

• Phase 1 : Comme prévu, SGD a atteint une erreur de reconstruction locale plus faible ($0,0312$) par rapport à Muon ($0,0447$) après 600 étapes, confirmant le sacrifice de la fidélité du gradient local par Muon.
• Phase 2 : Lors de l'optimisation de $W_2$ avec le $W_1$ gelé, le modèle entraîné avec Muon en Phase 1 a nécessité moins d'étapes (13 321) pour atteindre le seuil d'erreur cible par rapport au modèle entraîné avec SGD (13 436 étapes).
• Interprétation : La représentation apprise par Muon, bien que sous-optimale pour la cible locale, possédait un spectre singulier plus "plat" (plus grande platitude spectrale), résultant en un problème mieux conditionné pour la couche en aval.

Expériences d'ordonnancement $\tau$

• Muon a systématiquement atteint la perte de composition cible plus rapidement que SGD à travers tous les ordonnancements testés, y compris les mises à jour conjointes complètes ($\tau=0$) et les mises à jour alternées ($\tau > 0$).
• Analyse de la dynamique :
  • Pendant les "segments bleus" (mise à jour de $W_1$ uniquement), SGD a souvent réduit la perte plus rapidement au début, reflétant le coût de fidélité du gradient de Muon.
  • Pendant les "segments rouges" (mise à jour de $W_2$ uniquement), la trajectoire entraînée par Muon a permis une convergence plus rapide.
  • Crucialement, l'amélioration de $W_2$ dans la trajectoire de Muon a rendu les sous-problèmes de $W_1$ ultérieurs plus faciles à résoudre, créant une boucle de rétroaction qui compense le ralentissement local initial.

Contributions principales

1. Interprétation mécaniste : L'article fournit une explication simple et intuitive du succès de Muon : il agit comme une connexion résiduelle implicite. Il équilibre le besoin immédiat de descente de gradient avec le besoin à long terme de préservation de la représentation.
2. Nouvelle perspective de conception : Il déplace le paradigme de la conception des optimiseurs, qui ne doit plus juger les mises à jour uniquement par la diminution de la perte locale, mais aussi par la manière dont les mises à jour façonnent les représentations pour les couches en aval.
3. Preuve empirique : À travers des expériences linéaires contrôlées, les auteurs démontrent que sacrifier la vitesse d'optimisation locale peut engendrer une convergence de bout en bout plus rapide en améliorant le nombre de conditionnement du problème d'optimisation en aval.

Signification et affirmations

L'article affirme que cette interprétation complète les explications existantes (spectrales, de courbure, de région de confiance) plutôt que de les remplacer. En présentant Muon comme une connexion résiduelle implicite, les auteurs suggèrent que :

• Les optimiseurs doivent équilibrer les objectifs locaux et en aval : Un optimiseur efficace n'a pas besoin d'être jugé uniquement par la diminution immédiate de la perte sur un sous-problème par couche.
• Le contexte compte : Les bénéfices du comportement de type Muon (sacrifier la fidélité locale pour l'utilisabilité en aval) sont particulièrement précieux lorsque les couches en aval disposent d'une capacité et d'un temps d'entraînement suffisants pour exploiter des représentations mieux conditionnées, notamment dans les chaînes profondes de transformations composées.
• Directions futures : Cette perspective suggère des conceptions d'optimiseurs adaptatifs qui contrôlent la force de l'orthogonalisation ou de l'aplatissement spectral en fonction des spectres des couches, du conditionnement ou de la progression en aval.

Les auteurs maintiennent un ton modeste, reconnaissant que leurs expériences sont délibérément contrôlées et linéaires. Ils ne prétendent pas avoir prouvé ce mécanisme dans tous les réseaux non linéaires modernes, mais soutiennent qu'il offre un mécanisme plausible et intuitif pour expliquer les avantages empiriques observés de Muon.