Training-Loss Guarantees for Muon with Finite-Step Newton--Schulz Orthogonalization
cs.LG, cs.AI, math.OC
Submitted: 2026-10-02
Updated: 2026-10-02
Comments: 22 pages, 5 figures
License: http://creativecommons.org/licenses/by/4.0/
The gist: Existing convergence analyses of Muon either assume exact orthogonalization or analyze classical Newton--Schulz polynomials, and guarantee only stationarity, so it is unresolved what Muon's five
Abstract
Existing convergence analyses of Muon either assume exact orthogonalization or analyze classical Newton--Schulz polynomials, and guarantee only stationarity, so it is unresolved what Muon's five tuned Newton--Schulz steps preserve and whether that suffices to reach a prescribed neural-network training loss. We establish a finite-time training guarantee that accounts for both momentum accumulation before orthogonalization and the tuned finite-step update. For full-batch training of a sufficiently wide two-layer ReLU network with fixed random output weights and a positive-definite limiting neural tangent kernel, we prove that Muon reaches any target empirical squared loss epsilon>0 with high probability over initialization. For every momentum parameter μ in[0,1), a target-dependent constant learning rate proportional to (1-μ) sqrt epsilon yields a hitting-time bound of O((1-μ)-1 epsilon-1/2), with other problem parameters fixed. The sufficient width is independent of both target accuracy and momentum. The analysis shows that the tuned Newton--Schulz map preserves alignment with the momentum buffer while bounding the update's spectral norm. Control of gradient variation near initialization transfers this alignment to the current gradient, ensuring descent until the target is reached without requiring exact orthogonalization. Numerical experiments support these mechanisms at widths below the sufficient theoretical threshold: gradient-update alignment remains above the analytical reference, and all 30 runs across six widths and five student initializations on a fixed teacher-student dataset reach the target loss while maintaining kernel positivity.
Sources
- Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks
- Old Optimizer, New Norm: An Anthology
- Modular Duality in Deep Learning
- On the Convergence of Muon and Beyond
- Convergence guarantees for Muon: New parameter regimes and generalizations
- Gradient Descent Finds Global Minima of Deep Neural Networks
- Gradient Descent Provably Optimizes Over-parameterized Neural Networks
- Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow ReLU networks
- Convergence of Muon with Newton-Schulz
- Learning Overparameterized Neural Networks via Stochastic Gradient Descent on Structured Data
- Muon is Scalable for LLM Training
- Convergence Bound and Critical Batch Size of Muon Optimizer
- On the Convergence Analysis of Muon
- Kimi K2: Open Agentic Intelligence
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks