Mousse: 곡률 인지 사전 조건화를 통한 뮤온(Muon)의 기하학적 개선

Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning

최근 스펙트럼 최적화 분야에서 상당한 진전을 보인 뮤온(Muon)은 업데이트 단계를 슈티펠 다양체(Stiefel manifold)로 제한함으로써 학습 속도를 크게 향상시키고 일반화 성능을 개선할 수 있음을 보여주었습니다. 그러나 뮤온은 암묵적으로 등방적인 최적화 지형을 가정하며, 모든 고유 방향에 걸쳐 균일한 스펙트럴 업데이트 노름을 적용합니다. 우리는 이러한 "평등한" 제약 조건이 딥 신경망에서 최적이 아니라고 주장합니다. 왜냐하면 신경망의 곡률 스펙트럼은 일반적으로 매우 비대칭적이고 불안정하기 때문입니다. 이러한 지형에서 뮤온은 높은 곡률 방향에서 불안정성을 증폭시킬 위험이 있으며, 동시에 평탄한 방향에서의 필요한 진행을 제한할 수 있습니다. 본 연구에서는 **Mousse** (**M**uon **O**ptimization **U**tilizing **S**hampoo's **S**tructural **E**stimation)라는 새로운 최적화 알고리즘을 제안합니다. Mousse는 스펙트럴 방법의 구조적 안정성과 2차 사전 조건화의 기하학적 적응성을 결합합니다. Mousse는 뉴턴-슐츠 직교화를 모멘텀 행렬에 직접 적용하는 대신, 샴푸(Shampoo)에서 파생된 크로네커 분해 통계(Kronecker-factored statistics)에 의해 유도된 표백된 좌표계에서 작동합니다. 수학적으로, Mousse는 비등방적 신뢰 영역으로 제한된 스펙트럴 최속 하강(spectral steepest descent) 문제의 해로 공식화되며, 최적의 업데이트는 표백된 기울기의 극 분해(polar decomposition)를 통해 도출됩니다. 1억 6천만 개에서 8억 개 매개변수를 갖는 언어 모델에 대한 실험 결과에서, Mousse는 뮤온보다 일관되게 우수한 성능을 보이며, 계산 오버헤드가 거의 없는 상태에서 학습 단계를 약 12% 줄였습니다.

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 \textbf{Mousse} (\textbf{M}uon \textbf{O}ptimization \textbf{U}tilizing \textbf{S}hampoo's \textbf{S}tructural \textbf{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.