Por que o Muon supera o Adam: uma perspectiva de curvatura

Resumo

O Muon melhora a eficiência do treinamento em cerca de duas vezes em relação ao Adam no treinamento de modelos de linguagem de grande escala, mas a origem geométrica local dessa vantagem permanece obscura. Nosso trabalho dá um primeiro passo para desmistificar a superioridade do Muon sobre o Adam sob uma perspectiva de curvatura. Primeiro, aplicamos uma aproximação de Taylor de segunda ordem à paisagem de treinamento e mostramos que o Muon alcança um decréscimo maior de perda em um passo do que o Adam, com perda de validação equivalente. Os dois otimizadores têm ganhos de primeira ordem comparáveis, mas o Muon incorre consistentemente em uma penalidade de curvatura de segunda ordem menor. Segundo, decompomos essa penalidade de curvatura na norma de atualização ao quadrado e na Nitidez Direcional Normalizada (NDS, na sigla em inglês). Descobrimos que Muon e Adam têm normas de atualização comparáveis; portanto, a penalidade de curvatura menor do Muon é impulsionada por uma NDS mais baixa, e não pela escala de atualização. Terceiro, estudamos como os dados de treinamento e a estrutura do modelo moldam a vantagem na NDS do Muon. Usando dados de Gramática Livre de Contexto Probabilística de Zipf (PCFG, na sigla em inglês) com desequilíbrio controlado, mostramos que o desequilíbrio dos dados amplifica a vantagem na NDS do Muon sobre o Adam. Uma decomposição intra e entre camadas mostra ainda que, nos estágios intermediário e final do treinamento, a NDS mais baixa do Muon é mantida principalmente por uma curvatura intra-camada menor. Além das evidências empíricas, analisamos problemas quadráticos estilizados com curvatura heterogênea e alinhamento do gradiente em direção a modos de alta curvatura. Provamos que o Muon atinge uma NDS média menor do que a Descida do Gradiente (GD) ao equilibrar a energia de atualização entre grupos de curvatura; quando a heterogeneidade da curvatura é suficientemente forte, isso também resulta em uma perda quadrática local menor após o mesmo número de passos.

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.