Preconditioning Benefits of Spectral Orthogonalization in Muon
This work provides theoretical insights into a key optimizer for large language models, but it is incremental as it focuses on simplified cases rather than full-scale applications.
The paper tackles the problem of understanding the preconditioning benefits of spectral orthogonalization in the Muon optimizer by analyzing a simplified variant in matrix factorization and linear transformer in-context learning. It proves that this variant converges linearly with iteration complexities independent of condition number, outperforming gradient descent and 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.