🤖 machine learning

Aurora: A Leverage-Aware Spectral Optimizer

Le papier introduit Aurora, un optimiseur spectral sensible au levier qui impose l'uniformité des lignes dans les mises à jour des paramètres de matrice tout en préservant la géométrie du facteur polaire de Muon, empêchant ainsi la stagnation des neurones et atteignant des performances de pointe, particulièrement lors de l'entraînement de couches MLP très larges.

Auteurs originaux : Alec Dewulf, Dhruv Pai, Li Yang, Ashley Zhang, Ben Keigwin

Publié 2026-06-29

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Alec Dewulf, Dhruv Pai, Li Yang, Ashley Zhang, Ben Keigwin

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 formez une équipe massive de travailleurs (un réseau de neurones) pour résoudre des puzzles complexes. Dans cette équipe, il existe des départements spécifiques appelés couches MLP où le travail est divisé en stations individuelles. Chaque station est un « neurone » responsable d'une infime partie de la réponse finale.

Le document présente un nouveau manager nommé Aurora qui corrige une faille majeure dans la façon dont l'ancien manager, Muon, répartissait le travail.

Le Problème : La Boucle du « Travailleur Paresseux »

Sous l'ancien système (Muon), le manager regardait la performance de toute l'équipe et tentait d'équilibrer le travail. Cependant, pour certains types de tâches (représentées mathématiquement par des « matrices hautes »), Muon a accidentellement créé un cercle vicieux :

La Répartition Inégale : Certains travailleurs (neurones) recevaient des mises à jour massives et passionnantes de leurs compétences, tandis que d'autres recevaient des mises à jour minuscules, presque invisibles.
La Spirale de la Mort : Les travailleurs qui recevaient de petites mises à jour commençaient à s'effacer. Comme ils n'apprenaient pas, ils ne contribuaient plus à la production de l'équipe. Finalement, ils sont devenus des « neurones morts » : ils étaient toujours sur la liste de paie, mais ne faisaient rien.
La Boucle de Rétroaction : À mesure que ces travailleurs cessaient de contribuer, le manager (Muon) voyait moins de nécessité de les mettre à jour, ce qui leur donnait des mises à jour encore plus petites. C'était une boucle de négligence auto-renforcée.

Le document montre qu'avec Muon, un nombre important de ces travailleurs ont essentiellement « démissionné », surtout dans les premières couches du réseau.

L'Ancienne Solution : L'Approche de la « Force Brute »

Certains chercheurs ont tenté de corriger cela en forçant simplement chaque travailleur à avoir la même quantité d'énergie (normalisation par ligne). Ils regardaient l'équipe, voyaient quelqu'un de faible, et le boostaient simplement.

Cependant, le document soutient que cette approche brutale revenait à essayer de réparer une montre délicate avec un marteau-pilon. Bien qu'elle empêchait les travailleurs de mourir, elle déformait la forme du travail. Elle imposait à l'équipe une géométrie qui ne correspondait pas au flux naturel du puzzle, rendant toute l'équipe moins efficace. C'était un compromis : vous sauviez les travailleurs, mais vous brisiez le rythme du travail.

La Nouvelle Solution : Aurora

Aurora est un nouvel optimiseur qui résout ce problème sans briser le rythme. Voyez Aurora comme un maître chorégraphe qui comprend deux règles à la fois :

La Règle de la Géométrie : L'équipe doit se déplacer selon un motif orthogonal spécifique et élégant (comme une danse parfaitement synchronisée) pour résoudre le puzzle efficacement. C'est le « facteur polaire » que Muon maîtrisait bien.
La Règle de l'Équité : Chaque travailleur doit recevoir une quantité égale d'énergie et d'attention.

Aurora gère les deux simultanément. Au lieu de simplement booster les travailleurs faibles après coup, il calcule le mouvement parfait qui satisfait à la fois le motif de la danse et l'exigence d'équité en une seule étape.

Comment cela fonctionne en pratique :

Il examine les parties « hautes » du réseau (les projections de montée et de porte dans les couches MLP).
Il exécute un calcul itératif rapide (comme quelques tours de calcul mental) pour trouver la mise à jour parfaite qui maintient la perfection des pas de danse de l'équipe tout en veillant à ce que personne ne soit laissé dans l'ombre.
Il prévient totalement le problème des « neurones morts ».

Les Résultats

Les auteurs ont testé Aurora sur de grands modèles de langage (340 millions et 1,1 milliard de paramètres) et ont constaté :

Aucun Neurone Mort : Contrairement à Muon, Aurora a maintenu chaque neurone actif et contributif.
Une Meilleure Performance : Les modèles entraînés avec Aurora ont appris plus rapidement et ont obtenu de meilleurs scores aux tests de raisonnement et de connaissances (comme MMLU) que ceux entraînés avec Muon ou les correcteurs de la « force brute ».
Plus c'est Large, Mieux c'est : Le document a révélé que plus l'équipe était large (plus il y avait de neurones dans les couches MLP), plus l'avantage d'Aurora était grand. Cela suggère qu'Aurora est la clé pour entraîner des réseaux extrêmement larges et puissants sans gaspiller de ressources avec des travailleurs morts.

L'Essentiel

Aurora est une manière plus intelligente d'entraîner des modèles d'IA. Il corrige une faille cachée où les « neurones » de l'IA mouraient silencieusement, garantissant que chaque partie du réseau est utilisée efficacement, tout en préservant la structure mathématique du processus d'apprentissage. C'est comme passer d'un manager qui ignore accidentellement la moitié de son personnel à un manager qui s'assure que tout le monde danse en parfaite synchronisation.

Résumé Technique : Aurora – Un Optimiseur Spectral Sensible au Levier (Leverage-Aware)

Énoncé du Problème

L'article identifie une pathologie critique dans l'entraînement des grands modèles de langage utilisant l'optimiseur Muon, un populaire optimiseur spectral qui approxime les mises à jour de poids en utilisant le facteur polaire de la matrice de momentum du gradient. Plus précisément, pour les paramètres de matrices hautes (où le nombre de lignes $m$ dépasse le nombre de colonnes $n$ , comme dans les matrices de projection up et gate des couches MLP), les mises à jour de Muon peuvent présenter des normes de lignes hautement non uniformes.

Cette non-uniformité crée une boucle de rétroaction auto-renforcée :

Les neurones correspondant aux lignes avec de petites normes reçoivent des mises à jour perpétuellement faibles.
Ces neurones ne parviennent pas à contribuer de manière significative aux sorties du réseau, devenant ainsi effectivement « morts ».
L'article fournit une preuve mécaniste liant ce phénomène aux scores de levier (leverage scores), montrant que les faibles normes de lignes dans le tampon de momentum sont corrélées à de faibles scores de levier, que Muon préserve plutôt que de corriger.

Les auteurs observent que si les méthodes de normalisation de lignes existantes (ex. : NorMuon) atténuent ce problème en redimensionnant les lignes, elles le font d'une manière qui déforme la géométrie de la mise à jour de Muon, l'éloignant du véritable facteur polaire. Ce compromis entre l'imposition de normes de lignes uniformes et la préservation de la géométrie précise du facteur polaire est le problème central traité.

Méthodologie

1. Analyse de la Précision de la Normalisation et du Polaire

Les auteurs analysent d'abord la tension entre la précision du facteur polaire et la normalisation des lignes.

Erreur d'Approximation Polaire (PAE) : Ils démontrent qu'une plus grande précision dans le calcul du facteur polaire (via des itérations de Newton-Schulz comme CANS-12) est corrélée à une perte de validation plus faible.
Conflit Géométrique : Ils prouvent que pour les matrices hautes, il est mathématiquement impossible de satisfaire simultanément l'orthogonalité par colonnes (la propriété du facteur polaire) et l'unité des normes de lignes. La normalisation a posteriori force la mise à jour à s'éloigner du véritable facteur polaire, introduisant une erreur géométrique significative.
U-NorMuon : Comme étape intermédiaire, ils proposent U-NorMuon, qui applique la normalisation des lignes uniquement aux matrices hautes et supprime les tampons de momentum d'état pour la normalisation. Cela isole le bénéfice de la normalisation des lignes sans le surcoût des statistiques obsolètes, montrant qu'il surpasse le Muon standard et NorMuon.

2. L'Optimiseur Aurora

Pour résoudre le conflit entre la préservation de la géométrie polaire et l'imposition de normes de lignes uniformes, les auteurs proposent Aurora.

Formulation : Aurora formule la mise à jour comme un problème d'optimisation sous contraintes qui maximise la trace du produit scalaire du gradient sous deux contraintes :
1. Semi-orthogonalité à gauche : $U^\top U = \eta^2 I$ (préservant la structure du facteur polaire).
2. Égalité des normes de lignes : $\|U_{i,:}\|^2 = \eta^2 \frac{n}{m}$ pour toutes les lignes $i$ .
Riemannian-Aurora : Ils dérivent un solveur de référence théorique, Riemannian-Aurora, qui traite ces contraintes comme une variété jointe (l'intersection de la variété de Stiefel et de la variété des normes de lignes égales). Bien que théoriquement précis, il est informatiquement prohibitif en raison de la nécessité de résoudre de grands systèmes linéaires impliquant des contraintes de levier.
Aurora Pratique : Ils proposent une instanciation pratique et itérative qui alterne entre :
1. Normalisation des lignes : Redimensionnement des lignes pour égaliser les normes.
2. Approximation du Facteur Polaire : Application d'une itération de Newton-Schulz pour ré-orthogonaliser la matrice.
  Ce processus est amorti par un paramètre $\beta$ pour assurer la convergence. Les auteurs prouvent que, sous des conditions de régularité locale, cette projection alternée converge linéairement vers une solution satisfaisant les deux contraintes.

Principales Contributions

Analyse des Optimiseurs Spectraux Normalisés : L'article démontre que la normalisation des lignes a posteriori dégrade la précision du facteur polaire pour les matrices hautes. Il introduit U-NorMuon, une variante sans état qui restreint la normalisation aux matrices hautes, prouvant que l'uniformité des lignes et l'orthogonalisation précise sont simultanément désirables.
Étude de la Mort des Neurones : Les auteurs fournissent une explication mécaniste de la mort des neurones dans les modèles entraînés avec Muon, en la liant aux faibles scores de levier dans les paramètres de matrices hautes. Ils vérifient cela empiriquement lors de leurs propres entraînements et sur des modèles à poids ouverts (ex. : Rnj-1, Laguna-XS.2), trouvant des cohortes significatives de neurones morts dans les couches denses MLP.
Aurora et Riemannian-Aurora : Ils dérivent Aurora comme une méthode pour imposer des normes de lignes uniformes tout en respectant strictement le cadre théorique de Muon. Ils fournissent également Riemannian-Aurora comme un solveur de référence à coût élevé pour le problème à double contrainte.
Vérification Extensive du Pré-entraînement : Les auteurs valident Aurora à des échelles de 340M et 1.1B de paramètres, démontant une performance supérieure par rapport à Muon, NorMuon et U-NorMuon.

Résultats

Performance Aval (Downstream)

Scores de Benchmark : À l'échelle 1.1B, Aurora surpasse significativement la ligne de base Muon, obtenant un score supérieur de 9.1 points sur MMLU. Il atteint également le score moyen le plus élevé sur une suite de benchmarks de raisonnement et de connaissance (ARC, HellaSwag, etc.) aux échelles 340M et 1.1B.
Speedrun : Sur le speedrun modded-nanoGPT (piste optimiseur), Aurora combiné aux méthodes existantes (Contra-Muon, flooring) a établi un nouveau record (SOTA), atteignant la perte de validation cible 25 étapes plus rapidement que la ligne de base NorMuon.

Utilisation des Neurones

Scores de Levier : Les modèles entraînés avec Muon présentent une distribution bimodale des scores de levier, avec un mode significatif proche de zéro (neurones morts). En revanche, Aurora et U-NorMuon produisent des distributions de scores de levier isotropes, éliminant efficacement les neurones morts.
Vérification sur Modèles à Poids Ouverts : L'analyse des modèles à poids ouverts entraînés avec Muon (Rnj-1 et Laguna-XS.2) confirme la présence de neurones morts dans les premières couches denses (jusqu'à ~28 % dans la première couche de Laguna-XS.2), validant la généralisabilité du problème.

Efficacité et Mise à l'Échelle (Scaling)

Débit (Throughput) : Dans un environnement distribué, le surcoût computationnel d'Aurora est minimal car les multiplications de matrices supplémentaires peuvent être chevauchées avec la communication des gradients. Sur un seul GPU, le débit diminue légèrement, mais le surcoût en temps réel (wall-clock) dans l'entraînement distribué est modeste.
Scaling de la Largeur MLP : Les gains de performance d'Aurora par rapport à Muon progressent de manière monotone avec le facteur d'expansion du MLP. À mesure que les couches MLP deviennent plus larges, l'avantage d'Aurora augmente, suggérant qu'il est particulièrement efficace pour l'entraînement de réseaux très larges.

Signification et Revendications

L'article affirme qu'Aurora résout une pathologie spécifique, observée empiriquement, dans l'entraînement des modèles de langage — la mort des neurones dans les paramètres de matrices hautes — en dérivant un optimiseur qui respecte à la fois la géométrie théorique de la mise à jour de Muon et la nécessité pratique de normes de lignes uniformes.

Les auteurs insistent sur une philosophie de conception ascendante (bottom-up) : dériver des algorithmes basés sur les dynamiques d'entraînement observées (mort des neurones, déséquilibres de scores de levier) plutôt que sur des modèles purement théoriques. Ils soutiennent que la co-conception architecture-optimiseur, où les optimiseurs sont adaptés à des classes de paramètres spécifiques (ex. : matrices hautes vs larges), est une direction prometteuse qui surpasse les stratégies standards « taille unique ».

Enfin, l'article suggère qu'Aurora permet l'entraînement efficace de couches MLP très larges, une capacité qui pourrait être limitée dans l'entraînement actuel basé sur Muon en raison de l'accumulation de neurones morts. Ce travail positionne Aurora non pas seulement comme un boost de performance, mais comme une correction nécessaire des hypothèses géométriques sous-jacentes à l'optimisation spectrale actuelle dans les modèles de langage à grande échelle.