PolarGrad: A Class of Matrix-Gradient Optimizers from a Unifying Preconditioning Perspective

2025-05-27

AI Summary

To address slow convergence and training instability in deep learning optimization, this paper proposes PolarGrad, a structure-aware preconditioning framework based on the polar decomposition of gradient matrices. Unlike conventional vector-based preconditioners (e.g., Adam) and existing matrix-based methods (e.g., Shampoo, Muon), PolarGrad is the first to explicitly decouple two distinct preconditioning mechanisms—curvature anisotropy and gradient anisotropy—and unifies modeling of matrix-gradient structure via polar decomposition, incorporating nuclear-norm scaling and efficient numerical algorithms. Muon is rigorously recovered as a special case of PolarGrad, and empirical phenomena such as warmup are theoretically justified. Experiments demonstrate that PolarGrad achieves significantly faster convergence and enhanced training stability compared to Adam and Muon across diverse matrix optimization tasks and large language model pretraining.

Technology Category

• Machine Learning: Optimization
• Natural Language Processing: Learning & Optimization for NLP
• Computer Vision: Learning & Optimization for CV

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

The ever-growing scale of deep learning models and datasets underscores the critical importance of efficient optimization methods. While preconditioned gradient methods such as Adam and AdamW are the de facto optimizers for training neural networks and large language models, structure-aware preconditioned optimizers like Shampoo and Muon, which utilize the matrix structure of gradients, have demonstrated promising evidence of faster convergence. In this paper, we introduce a unifying framework for analyzing"matrix-aware"preconditioned methods, which not only sheds light on the effectiveness of Muon and related optimizers but also leads to a class of new structure-aware preconditioned methods. A key contribution of this framework is its precise distinction between preconditioning strategies that treat neural network weights as vectors (addressing curvature anisotropy) versus those that consider their matrix structure (addressing gradient anisotropy). This perspective provides new insights into several empirical phenomena in language model pre-training, including Adam's training instabilities, Muon's accelerated convergence, and the necessity of learning rate warmup for Adam. Building upon this framework, we introduce PolarGrad, a new class of preconditioned optimization methods based on the polar decomposition of matrix-valued gradients. As a special instance, PolarGrad includes Muon with updates scaled by the nuclear norm of the gradients. We provide numerical implementations of these methods, leveraging efficient numerical polar decomposition algorithms for enhanced convergence. Our extensive evaluations across diverse matrix optimization problems and language model pre-training tasks demonstrate that PolarGrad outperforms both Adam and Muon.

Problem

• Developing efficient optimization methods for large-scale deep learning models
• Analyzing matrix-aware preconditioned methods for neural network training
• Introducing PolarGrad to improve convergence in language model pre-training

Innovation

• Unifying framework for matrix-aware preconditioned methods
• PolarGrad optimizers using polar decomposition
• Efficient numerical polar decomposition algorithms