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

TL;DR

PRISM introduces a distribution-free framework that accelerates matrix-function iterations by adaptively fitting spectrum-aware polynomials at each step and using randomized sketching to keep costs low. It provides a unified template (Part I: basic setup; Part II: spectrum-adaptive acceleration) that applies to Newton-Schulz, matrix sign, square roots, and polar decomposition, without requiring prior spectral bounds. The method delivers theoretical guarantees and practical speedups, demonstrated by accelerating neural-network optimizers (Shampoo and Muon) on realistic models and data, including MP- and heavy-tailed spectra. By enabling instance-specific polynomial updates with low overhead, PRISM turns spectral adaptivity into a robust, GPU-friendly primitive for large-scale matrix computations in ML pipelines.

Abstract

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.

Figures (11)

• Figure 1.1: Speedup in GPU time over the classical Newton-Schulz for polar decomposition (left) and square root (right). We keep $\sigma_{\max}=1$ and vary $\sigma_{\min} \in [10^{-12}, 1/2]$. The PolarExpress variant we use is optimized for $\sigma_{\min} = 10^{-3}$ for polar decomposition (and hence it is optimized for $\sigma_{\min} = 10^{-6}$ for square root). All algorithms are run until convergence. In both cases, we see a performance degradation as $\sigma_{\min}$ deviates from the one PolarExpress is optimized for. PRISM (this work) does not require or assume $\sigma_{\min}$ and has a stable speedup across the entire range.
• Figure 4.1: Better polynomial approximation leads to faster convergence. Left: Approximating $f(\xi)$ using its Taylor approximation $f_1(\xi)$ around $\xi=0$ versus the alternative $g_1(\xi;1)$. Right: The initial convergence behavior in residual error $\xi_k = 1-x_k^2$ using the standard and "accelerated" Newton-Schulz, respectively, for $x_0 = 10^{-6}$. Using $g_1(\xi;1)$ leads to an exponential speedup.
• Figure 6.1: Convergence of degree-5 polynomial methods for othogonalizing a Gaussian random matrix ${\bm{A}} \in \mathbb{R}^{n \times m}$ with varying aspect ratio $\gamma = n/m$. The figures from left to right show the Frobenius norm error $\|{\bm{I}} - {\bm{X}}_k^T{\bm{X}}_k\|_F$ for $\gamma = 1, 4, 50$, respectively. The last figure on the right shows the $\alpha_k$'s computed by (\ref{['eq:alpha_sketch']}) in PRISM for different aspect ratios at each iteration.
• Figure 6.2: Convergence of degree-5 polynomial methods for othogonalizing random matrices generated by HTMP hodgkinson2025models with different parameter $\kappa$. Smaller $\kappa$ indicates a heavier tail in the spectra. The figures from left to right show the Frobenius norm error $\|{\bm{I}}-{\bm{X}}_k^T{\bm{X}}_k\|_F$ for $\kappa = 0.1, 0.5, 100$, respectively. The rightmost figure shows the $\alpha_k$'s computed by (\ref{['eq:alpha_sketch']}) in PRISM.
• Figure 6.3: Improvement to the Shampoo optimizer in terms of training speed. We compare three methods to compute the inverse root preconditioner inside Shampoo. Left: ResNet-20 on CIFAR10. Right: ResNet-32 on CIFAR100.

Theorems & Definitions (10)

• Theorem 4.1
• Theorem 4.2
• Theorem 5.1: higham1997stable
• Theorem 5.2: higham2004computing
• Lemma B.1
• Claim 1
• Claim 2
• Claim 3
• Claim 4
• Claim 5