To Use or not to Use Muon: How Simplicity Bias in Optimizers Matters

Cet article examine les inconvénients potentiels de l'optimiseur Muon, en soulignant que son accélération du processus d'entraînement s'accompagne de la suppression d'un biais de simplicité naturel présent dans des méthodes comme SGD, ce qui peut nuire à la capacité du modèle à identifier des structures sous-jacentes communes et augmenter le risque de surapprentissage sur des caractéristiques spurious.

L'essentiel

🚀 Le Dilemme : La Vitesse contre la Sagesse

Imaginez que vous apprenez à jouer du piano. Vous avez deux méthodes pour apprendre un morceau complexe :

1. La méthode "SGD" (l'ancienne école) : Vous commencez par les notes les plus simples et les plus importantes. Vous maîtrisez la mélodie de base, puis vous ajoutez petit à petit les accords, puis les ornements. C'est lent, mais vous comprenez la structure de la musique.
2. La méthode "Muon" (la nouvelle star) : C'est un super-élève qui apprend tout d'un coup ! Il joue toutes les notes, les accords et les ornements simultanément, à une vitesse fulgurante. Il termine le morceau trois fois plus vite que vous.

Le problème ? Parce que Muon va si vite et apprend tout en même temps, il risque de ne jamais vraiment comprendre pourquoi la mélodie fonctionne. Il risque de mémoriser le morceau note par note sans en saisir l'âme, et de se tromper dès qu'on change légèrement la partition.

C'est exactement ce que disent les chercheurs Sara Dragutinović et Rajesh Ranganath dans leur article. Ils nous disent : "Attention à la vitesse ! Parfois, aller plus lentement permet de mieux comprendre."

---

🧠 Le Secret de la "Biais de Simplicité"

Pourquoi la méthode lente (SGD) est-elle parfois meilleure ? Parce qu'elle possède ce qu'on appelle un "biais de simplicité".

• L'analogie du détective :
  Imaginez que vous essayez de résoudre un mystère.
  • SGD est un détective méthodique. Il commence par les indices les plus évidents (le "qui", le "quoi"). Une fois qu'il a résolu cette partie, il passe aux détails plus fins. Il construit sa théorie étape par étape.
  • Muon est un détective qui regarde tous les indices en même temps avec des jumelles ultra-puissantes. Il voit tout instantanément.

Le problème avec Muon, c'est qu'en voulant tout voir tout de suite, il peut se laisser piéger par des fausses pistes (des détails qui semblent importants mais qui ne le sont pas).

⚠️ Les Deux Pièges où Muon Trébuche

Les chercheurs ont testé Muon et SGD dans deux situations spécifiques pour voir qui gagnait.

1. Le Piège de la "Mémoire vs Compréhension" (L'exemple du Routage)

Imaginez un serveur dans un restaurant qui doit apprendre à servir différents plats à différentes tables.

• La tâche : Apprendre que "le plat A va toujours à la table 1, le plat B à la table 2", peu importe si le client arrive par la porte de gauche ou de droite.
• Ce que fait SGD : Il comprend la règle générale. Il apprend que "Plat A = Table 1". Si un nouveau client arrive par une porte qu'il n'a jamais vue, il sait quand même où servir le plat.
• Ce que fait Muon : Il mémorise chaque combinaison spécifique qu'il a vue. "Si le client entre par la porte gauche et commande le plat A, alors table 1". Mais si le client entre par une porte nouvelle, Muon est perdu. Il a appris par cœur, pas par compréhension.

Leçon : Muon est excellent pour mémoriser des données, mais moins bon pour trouver la règle cachée qui permet de s'adapter à de nouvelles situations.

2. Le Piège de la "Fausse Indice" (L'exemple du MNIST)

Imaginez que vous devez reconnaître des chiffres écrits à la main (comme sur un formulaire).

• La triche : Dans les images d'entraînement, on a ajouté un petit point rouge dans le coin de chaque image. Pour le chiffre "1", le point est en haut à gauche. Pour le "2", il est en bas à droite.
• Ce que fait SGD : Au début, il regarde le chiffre lui-même (la forme). C'est plus difficile, donc il progresse lentement. Mais il finit par apprendre la vraie forme du chiffre.
• Ce que fait Muon : Comme il apprend tout en même temps, il remarque très vite que le point rouge est un indice facile. Il se dit : "Ah ! Si je vois le point rouge en haut à gauche, c'est un 1 !" Il triche.
• Le résultat : Si on lui montre un chiffre "1" sans le point rouge, Muon se trompe car il ne regarde plus le chiffre, mais le point. SGD, lui, a appris à lire le chiffre.

Leçon : Muon est trop rapide pour ignorer les "trucs" faciles (les corrélations fortuites), ce qui le rend fragile quand les données changent.

---

💡 La Conclusion : Pourquoi cela nous concerne ?

Pendant longtemps, tout le monde utilisait SGD (la méthode lente) ou Adam (une méthode intermédiaire). Récemment, Muon est arrivé sur le marché en promettant d'entraîner les intelligences artificielles beaucoup plus vite. Tout le monde s'est précipité pour l'utiliser.

Cet article est un avertissement bienveillant. Il dit :

"Oui, Muon est rapide. Oui, il gagne des compétitions de vitesse. Mais en allant si vite, il perd une qualité précieuse : la capacité de découvrir les structures simples et profondes qui rendent l'intelligence artificielle intelligente et robuste."

En résumé :

• Si vous voulez juste entraîner un modèle sur des données fixes et que la vitesse est la seule priorité, Muon est génial.
• Si vous voulez un modèle qui comprend le monde, généralise bien, et ne se fait pas piéger par des détails trompeurs, SGD (ou des méthodes similaires) reste souvent le meilleur choix.

La morale de l'histoire : En ingénierie, comme en cuisine, ce n'est pas toujours le plat qui est prêt le plus vite qui est le meilleur. Parfois, il faut laisser le temps à la recette de se développer pour en extraire toute la saveur. Ne choisissez pas votre outil de formation uniquement pour sa vitesse ; regardez aussi comment il apprend.

Résumé technique

Titre : Utiliser ou non Muon : L'importance du biais de simplicité dans les optimiseurs

1. Problématique

Depuis longtemps, l'optimiseur Adam est le choix par défaut pour l'entraînement des réseaux de neurones profonds. Récemment, un nouvel optimiseur nommé Muon (MomentUm Orthogonalized by Newton-Schulz) a gagné en popularité grâce à sa vitesse d'entraînement supérieure, devenant même un nouveau standard dans certains contextes (ex. : NanoChat, NanoGPT Speedrun).

Cependant, la communauté manque de compréhension théorique sur les biais inductifs introduits par Muon. Alors que la littérature se concentre principalement sur les gains de vitesse (convergence plus rapide), ce papier soulève une question critique : la vitesse de Muon s'accompagne-t-elle d'un coût caché en termes de qualité de la solution finale et de généralisation ? Les auteurs suggèrent que Muon pourrait éliminer un « biais de simplicité » naturellement présent dans des méthodes plus anciennes comme la Descente de Gradient (SGD), ce qui pourrait nuire à la capacité du modèle à découvrir des structures sous-jacentes communes et à éviter le surapprentissage de caractéristiques spurious (artificielles).

2. Méthodologie

Pour analyser les mécanismes sous-jacents, les auteurs adoptent une approche combinant théorie et expérimentation :

• Analyse Théorique (Réseaux Linéaires Profonds) :

  • Ils étudient des réseaux linéaires profonds à deux couches ($\hat{y} = VUx$) pour isoler la dynamique d'apprentissage.
  • Ils introduisent une version théorique simplifiée de Muon, appelée Spectral Gradient Descent (Spectral GD). Cette version utilise une décomposition en valeurs singulières (SVD) exacte et orthogonale les mises à jour, sans utiliser de momentum ni d'approximations de Newton-Schulz (pour la tractabilité mathématique).
  • Ils comparent les trajectoires d'apprentissage de la Descente de Gradient (GD) classique et du Spectral GD sur le paysage de perte.

• Expérimentations Empiriques :

  • Tâche de « Routage » (Shared Representations) : Un scénario multi-modal où le modèle doit apprendre une structure commune sous-jacente à partir de différentes sources d'entrée, mais où seules certaines paires entrée-sortie sont observées durant l'entraînement.
  • Corrélations Spurious (MNIST) : Utilisation du dataset MNIST avec l'ajout d'un pixel « spurious » (artificiel) dont l'intensité varie par classe. L'objectif est de voir si l'optimiseur apprend la forme du chiffre (généralisable) ou se fie au pixel artificiel (spurious).
  • Comparaison directe entre Muon (version complète), SGD et Adam.

3. Contributions Clés et Résultats Théoriques

A. Perte du Biais de Simplicité (Simplicity Bias)

La contribution théorique majeure réside dans la caractérisation des trajectoires d'apprentissage :

• GD (Descente de Gradient) : Suit une dynamique « de selle à selle » (saddle-to-saddle). Elle apprend les composantes principales de la solution séquentiellement, en commençant par les plus grandes valeurs singulières. Cela crée un biais de simplicité : le modèle augmente progressivement sa complexité (rang de la solution), agissant comme une régularisation implicite.
• Spectral GD (et Muon) : Apprend toutes les composantes principales simultanément et à la même vitesse. Il ne passe pas par les points de selle intermédiaires, ce qui explique sa rapidité. Cependant, cela signifie qu'il perd le biais de simplicité : il apprend tout en même temps, y compris le bruit et les détails fins, sans hiérarchiser les structures importantes.

B. Conséquences sur la Généralisation

La perte de ce biais de simplicité a des conséquences négatives démontrées expérimentalement :

1. Échec à apprendre les représentations partagées : Dans la tâche de routage, le SGD réussit à généraliser à des paires entrée-sortie non vues durant l'entraînement en découvrant la structure commune sous-jacente (faible rang). À l'inverse, Spectral GD (et Muon) mémorise les paires vues durant l'entraînement mais échoue à généraliser, conduisant à une solution de rang élevé et disjointe.
2. Sensibilité aux caractéristiques spurious : Dans l'expérience MNIST, le SGD apprend d'abord la structure dominante (la forme du chiffre) avant d'apprendre les caractéristiques spurious. Muon, apprenant tout simultanément, commence à utiliser les caractéristiques spurious beaucoup plus tôt.
   • Résultat : Si l'on utilise l'arrêt anticipé (early stopping), SGD obtient une meilleure précision sur les données sans caractéristiques spurious.
   • Nuance : Si la caractéristique spurious devient très dominante (intensité élevée), Muon peut surpasser SGD, car il n'a pas besoin d'attendre que la structure dominante soit apprise avant d'ajuster les autres composantes.

4. Signification et Implications

• Vitesse vs Qualité : Le papier démontre que la supériorité de Muon en termes de vitesse de convergence n'est pas universellement bénéfique. Elle peut être contre-productive lorsque la tâche nécessite de découvrir une structure latente simple ou de généraliser hors distribution.
• Rôle des Biais Inductifs : Les optimiseurs ne sont pas de simples accélérateurs ; ils imposent des biais inductifs fondamentaux. Le biais de simplicité de SGD, souvent perçu comme un frein à la vitesse, est en réalité un mécanisme crucial pour la robustesse et la généralisation.
• Avertissement pour la conception d'optimiseurs : Les auteurs avertissent que le développement de nouveaux optimiseurs (comme Muon) doit évaluer non seulement la vitesse, mais aussi les trajectoires d'apprentissage et les biais induits. Une optimisation trop « gourmande » (apprenant tout en même temps) peut mener à des solutions qui mémorisent le bruit plutôt que de comprendre la logique du problème.

Conclusion

Ce papier remet en question l'adoption aveugle de Muon comme nouvel optimiseur par défaut. Il propose que le choix de l'optimiseur dépende de la nature du problème : Muon est excellent pour les tâches où l'apprentissage uniforme et rapide de toutes les modalités est souhaitable (ex. : données déséquilibrées), mais SGD reste supérieur pour les tâches nécessitant la découverte de structures causales, la généralisation hors distribution et la résistance aux corrélations spurious. L'étude invite à une évaluation plus nuancée des optimiseurs, au-delà de la simple métrique de vitesse.