Preconditioning Benefits of Spectral Orthogonalization in Muon

TL;DR

This paper analyzes a simplified version of the Muon optimizer, demonstrating its spectral orthogonalization preconditioning benefits through theoretical convergence guarantees and insights into its dynamics.

Contribution

It provides the first rigorous analysis of Muon's spectral orthogonalization, showing linear convergence independent of condition number in matrix problems.

Findings

Simplified Muon converges linearly in matrix factorization and in-context learning.

Spectral orthogonalization decouples dynamics into independent scalar sequences.

Muon's preconditioning effect improves convergence over gradient descent and Adam.

Abstract

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.