Почему Muon превосходит Adam: с точки зрения кривизны

Аннотация

Muon повышает эффективность обучения примерно в два раза по сравнению с Adam при обучении больших языковых моделей, однако локальный геометрический источник этого преимущества остаётся неясным. Наша работа делает первый шаг к демистификации превосходства Muon над Adam с точки зрения кривизны. Во-первых, мы применяем аппроксимацию второго порядка Тейлора к ландшафту обучения и показываем, что Muon обеспечивает большее уменьшение потерь за один шаг, чем Adam, при одинаковых валидационных потерях. Оба оптимизатора имеют сопоставимые выигрыши первого порядка, но Muon последовательно несёт меньший штраф за кривизну второго порядка. Во-вторых, мы разлагаем этот штраф за кривизну на квадрат нормы обновления и нормализованную направленную остроту (Normalized Directional Sharpness, NDS). Мы обнаруживаем, что нормы обновления у Muon и Adam сравнимы, поэтому меньший штраф за кривизну у Muon обусловлен более низкой NDS, а не масштабом обновления. В-третьих, мы изучаем, как обучающие данные и структура модели формируют преимущество Muon по NDS. Используя данные, сгенерированные вероятностной контекстно-свободной грамматикой Зипфа (Zipf-PCFG) с контролируемым дисбалансом, мы показываем, что дисбаланс данных усиливает преимущество Muon по NDS перед Adam. Дальнейшее разложение на внутрислойную и межслойную составляющие показывает, что на средних и поздних этапах обучения более низкая NDS Muon в основном поддерживается меньшей внутрислойной кривизной. Помимо эмпирических данных, мы анализируем стилизованные квадратичные задачи с гетерогенной кривизной и выравниванием градиентов по направлению к модам с высокой кривизной. Мы доказываем, что Muon достигает меньшей средней NDS, чем градиентный спуск (GD), за счёт балансировки энергии обновления между группами кривизны; когда гетерогенность кривизны достаточно сильна, это также приводит к меньшим локальным квадратичным потерям после одинакового числа шагов.

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.