PRISM: Distribution-free Adaptive Computation of Matrix Functions for Accelerating Neural Network Training
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866912860136800256 |
|---|---|
| author | Yang, Shenghao Wang, Zhichao Balabanov, Oleg Erichson, N. Benjamin Mahoney, Michael W. |
| author_facet | Yang, Shenghao Wang, Zhichao Balabanov, Oleg Erichson, N. Benjamin Mahoney, Michael W. |
| contents | Matrix functions such as square root, inverse roots, and orthogonalization play a central role in preconditioned gradient methods for neural network training. This has motivated the development of iterative algorithms that avoid explicit eigendecompositions and rely primarily on matrix multiplications, making them well suited for modern GPU accelerators. We present PRISM (Polynomial-fitting and Randomized Iterative Sketching for Matrix functions computation), a general framework for accelerating iterative algorithms for computing matrix functions. PRISM combines adaptive polynomial approximation with randomized sketching: at each iteration, it fits a polynomial surrogate to the current spectrum via a sketched least-squares problem, adapting to the instance at hand with minimal overhead. We apply PRISM to accelerate Newton-Schulz-like iterations for matrix square roots and orthogonalization, which are core primitives in machine learning. Unlike prior methods, PRISM requires no explicit spectral bounds or singular value estimates; and it adapts automatically to the evolving spectrum. Empirically, PRISM accelerates training when integrated into Shampoo and Muon optimizers. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2601_22137 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | PRISM: Distribution-free Adaptive Computation of Matrix Functions for Accelerating Neural Network Training Yang, Shenghao Wang, Zhichao Balabanov, Oleg Erichson, N. Benjamin Mahoney, Michael W. Machine Learning Artificial Intelligence Numerical Analysis Optimization and Control Matrix functions such as square root, inverse roots, and orthogonalization play a central role in preconditioned gradient methods for neural network training. This has motivated the development of iterative algorithms that avoid explicit eigendecompositions and rely primarily on matrix multiplications, making them well suited for modern GPU accelerators. We present PRISM (Polynomial-fitting and Randomized Iterative Sketching for Matrix functions computation), a general framework for accelerating iterative algorithms for computing matrix functions. PRISM combines adaptive polynomial approximation with randomized sketching: at each iteration, it fits a polynomial surrogate to the current spectrum via a sketched least-squares problem, adapting to the instance at hand with minimal overhead. We apply PRISM to accelerate Newton-Schulz-like iterations for matrix square roots and orthogonalization, which are core primitives in machine learning. Unlike prior methods, PRISM requires no explicit spectral bounds or singular value estimates; and it adapts automatically to the evolving spectrum. Empirically, PRISM accelerates training when integrated into Shampoo and Muon optimizers. |
| title | PRISM: Distribution-free Adaptive Computation of Matrix Functions for Accelerating Neural Network Training |
| topic | Machine Learning Artificial Intelligence Numerical Analysis Optimization and Control |
| url | https://arxiv.org/abs/2601.22137 |