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

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.

💡 Research Summary

The paper tackles the pressing need for more efficient optimizers as deep learning models scale to billions of parameters and massive datasets. While Adam and its variants dominate current practice, they treat all parameters as a single flattened vector and apply diagonal preconditioning that approximates the inverse square root of the Hessian. This “curvature preconditioning” works well for scalar and vector parameters but ignores the inherent matrix or tensor structure of many layers (e.g., attention weight matrices, convolution kernels).

The authors introduce a unifying preconditioning framework that distinguishes two fundamentally different strategies: (1) curvature‑based preconditioning, which reduces the condition number of the Hessian and is appropriate for vector‑valued parameters, and (2) gradient‑based preconditioning, which directly improves the condition number of the update direction (gradient or momentum) for matrix‑valued parameters. They argue that the latter is essential because a poorly conditioned gradient can dramatically slow convergence or cause instability, especially in the early phases of training large language models.

Building on this insight, the paper revisits Muon, a recent structure‑aware optimizer that orthogonalizes matrix gradients via polar decomposition. The authors show that Muon can be interpreted as a gradient‑preconditioned method: the polar factor U (the orthogonal component) has condition number 1, thus providing the optimal preconditioning for the matrix update. However, Muon’s original implementation relies on a Newton–Schulz iteration for the polar decomposition, which can be numerically fragile and requires careful tuning.

From this analysis emerges PolarGrad, a new family of matrix‑gradient optimizers. The core idea is to compute the polar decomposition G = U P of the current gradient G, then scale the update by the nuclear norm ‖P‖_* (the sum of singular values). The update rule becomes
ΔW = −η · ‖P‖_*⁻¹ · U · M,
where M is a momentum term (e.g., Adam’s first‑moment estimate) and η is the learning rate. When the scaling factor is set to 1, PolarGrad reduces exactly to Muon, showing that Muon is a special case of the broader PolarGrad framework.

To address the computational bottleneck of polar decomposition, the authors adopt two state‑of‑the‑art algorithms: QDWH (Quasi‑Diagonal‑Weighted‑Halley) and ZOLO‑PD (Zero‑Order‑Lipschitz‑Optimized Polar Decomposition). Both achieve rapid convergence with only a few iterations, are highly parallelizable on GPUs, and incur only a modest 10–20 % overhead compared with Adam while delivering substantially better conditioning of the update direction.

The experimental section is extensive. First, synthetic matrix tasks (random regression and low‑rank factorization) demonstrate that PolarGrad reaches target loss in fewer epochs and with lower final error than both Adam and Muon. Second, large‑scale language‑model pre‑training experiments are conducted on a 1.3 B‑parameter GPT‑style model and a 15 B‑parameter Mixture‑of‑Experts model. Across these benchmarks, PolarGrad consistently achieves lower perplexity (≈2 % improvement over Adam, ≈0.5–1 % over Muon) and converges faster, especially in the first 10 % of training steps. Notably, PolarGrad remains stable without learning‑rate warm‑up, a requirement that is still essential for Adam in these settings. Ablation studies confirm that both the nuclear‑norm scaling and the use of a high‑quality polar decomposition are critical: removing the scaling or replacing the polar step with a simple SVD‑based orthogonalization degrades performance noticeably.

The paper concludes that recognizing the distinction between vector‑wise curvature preconditioning and matrix‑wise gradient preconditioning provides a clearer theoretical explanation for several empirical phenomena: Adam’s training instabilities on matrices, Muon’s speed‑up, and the need for warm‑up. PolarGrad operationalizes this insight into a practical optimizer that outperforms the current state‑of‑the‑art on both synthetic and real‑world large‑scale tasks. The authors release an open‑source implementation for both PyTorch and JAX, and outline future directions such as extending the framework to tensor‑level preconditioning, automated selection of preconditioning strategies per layer, and applying PolarGrad to multimodal foundation models. Overall, PolarGrad represents a significant step toward optimizers that respect the algebraic structure of deep‑learning parameters, delivering faster and more stable training for the next generation of massive models.