Muon+: Towards Better Muon via One Additional Normalization Step

Ce papier propose Muon+, une amélioration simple et efficace de l'optimiseur Muon intégrant une étape de normalisation supplémentaire, qui démontre des performances supérieures en termes de perplexité lors de l'entraînement préliminaire de modèles de langage à grande échelle.

L'essentiel

🚀 MUON+ : L'optimiseur qui apprend à "respirer" pour mieux apprendre

Imaginez que vous entraînez un cerveau artificiel (un modèle d'IA) pour qu'il écrive des textes, réponde à des questions ou crée des histoires. Pour apprendre, ce cerveau doit ajuster des milliards de petits boutons (les paramètres) en fonction de ses erreurs.

C'est là qu'intervient l'optimiseur. C'est le "professeur" qui dit au cerveau : "Tu as fait une erreur ici, tourne ce bouton un tout petit peu vers la gauche."

1. Le Problème : Le Professeur Muon (La version précédente)

Récemment, un nouveau professeur nommé Muon est arrivé en classe. Il était très doué. Au lieu de simplement dire "tourne à gauche", il utilisait une technique mathématique sophistiquée (l'orthogonalisation) pour s'assurer que les corrections ne se contredisaient pas entre elles.

• L'analogie : Imaginez un groupe de danseurs. Si chacun bouge dans une direction aléatoire, ils se cognent et la chorégraphie est chaotique. Muon, c'est comme un chef de danse qui force chaque danseur à bouger dans une direction parfaitement perpendiculaire aux autres. Cela évite les collisions et rend la danse plus fluide.

Cependant, même avec ce chef de danse, il y avait un petit problème : parfois, les mouvements restaient un peu trop "lourds" ou déséquilibrés.

2. La Solution : MUON+ (Le petit coup de pouce)

Les auteurs de cet article ont dit : "Attendez, Muon est génial, mais il manque une petite étape pour être parfait."

Ils ont ajouté une seule étape supplémentaire : la normalisation.

• L'analogie du chef d'orchestre :
  Imaginez que Muon a déjà fait en sorte que les musiciens jouent dans des directions différentes (orthogonalisation). Mais certains jouent fort (trop de volume) et d'autres jouent très doucement. Le résultat est déséquilibré.

  MUON+ ajoute un ingénieur du son juste après la direction de Muon. Son travail ? Regarder chaque musicien et dire : "Toi, baisse un peu le volume. Toi, monte un peu. On veut que tout le monde soit au même niveau d'intensité."

  C'est cette étape de "mise à niveau" (normalisation) qui fait toute la différence.

3. Les Résultats : Pourquoi c'est génial ?

Les chercheurs ont testé cette nouvelle méthode (MUON+) sur des modèles de tailles très différentes, du petit (taille d'un chat) au géant (taille d'un éléphant), et même sur des entraînements très longs.

Voici ce qu'ils ont découvert :

• Plus rapide et plus stable : Avec MUON+, le modèle apprend plus vite et fait moins d'erreurs (ce qu'on appelle la "perplexité" diminue).
• Robuste : Même si on change un peu la façon d'enseigner (le taux d'apprentissage), MUON+ continue de bien fonctionner, alors que l'ancien Muon pouvait parfois s'effondrer.
• Économique : Comme le modèle apprend mieux, on a besoin de moins de calculs pour atteindre le même résultat. C'est comme apprendre à conduire en moins de temps avec moins d'essence.

4. En résumé, c'est quoi la magie ?

L'article nous apprend quelque chose de très important : la structure compte plus que la complexité.

On pensait qu'il fallait des formules mathématiques ultra-complexes pour améliorer l'IA. En réalité, il suffisait d'ajouter une petite règle de bon sens : "Après avoir organisé les mouvements, assure-toi qu'ils sont tous de la même taille."

C'est un peu comme si vous prépariez un gâteau. Vous avez déjà les bons ingrédients (Muon). Mais si vous ne mélangez pas bien la pâte pour que tout soit homogène (la normalisation de MUON+), le gâteau ne sera pas aussi bon. MUON+, c'est simplement le fouet qui rend le gâteau parfait.

Le mot de la fin :
MUON+ ne remplace pas la révolution Muon, il la complète avec une touche de simplicité. C'est la preuve que parfois, pour faire avancer l'intelligence artificielle, il ne faut pas toujours construire des fusées plus grosses, mais juste ajuster un petit boulon au bon endroit.

Résumé technique

1. Problématique

L'entraînement de modèles de langage fondation (LLM) à grande échelle (comme GPT, LLaMA, Gemini) devient prohibitif en termes de coûts de calcul. Bien que les optimiseurs classiques comme Adam et AdamW dominent encore le domaine, des recherches récentes se concentrent sur l'amélioration de l'efficacité de l'entraînement.

L'optimiseur Muon a émergé comme une solution prometteuse pour l'entraînement préliminaire (pre-training) de grands modèles. Son principe repose sur l'orthogonalisation de la matrice de moment (momentum) via des itérations de Newton-Schulz, visant à contrer l'effondrement du rang des gradients (gradient rank collapse). Bien que Muon ait démontré une grande évolutivité et soit utilisé par des modèles de pointe (Kimi, GLM), les auteurs constatent qu'il existe une marge de manœuvre pour améliorer sa stabilité et ses performances finales.

Le problème central abordé est l'optimisation de la stabilité et de la convergence de Muon lors de l'entraînement à très grande échelle, en particulier dans des régimes de calcul optimal et lors d'entraînements prolongés (overtraining) avec des ratios tokens/paramètres (T2P) élevés.

2. Méthodologie : MUON+

Les auteurs proposent MUON+, une amélioration simple mais efficace de l'optimiseur Muon. L'idée centrale est d'introduire une étape de normalisation supplémentaire immédiatement après l'orthogonalisation de la matrice de moment, mais avant la mise à jour des poids.

Algorithme de mise à jour

Contrairement à Muon standard qui met à jour les poids $W_t$ directement à partir de la matrice orthogonalisée $O_t$, MUON+ applique un opérateur de normalisation $Norm^{(d)}$ :

1. Calcul du moment : $M_t = \mu M_{t-1} + (1 - \mu)G_t$
2. Orthogonalisation : $O_t = \text{Ortho}(M_t)$ (via Newton-Schulz)
3. Normalisation (Nouvelle étape) : $\tilde{O}_t = \text{Norm}^{(d)}(O_t)$
4. Mise à jour des poids : $W_t = W_{t-1} - \eta \cdot \sqrt{m/n} \cdot \tilde{O}_t$

L'opérateur de normalisation $\text{Norm}^{(d)}$ peut être appliqué selon différentes directions :

• Colonne (col) : Normalisation $\ell_2$ des colonnes.
• Ligne (row) : Normalisation $\ell_2$ des lignes.
• Combiné : Séquences de normalisations (ex: col_row ou row_col).

Les auteurs montrent que cette normalisation structurelle des mises à jour orthogonales est le moteur principal de l'amélioration, plus que les mécanismes adaptatifs complexes (comme l'adaptation du second moment) présents dans d'autres variantes récentes.

3. Contributions Clés

1. Proposition de MUON+ : Introduction d'une étape de normalisation post-orthogonalisation simple, augmentant la stabilité de l'optimisation sans complexité computationnelle significative.
2. Validation à grande échelle : Évaluation exhaustive sur une large gamme de modèles (de 60M à 1 milliard de paramètres) et d'architectures (GPT et LLaMA).
3. Étude des régimes d'entraînement :
   • Optimisation du calcul (Compute-optimal) : Entraînement avec un ratio T2P d'environ 20.
   • Sur-entraînement (Overtraining) : Extension du ratio T2P à un niveau industriel d'environ 200, démontrant la robustesse de la méthode sur des durées d'entraînement très longues.
4. Analyse par ablation : Identification que la normalisation elle-même est le facteur dominant de la performance, surpassant les ajustements de second moment ou les formulations de variétés complexes dans ce contexte.
5. Robustesse aux hyperparamètres : Démonstration que MUON+ est moins sensible au choix du taux d'apprentissage (learning rate) que Muon standard, en particulier pour les grands modèles.

4. Résultats Expérimentaux

Les expériences ont été menées sur le jeu de données FineWeb avec une précision mixte (bfloat16) sur des GPU H100/A100.

Performance sur les modèles GPT (130M - 774M)

• MUON+ surpasse systématiquement Muon sur tous les modèles.
• GPT-Small (124M) : Réduction de la perplexité de validation de 2.02 (29.66 $\to$ 27.64).
• GPT-Base (362M) : Réduction de 1.72 (21.70 $\to$ 19.98).
• GPT-Large (774M) : Réduction de 0.91 (17.82 $\to$ 16.91).

Performance sur les modèles LLaMA (60M - 1B)

• MUON+ bat à la fois AdamW et Muon standard.
• LLaMA-1B : La perplexité passe de 10.68 (Muon) à 10.31 (MUON+).
• Les gains sont constants à travers toutes les échelles de modèles.

Robustesse et Sur-entraînement (Overtraining)

• Dans des scénarios de sur-entraînement (72 milliards de tokens, ratio T2P $\approx$ 200), MUON+ maintient une perplexité inférieure et une stabilité d'optimisation supérieure.
  • GPT-Base (72B tokens) : 16.97 (Muon) $\to$ 15.84 (MUON+).
  • LLaMA-350M (72B tokens) : 11.48 (Muon) $\to$ 11.03 (MUON+).
• Les courbes de perte montrent que l'écart de performance se maintient tout au long de l'entraînement, indiquant que la méthode ne se dégrade pas dans les phases tardives.

Ablations et Sensibilité

• Directions de normalisation : Les combinaisons col_row et row_col offrent les meilleures performances. La normalisation par ligne (row) tend à être légèrement supérieure à la normalisation par colonne (col).
• Taux d'apprentissage : MUON+ présente une dégradation de performance beaucoup plus faible que Muon lorsque le taux d'apprentissage est sous-optimal (trop élevé), ce qui facilite le réglage des hyperparamètres.
• Méthodes d'approximation : La méthode fonctionne bien avec différentes approximations de la décomposition SVD (Newton-Schulz, PolarExpress, etc.), confirmant son indépendance vis-à-vis de la méthode d'orthogonalisation spécifique.

5. Signification et Impact

Ce travail démontre que la complexité algorithmique n'est pas toujours nécessaire pour améliorer les optimiseurs de grands modèles. En ajoutant une simple étape de normalisation, MUON+ améliore significativement la qualité finale des modèles et la stabilité de l'entraînement.

• Efficacité : La méthode est applicable sans coût computationnel majeur, rendant l'entraînement de modèles de plus en plus grands plus efficace.
• Généralité : Les résultats sont cohérents sur différentes architectures (GPT, LLaMA) et tailles, suggérant que la normalisation structurelle des mises à jour orthogonales est un principe fondamental bénéfique pour l'entraînement des LLM.
• Industrialisation : La capacité de MUON+ à performer dans des régimes de sur-entraînement (T2P $\approx$ 200) en fait un candidat idéal pour les pipelines de pré-entraînement industriels où l'on cherche à maximiser l'utilisation des données.

En résumé, MUON+ offre une voie simple et robuste pour pousser les limites de l'efficacité de l'entraînement des modèles de langage, validant l'hypothèse que la normalisation des mises à jour orthogonales est un levier critique souvent sous-estimé.