Technical Summary: PRISM

Problem Statement

Matrix functions, such as square roots, inverse roots, and orthogonalization (polar decomposition), are fundamental to preconditioned gradient methods in neural network training (e.g., Shampoo, Muon). While iterative algorithms relying on General Matrix Multiplications (GEMMs) are preferred for GPU acceleration over methods requiring explicit Singular Value Decomposition (SVD) or eigenvalue estimation, existing state-of-the-art approaches suffer from significant limitations:

1. Lack of Adaptivity: Recent methods (e.g., PolarExpress, CANS) often rely on fixed polynomial updates optimized for a predefined spectral interval (e.g., $[\ell, u]$). If the actual singular value distribution of the input matrix deviates from this interval, convergence degrades dramatically.
2. Dependency on Prior Knowledge: These methods typically require prior knowledge of the singular value distribution or tight bounds on the smallest and largest singular values, which are often unavailable or as costly to estimate as the matrix function itself.
3. Narrow Scope: Many recent acceleration techniques are tailored to specific use cases (e.g., only polar decomposition in the Muon optimizer) rather than providing a general framework for various matrix functions.

Methodology: The PRISM Framework

The authors propose PRISM (Polynomial-fitting and Randomized Iterative Sketching for Matrix functions computation), a general meta-algorithm designed to accelerate iterative matrix function computations without requiring explicit spectral bounds. PRISM operates in two distinct parts:

Part I: Basic Setup (Iterative Polynomial Approximation)

PRISM frames the computation of a matrix function $T(A)$ as an iterative polynomial approximation problem.

• For a scalar $x$ approximating $T(a)$, the problem is reformulated as approximating a residual function $f(\xi)$ where $\xi$ measures the distance to the solution.
• Standard iterative methods (like Newton-Schulz) are derived by using a fixed-order Taylor polynomial $f_d(\xi)$ around $\xi=0$.
• The matrix iteration takes the form $X_{k+1} = X_k f_d(R_k)$, where $R_k$ is the residual matrix (e.g., $I - X_k^2$ for sign functions).

Part II: Acceleration (Adaptive Fitting and Sketching)

To overcome the slow initial convergence of fixed Taylor polynomials when the spectrum of $R_k$ is not close to 0, PRISM introduces two mechanisms:

1. Adaptive Polynomial Fitting:
   Instead of using a fixed Taylor polynomial, PRISM replaces the highest-order coefficient with a learnable parameter $\alpha$. The polynomial is updated at each iteration to fit the current spectrum of the residual matrix $R_k$.

   • The update polynomial is defined as $g_d(\xi; \alpha) = f_{d-1}(\xi) + \alpha \xi^d$.
   • The optimal $\alpha_k^*$ is found by minimizing the residual norm $\|I - X_{k+1}^2\|_F^2$ (or equivalent for other functions) with respect to $\alpha$. This effectively fits the polynomial to the eigenvalues of $R_k$ in a least-squares sense.
   • To ensure stability, $\alpha_k$ is constrained to a fixed interval $[\ell, u]$ (e.g., $[1/2, 1]$ for $d=1$), which guarantees convergence properties independent of the input spectrum.

2. Randomized Sketching:
   Computing the exact optimal $\alpha_k^*$ requires evaluating traces of high powers of the residual matrix, which is computationally expensive ($O(n^3)$ or higher).

   • PRISM employs randomized sketching (using Oblivious Subspace Embeddings, e.g., Gaussian matrices) to approximate the least-squares problem.
   • By solving the minimization problem on a sketched version of the residual, the cost of finding $\alpha_k$ is reduced to $O(n^2 \log n)$ (or $O(n^2 p)$ where $p \ll n$), which is negligible compared to the $O(n^3)$ cost of the matrix multiplications in the main iteration.
   • Theoretical guarantees (Theorem 2) show that with a sketch dimension $p = O(\log n)$, the convergence rate remains essentially identical to the exact minimization.

Key Contributions

1. General Acceleration Framework: PRISM provides a unified approach applicable to a wide range of iterative algorithms, including Newton-Schulz (for square roots, inverse roots, sign, and polar decomposition), Chebyshev methods, and Inverse Newton methods.
2. Distribution-Free Adaptivity: Unlike prior work, PRISM dynamically adapts to the evolving spectrum of the current iterate without requiring prior knowledge of singular value bounds or distributions.
3. Efficient Implementation: The integration of randomized sketching ensures that the overhead of adaptive polynomial fitting is minimal, preserving the GPU-friendly nature of GEMM-dominated iterations.
4. Theoretical Guarantees: The paper provides convergence proofs showing that PRISM-accelerated iterations maintain quadratic (or higher-order) convergence rates even with approximate sketching, provided the sketch dimension is sufficiently large.

Experimental Results

The authors validate PRISM through both synthetic benchmarks and integration into neural network optimizers:

• Convergence Speed: On matrices with varying aspect ratios and spectral distributions (including Gaussian random matrices and Heavy-Tailed Marchenko-Pastur distributions mimicking pre-trained models), PRISM consistently outperforms classical Newton-Schulz and fixed-interval methods like PolarExpress.
  • PolarExpress often slows down or fails to converge when the actual singular value range deviates from its optimized interval.
  • PRISM maintains stable speedups across the entire range of $\sigma_{min}$, adapting the coefficient $\alpha_k$ automatically.
• Neural Network Training:
  • Shampoo Optimizer: Integrating PRISM to compute inverse square root preconditioners for ResNet-20/32 on CIFAR datasets resulted in faster training times compared to eigen-decomposition and PolarExpress baselines.
  • Muon Optimizer: Applying PRISM to the polar decomposition of gradient matrices in GPT-2 training (FineWeb dataset) demonstrated improved validation losses and training efficiency compared to AdamW, PolarExpress, and fixed-degree Newton-Schulz variants.

Significance and Claims

The paper claims that PRISM offers a practical and general route to instance-adaptive, GPU-friendly matrix function computation. Its primary significance lies in transforming "spectral adaptivity"—previously a source of tuning burden requiring manual spectral bounds—into a reliable, automated algorithmic primitive. By decoupling acceleration from prior spectral knowledge, PRISM enables robust performance across diverse spectral regimes common in machine learning (e.g., heavy-tailed distributions in pre-trained models), making it a versatile tool for accelerating large-scale neural network training.