Nonsmooth Optimization via Orthogonalized Momentum

Lexiao Lai; Tianyi Lin; Jiayu Zhang

math.OCSep 12, 2026

Nonsmooth Optimization via Orthogonalized Momentum

Authors: Lexiao Lai, Tianyi Lin, Jiayu Zhang

Organizations: Department of Mathematics, The University of Hong Kong · Department of Industrial Engineering and Operations Research Columbia University

Abstract

Modern real application problems involve matrix-valued parameters, yet conventional optimizers treat them as vectors, thereby motivating matrix-aware methods that exploit input-output geometry, such as Muon which orthogonalizes the momentum matrices before parameter updates. Its empirical success raises a conceptual question: can orthogonalized momentum remain effective beyond smooth optimization? This paper studies this question for locally Lipschitz functions using a generalized derivative framework compatible with backpropagation. Our first contribution is to identify a key limitation: for every fixed momentum factor $β\in[0,1)$ , Muon can fail to approach the global optimal solution of a convex Lipschitz objective from almost every initialization, when step sizes adapt to the full gradient history. The failure can occur even along bounded iterates. Our example is inspired by the one of Parshakova et al. which only covers $β\in[0,\frac{1}{2})$ . Then, we show that the obstruction lies in fixed momentum rather than orthogonalization. Indeed, when the momentum factor is adaptive and approaches 1 together with a vanishing step size, Muon recovers asymptotic convergence for nonconvex nonsmooth optimization under the boundedness and regularity conditions. Moreover, we propose MAGD, which combines orthogonalized momentum with gradient, weighted based on their relative progress. MAGD retains asymptotic convergence in nonconvex settings and achieves an $O(\min\{m,n\}ε^{-2})$ rate in convex settings. A lower bound shows the optimal dimension dependence. Experiments on synthetic problems, image classification, and LLM pretraining show MAGD is a simple and practical alternative to Muon. Together, our results characterize when orthogonalized momentum fails without smoothness and how it can be made reliable and we hope that the analysis may be useful more broadly.

Figures & tables

Panel of Figure 1 — Figure 1: Reduction of a convex nonsmooth optimization algorithm $\mathsf{A}$ to a string-guessing algorithm $\mathsf{B}$

$\delta$	Muon (fixed $\beta_{t}=0.7$ )	Muon ( $\beta_{t}=\max\{1-256\eta_{t},0\}$ )	GD	MAGD (fixed $\beta_{t}^{\rm{M}}=0.7$ )
$0$	$0\%$	$85\%$	$100\%$	$100\%$
$10^{-4}$	$55\%$	$100\%$	$100\%$	$100\%$
$10^{-3}$	$100\%$	$100\%$	$100\%$	$100\%$
$10^{-2}$	$100\%$	$100\%$	$100\%$	$100\%$
$10^{-1}$	$100\%$	$100\%$	$100\%$	$100\%$

Table 1: Percentage of the 100 initializations for which the objective gap is below

10^{-6}

Spectrum	Construction
max_spiked	One eigenvalue at $s_{\min}$ , all others at $s_{\max}$ .
min_spiked	One eigenvalue at $s_{\max}$ , all others at $s_{\min}$ .
uniform	$s_{i}$ sampled uniformly on $[s_{\min},s_{\max}]$ .
gaussian	$s_{i}$ sampled from a clipped Gaussian centered in $[s_{\min},s_{\max}]$ .
linear_decay_to_max	$s_{i}=s_{\max}-(s_{\max}-s_{\min})\sqrt{u_{i}}$ , $u_{i}\sim\mathrm{Unif}[0,1]$ .
geometric_decay_to_max	$s_{i}=\max\{s_{\min},s_{\max}\cdot 0.9^{i-1}\}$ .

Table 2: Spectrum families used in the quadratic optimization experiment.

Figure 2: Orders of magnitude of loss decrease after 500 iterations on the quadratic functions.

Method	Val. Loss	Val. Accuracy	Test Loss	Test Accuracy
MAGD (NS3)	$0.4312\pm 0.0002$	$0.9253\pm 0.0002$	$0.4271\pm 0.0003$	$0.9277\pm 0.0002$
Muon (NS3)	$0.4313\pm 0.0003$	$0.9253\pm 0.0002$	$0.4271\pm 0.0003$	$0.9279\pm 0.0002$
MAGD (NS5)	$0.4310\pm 0.0003$	$0.9260\pm 0.0002$	$0.4271\pm 0.0002$	$0.9282\pm 0.0002$
Muon (NS5)	$0.4310\pm 0.0003$	$0.9259\pm 0.0002$	$0.4273\pm 0.0002$	$0.9282\pm 0.0001$
MAGD (PolarExpress)	$0.4291\pm 0.0002$	$0.9266\pm 0.0002$	$0.4249\pm 0.0002$	$0.9288\pm 0.0002$
Muon (PolarExpress)	$0.4294\pm 0.0003$	$0.9266\pm 0.0002$	$0.4253\pm 0.0003$	$0.9286\pm 0.0002$

Table 3: Results are the mean

\pm

standard error over 100 seeds for both validation and test. NS3 and NS5 mean three and five Newton-Schulz steps, respectively. PolarExpress uses five steps.

Figure 3: Performance of Muon and MAGD on nanochat LLM pretraining.

Muon Optimizer

Muon optimizer uses momentum and matrix orthogonalization to shape neural-network updates, especially in large language-model training.

See all 143 papers

Aug 26, 2026

Muon with Finite Newton-Schulz: The Smoothing Benefit in Nonsmooth Nonconvex Optimization

Muon has emerged as a strong optimizer for the matrix-valued parameters in large language model pretraining, approximately orthogonalizing its momentum with a few Newton-Schulz iterations. Existing theory either replaces this iteration with the exact polar factor it approximates, or treats its finite depth as an approximation error, and thus the iteration Muon actually runs can only hurt the guarantees. We show that finite Newton-Schulz can instead be beneficial for nonsmooth nonconvex optimization. To this end, we analyze Muon through the online-to-nonconvex conversion, which views the update rule as an online learner and converts its regret bound into a stationarity guarantee. The finite Newton-Schulz iteration smooths the discontinuous polar map into a Lipschitz map of the singular values, and Muon with finite Newton-Schulz can be regarded as an online learner with a smoothed spectral potential. This smoothing is exactly what the conversion needs: we prove that a Newton-Schulz depth growing only logarithmically in the target accuracy suffices for convergence to stationary points in nonsmooth nonconvex optimization, whereas Muon with the exact-polar update may fail to converge. The resulting sample complexity bounds match the best-known guarantees for nonsmooth nonconvex optimization and are optimal for smooth nonconvex optimization up to problem-dependent factors. The argument extends beyond Newton-Schulz to general spectral maps with the same smoothing property.

Mingyi Li, Taira Tsuchiya

The University of Tokyo · The University of Tokyo and RIKEN

Muon Optimizer Newton-Schulz Iteration

Ask PDF

Jun 7, 2026

OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality

Orthogonalized momentum updates, as used in Muon-style optimizers, have recently shown strong empirical stability in large-scale deep learning. However, most current orthogonalized methods are still paired with fixed, externally scheduled, or otherwise open-loop magnitude rules, so their scale is not directly calibrated from the realized optimization trajectory. Motivated by the closed-loop perspective behind Lipschitz-free and noise-adaptive methods, we propose OptMuon, a family of adaptive momentum orthogonalization methods for stochastic nonconvex optimization. OptMuon combines Muon-style polar-factor directions with a trajectory-dependent AdaGrad-Norm-type coefficient schedule, so that the update magnitude is determined by the observed gradient and momentum history rather than by a prescribed Lipschitz-dependent rule. The schedule does not use the smoothness constant, the variance level, or the bounded-gradient constant in parameter selection, and its running-maximum correction prevents isolated gradient spikes from causing excessive coefficient collapse. Under lower-boundedness, unbiased stochastic gradients with bounded variance, smoothness, and an almost-sure bounded stochastic-gradient condition, we prove two complementary expected-stationarity guarantees. OptMuon-A achieves the noise-adaptive rate

\tilde{\mathcal O}(T^{-1/2}+σ^{1/2}T^{-1/4})

under average smoothness, while OptMuon-I achieves

\tilde{\mathcal O}(T^{-1/2}+σ^{1/3}T^{-1/3})

under individual smoothness. In the zero-noise regime, both bounds automatically reduce to a nearly optimal deterministic first-order rate

\tilde{\mathcal O}(T^{-1/2})

without manual hyperparameter retuning. These results show that closed-loop scalar adaptation can be combined with Muon-style momentum orthogonalization while retaining noise adaptivity and zero-noise optimality up to logarithmic factors.

Ganzhao Yuan

Faculty of Computer Science and Artificial Intelligence Shenzhen University of Advanced Technology (SUAT), China

Momentum Methods Nonconvex Stochastic Optimization Muon Optimizer Stochastic Optimization Convergence Adaptive Gradient Methods

Ask PDF

May 26, 2026

Convergence of Spectral Descent for Non-smooth Optimization

The Muon optimizer has recently demonstrated remarkable empirical success in training large language models. However, the theoretical understanding of its mechanisms remains limited. Current convergence guarantees for Muon rely heavily on smoothness assumptions, leaving its non-smooth convergence behavior largely unexplored. In this work, we take a step toward bridging this gap by investigating Spectral Descent (SD), a simplified variant of Muon, together with its truncated counterpart, Truncated Spectral Descent (TSD). Under convexity, Lipschitz continuity, and sharpness conditions, we establish global linear convergence for both SD and TSD in non-smooth convex formulations. We also study regularized variants equipped with decoupled weight decay and derive sublinear convergence guarantees through their connection with Frank-Wolfe methods. Finally, we apply our theoretical framework to robust low-rank matrix recovery under mixed sparse and dense noise regimes and provide rigorous recovery guarantees. Numerical experiments support the theoretical findings and demonstrate the effectiveness of Muon-type methods for non-smooth optimization.

Yixuan Yang, Yuqing He, Song Li

School of Mathematical Sciences, Zhejiang University, Hangzhou, China

Muon Optimizer Nonsmooth Convex Optimization

Ask PDF

Abstract

Muon Optimizer

Explore similar work

Muon with Finite Newton-Schulz: The Smoothing Benefit in Nonsmooth Nonconvex Optimization

OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality

Convergence of Spectral Descent for Non-smooth Optimization