Understanding Gradient Orthogonalization for Deep Learning via Non-Euclidean Trust-Region Optimization

2025-03-16

AI Summary

This work addresses the lack of a unified theoretical foundation for gradient orthogonalization methods in deep neural network training. Methodologically, it introduces a novel perspective grounded in non-Euclidean trust-region optimization: gradient orthogonalization is formalized as a first-order trust-region method with the matrix spectral norm as the metric, and its convergence is established under star-convexity assumptions—the first such result. Technically, it provides a unified analysis of momentum-based stochastic non-Euclidean trust-region optimization and rigorously proves that Muon is a special case of this framework. The iteration complexity of normalized SGD with momentum and Muon is improved to $O(varepsilon^{-3})$, surpassing the prior $O(varepsilon^{-3.5})$. Crucially, this work delivers the first unified theoretical justification for orthogonalization algorithms—including Orthogonal-SGDM and Muon—and reveals that Muon’s superiority stems from its implicit satisfaction of the spectral-norm constraint.

Technology Category

• Machine Learning: Optimization
• Search and Optimization: Non-convex Optimization
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Security and Privacy: Large-scale security measurements

Abstract

Optimization with matrix gradient orthogonalization has recently demonstrated impressive results in the training of deep neural networks (Jordan et al., 2024; Liu et al., 2025). In this paper, we provide a theoretical analysis of this approach. In particular, we show that the orthogonalized gradient method can be seen as a first-order trust-region optimization method, where the trust-region is defined in terms of the matrix spectral norm. Motivated by this observation, we provide the first theoretical analysis of the stochastic non-Euclidean trust-region gradient method with momentum, which recovers the Muon optimizer (Jordan et al., 2024) as a special case. In addition, we establish the convergence of the normalized SGD with momentum (Cutkosky and Mehta, 2020) in the constrained and composite setting, show that its iteration complexity of finding an $varepsilon$-accurate solution can be improved from $mathcal{O}(varepsilon^{-3.5})$ to $mathcal{O}(varepsilon^{-3})$ under the star-convexity assumption, and obtain similar results for the Muon algorithm. Finally, our theoretical findings provide an explanation for the practical superiority of Muon compared to the Orthogonal-SGDM algorithm of Tuddenham et al. (2022).

Problem

• Analyzes gradient orthogonalization in deep learning optimization.
• Theorizes non-Euclidean trust-region gradient method with momentum.
• Improves convergence rates for stochastic gradient descent methods.

Innovation

• Non-Euclidean trust-region optimization for deep learning
• Matrix gradient orthogonalization improves training efficiency
• Muon optimizer outperforms Orthogonal-SGDM in practice