Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Preconditioning Benefits of Spectral Orthogonalization in Muon

Jianhao Ma, Yu Huang, Yuejie Chi, Yuxin Chen

TL;DR

This paper analyzes Muon, a spectrum-aware optimizer that performs gradient orthogonalization via the matrix sign function, to understand its preconditioning mechanism. By focusing on two fundamental matrix problems—symmetric matrix factorization and in-context learning with linear transformers—it shows that simplified Muon achieves linear, condition-number-free convergence, with iteration counts scaling only with $\, ext{log}(1/\varepsilon)$ rather than the problem's condition number. The key insight is that Muon's gradient spectrum normalization decouples into independent scalar sequences in the spectral domain, yielding robust convergence even under overparameterization. The results are complemented by numerical experiments, showing Muon's superiority over GD and SignGD in ill-conditioned and over-parameterized regimes and suggesting broad applicability of spectrum-aware preconditioning beyond the studied settings.

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.

Preconditioning Benefits of Spectral Orthogonalization in Muon

TL;DR

This paper analyzes Muon, a spectrum-aware optimizer that performs gradient orthogonalization via the matrix sign function, to understand its preconditioning mechanism. By focusing on two fundamental matrix problems—symmetric matrix factorization and in-context learning with linear transformers—it shows that simplified Muon achieves linear, condition-number-free convergence, with iteration counts scaling only with rather than the problem's condition number. The key insight is that Muon's gradient spectrum normalization decouples into independent scalar sequences in the spectral domain, yielding robust convergence even under overparameterization. The results are complemented by numerical experiments, showing Muon's superiority over GD and SignGD in ill-conditioned and over-parameterized regimes and suggesting broad applicability of spectrum-aware preconditioning beyond the studied settings.

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.
Paper Structure (66 sections, 26 theorems, 226 equations, 3 figures, 1 table)

This paper contains 66 sections, 26 theorems, 226 equations, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. The Muon algorithm and prior theory
  3. This paper: preconditioning with Muon
  4. Additional related work
  5. Notation
  6. Main results: two case studies
  7. Matrix factorization
  8. Convergence guarantees for Muon
  9. Comparisons with other optimizers
  10. Intuition
  11. Decoupling of Muon dynamics into independent scalar sequences.
  12. Why do SignGD and Adam fail?
  13. Numerical experiments
  14. In-context learning with linear transformers
  15. Motivation: in-context learning of a single-layer linear transformer.
  16. ...and 51 more sections

Key Result

Theorem 1

Suppose that $\lambda_{\max}^{\star}\geq \lambda_1^{\star}\geq \dots \geq \lambda_r^{\star}> 0$, and consider any $0<\varepsilon<\lambda_{\max}^{\star}$.

Figures (3)

  • Figure 1: Numerical convergence behavior of Muon, SignGD, and GD on matrix factorization tasks under varying condition numbers and search ranks.
  • Figure 2: Numerical convergence behavior of Muon, SignGD, and GD on in-context learning problems with one-layer linear transformers under varying condition numbers.
  • Figure 3: Numerical comparison of the preconditioners of Muon and ScaledGD for matrix factorization at various training steps along a Muon trajectory.

Theorems & Definitions (38)

  • Theorem 1
  • Remark 1
  • Theorem 2
  • Theorem 3
  • Theorem 4
  • Lemma 1
  • Lemma 2: Convergence of scalar Muon
  • proof : Proof of \ref{['lem::never-reach-zero']}
  • proof : Proof of \ref{['lem::1d-muon']}
  • Lemma 3
  • ...and 28 more