Muon 위에서의 이동: Muon 옵티마이저의 Hamiltonian 확률 구배 흐름 관점

우리는 이상적인 Muon 옵티마이저(optimizer)의 해석적으로 매끄러운 버전인 정규화된 Muon(regularized Muon)에 의해 유도된, 행렬 값 파라미터(matrix-valued parameters) 상에 정의된 확률 측도(probability measures) 공간에서의 구배 흐름(gradient flow)을 개발합니다. 핵심적인 관찰은 정규화된 직교화 사상(orthogonalization map)이 핵 노름(nuclear norm)의 매끄러운 Fenchel-dual smoothing의 구배라는 점입니다. 이는 (정규화된) Muon 업데이트를 업데이트 변수에서의 mirror/prox 단계로 식별하며, 여기서 모멘텀(momentum)은 dual 좌표로 작용합니다. 우리는 이 구조를 사용하여 Muon을 단일 행렬 파라미터에서 신경망 학습의 평균장(mean-field) 기술에 의해 동기 부여된 $J(ρ)=Rig(\int F d ρ\big)$ 형태의 유한 입자 확률 목적 함수(finite-particle probability objectives)로 확장하며, 관성적 연속 시간 극한(inertial continuous-time limit)을 도출합니다. 이 구조를 사용하여, 단계 크기(step size)와 모멘텀의 관성적 스케일링 하에서의 유한 입자 연속 시간 극한을 도출하고, 그 다음 파라미터-모멘텀 쌍에 대한 확률 법칙 상의 위상 공간(phase-space) 평균장 방정식으로 넘어갑니다. 결과적인 흐름은 정규화된 Muon mirror 포텐셜에 의해 유도된 운동 에너지(kinetic energy)를 갖는 감쇠된 Hamiltonian 확률 역학(damped Hamiltonian probability dynamics)임을 보일 수 있습니다. 우리는 Hamiltonian 에너지(Hamiltonian energy)가 단조적으로 감소함을 보여주는 정확한 Hamiltonian 소산 항등식(Hamiltonian dissipation identity)을 증명합니다. 목표 목적 함수 자체가 관성적 Muon 역학을 따라 단조적일 필요는 없지만, 추가적인 구배 지배(gradient-dominance), 유계 모멘텀(bounded-momentum), 그리고 곡률/정렬(curvature/alignment) 가정 하에서, 우리는 목적 함수 격차(objective gap)에 대한 연속 및 이산 시간 지수 수렴 속도(exponential convergence rates)를 얻습니다. 우리는 또한 평균장 극한 방정식의 적절성(well-posedness)을 연구하고 상호작용하는 입자 시스템에 대한 혼돈 전파(propagation of chaos) 보장을 확립합니다. 마지막으로, 우리는 이 정식화를 곱 행렬 공간(product matrix spaces) 상의 Hilbert 값 특징 맵(Hilbert-valued feature maps)으로 확장하여, 매끄러운 Transformer Mixture-of-Experts 모델에 적용 가능한 블록 단위 Muon 확률 흐름(blockwise Muon probability flow)을 산출합니다.