Résumé technique : Réévaluation de Muon pour la factorisation de matrices

Énoncé du problème

Muon (MomentUm Orthogonalized by Newton–Schulz) est récemment apparu comme un optimiseur puissant pour l'apprentissage profond à grande échelle, particulièrement pour l'entraînement des grands modèles de langage (LLM), où il surpasse apparemment les méthodes adaptatives standards comme Adam et AdamW. La littérature théorique interprète Muon comme effectuant une descente de gradient la plus raide sous la norme spectrale, en exploitant une orthogonalisation approximative pour remodeler les mises à jour de gradient.

Cependant, une lacune critique subsiste : il n'est pas clair si le succès empirique de Muon provient de sa règle de mise à jour intrinsèque ou s'il s'agit d'un artefact lié à l'échelle, à l'architecture et aux distributions de données spécifiques des réseaux profonds modernes. Les analyses théoriques existantes reposent souvent sur des cadres idéalisés ou des propriétés locales, échouant à spécifier les conditions concrètes sous lesquelles Muon devrait théoriquement ou empiriquement surpasser les optimiseurs classiques.

Ce document traite cette incertitude en isolant l'optimiseur des facteurs de confusion (architecture, échelle, données) en évaluant Muon sur la factorisation de matrices de rang faible. Ce problème est choisi pour sa simplicité, son paysage d'optimisation bien compris et sa structure spectrale explicite, permettant une évaluation systématique et exhaustive de la manière dont la nature spectrale de Muon interagit avec la courbure et les structures de valeurs singulières.

Méthodologie

Les auteurs mènent une comparaison contrôlée et systématique de Muon par rapport aux bases adaptatives (AdamW, Gradient Descent pur, et SignGD) à travers cinq contextes distincts de problèmes de récupération de matrices. Le protocole d'évaluation est conçu pour éliminer le biais de la sélection des hyperparamètres, que les auteurs identifient comme une source primaire de conclusions trompeuses dans les travaux antérieurs.

Conception expérimentale

1. Variantes de problèmes :

   • Factorisation de matrice symétrique de rang faible : Récupération d'une matrice semi-définie positive de rang $r$ à partir de $UU^\top$.
   • Complétion de matrice : Récupération d'une matrice à partir d'un sous-ensemble d'entrées observées (taux d'observation de 20 %).
   • Factorisation de matrices non négatives (NMF) : Factorisation contrainte où la cible et les facteurs sont non négatifs.
   • Factorisation Tensor-Train : Une configuration de factorisation bilinéaire appliquée à des cibles PSD symétriques, servant de modèle minimal pour les paramétrisations de produits dans l'apprentissage profond.
   • Factorisation par noyau Gaussien : Un paysage lisse et bien conditionné utilisant une matrice de noyau RBF.

2. Variables contrôlées :

   • Conditionnement ($\kappa$) : Les expériences balayent des nombres de conditionnement allant de bien conditionnés ($\kappa=1$) à sévèrement mal conditionnés ($\kappa=625$ ou $10^4$).
   • Forme spectrale : La distribution des valeurs propres est variée (uniforme, décroissante, à queue lourde, regroupée) tout en maintenant le nombre de conditionnement fixe pour tester la sensibilité à la densité spectrale.
   • Paramétrage : Les expériences couvrent les régimes sous-paramétrés, exactement paramétrés et sur-paramétrés.

3. Protocole d'évaluation :

   • Réglage des hyperparamètres : Au lieu d'utiliser un taux d'apprentissage "par défaut" unique, les auteurs effectuent des balayages exhaustifs de taux d'apprentissage (grilles logarithmiques) pour chaque optimiseur et configuration de problème.
   • Métrique : La récupération est mesurée via l'erreur quadratique moyenne normalisée (NMSE).
   • Comparaison : Les optimiseurs sont comparés sur la base de leur "plage de taux d'apprentissage stable" (la bande de tailles de pas où la convergence vers le plancher de la perte se produit) et de leur meilleur taux de perte final ajusté.

Contributions clés et résultats

1. Sensibilité aux hyperparamètres et réglage

Le document démontre que le classement des optimiseurs est hautement sensible à la sélection du taux d'apprentissage.

• Par défaut vs Ajusté : Les classements dérivés d'un taux d'apprentissage par défaut unique (ex. $10^{-3}$) contredisent fréquemment les classements obtenus après un ajustage individuel.
• Absence de dominance universelle : Lorsqu'une méthode est individuellement ajustée, aucun optimiseur ne domine tous les contextes.
  • AdamW et GD égalent ou dépassent souvent Muon sur les tâches de factorisation de rang faible standard et de complétion de matrice.
  • Muon présente un avantage constant principalement dans la Factorisation de Matrices Non Négatives (NMF) et certains problèmes de complétion hautement mal conditionnés.

2. Le rôle du conditionnement et de la structure spectrale

• Le conditionnement réordonne les méthodes : À mesure que le nombre de conditionnement ($\kappa$) augmente, le taux d'apprentissage optimal se déplace, et la performance relative des optimiseurs change. Par exemple, GD performe bien dans les régimes bien conditionnés mais échoue dans les cas de mal conditionnement extrême, là où AdamW excelle.
• La forme spectrale importe : Même à un nombre de conditionnement extrême et fixe ($\kappa=10^4$), la distribution intérieure des valeurs propres impacte significativement la performance. AdamW s'adapte efficacement aux spectres à queue lourde ou regroupés, tandis que Muon sous-performe dans la plupart des formes spectrales dans ces conditions extrêmes.

3. Performance spécifique en NMF

Muon conserve un avantage distinct dans les tâches de NMF. Les auteurs suggèrent que cela est dû au fait que les mises à jour orthogonalisées de Muon découragent les facteurs redondants et favorisent des représentations plus diverses, un bénéfice qui n'est pas observé dans les tâches de factorisation non contraintes ou symétriques.

4. Paysages bien comportés

Sur la tâche de factorisation par noyau Gaussien (un paysage lisse et bien conditionné), le mécanisme d'orthogonalisation complexe de Muon n'apporte aucun avantage par rapport à un AdamW bien ajusté. Les deux convergent vers le même plancher global, renforçant l'idée que les avantages de Muon ne sont pas universels mais dépendent du contexte.

Signification et affirmations

Le document soutient que la supériorité perçue de Muon dans l'apprentissage profond à grande échelle peut être confondue avec d'autres facteurs tels que l'échelle du modèle, l'architecture et les caractéristiques des données. En isolant l'optimiseur sur des problèmes contrôlés, les auteurs offrent une image plus nuancée des capacités de Muon :

• Utilité dépendante du contexte : Muon n'est pas un optimiseur universellement supérieur. Ses bénéfices sont liés à des structures de problèmes spécifiques, particulièrement celles impliquant des contraintes de non-négativité ou des dégénérescences spectrales spécifiques trouvées dans la NMF.
• Standards d'évaluation : Ce travail souligne la nécessité d'évaluer les optimiseurs modernes à travers des plages d'hyperparamètres ajustées et des régimes de conditionnement divers, plutôt que de se fier à des paramètres par défaut ou à des benchmarks de bout en bout isolés.
• Implications théoriques : Les résultats remettent en question l'hypothèse selon laquelle la descente de gradient la plus raide sous la norme spectrale (l'interprétation théorique de Muon) est intrinsèquement supérieure à l'adaptativité par coordonnée (AdamW) dans tous les contextes structurés. Au lieu de cela, la "meilleure" approche dépend fortement des caractéristiques géométriques et spectrales spécifiques de la fonction objectif.

En conclusion, les auteurs ne prétendent pas que Muon est inefficace, mais plutôt que ses avantages sont dépendants du problème. L'étude préconise une approche de benchmarking plus rigoureuse et contrôlée pour comprendre quand et pourquoi l'orthogonalisation sensible au spectre aide l'optimisation.