Warum Muon Adam übertrifft: Eine Krümmungsperspektive

Zusammenfassung

Muon verbessert die Trainingseffizienz beim Training großer Sprachmodelle im Vergleich zu Adam um etwa das Zweifache, doch der lokale geometrische Ursprung dieses Vorteils bleibt unklar. Unsere Arbeit unternimmt einen ersten Schritt zur Entmystifizierung der Überlegenheit von Muon gegenüber Adam aus einer Krümmungsperspektive. Erstens wenden wir eine Taylor-Approximation zweiter Ordnung auf die Trainingslandschaft an und zeigen, dass Muon bei gleichem Validierungsverlust eine größere Einschritt-Verlustabnahme erzielt als Adam. Die beiden Optimierer weisen vergleichbare Gewinne erster Ordnung auf, doch Muon zieht sich durchgängig eine geringere Krümmungsstrafe zweiter Ordnung zu. Zweitens zerlegen wir diese Krümmungsstrafe in die quadratische Aktualisierungsnorm und die Normalisierte Richtungsschärfe (NDS). Wir stellen fest, dass Muon und Adam vergleichbare Aktualisierungsnormen aufweisen, sodass Muons geringere Krümmungsstrafe durch eine niedrigere NDS und nicht durch die Aktualisierungsskala getrieben wird. Drittens untersuchen wir, wie Trainingsdaten und Modellstruktur Muons NDS-Vorteil formen. Anhand von Zipf-probabilistischen kontextfreien Grammatikdaten (PCFG) mit kontrolliertem Ungleichgewicht zeigen wir, dass Datenungleichgewicht Muons NDS-Vorteil gegenüber Adam verstärkt. Eine Intra-/Inter-Layer-Zerlegung zeigt zudem, dass Muons niedrigere NDS in der mittleren und späten Trainingsphase hauptsächlich durch eine geringere schichtinterne Krümmung aufrechterhalten wird. Über empirische Belege hinaus analysieren wir stilisierte quadratische Probleme mit heterogener Krümmung und Gradientenausrichtung auf Modi hoher Krümmung. Wir beweisen, dass Muon eine geringere durchschnittliche NDS als GD erreicht, indem es die Aktualisierungsenergie über Krümmungsgruppen ausgleicht; bei ausreichend starker Krümmungsheterogenität führt dies auch nach der gleichen Anzahl von Schritten zu einem geringeren lokalen quadratischen Verlust.

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.