Por qué Muon supera a Adam: Una perspectiva de curvatura

Resumen

Muon mejora la eficiencia de entrenamiento en aproximadamente el doble en comparación con Adam en el entrenamiento de modelos de lenguaje grandes, pero la fuente geométrica local de esta ventaja sigue sin estar clara. Nuestro trabajo da un primer paso hacia desmitificar la superioridad de Muon sobre Adam desde una perspectiva de curvatura. Primero, aplicamos una aproximación de Taylor de segundo orden al paisaje de entrenamiento y mostramos que Muon logra una mayor disminución de pérdida en un paso que Adam con una pérdida de validación equivalente. Ambos optimizadores tienen ganancias de primer orden comparables, pero Muon incurre consistentemente en una penalización de curvatura de segundo orden menor. Segundo, descomponemos esta penalización de curvatura en la norma cuadrada de actualización y la Nitidez Direccional Normalizada (NDS). Encontramos que Muon y Adam tienen normas de actualización comparables, por lo que la menor penalización de curvatura de Muon se debe a una NDS más baja, no a la escala de actualización. Tercero, estudiamos cómo los datos de entrenamiento y la estructura del modelo moldean la ventaja de NDS de Muon. Utilizando datos de Gramática Libre de Contexto Probabilística de Zipf (PCFG) con desequilibrio controlado, mostramos que el desequilibrio de datos amplifica la ventaja de NDS de Muon sobre Adam. Una descomposición intra/inter capas revela además que, en las etapas media y tardía del entrenamiento, la menor NDS de Muon se mantiene principalmente mediante una curvatura intra-capa más pequeña. Más allá de la evidencia empírica, analizamos problemas cuadráticos estilizados con curvatura heterogénea y alineación del gradiente hacia modos de alta curvatura. Demostramos que Muon alcanza una NDS promedio menor que el Descenso por Gradiente (GD) al equilibrar la energía de actualización entre grupos de curvatura; cuando la heterogeneidad de curvatura es suficientemente fuerte, esto también produce una pérdida cuadrática local menor después del mismo número de pasos.

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.