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

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 develop the stochastic non-Euclidean trust-region gradient method with momentum, which recovers the Muon optimizer (Jordan et al., 2024) as a special case, along with normalized SGD and signSGD with momentum (Cutkosky and Mehta, 2020; Sun et al., 2023). In addition, we prove state-of-the-art convergence results for the proposed algorithm in a range of scenarios, which involve arbitrary non-Euclidean norms, constrained and composite problems, and non-convex, star-convex, first- and second-order smooth functions. Finally, our theoretical findings provide an explanation for several practical observations, including the practical superiority of Muon compared to the Orthogonal-SGDM algorithm of Tuddenham et al. (2022) and the importance of weight decay in the training of large-scale LLMs.

• The paper establishes a theoretical link between orthogonalized gradients and non-Euclidean trust-region optimization, underpinning the Muon optimizer’s performance.
• It rigorously proves convergence guarantees, enhancing iteration complexity from O(ε⁻³.⁵) to O(ε⁻³) in training deep neural networks.
• The study extends normalized SGD to composite and star-convex scenarios, providing actionable insights for advanced deep learning optimization.

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

Introduction to the Paper

The paper "Understanding Gradient Orthogonalization for Deep Learning via Non-Euclidean Trust-Region Optimization" explores an innovative method of optimizing deep learning models using orthogonalized gradients. The research situates itself within the broader context of adaptive gradient optimization algorithms, which are pivotal in training neural networks. This paper explores the theoretical underpinning of integrating matrix gradient orthogonalization with trust-region methods, revealing its efficacy in training deep neural networks.

Main Contributions

This research primarily contributes by establishing the theoretical connection between orthogonalized gradients and non-Euclidean trust-region optimization. This insight importantly underlies the Muon optimizer, demonstrating why it performs better than the Orthogonal-SGDM algorithm.

1. Theoretical Analysis of Orthogonalized Gradients: The paper interprets the orthogonalized gradient method as a trust-region approach based on the matrix spectral norm. This presents a fresh perspective compared to the interpretation of traditional methods such as Shampoo.
2. Convergence Guarantees: It provides a rigorous convergence analysis of the Muon optimizer. These results validate the performance improvements observed in practical implementations compared to other optimizers like Orthogonal-SGDM.
3. Generalized Framework and Applications: The research extends previous work on normalized SGD to account for composite and non-Euclidean setting optimizations, providing a comprehensive framework that includes constraints and star-convexity scenarios.

Methodology and Theoretical Findings

• Orthogonalized Gradient as a Trust-Region Method: The study identified that the orthogonalized gradient update can be understood as a non-Euclidean trust-region method. This methodological innovation simplifies understanding its convergence properties and underpins the efficacy of the Muon algorithm.
• Algorithmic Interpretation: The non-Euclidean trust-region gradient method with momentum, a variant introduced in this paper, is shown to recover various optimizers, including signSGD and Muon. It is uniquely robust across different norms, including spectral and nuclear norms.
• Improved Iteration Complexity: By leveraging a star-convexity assumption, the paper demonstrates an improved iteration complexity from $\mathcal{O}(\varepsilon^{-3.5})$ to $\mathcal{O}(\varepsilon^{-3})$. This enhancement over the normalized SGD results indicates efficiency in convergence for specific neural network problems.

Practical Implications and Future Directions

The elucidation of the link between orthogonalized gradients and trust-region methods has significant implications for optimizing deep neural networks. In particular, the insights support the refinement and deployment of the Muon optimizer in practical deep learning applications, potentially leading to more efficient training regimes for LLMs.

Additionally, the paper prompts further exploration of trust-region methods in first-order optimization settings. Future research could investigate extending these foundational results to broader classes of problems, potentially uncovering further optimizations and applications in deep learning and other machine learning domains.

Conclusion

This paper provides a robust theoretical framework for understanding the role of orthogonalized gradients within trust-region optimization contexts. Its convergence analyses and algorithmic insights not only advance theoretical knowledge but also have practical implications for optimizing complex neural networks. The work signifies a substantive step in combining gradient orthogonalization with efficient optimization strategies, paving the way for further advancements in adaptive optimization techniques.