🤖 machine learning

Muown: Row-Norm Control for Muon Optimization

Le papier présente Muown, un optimiseur plug-and-play qui contrôle explicitement les vecteurs de magnitude des lignes pour prévenir la dérive de la norme spectrale dans Muon, améliorant ainsi la stabilité de l'entraînement, réduisant la sensibilité à la décroissance des poids et atteignant une perplexité supérieure à travers différentes échelles de modèles.

Auteurs originaux : Kai Lion, Florian Hübler, Bingcong Li, Antonio Orvieto, Niao He

Publié 2026-05-12

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Kai Lion, Florian Hübler, Bingcong Li, Antonio Orvieto, Niao He

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous entraînez un cerveau numérique géant (un Modèle de Langage à Grande Échelle) à écrire des histoires, à résoudre des problèmes de mathématiques ou à discuter avec vous. Pour ce faire, le cerveau doit apprendre en ajustant des millions de petits boutons (poids) à l'intérieur de ses circuits.

Pendant longtemps, la meilleure façon de tourner ces boutons était une méthode appelée AdamW. Récemment, une nouvelle méthode plus rapide, appelée Muon, est arrivée. C'était comme passer d'un vélo à une voiture de sport ; elle apprenait plus vite et mieux. Mais il y avait un piège : la voiture de sport avait tendance à accélérer hors de contrôle.

Le Problème : Le Moteur « Dérivant »

L'article identifie un problème spécifique avec Muon. À mesure que le modèle s'entraîne, la « taille » de ses connexions internes (spécifiquement, la force des lignes de ses matrices de poids) commence à croître de manière incontrôlable.

Pensez-y comme à un ballon que l'on gonfle.

Muon est excellent pour déterminer dans quelle direction souffler le ballon (la direction).
Cependant, il continue de pomper plus d'air que nécessaire, ce qui fait que le ballon grossit de plus en plus jusqu'à ce qu'il éclate (erreurs numériques) ou devienne instable.
La solution standard consistait à attacher une ficelle autour du ballon et à la serrer chaque fois qu'il devenait trop gros (Décroissance des Poids). Mais les auteurs ont découvert que c'était comme utiliser un marteau-pilon pour casser une noix : cela ralentissait le processus d'apprentissage car cela comprimait tout le ballon, même les parties qui allaient bien.

Le Diagnostic : Deux Parties au Problème

Les auteurs, dirigés par Kai Lion, ont décidé d'ouvrir le ballon pour voir ce qui grandissait réellement. Ils ont réalisé que la « taille » de la connexion est composée de deux parties distinctes :

La Magnitude (Le Volume) : La taille globale de la ligne de nombres.
La Cohérence (La Forme) : La manière dont les nombres à l'intérieur de cette ligne sont disposés les uns par rapport aux autres.

Ils ont découvert que seule la Magnitude dérivait hors de contrôle. La Forme (Cohérence) se comportait en fait parfaitement bien et restait stable. Le ballon ne changeait pas de forme ; il grossissait simplement.

La Solution : Muown (La Pompe Intelligente)

Les auteurs ont créé un nouvel optimiseur appelé Muown.

Au lieu de traiter toute la connexion comme un gros bloc unique, Muown divise la tâche en deux missions séparées :

L'Équipe Direction (Muon) : Cette équipe continue de faire ce que Muon fait le mieux : déterminer la direction parfaite pour mettre à jour les poids. Ils laissent cette partie inchangée.
L'Équipe Volume (Nouvelle) : C'est la nouvelle partie. Ils prennent la « Magnitude » (le volume) et la traitent comme une variable séparée et explicite. Ils donnent à ce volume un ensemble spécial de règles (une géométrie mathématique spécifique appelée $\ell_\infty$ ) pour l'empêcher de croître trop vite.

L'Analogie :
Imaginez que vous conduisez une voiture où le volant et la pédale d'accélération sont collés ensemble. Si vous tournez le volant, la pédale d'accélération est aussi enfoncée, ce qui fait accélérer la voiture de manière incontrôlable.

Muon était comme cette voiture bloquée.
Muown est comme le découplage du volant de la pédale d'accélération. Vous pouvez toujours diriger parfaitement (en utilisant la logique de Muon), mais vous avez maintenant un régulateur de vitesse intelligent et séparé pour la pédale d'accélération (la magnitude) qui maintient la vitesse exactement là où elle doit être sans vous ralentir.

Les Résultats

L'article a testé cela sur des modèles allant de petits (124 millions de paramètres) à très grands (2,7 milliards de paramètres). Voici ce qu'ils ont trouvé :

Meilleures Performances : Muown a produit systématiquement de meilleurs résultats (une « perplexité » plus faible, qui est une mesure de la confusion du modèle) que Muon, AdamW et autres concurrents.
Plus Forgivant : Avec Muon, il fallait être très prudent avec vos paramètres (taux d'apprentissage). Si vous en choisissiez un mauvais, le modèle échouait. Muown est beaucoup plus robuste ; il fonctionne bien sur une large gamme de paramètres, comme une voiture avec une « zone idéale » de conduite plus large.
Aucun Coût Supplémentaire : Habituellement, ajouter un nouveau système de contrôle ralentit une voiture. Mais parce que les auteurs ont intégré ce nouveau « contrôle de volume » directement dans le flux de travail existant du moteur, cela n'a ajouté presque aucun temps supplémentaire au processus d'entraînement.
Stabilité : Le « ballon » a cessé de dériver. La norme spectrale (la taille des connexions) est restée stable, empêchant les erreurs numériques qui tourmentaient le Muon original.

Résumé

L'article soutient que l'instabilité de l'optimiseur Muon populaire n'était pas un défaut de sa capacité à trouver la direction, mais plutôt un manque de contrôle sur le « volume » de ses mises à jour. En séparant le « volume » de la « direction » et en les contrôlant avec des outils différents et spécialisés, Muown maintient l'entraînement rapide et stable sans avoir besoin de corrections lourdes comme la décroissance des poids. C'est un remplacement direct qui rend le processus d'entraînement plus fluide, plus rapide et plus fiable.

Résumé technique : Muown - Contrôle de la norme des lignes pour l'optimisation Muon

Énoncé du problème

Muon est apparu comme un optimiseur compétitif pour l'entraînement préalable des grands modèles de langage (LLM), surpassant souvent AdamW en exploitant la structure matricielle des couches linéaires des Transformers. Cependant, des observations récentes indiquent que le comportement de Muon à grande échelle est sensible à la décroissance des poids (weight decay) et sujet à des instabilités numériques, telles que des débordements en bfloat16 et des logits d'attention explosifs.

Un problème critique identifié est la dérive constante vers le haut de la norme spectrale ( $\|W_t\|_{S_\infty}$ ) des matrices de poids pendant l'entraînement lorsque Muon est utilisé sans décroissance des poids découplée. Cette dérive viole la condition spectrale requise pour un apprentissage de caractéristiques stable ( $\|W\|_{S_\infty} = \Theta(\sqrt{m/n})$ ). Bien que la décroissance des poids découplée puisse atténuer ce phénomène en rétrécissant uniformément la matrice de poids, il s'agit d'une intervention non ciblée qui met à l'échelle l'ensemble du spectre indépendamment des directions spécifiques qui entraînent la croissance, ralentissant potentiellement la convergence lors des premières phases d'entraînement. L'article vise à identifier le moteur empirique spécifique de cette dérive de la norme spectrale et à proposer un mécanisme de contrôle plus ciblé.

Méthodologie : Muown

Les auteurs introduisent Muown, un remplacement direct de Muon qui résout la dérive de la norme spectrale grâce à une stratégie de reparamétrisation plutôt que par des contraintes agrégées.

1. Décomposition de la norme spectrale

L'idée centrale est une décomposition de la norme spectrale en deux facteurs (Proposition 1) :
$\|W_t\|_{S_\infty}^2 = \|g_t\|_\infty^2 \cdot \lambda_{\max}(P_t C_t P_t)$
Où :

$g_t = \|W_t\|_{row}$ est le vecteur des magnitudes des lignes (normes euclidiennes des lignes).
$D_t$ est la matrice normalisée par ligne ( $W_t = \text{Diag}(g_t)D_t$ ).
$C_t = D_t D_t^\top$ est la matrice de Gram des lignes de norme unitaire.
$P_t$ est une matrice diagonale des magnitudes de lignes normalisées.
$\lambda_{\max}(P_t C_t P_t)$ représente le facteur de cohérence des lignes.

L'analyse empirique sur des modèles de 500M paramètres révèle que le facteur de magnitude des lignes ( $\|g_t\|_\infty$ ) est le moteur principal de la dérive de la norme spectrale, évoluant à pas égal avec $\|W_t\|_{S_\infty}$ . En revanche, le facteur de cohérence des lignes reste bien comporté et borné tout au long de l'entraînement.

2. Reparamétrisation et optimisation

Sur la base de ce diagnostic, Muown reparamétrise la matrice de poids $W$ comme suit :
$W(g, R) = \text{Diag}\left(\frac{g}{\|R\|_{row}}\right) R$
Ici, $g \in \mathbb{R}^m$ représente le vecteur explicite des magnitudes des lignes, et $R \in \mathbb{R}^{m \times n}$ représente la composante directionnelle.

Composante directionnelle ( $R$ ) : Muon est appliqué sans modification à $R$ , la traitant selon la géométrie de la norme spectrale (descente spectrale la plus raide).
Composante magnitude ( $g$ ) : Les magnitudes des lignes $g$ sont promues en variables explicites de l'optimiseur. Elles sont mises à jour selon la géométrie $\ell_\infty$ , qui est la norme duale naturelle pour la norme d'opérateur des matrices diagonales.
Règle de mise à jour : La mise à jour de $g$ utilise un optimiseur Adam (agissant comme un SignSGD lissé sous la géométrie $\ell_\infty$ ) pour gérer le bruit stochastique, tandis que la mise à jour de $R$ suit la procédure Muon standard avec une orthogonalisation Newton-Schulz.

Cette approche est mise en œuvre implicitement au sein de l'état de l'optimiseur, laissant le passage avant de l'architecture du modèle inchangé.

Contributions clés

Diagnostic de la dérive spectrale : L'article décompose la norme spectrale pour identifier la magnitude maximale des lignes ( $\|g_t\|_\infty$ ) comme le moteur empirique de la croissance de la norme spectrale sous Muon, la distinguant du facteur de cohérence des lignes stable.
Algorithme Muown : Un nouvel optimiseur qui traite les magnitudes des lignes comme des variables entraînables explicites mises à jour sous la géométrie $\ell_\infty$ , tout en préservant les mises à jour directionnelles de Muon. Cela évite le rétrécissement indiscriminé de la décroissance des poids et les priors d'orthogonalité des reparamétrisations spectrales.
Garanties de convergence : Les auteurs prouvent que Muown atteint des taux de convergence optimaux non convexes dans les régimes déterministe et stochastique. L'analyse est menée dans une norme duale mixte ( $\ell_1$ pour les magnitudes, Schatten-1 pour les directions) alignée sur les géométries de mise à jour. La borne théorique suggère une réduction du coefficient de bruit stochastique effectif par rapport à Muon.
Validation empirique : Des expériences extensives sur FineWeb-Edu à travers des tailles de modèles allant de 124M à 2,7B de paramètres démontrent des améliorations constantes.

Résultats expérimentaux

Améliorations de la perplexité : Muown atteint systématiquement une perplexité inférieure à Muon, SOAP, AdamW et Lion à toutes les échelles de modèles testées (124M à 2,7B) et pour tous les taux d'apprentissage.
Robustesse aux hyperparamètres :
- Taux d'apprentissage : Muown élargit le plateau des taux d'apprentissage quasi-optimaux à travers différentes largeurs de modèles, réduisant la sensibilité au réglage du taux d'apprentissage par échelle.
- Décroissance des poids : Muown atténue le besoin de réglage de la décroissance des poids ; il fonctionne bien sans décroissance des poids ( $\lambda=0$ ), alors que Muon nécessite un réglage minutieux (optimal $\lambda=0.1$ ) pour éviter la divergence ou des performances sous-optimales.
Stabilité : Muown empêche la dérive systématique de la norme spectrale observée dans Muon. Dans les régimes de taux d'apprentissage élevés (par exemple $4 \times 10^{-3}$ ), Muon diverge souvent, tandis que Muown reste stable et quasi-optimal.
Propriétés des gradients : L'analyse montre que Muown maintient un rang effectif normalisé plus élevé des gradients de mini-lots et présente un coefficient de bruit stochastique plus faible tout au long de l'entraînement par rapport à Muon.
Surcharge computationnelle : La méthode introduit une surcharge de temps d'étape négligeable. Lorsque l'optimiseur est fragmenté (sharded) sur les rangs (environnement distribué), les temps d'étape pour Muown et Muon deviennent indiscernables. La surcharge mémoire est minimale (linéaire par rapport au nombre de lignes), ajoutant seulement quelques mégaoctets par couche.

Importance et revendications

L'article revendique que la dérive de la norme spectrale observée dans Muon n'est pas une propriété irréductible de l'optimiseur, mais une conséquence d'une croissance incontrôlée des magnitudes des lignes. En isolant ce moteur et en le promouvant en variable explicite de l'optimiseur, Muown atteint :

Convergence améliorée : Une meilleure perplexité et stabilité sans besoin d'une décroissance des poids agressive.
Solidité théorique : Des taux de convergence optimaux dans un cadre sensible à la géométrie qui prend explicitement en compte la nature mixte des mises à jour.
Efficacité pratique : Un remplacement direct qui correspond au temps d'exécution de Muon dans les environnements distribués tout en offrant une robustesse supérieure aux choix d'hyperparamètres.

Les auteurs positionnent Muown comme une solution qui comble le fossé entre les exigences théoriques de stabilité de l'apprentissage de caractéristiques et les besoins pratiques de l'entraînement à grande échelle des LLM, offrant une alternative plus ciblée à la décroissance des poids découplée.