The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm

• URL: http://arxiv.org/abs/2505.16932v3
• Date: Thu, 25 Sep 2025 20:39:00 GMT
• Title: The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm
• Abstract summary: We introduce Polar Express, a new method for computing the polar decomposition.<n>We prove that this strategy minimizes error in a worst-case sense, allowing Polar Express to converge as rapidly as possible.<n>When integrated into the Muon training framework, our method leads to consistent improvements in validation loss.
• Score: 24.020943794160555
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Computing the polar decomposition and the related matrix sign function has been a well-studied problem in numerical analysis for decades. Recently, it has emerged as an important subroutine within the Muon algorithm for training deep neural networks. However, the requirements of this application differ sharply from classical settings: deep learning demands GPU-friendly algorithms that prioritize high throughput over high precision. We introduce Polar Express, a new method for computing the polar decomposition. Like Newton-Schulz and other classical polynomial methods, our approach uses only matrix-matrix multiplications, making it very efficient on GPUs. Inspired by earlier work of Chen & Chow and Nakatsukasa & Freund, Polar Express adapts the update rule at each iteration by solving a minimax optimization problem. We prove that this strategy minimizes error in a worst-case sense, allowing Polar Express to converge as rapidly as possible both in the early iterations and asymptotically. We also address finite-precision issues, making it practical to use in bfloat16. When integrated into the Muon training framework, our method leads to consistent improvements in validation loss when training a GPT-2 model on one billion tokens from the FineWeb dataset, outperforming recent alternatives across a range of learning rates.

Related papers

• Federated stochastic bilevel optimization with fully first-order gradients [57.1147486991903]
  Federated matrices bilevel optimization has been actively studied in recent years due to its widespread applications in machine learning.<n>Most existing federated bilevel optimization algorithms require the computation of second-order Hessian and Jacobian matrices.<n>We propose a novel federated convergence variance-reduced bilevel descent algorithm that relies solely on first-order oracles.
  arXiv  Detail & Related papers  (2026-09-14T21:12:59Z)
• GPU-Enabled Large-Scale Optimization Using Randomized Linear Algebra [13.9407640377773]
  This paper introduces rlaopt, a PyTorch-based package for large-scale optimization and scientific computing using randomized numerical linear algebra (RandNLA)<n>Rlaopt addresses this gap by providing GPU-enabled solvers for positive-definite linear systems and convex empirical risk minimization with constraints and regularizers.<n>Rlaopt automatically checks compatibility with the selected solver and performs the required problem decomposition.
  arXiv  Detail & Related papers  (2026-09-08T02:14:05Z)
• Inertial Quadratic Majorization Minimization with Application to Kernel Regularized Learning [1.0282274843007797]
  We introduce the Quadratic Majorization Minimization with Extrapolation (QMME) framework and establish its sequential convergence properties.<n>To demonstrate practical advantages, we apply QMME to large-scale kernel regularized learning problems.
  arXiv  Detail & Related papers  (2025-07-06T05:17:28Z)
• Towards Learning High-Precision Least Squares Algorithms with Sequence Models [42.217390215093516]
  We show that softmax Transformers struggle to perform high-precision multiplications.<n>We identify limitations present in existing architectures and training procedures.<n>For the first time, we demonstrate the ability to train to near machine precision.
  arXiv  Detail & Related papers  (2025-03-15T23:25:11Z)
• WENDy for Nonlinear-in-Parameters ODEs [1.9573380763700712]
  The Weak-form Estimation of Non-linear Dynamics (WEN) is extended to accommodate systems of ordinary differential equations that are nonlinear-ins.<n>We present results on a suite of benchmark systems to demonstrate the practical benefits of our approach.
  arXiv  Detail & Related papers  (2025-02-13T01:40:21Z)
• Training Deep Learning Models with Norm-Constrained LMOs [56.00317694850397]
  We propose a new family of algorithms that uses the linear minimization oracle (LMO) to adapt to the geometry of the problem.<n>We demonstrate significant speedups on nanoGPT training using our algorithm, Scion, without any reliance on Adam.
  arXiv  Detail & Related papers  (2025-02-11T13:10:34Z)
• Newton Losses: Using Curvature Information for Learning with Differentiable Algorithms [80.37846867546517]
  We show how to train eight different neural networks with custom objectives. We exploit their second-order information via their empirical Fisherssian matrices. We apply Loss Lossiable algorithms to achieve significant improvements for less differentiable algorithms.
  arXiv  Detail & Related papers  (2024-10-24T18:02:11Z)
• Optimizing Tensor Computation Graphs with Equality Saturation and Monte Carlo Tree Search [0.0]
  We present a tensor graph rewriting approach that uses Monte Carlo tree search to build superior representation. Our approach improves the inference speedup of neural networks by up to 11% compared to existing methods.
  arXiv  Detail & Related papers  (2024-10-07T22:22:02Z)
• Stochastic Optimization for Non-convex Problem with Inexact Hessian Matrix, Gradient, and Function [99.31457740916815]
  Trust-region (TR) and adaptive regularization using cubics have proven to have some very appealing theoretical properties. We show that TR and ARC methods can simultaneously provide inexact computations of the Hessian, gradient, and function values.
  arXiv  Detail & Related papers  (2023-10-18T10:29:58Z)
• Learning the Positions in CountSketch [49.57951567374372]
  We consider sketching algorithms which first compress data by multiplication with a random sketch matrix, and then apply the sketch to quickly solve an optimization problem. In this work, we propose the first learning-based algorithms that also optimize the locations of the non-zero entries.
  arXiv  Detail & Related papers  (2023-06-11T07:28:35Z)
• Efficient Graph Laplacian Estimation by Proximal Newton [12.05527862797306]
  A graph learning problem can be formulated as a maximum likelihood estimation (MLE) of the precision matrix. We develop a second-order approach to obtain an efficient solver utilizing several algorithmic features.
  arXiv  Detail & Related papers  (2023-02-13T15:13:22Z)
• Adaptive Federated Minimax Optimization with Lower Complexities [82.51223883622552]
  We propose an efficient adaptive minimax optimization algorithm (i.e., AdaFGDA) to solve these minimax problems. It builds our momentum-based reduced and localSGD techniques, and it flexibly incorporate various adaptive learning rates.
  arXiv  Detail & Related papers  (2022-11-14T12:32:18Z)
• The Curse of Unrolling: Rate of Differentiating Through Optimization [35.327233435055305]
  Un differentiation approximates the solution using an iterative solver and differentiates it through the computational path. We show that we can either 1) choose a large learning rate leading to a fast convergence but accept that the algorithm may have an arbitrarily long burn-in phase or 2) choose a smaller learning rate leading to an immediate but slower convergence.
  arXiv  Detail & Related papers  (2022-09-27T09:27:29Z)
• Accelerated Proximal Alternating Gradient-Descent-Ascent for Nonconvex Minimax Machine Learning [12.069630105460766]
  An Alternating Table-descentascent (AltGDA) is an computation optimization algorithm that has been widely used for training in various machine learning applications. In this paper, we develop a single-loop fast computation-of-the-loop gradient-of-the-loop algorithm to solve non minimax optimization problems.
  arXiv  Detail & Related papers  (2021-12-22T04:33:27Z)
• Newton-LESS: Sparsification without Trade-offs for the Sketched Newton Update [88.73437209862891]
  In second-order optimization, a potential bottleneck can be computing the Hessian matrix of the optimized function at every iteration. We show that the Gaussian sketching matrix can be drastically sparsified, significantly reducing the computational cost of sketching. We prove that Newton-LESS enjoys nearly the same problem-independent local convergence rate as Gaussian embeddings.
  arXiv  Detail & Related papers  (2021-07-15T17:33:05Z)
• Minimax Optimization with Smooth Algorithmic Adversaries [59.47122537182611]
  We propose a new algorithm for the min-player against smooth algorithms deployed by an adversary. Our algorithm is guaranteed to make monotonic progress having no limit cycles, and to find an appropriate number of gradient ascents.
  arXiv  Detail & Related papers  (2021-06-02T22:03:36Z)
• SHINE: SHaring the INverse Estimate from the forward pass for bi-level optimization and implicit models [15.541264326378366]
  In recent years, implicit deep learning has emerged as a method to increase the depth of deep neural networks. The training is performed as a bi-level problem, and its computational complexity is partially driven by the iterative inversion of a huge Jacobian matrix. We propose a novel strategy to tackle this computational bottleneck from which many bi-level problems suffer.
  arXiv  Detail & Related papers  (2021-06-01T15:07:34Z)
• Recent Theoretical Advances in Non-Convex Optimization [56.88981258425256]
  Motivated by recent increased interest in analysis of optimization algorithms for non- optimization in deep networks and other problems in data, we give an overview of recent results of theoretical optimization algorithms for non- optimization.
  arXiv  Detail & Related papers  (2020-12-11T08:28:51Z)
• Gradient Free Minimax Optimization: Variance Reduction and Faster Convergence [120.9336529957224]
  In this paper, we denote the non-strongly setting on the magnitude of a gradient-free minimax optimization problem. We show that a novel zeroth-order variance reduced descent algorithm achieves the best known query complexity.
  arXiv  Detail & Related papers  (2020-06-16T17:55:46Z)
• Hardness of Random Optimization Problems for Boolean Circuits, Low-Degree Polynomials, and Langevin Dynamics [78.46689176407936]
  We show that families of algorithms fail to produce nearly optimal solutions with high probability. For the case of Boolean circuits, our results improve the state-of-the-art bounds known in circuit complexity theory.
  arXiv  Detail & Related papers  (2020-04-25T05:45:59Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.