[논문 리뷰] Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning

Yechen Zhang, Shuhao Xing외 6명|arXiv (Cornell University)|2026. 03. 10.

Stochastic Gradient Optimization Techniques인용 수 0

한 줄 요약

Mousse가 Shampoo에 의해 유도된 흰색 공간에서 Newton-Schulz 직교화를 적용하여 Muon을 향상시키고, 160M–800M 매개변수 언어 모델 전반에서 ~12% 더 빠른 수렴을 달성하며 오버헤드는 미미합니다.

ABSTRACT

Recent advances in spectral optimization, notably Muon, have demonstrated that constraining update steps to the Stiefel manifold can significantly accelerate training and improve generalization. However, Muon implicitly assumes an isotropic optimization landscape, enforcing a uniform spectral update norm across all eigen-directions. We argue that this "egalitarian" constraint is suboptimal for Deep Neural Networks, where the curvature spectrum is known to be highly heavy-tailed and ill-conditioned. In such landscapes, Muon risks amplifying instabilities in high-curvature directions while limiting necessary progress in flat directions. In this work, we propose extbf{Mousse} ( extbf{M}uon extbf{O}ptimization extbf{U}tilizing extbf{S}hampoo's extbf{S}tructural extbf{E}stimation), a novel optimizer that reconciles the structural stability of spectral methods with the geometric adaptivity of second-order preconditioning. Instead of applying Newton-Schulz orthogonalization directly to the momentum matrix, Mousse operates in a whitened coordinate system induced by Kronecker-factored statistics (derived from Shampoo). Mathematically, we formulate Mousse as the solution to a spectral steepest descent problem constrained by an anisotropic trust region, where the optimal update is derived via the polar decomposition of the whitened gradient. Empirical results across language models ranging from 160M to 800M parameters demonstrate that Mousse consistently outperforms Muon, achieving around $\sim$12\% reduction in training steps with negligible computational overhead.

연구 동기 및 목표

• Muon의 등방성 스펙트럴 제약과 높은 이방성 신경 곡률 사이의 불일치를 동기화한다.
• 스펙트럼 제약과 함께 2차 전처리기을 통합하는 곡률 인지 스펙트럼 최적화 프레임워크를 제안한다.
• Mousse 업데이트를 안정적으로 구현하고 효율적으로 수행하는 기술을 개발하고 분석한다.
• 160M에서 800M 매개변수를 가진 언어 모델에서 견고성과 효율성 향상을 시연한다.

제안 방법

• Muon을 op-norm 제약이 있는 스펙트럴 급경사 하 문제로 프레이밍한다.
• Shampoo의 Kronecker-팩터링 곡률 통계로 그래디언트를 화이트닝하여 구 형태를 구면 기하로 유도한다.
• whitened 좌표에서 Newton-Schulz 직교화를 적용하여 Stiefel 매니폴드에서 업데이트를 얻는다.
• 제약 최소화를 전처치된 그래디언트를 통해 해결하고, L^{-1/4}와 R^{-1/4} 전처치기를 포함하는 업데이트 공식을 얻는다.
• 안정성 기법: 추적 정규화(Trace Normalization)와 스펙트럴 템퍼링(alpha = 0.125)을 도입하여 조건 수와 곡률 강도를 제어한다.
• 계산량과 메모리 오버헤드를 줄이기 위한 단면(preconditioner) 변형을 제공한다.

실험 결과

연구 질문

• RQ1 비등방성 곡률 인지 기하가 Muon의 등방성 제약과 비교하여 스펙트럴 업데이트에 어떤 영향을 미치는가?
• RQ2 Shampoo 통계로 화이트닝이 곡률 인지 프레임워크 내에서 효과적인 Newton-Schulz 업데이트를 가능하게 하는가?
• RQ3 깊은 네트워크에서 곡률 인지 스펙트럴 최적화를 견고하게 만드는 필수 안정성 기법은 무엇인가?
• RQ4 Mousse를 사용한 대형 언어 모델 학습에서 수렴 속도 및 샘플 효율성의 실험적 이득은 무엇인가?
• RQ5 단면 전처치기가 성능이나 안정성을 희생하지 않고도 실행 가능한가?

주요 결과

• Mousse는 800M 매개변수 모델에서 Muon에 비해 목표 손실에 도달하는 학습 스텝 수를 약 12% 감소시킨다.
• Mousse는 Muon과 거의 동일한 벽시계 시간 학습을 유지하며 오버헤드가 미미하다.
• Mousse는 160M에서 800M 모델 크기에 걸쳐 최종 검증 손실을 더 낮춘다.
• Trace Normalization과 Spectral Tempering은 안정성과 효과적인 곡률 보정에 필수적이다.
• 단면(preconditioner) 방식은 전체 Kronecker 기반 접근법에 비해 계산 및 메모리 비용을 줄이며 성능에 비례적으로 근접하다.