Muon은 느린 saddle-to-saddle 역학 없이 행렬 분해에서 균형 잡힌 해를 학습한다

행렬 분해 (Matrix factorization, 즉 $\min_{\mathbf{P},\mathbf{Q}} |\mathbf{M}^\star - \mathbf{P}^\top\mathbf{Q}|_{\mathrm{F}}^2$ 형태의 문제)는 비선형 파라미터 역학 (nonlinear parameter dynamics)과 표현 학습 (representation learning)을 모두 보여주는 최소한의 학습 문제입니다. 이 설정에서, 우리는 Muon 옵티마이저 (optimizer) 하에서의 파라미터 궤적 (parameter trajectories)이 경사 하강법 (gradient descent)의 궤적과 어떻게 다른지 연구합니다. 우리는 세 가지 주요한 역학적 차이점을 식별했습니다: 1) Muon은 작은 초기화 (small initialization)로부터 발생하는 느린 saddle-to-saddle 역학을 피합니다. 대신 Muon은 $\mathbf{M}^\star$의 모든 상위 모드 (top modes)를 동일한 속도로 학습하며, 더 작은 모드들이 먼저 수렴합니다. 2) Muon은 학습률 (learning rate)이 국소 손실 날카로움 (local loss sharpness)에 의해 설정된 임계값 (critical threshold)을 초과하더라도 안정성을 유지합니다. 이를 통해 학습률을 문제의 조건수 (condition number)로부터 자유롭게 하여, 지수적 학습률 어닐링 (exponential learning rate annealing)을 통한 빠른 수렴을 가능하게 합니다. 3) 가중치 (weights)가 서로 및 타겟과 정렬(aligned)되면, Muon 흐름 (Muon flow)은 행렬 양 $\sqrt{\mathbf{P}^\top \mathbf{P}}-\sqrt{\mathbf{Q}^\top \mathbf{Q}}$를 보존하는 반면, 경사 흐름 (gradient flow)은 행렬 $\mathbf{P}^\top\mathbf{P} - \mathbf{Q}^\mathbf{Q}$를 보존하는 것으로 알려져 있습니다. 서로 다른 보존량 (conserved quantities)을 가짐에도 불구하고, 두 옵티마이저 모두 소멸하는 초기화 (vanishing initialization)로부터 이른바 extit{균형 잡힌 (balanced)} 해를 찾아냅니다. 작은 무작위 초기화 (small random initialization)로부터 학습할 때, 가중치는 학습 초기에 자발적으로 정렬됩니다. 우리는 단순한 설정에서 정렬 속도 (alignment rates)를 도출하였으며, 이것이 일반적인 상황에서의 경험적 정렬 속도를 예측함을 보여줍니다. 마지막으로, 우리는 단 두 번의 최적화 단계만으로 거의 완벽한 정렬을 달성하는 학습률 스케줄 (learning rate schedule)을 구축하기 위해 Muon의 구조적 특성을 활용합니다.