The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm
Fuente:
arXiv
Guardado en:
| Autores principales: | , , , |
|---|---|
| Formato: | Preprint |
| Publicado: |
2025
|
| Materias: | |
| Acceso en línea: | |
| Etiquetas: |
Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
|
| _version_ | 1866910189917044736 |
|---|---|
| author | Amsel, Noah Persson, David Musco, Christopher Gower, Robert M. |
| author_facet | Amsel, Noah Persson, David Musco, Christopher Gower, Robert M. |
| contents | 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 optimizer 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 Muon, our method yields consistent improvements in validation loss for a GPT-2 model trained on one to ten billion tokens from the FineWeb dataset, outperforming recent alternatives across a range of learning rates. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2505_16932 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm Amsel, Noah Persson, David Musco, Christopher Gower, Robert M. Machine Learning Artificial Intelligence Computation and Language Numerical Analysis Optimization and Control 65F30, 68T07, 68N19 G.1.3; I.2.6; F.2.1; G.1.6 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 optimizer 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 Muon, our method yields consistent improvements in validation loss for a GPT-2 model trained on one to ten billion tokens from the FineWeb dataset, outperforming recent alternatives across a range of learning rates. |
| title | The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm |
| topic | Machine Learning Artificial Intelligence Computation and Language Numerical Analysis Optimization and Control 65F30, 68T07, 68N19 G.1.3; I.2.6; F.2.1; G.1.6 |
| url | https://arxiv.org/abs/2505.16932 |