Адам улучшает Muon: адаптивная оценка моментов с ортогонализацией импульса

Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

February 19, 2026

Авторы: Minxin Zhang, Yuxuan Liu, Hayden Scheaffer

cs.AI

Аннотация

Эффективная стохастическая оптимизация обычно объединяет направление обновления, хорошо зарекомендовавшее себя в детерминированном режиме, с механизмом адаптации к стохастическим возмущениям. В то время как Adam использует адаптивные оценки моментов для обеспечения стабильности, Muon задействует матричную структуру слоев весов посредством ортогонализированного момента, демонстрируя превосходную производительность при обучении больших языковых моделей. Мы предлагаем новый оптимизатор и его диагональное расширение — NAMO и NAMO-D, которые представляют собой первую принципиальную интеграцию ортогонализированного момента с адаптацией к шуму по норме в стиле Adam. NAMO масштабирует ортогонализированный момент, используя единственный адаптивный шаг обучения, сохраняя ортогональность и превосходя Muon при пренебрежимо малых дополнительных затратах. NAMO-D, напротив, умножает ортогонализированный момент справа на диагональную матрицу с ограниченными элементами. Такая конструкция позволяет осуществлять поэлементную (по нейронам) адаптацию к шуму и согласуется с распространенной структурой гессиана, близкой к блочно-диагональной. В рамках стандартных предположений мы устанавливаем оптимальные скорости сходимости для обоих алгоритмов в детерминированной постановке и показываем, что в стохастической постановке их гарантии сходимости адаптируются к уровню шума стохастических градиентов. Эксперименты по предварительному обучению моделей GPT-2 демонстрируют улучшенную производительность как NAMO, так и NAMO-D по сравнению с базовыми алгоритмами AdamW и Muon, причем NAMO-D достигает дополнительного преимущества перед NAMO за счет гиперпараметра ограничения, который балансирует конкурирующие цели сохранения хорошо обусловленного направления обновления и использования детализированной адаптации к шуму.

English

Efficient stochastic optimization typically integrates an update direction that performs well in the deterministic regime with a mechanism adapting to stochastic perturbations. While Adam uses adaptive moment estimates to promote stability, Muon utilizes the weight layers' matrix structure via orthogonalized momentum, showing superior performance in large language model training. We propose a new optimizer and a diagonal extension, NAMO and NAMO-D, providing the first principled integration of orthogonalized momentum with norm-based Adam-type noise adaptation. NAMO scales orthogonalized momentum using a single adaptive stepsize, preserving orthogonality while improving upon Muon at negligible additional cost. NAMO-D instead right-multiplies orthogonalized momentum by a diagonal matrix with clamped entries. This design enables neuron-wise noise adaptation and aligns with the common near block-diagonal Hessian structure. Under standard assumptions, we establish optimal convergence rates for both algorithms in the deterministic setting and show that, in the stochastic setting, their convergence guarantees adapt to the noise level of stochastic gradients. Experiments on pretraining GPT-2 models demonstrate improved performance of both NAMO and NAMO-D compared to the AdamW and Muon baselines, with NAMO-D achieving further gains over NAMO via an additional clamping hyperparameter that balances the competing goals of maintaining a well-conditioned update direction and leveraging fine-grained noise adaptation.