🤖 machine learning

MONA: Muon Optimizer with Nesterov Acceleration for Scalable Language Model Training

L'article présente MONA, un nouvel optimiseur qui intègre une accélération de style Nesterov dans le cadre Muon pour échapper aux minima locaux aigus et atteindre une convergence et des performances en aval de pointe sur des modèles de langage allant de 1 à 68 milliards de paramètres.

Auteurs originaux : Jiacheng Li, Jianchao Tan, Hongtao Xu, Jiaqi Zhang, Yifan Lu, Yerui Sun, Yuchen Xie, Xunliang Cai

Publié 2026-05-27

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Jiacheng Li, Jianchao Tan, Hongtao Xu, Jiaqi Zhang, Yifan Lu, Yerui Sun, Yuchen Xie, Xunliang Cai

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 essayez d'enseigner à un robot géant et surdoué comment parler la langue humaine. Pour ce faire, vous devez ajuster des millions de petits boutons à l'intérieur de son cerveau. Le processus consistant à trouver les réglages parfaits pour ces boutons s'appelle « l'entraînement », et l'outil que vous utilisez pour tourner ces boutons s'appelle un optimiseur.

Pendant longtemps, l'outil standard pour ce travail s'appelait AdamW. Il est fiable, comme un randonneur constant et prudent qui fait de petits pas et consulte la carte en permanence. Récemment, un nouvel outil appelé Muon est arrivé. Muon est comme un randonneur qui observe l'ensemble de la chaîne de montagnes d'un seul coup d'œil, traitant des groupes de boutons comme une seule unité. Cela rend Muon plus rapide et plus efficace, mais il présente toujours un défaut : il peut rester coincé dans une « vallée abrupte ».

Le problème : Le piège de la vallée abrupte

Imaginez que le paysage du processus d'apprentissage du robot soit un terrain vallonné.

Les vallées plates sont de bons endroits pour s'installer ; elles signifient que le robot a trouvé une solution stable et générale qui fonctionne bien dans de nombreuses situations.
Les vallées abruptes sont des fosses étroites et profondes. Si le robot tombe dedans, il peut sembler qu'il a trouvé le point le plus bas, mais il est en réalité coincé dans un endroit très spécifique et fragile. Si vous le poussez légèrement, il sort de son endroit « parfait » et performe mal.

AdamW et Muon peuvent tous deux accidentellement tomber dans ces vallées abruptes et y rester coincés. Ils ne disposent pas d'un moyen de « sentir » la forme du sol à l'avance pour savoir si un endroit est trop abrupt.

La solution : MONA (Muon avec accélération de Nesterov)

Les auteurs de cet article ont créé un nouvel outil appelé MONA. Imaginez MONA comme Muon équipé d'une paire de lunettes à sonar.

Voici comment cela fonctionne en utilisant une analogie simple :

L'ancienne méthode (Muon) : Muon regarde la pente exactement là où il se tient et descend. C'est excellent pour avancer efficacement, mais il ne sait pas si le sol devant est une pente douce ou un bord de falaise.
La nouvelle méthode (MONA) : Avant que Muon ne fasse un pas, MONA examine comment la pente a changé entre le dernier pas et le pas actuel.
- Si la pente a changé très soudainement (comme heurter une falaise abrupte), MONA sait : « Ouh là, c'est une vallée abrupte ! Poussons plus fort pour rebondir et sortir d'ici. »
- Si la pente a changé très lentement (une vallée douce et plate), MONA dit : « Cela semble sûr. Installons-nous ici. »

MONA ajoute un « terme d'accélération » spécial aux mathématiques. Il calcule la différence entre la direction actuelle et la direction précédente. Cette différence agit comme un capteur qui détecte l'« abruptitude » du terrain. Si le terrain est abrupt, il pousse le robot hors du fossé. S'il est plat, il laisse le robot se reposer.

Que ont-ils découvert ?

Les chercheurs ont testé cet nouvel outil sur trois tailles différentes de modèles de langage (petit, moyen et énorme), allant de 1 milliard à 68 milliards de paramètres. Ils les ont entraînés sur d'énormes quantités de texte (jusqu'à 1 trillion de mots).

Meilleur entraînement : Dans chaque test, MONA a aidé les modèles à apprendre plus vite et à atteindre un état final meilleur que l'ancien standard (AdamW) et le nouveau Muon.
Robots plus intelligents : Les modèles entraînés avec MONA étaient meilleurs dans des tâches générales, des problèmes de mathématiques et l'écriture de code. Par exemple, sur un modèle de 68 milliards de paramètres, MONA a obtenu les meilleurs scores sur les benchmarks pour les mathématiques et la programmation.
Ajustement fin : Même après l'entraînement initial, lorsqu'ils ont donné aux modèles un entraînement supplémentaire spécifique (comme leur apprendre à écrire spécifiquement du code), les modèles qui avaient commencé avec MONA ont mieux performé que ceux qui avaient commencé avec Muon.

Rendre cela pratique (MONA-Lite)

Les auteurs ont également réalisé que l'ajout de ce système « sonar » utilisait un peu plus de mémoire informatique. Pour corriger cela, ils ont créé une version allégée appelée MONA-Lite.

Ils ont utilisé une astuce pour compresser les données de mémoire (comme passer de la vidéo haute définition à la définition standard) et une méthode pour calculer les changements de pente à la volée sans stocker de fichiers supplémentaires.
Cela a réduit la mémoire supplémentaire nécessaire d'environ 75 %, rendant son utilisation facile même sur des ordinateurs à mémoire limitée, sans perdre les avantages.

Résumé

En bref, MONA est une mise à niveau de l'optimiseur Muon. Il conserve la vitesse et l'efficacité de Muon mais ajoute un « capteur de courbure » qui aide l'IA à éviter de rester coincée dans de mauvais endroits abrupts pendant l'apprentissage. Cela se traduit par des modèles de langage plus intelligents et plus capables, meilleurs en mathématiques, en programmation et en raisonnement général, tout en étant suffisamment efficaces pour fonctionner sur du matériel standard.

Résumé Technique : MONA – Optimiseur Muon avec Accélération Nesterov

1. Énoncé du Problème

L'entraînement de grands modèles de langage (LLM) exige de plus en plus des optimiseurs dotés d'une efficacité d'échantillonnage supérieure. Bien que AdamW ait été la norme pendant plus d'une décennie, des alternatives récentes comme Muon ont émergé. Muon améliore les méthodes du premier ordre en traitant les matrices de poids comme des unités géométriques uniques et en appliquant une orthogonalisation de matrice (via l'itération de Newton-Schulz) au tampon de momentum. Cette approche s'aligne sur la descente la plus raide sous la norme spectrale, offrant une efficacité computationnelle d'environ $2\times$ supérieure à AdamW dans l'entraînement à grande échelle de mélanges d'experts (MoE).

Cependant, Muon partage une limitation fondamentale avec toutes les méthodes de gradient du premier ordre : il manque de mécanismes explicites pour explorer le paysage de perte au-delà des informations de gradient locales. Par conséquent, Muon peut rester piégé dans des minima locaux aigus, en particulier dans les scénarios d'entraînement avec de grands lots. Les minima aigus sont associés à une mauvaise généralisation, tandis que les minima plus plats tendent à produire de meilleures performances en aval. Bien que l'optimiseur ALTO ait introduit un terme d'accélération basé sur la moyenne mobile exponentielle (EMA) des différences de gradient pour échapper aux minima aigus, sa conception en tant qu'adaptateur général pour Adam et LAMB présente des défis d'intégration avec le processus d'orthogonalisation hautement non linéaire de Muon.

2. Méthodologie : MONA

Les auteurs proposent MONA (Muon Optimizer with Nesterov Acceleration), qui intègre de manière transparente un terme d'accélération conscient de la courbure directement dans le cadre d'orthogonalisation de Muon.

Algorithme Central

MONA modifie le pipeline Muon standard en introduisant un terme d'accélération avant l'accumulation de momentum et l'orthogonalisation. Le processus de mise à jour pour un paramètre matriciel $W$ à l'itération $k$ se déroule comme suit :

Calcul de la Différence de Gradient : Calculer la différence entre le gradient stochastique actuel $G_k$ et le gradient précédent $G_{k-1}$ :
$D_k = G_k - G_{k-1}$
Mise à Jour du Tampon d'Accélération : Maintenir une moyenne mobile exponentielle de ces différences :
$A_k = \beta_a A_{k-1} + (1 - \beta_a) D_k$
Ici, $\beta_a$ contrôle la mémoire de l'accélération.
Gradient Conscient de la Courbure : Construire un gradient augmenté $\tilde{G}_k$ en ajoutant le terme d'accélération mis à l'échelle au gradient brut :
$\tilde{G}_k = G_k + \alpha A_k$
Le coefficient $\alpha$ met l'accélération à l'échelle. La motivation théorique suggère que $G_k - G_{k-1}$ approxime $H_k(\theta_k - \theta_{k-1})$ , où $H_k$ est le Hessien. Ainsi, le terme pointe loin des minima aigus (où les changements de gradient sont importants) et reste faible dans les régions plates.
Orthogonalisation Muon : Le gradient augmenté est injecté dans les étapes standard de momentum et d'orthogonalisation de Muon :
$M_k = \mu M_{k-1} + \tilde{G}_k$
$O_k = \text{Newton-Schulz}(M_k)$
$W_{k+1} = W_k - \eta (\gamma O_k + \lambda W_k)$

Raison d'être de la Conception

Appliquer le terme d'accélération avant l'orthogonalisation est crucial. S'il était appliqué après, il perturberait la structure orthogonale sur laquelle Muon s'appuie pour la régularisation de norme spectrale. En transformant d'abord le gradient brut en une direction consciente de la courbure, le tampon de momentum capture ces directions enrichies, qui sont ensuite traitées via le pipeline standard de Muon.

Pour les paramètres à valeurs vectorielles (par exemple, les embeddings, les biais), MONA utilise AdamW par défaut, respectant ainsi les conventions existantes de Muon.

3. Analyse Théorique

L'article fournit une analyse de convergence pour MONA sous des hypothèses standard (lissage L, gradients non biaisés avec variance bornée, et alignement directionnel attendu).

Vitesse de Convergence : Les auteurs prouvent que MONA atteint une vitesse de convergence de $O(K^{-1/2})$ dans des contextes non convexes, correspondant aux taux asymptotiques de SGD et AdamW. Le terme d'accélération affecte les constantes dans la borne mais pas le taux asymptotique.
Échappement aux Minima Aigus : L'analyse formalise comment le terme d'accélération résiste sélectivement aux minima aigus. Près d'un point stationnaire, l'accélération attendue $E[A_k]$ est proportionnelle à la somme pondérée par le Hessien des directions passées. Pour les directions avec de grandes valeurs propres (minima aigus), l'amplitude de l'accélération est grande, repoussant l'optimiseur. Pour les petites valeurs propres (directions plates), l'accélération est faible, permettant une convergence stable. Ce mécanisme préserve la régularisation de norme spectrale de Muon tout en ajoutant une conscience de la courbure.

4. Résultats Empiriques

MONA a été évalué sur trois échelles de modèles de préentraînement MoE (1B, 6B et 68B paramètres) et une étape de fine-tuning supervisé (SFT).

Performance de Préentraînement

Modèles : MOE-1B-A0d2B (400B tokens), MOE-6B-A0d5B (1,2T tokens) et MOE-68B-A3B (700B tokens).
Métriques : La perte de validation a été surveillée sur les domaines du code, du raisonnement mathématique, de l'anglais général et des textes académiques chinois.
Constats : MONA a constamment obtenu une perte de validation inférieure à la fois Muon et AdamW sur les trois échelles de modèles et tous les domaines. L'écart de performance était particulièrement évident sur le plus grand modèle de 68B, où MONA a démontré une convergence plus stable.

Benchmarks en Aval (MOE-68B-A3B)

Le modèle de 68B entraîné avec MONA a été évalué sur une suite complète de benchmarks :

Capacité Générale : MONA a obtenu le score moyen le plus élevé (0,4557) sur des benchmarks incluant MMLU, MMLU-Pro, CMMLU, CEVAL et BBH, surpassant Muon (0,4478) et AdamW (0,4382).
Raisonnement Mathématique et Génération de Code : MONA a montré des performances supérieures sur GSM8K, GPQA et divers benchmarks de génération de code (HumanEval+, MBPP+, BigCodeBench, LiveCodeBench). Notamment, MONA a apporté des gains significatifs sur les tâches nécessitant des représentations de code transférables.
Fine-Tuning Supervisé (SFT) : Lorsque le checkpoint préentraîné avec MONA a subi un SFT sur des données de code, il a maintenu une perte d'entraînement plus faible tout au long du processus par rapport à la ligne de base préentraînée avec Muon. Les évaluations post-SFT sur la suite BigCode ont confirmé que les modèles initialisés avec MONA ont obtenu des scores plus élevés sur 6 tâches sur 7 par rapport aux modèles initialisés avec Muon.

Efficacité et Déploiement

Pour répondre à la surcharge mémoire liée au stockage du gradient précédent ( $G_{k-1}$ ) et du tampon d'accélération ( $A_k$ ), les auteurs ont introduit MONA-Lite :

Quantification : Stockage des états auxiliaires en bfloat16 (BF16) au lieu de float32.
Calcul en Flux : Calcul des différences de gradient in situ pour éliminer le besoin d'un tampon dédié $G_{k-1}$ .
Résultat : Ces stratégies réduisent la surcharge mémoire supplémentaire d'environ 75 % sans sacrifier la qualité de l'entraînement. L'analyse de la surcharge computationnelle a montré que MONA est seulement ~1 % plus lent au niveau de l'étape de l'optimiseur, sans différence pratique dans le temps d'itération de bout en bout par rapport à Muon.

5. Signification et Revendications

L'article revendique que MONA comble avec succès le fossé entre l'efficacité géométrique de Muon et le besoin d'une accélération consciente de la courbure. Sa signification principale réside dans :

Intégration Algorithmique : Démontrer que des termes d'accélération peuvent être intégrés dans des cadres d'orthogonalisation de matrices sans perturber leurs propriétés géométriques.
Évolutivité : Valider la méthode sur une large gamme de tailles de modèles (de 1B à 68B paramètres) et de nombres de tokens (jusqu'à 1 trillion de tokens), montrant des améliorations constantes par rapport à AdamW et Muon.
Généralisation : Fournir des preuves empiriques que l'échappement aux minima aigus lors du préentraînement conduit à une meilleure généralisation dans les tâches en aval, en particulier dans le raisonnement complexe et la génération de code.
Praticité : Offrir une variante économe en mémoire (MONA-Lite) qui rend l'optimiseur viable pour des environnements aux ressources contraintes.

Les auteurs concluent que MONA représente un optimiseur de pointe pour l'entraînement de LLM à grande échelle, atteignant une convergence et des performances en aval supérieures tout en maintenant l'efficacité computationnelle du cadre Muon.