뮤온에서 스펙트럴 직교화의 사전 조건화 효과

Preconditioning Benefits of Spectral Orthogonalization in Muon

뮤온 최적화 알고리즘은 그래디언트의 스펙트럴 직교화를 활용하는 행렬 구조 기반 알고리즘으로, 대규모 언어 모델의 사전 학습 분야에서 중요한 진전을 이루었습니다. 그러나 뮤온의 작동 원리, 특히 그래디언트 직교화의 역할은 아직 명확하게 이해되지 못하고 있으며, 구체적인 응용 분야에서 뮤온의 장점을 엄밀하게 설명하는 연구는 매우 드뭅니다. 본 연구에서는 행렬 분해와 선형 트랜스포머의 문맥 내 학습이라는 두 가지 사례 연구를 통해 뮤온의 단순화된 변형의 효과를 분석합니다. 두 문제 모두에서, 단순화된 뮤온이 반복 복잡도와 상관없이 선형적으로 수렴하며, 그래디언트 강하법 및 Adam 알고리즘보다 우수한 성능을 보인다는 것을 증명합니다. 분석 결과, 뮤온의 동역학은 스펙트럴 영역에서 서로 독립적인 스칼라 시퀀스로 분리되며, 각 시퀀스는 유사한 수렴 특성을 보입니다. 본 연구의 이론적 분석은 스펙트럴 직교화에 의해 유도되는 사전 조건화 효과를 명확히 제시하며, 뮤온이 이러한 행렬 최적화 문제에서 효과적인 이유와 그 잠재적 활용 가능성에 대한 통찰력을 제공합니다.

The Muon optimizer, a matrix-structured algorithm that leverages spectral orthogonalization of gradients, is a milestone in the pretraining of large language models. However, the underlying mechanisms of Muon -- particularly the role of gradient orthogonalization -- remain poorly understood, with very few works providing end-to-end analyses that rigorously explain its advantages in concrete applications. We take a step by studying the effectiveness of a simplified variant of Muon through two case studies: matrix factorization, and in-context learning of linear transformers. For both problems, we prove that simplified Muon converges linearly with iteration complexities independent of the relevant condition number, provably outperforming gradient descent and Adam. Our analysis reveals that the Muon dynamics decouple into a collection of independent scalar sequences in the spectral domain, each exhibiting similar convergence behavior. Our theory formalizes the preconditioning effect induced by spectral orthogonalization, offering insight into Muon's effectiveness in these matrix optimization problems and potentially beyond.