# PRISM: Distribution-free Adaptive Computation of Matrix Functions for Accelerating Neural Network Training

arXiv: https://arxiv.org/abs/2601.22137 · HTML: https://mlsanity.com/paper/2601.22137/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 70 references that are in the map.
- First posted on arXiv 29 Jan 2026.

## Builds on (most cited first, 30 of 35)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Universal Style Transfer via Feature Transforms](https://mlsanity.com/paper/1705.08086/) (2017)
- [PyHessian: Neural Networks Through the Lens of the Hessian](https://mlsanity.com/paper/1912.07145/) (2019)
- [Sketching as a Tool for Numerical Linear Algebra](https://mlsanity.com/paper/1411.4357/) (2014)
- [Implicit Self-Regularization in Deep Neural Networks: Evidence from Random Matrix Theory and Implications for Learning](https://mlsanity.com/paper/1810.01075/) (2018)
- [ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning](https://mlsanity.com/paper/2006.00719/) (2020)
- [The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm](https://mlsanity.com/paper/2505.16932/) (2025)
- [Traditional and Heavy-Tailed Self Regularization in Neural Network Models](https://mlsanity.com/paper/1901.08276/) (2019)
- [Fantastic Pretraining Optimizers and Where to Find Them](https://mlsanity.com/paper/2509.02046/) (2025)
- [Decorrelated Batch Normalization](https://mlsanity.com/paper/1804.08450/) (2018)
- [Practical Efficiency of Muon for Pretraining](https://mlsanity.com/paper/2505.02222/) (2025)
- [NVIDIA Tensor Core Programmability, Performance & Precision](https://mlsanity.com/paper/1803.04014/) (2018)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [Gluon: Making Muon & Scion Great Again! (Bridging Theory and Practice of LMO-based Optimizers for LLMs)](https://mlsanity.com/paper/2505.13416/) (2025)
- [Dion: Distributed Orthonormalized Updates](https://mlsanity.com/paper/2504.05295/) (2025)
- [Hutch++: Optimal Stochastic Trace Estimation](https://mlsanity.com/paper/2010.09649/) (2020)
- [A New Perspective on Shampoo's Preconditioner](https://mlsanity.com/paper/2406.17748/) (2024)
- [Heavy-Tailed Universality Predicts Trends in Test Accuracies for Very Large Pre-Trained Deep Neural Networks](https://mlsanity.com/paper/1901.08278/) (2019)
- [Deep CNNs Meet Global Covariance Pooling: Better Representation and Generalization](https://mlsanity.com/paper/1904.06836/) (2019)
- [Inference and Uncertainty Quantification for Noisy Matrix Completion](https://mlsanity.com/paper/1906.04159/) (2019)
- [Randomized Numerical Linear Algebra : A Perspective on the Field With an Eye to Software](https://mlsanity.com/paper/2302.11474/) (2023)
- [Isotropic Curvature Model for Understanding Deep Learning Optimization: Is Gradient Orthogonalization Optimal?](https://mlsanity.com/paper/2511.00674/) (2025)
- [When do spectral gradient updates help in deep learning?](https://mlsanity.com/paper/2512.04299/) (2025)
- [Spectral Evolution and Invariance in Linear-width Neural Networks](https://mlsanity.com/paper/2211.06506/) (2022)
- [Fast Matrix Square Roots with Applications to Gaussian Processes and Bayesian Optimization](https://mlsanity.com/paper/2006.11267/) (2020)
- [Why Approximate Matrix Square Root Outperforms Accurate SVD in Global Covariance Pooling?](https://mlsanity.com/paper/2105.02498/) (2021)
- [LiMuon: Light and Fast Muon Optimizer for Large Models](https://mlsanity.com/paper/2509.14562/) (2025)
- [Lectures on Randomized Numerical Linear Algebra](https://mlsanity.com/paper/1712.08880/) (2017)

## Built on by (most cited first, 4 of 4)

- [The Newton-Muon Optimizer](https://mlsanity.com/paper/2604.01472/) (2026)
- [Learning to Discover Iterative Spectral Algorithms](https://mlsanity.com/paper/2602.09530/) (2026)
- [AMO: Operator-level Adaptive Muon Orthogonalization](https://mlsanity.com/paper/2605.17806/) (2026)
- [Cost-free Spectral Estimation for Adaptive Newton--Schulz in Matrix Optimizers](https://mlsanity.com/paper/2609.33047/) (2026)

## Abstract (arXiv)

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.