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

TL;DR

PolarGrad introduces a unifying preconditioning framework for matrix-valued gradients, distinguishing gradient anisotropy from curvature anisotropy and using polar decomposition to drive updates. The method generalizes Muon with a nuclear-norm scaling term and is analyzed under exact and inexact polar factors, yielding convergence guarantees and showing when gradient-based rates outperform Hessian-based rates. Empirical results across matrix regression, matrix completion, and language-model pretraining demonstrate that PolarGrad can outperform Adam and Muon, with improved stability and convergence in many matrix-parameter settings. This work provides a principled approach to matrix-parameter optimization, enabling more scalable and robust training of large models by respecting matrix structure and gradient magnitudes.

Abstract

The ever-growing scale of deep learning models and training data 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.

Figures (23)

• Figure 1: Losses, residual and gradient condition numbers of matrix quadratic regression.
• Figure 2: Gradient nuclear norms of matrix quadratic regression (1st seed).
• Figure 3: Losses, gradient condition numbers and nuclear norms of matrix logistic regression.
• Figure 4: Losses and gradient condition numbers of low-rank matrix completion.
• Figure 5: Training losses and gradient condition numbers of Qwen2.5 pre-training: AdamW---AdamW for all parameters; Muon$+$AdamW (PolarSGDM)---Muon for hidden layers and AdamW (PolarSGDM) for embedding and head layers.

Theorems & Definitions (42)

• Definition 3.1: Polar decomposition
• Definition 3.2: Null-gradient consistency
• Definition 3.3: $L$-Lipschitz smoothness
• Proposition 3.1: $\mu$-strong convexity
• Theorem 3.2: PolarGrad
• Theorem 3.3: PolarSGD
• Theorem 3.4: Matrix sign descent and matrix signSGD
• Theorem 3.5: PolarGrad with general inexact polar oracles
• Theorem 3.6: PolarSGD with general inexact polar oracles
• Theorem 3.7: Matrix sign descent and matrix signSGD with general inexact polar oracles