Waarom Muon Adam overtreft: een krommingsperspectief

Samenvatting

Muon verbetert de trainingsefficiëntie bij het trainen van grote taalmodellen met ongeveer een factor twee ten opzichte van Adam, maar de lokale geometrische bron van dit voordeel blijft onduidelijk. Ons werk zet een eerste stap in het ontrafelen van Muons superioriteit ten opzichte van Adam vanuit een krommingsperspectief. Ten eerste passen we een tweede-orde Taylor-benadering toe op het trainingslandschap en tonen aan dat Muon een grotere éénstapsverliesafname behaalt dan Adam bij een overeenkomend validatieverlies. De twee optimalisatoren hebben vergelijkbare eerste-ordewinsten, maar Muon ondervindt consequent een kleinere tweede-orde krommingspenalty. Ten tweede ontleden we deze krommingspenalty in het kwadraat van de updatenorm en de genormaliseerde richtingsscherpte (Normalized Directional Sharpness, NDS). We vinden dat Muon en Adam vergelijkbare updatenormen hebben, dus Muons kleinere krommingspenalty wordt gedreven door een lagere NDS, niet door een kleinere updategrootte. Ten derde onderzoeken we hoe trainingsgegevens en modelstructuur Muons NDS-voordeel vormgeven. Met behulp van Zipf-probabilistische contextvrije grammatica (PCFG)-data met gecontroleerde onevenwichtigheid tonen we aan dat data-onevenwichtigheid Muons NDS-voordeel ten opzichte van Adam versterkt. Een ontleding in termen van binnenlaagse en tussenlaagse kromming laat verder zien dat Muons lagere NDS in de midden- en late stadia van de training voornamelijk in stand wordt gehouden door kleinere kromming binnen de lagen. Naast empirisch bewijs analyseren we gestileerde kwadratische problemen met heterogene kromming en gradiëntuitlijning richting hoogkrommingsmodi. We bewijzen dat Muon een kleinere gemiddelde NDS behaalt dan GD door de updatenergie over krommingsgroepen in evenwicht te brengen; wanneer de krommingsheterogeniteit voldoende sterk is, levert dit ook een lager lokaal kwadratisch verlies op na hetzelfde aantal stappen.

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.