Beyond the Ideal: Analyzing the Inexact Muon Update

TL;DR

This paper analyzes the practical inexact orthogonalization in Muon optimizer, providing explicit bounds on performance degradation and revealing how approximation errors influence optimal step size and momentum, with experimental validation on NanoGPT.

Contribution

First theoretical analysis of inexact Muon updates within the LMO framework, linking approximation errors to optimization hyperparameters and demonstrating their impact through experiments.

Findings

Inexact orthogonalization affects Muon's efficiency and requires co-tuning of hyperparameters.

Explicit bounds quantify performance loss due to approximation errors.

Experimental results on NanoGPT confirm the theoretical predictions.

Abstract

The Muon optimizer has rapidly emerged as a powerful, geometry-aware alternative to AdamW, demonstrating strong performance in large-scale training of neural networks. However, a critical theory-practice disconnect exists: Muon's efficiency relies on fast, approximate orthogonalization, yet all prior theoretical work analyzes an idealized, computationally intractable version assuming exact SVD-based updates. This work moves beyond the ideal by providing the first analysis of the inexact orthogonalized update at Muon's core. We develop our analysis within the general framework of Linear Minimization Oracle (LMO)-based optimization, introducing a realistic additive error model to capture the inexactness of practical approximation schemes. Our analysis yields explicit bounds that quantify performance degradation as a function of the LMO inexactness/error. We reveal a fundamental coupling between this inexactness and the optimal step size and momentum: lower oracle precision requires a smaller step size but larger momentum parameter. These findings elevate the approximation procedure (e.g., the number of Newton-Schulz steps) from an implementation detail to a critical parameter that must be co-tuned with the learning schedule. NanoGPT experiments directly confirm the predicted coupling, with optimal learning rates clearly shifting as approximation precision changes.