Nonsmooth Optimization via Orthogonalized Momentum
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 β∈[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 β∈[0,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.
Related
- Infeasible Momentum Orthogonalized Methods for Optimization on the Stiefel Manifold
- Sharpness-aware Orthogonalized Momentum Algorithms
- A Two-Stage Optimization Framework for Real-World Large-Scale University Course Timetabling Based on Orthogonalized Hierarchical Modeling and Conflict-Driven ALNS Repair
- Linear convergence rate analysis of Proximal-Tracking for nonsmooth distributed optimization
- Exact penalty methods with adaptive discretization for nonsmooth semi-infinite optimization
BibTeX
@misc{lai2026nonsmooth,
title = {{Nonsmooth Optimization via Orthogonalized Momentum}},
author = {Lexiao Lai and Tianyi Lin and Jiayu Zhang},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2609.13677},
eprint = {2609.13677},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.13677},
}
RIS
TY - UNPB
TI - Nonsmooth Optimization via Orthogonalized Momentum
AU - Lexiao Lai
AU - Tianyi Lin
AU - Jiayu Zhang
PY - 2026
DA - 2026/10/09
PB - arXiv
DO - 10.48550/arxiv.2609.13677
UR - https://arxiv.org/abs/2609.13677
AB - 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 β∈[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 β∈[0,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.
ER -