Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Beyond the Ideal: Analyzing the Inexact Muon Update

Egor Shulgin, Sultan AlRashed, Francesco Orabona, Peter Richtárik

TL;DR

This work addresses the gap between Muon's practical, approximate orthogonalization and its theoretical analysis by modeling the orthogonalization step as an inexact Linear Minimization Oracle. It develops a general framework for inexact LMO-based optimization, deriving explicit convergence bounds in both deterministic and stochastic settings and revealing a fundamental coupling: as LMO inexactness $\delta$ grows, the optimal step size $\gamma$ decreases while the momentum $\alpha$ should increase. The authors extend the results to $(L^0,L^1)$-smoothness and layer-wise settings, and validate the theory with nanoGPT on FineWeb and CIFAR-10 experiments, showing degraded performance with lower LMO precision and confirming the predicted hyperparameter shifts. The practical impact is clear: achieving better approximate orthogonalization (e.g., via PolarExpress or faster Newton-Schulz variants) translates directly into faster convergence and more robust hyperparameter tuning in large-scale neural network training.

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.

Beyond the Ideal: Analyzing the Inexact Muon Update

TL;DR

This work addresses the gap between Muon's practical, approximate orthogonalization and its theoretical analysis by modeling the orthogonalization step as an inexact Linear Minimization Oracle. It develops a general framework for inexact LMO-based optimization, deriving explicit convergence bounds in both deterministic and stochastic settings and revealing a fundamental coupling: as LMO inexactness grows, the optimal step size decreases while the momentum should increase. The authors extend the results to -smoothness and layer-wise settings, and validate the theory with nanoGPT on FineWeb and CIFAR-10 experiments, showing degraded performance with lower LMO precision and confirming the predicted hyperparameter shifts. The practical impact is clear: achieving better approximate orthogonalization (e.g., via PolarExpress or faster Newton-Schulz variants) translates directly into faster convergence and more robust hyperparameter tuning in large-scale neural network training.

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.
Paper Structure (44 sections, 17 theorems, 151 equations, 5 figures, 1 table, 1 algorithm)

This paper contains 44 sections, 17 theorems, 151 equations, 5 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Contributions
  3. Related work
  4. From Idealized Theory to Practical Implementation
  5. The idealized update
  6. The practical imperative of approximation
  7. Modeling inexactness
  8. Main Theoretical Results
  9. Deterministic case
  10. Stochastic case
  11. Extensions and generalizations
  12. Beyond standard smoothness.
  13. Layer-wise analysis.
  14. Experiments
  15. nanoGPT on FineWeb
  16. ...and 29 more sections

Key Result

Theorem 1

Let Assumption assump:smoothness hold. Let the sequence $\{x^k\}_{k=0}^{K-1}$ be generated by the update rule eq:deterministic_method with step sizes $\gamma_k > 0$. Assume the inexact LMO satisfies Assumption assump:inexact_lmo_main with $\delta_k < 1$ for all $k$ and let $\Delta^0 := f(x^0) - f^*$

Figures (5)

  • Figure 1: NanoGPT on FineWeb: convergence (left) and hyperparameter sensitivity (right) as PolarExpress (PolarExp) iterations vary.
  • Figure 2: CIFAR-10 training loss and test accuracy across different Newton-Schulz (NS) precision levels.
  • Figure 3: Supplementary results for nanoGPT on FineWeb: convergence behavior (left) and runtime scaling (right) as the number of PolarExpress (PolarExp) iterations varies.
  • Figure 4: LMO error evolution across training steps for different numbers of PolarExpress (PolarExp) iterations. Higher iteration counts maintain lower orthogonalization error throughout training.
  • Figure 5: Full final test accuracy heatmaps on CIFAR-10 across a grid of step sizes and momentum values for different numbers of Newton-Schulz ("Newton") iterations. Brighter colors indicate higher final test accuracy. The region of high performance is visibly larger and more stable for higher numbers of Newton-Schulz iterations.

Theorems & Definitions (33)

  • Theorem 1: General Result
  • Corollary 1: Constant Parameters
  • Corollary 2: Optimal Parameters
  • Lemma 1
  • Theorem 2
  • Corollary 3
  • proof
  • proof
  • proof
  • Corollary 4: Iteration Complexity with Constant Step Size
  • ...and 23 more