Pourquoi Muon surpasse Adam : une perspective de courbure

Résumé

Muon améliore l'efficacité d'entraînement par rapport à Adam dans l'entraînement des grands modèles de langage d'environ deux fois, mais la source géométrique locale de cet avantage reste floue. Notre travail constitue un premier pas vers la démystification de la supériorité de Muon sur Adam d'un point de vue de la courbure. Premièrement, nous appliquons une approximation de Taylor du second ordre au paysage d'entraînement et montrons que Muon obtient une plus grande diminution de la perte en une étape qu'Adam à perte de validation appariée. Les deux optimiseurs ont des gains de premier ordre comparables, mais Muon encourt systématiquement une plus faible pénalité de courbure de second ordre. Deuxièmement, nous décomposons cette pénalité de courbure en la norme de mise à jour au carré et la Netteté Directionnelle Normalisée (NDS). Nous constatons que Muon et Adam ont des normes de mise à jour comparables, donc la plus faible pénalité de courbure de Muon est due à une NDS plus faible, et non à l'échelle de mise à jour. Troisièmement, nous étudions comment les données d'entraînement et la structure du modèle façonnent l'avantage de Muon en NDS. En utilisant des données Zipf-PCFG (grammaire hors-contexte probabiliste) avec déséquilibre contrôlé, nous montrons que le déséquilibre des données amplifie l'avantage de Muon en NDS par rapport à Adam. Une décomposition intra-/inter-couche montre en outre qu'aux stades intermédiaires et tardifs de l'entraînement, la NDS plus faible de Muon est principalement soutenue par une plus faible courbure intra-couche. Au-delà des preuves empiriques, nous analysons des problèmes quadratiques stylisés avec une courbure hétérogène et un alignement du gradient vers les modes de haute courbure. Nous prouvons que Muon atteint une NDS moyenne plus faible que GD (descente de gradient) en équilibrant l'énergie de mise à jour entre les groupes de courbure ; lorsque l'hétérogénéité de la courbure est suffisamment forte, cela produit également une perte quadratique locale plus faible après le même nombre d'étapes.

English

Muon improves training efficiency over Adam in large language-model training by about two times, but the local geometric source of this advantage remains unclear. Our work takes a first step toward demystifying Muon's superiority over Adam from a curvature perspective. First, we apply a second-order Taylor approximation to the training landscape and show that Muon achieves a larger one-step loss decrease than Adam at matched validation loss. The two optimizers have comparable first-order gains, but Muon consistently incurs a smaller second-order curvature penalty. Second, we decompose this curvature penalty into the squared update norm and Normalized Directional Sharpness (NDS). We find that Muon and Adam have comparable update norms, so Muon's smaller curvature penalty is driven by lower NDS, not update scale. Third, we study how training data and model structure shape Muon's NDS advantage. Using Zipf-Probabilistic Context-Free Grammar (PCFG) data with controlled imbalance, we show that data imbalance amplifies Muon's NDS advantage over Adam. A within-/cross-layer decomposition further shows that, in the middle and late stages of training, Muon's lower NDS is mainly sustained by smaller within-layer curvature. Beyond empirical evidence, we analyze stylized quadratic problems with heterogeneous curvature and gradient alignment toward high-curvature modes. We prove that Muon attains a smaller average NDS than GD by balancing update energy across curvature groups; when curvature heterogeneity is sufficiently strong, this also yields lower local quadratic loss after the same number of steps.