Muon을 위한 자유로운 헤비 테일(Heavy-Tailed) 점심: 경험적 성공에 대한 이론적 정당화

Muon 및 Scion과 같이 행렬 값 업데이트(matrix-valued updates)를 사용하는 비유클리드 최적화(Non-Euclidean optimisation) 방법들은 최근 Transformer 모델 학습에서 강력한 경험적 성능을 보여주었으나, 유클리드 방법(Euclidean methods) 대비 이들의 이론적 이점은 여전히 명확히 이해되지 않고 있습니다. 우리는 확률적 경사(stochastic gradients)가 유계된 $p$-차 중심 모멘트(bounded $p$-th central moments, $p imes (1,2]$)를 갖는 헤비 테일(heavy-tailed) 비볼록(non-convex) 영역에서 이 격차를 다룹니다. 우리는 특정 비유클리드 방법들이 더 강력한 정체성 측정치(stationarity measures) 하에서 최적의 샘플 복잡도(sample complexity)를 달성하는 반면, 유클리드 방법들은 차원 의존적인(dimension-dependent) 추가 비용을 발생시킨다는 것을 보여줍니다. 그 결과, $m imes n$ 행렬에 대해 Muon은 유클리드 방법과 달리 추가적인 차원 의존성 없이 헤비 테일 노이즈를 흡수하며, 핵 노름(nuclear norm) 내에서 $\mathcal{O}\left(\min{m, n} \frac{Δ_1 L}{\varepsilon^2} \left(\frac {\sigma}\varepsilon \right)^{\frac p {p-1}}\right)$ 샘플 이내에 $\varepsilon$-정체 지점($\varepsilon$-stationary point)을 찾아냅니다. 나아가 우리는 이 샘플 복잡도가 핵 노름 정체성(nuclear-norm stationarity) 하에서 모든 1차 방법(first-order methods)에 대해 차원 의존성을 포함하여 최적임을 증명합니다. 대규모 언어 모델(LLM)에 대한 실험은 우리의 이론을 뒷받침합니다. 놀랍게도, 우리의 결과는 Muon의 스펙트럼 기하학(spectral geometry)을 넘어선 다른 Schatten 기하학(Schatten geometries)들도 특정 설정에서 경쟁력 있는 성능을 발휘할 수 있음을 시사합니다.