Adam이 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.