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.

翻译：矩阵函数（如平方根、逆根和正交化）在神经网络训练的预处理梯度方法中起着核心作用。这推动了迭代算法的发展，这些算法避免了显式的特征分解，主要依赖矩阵乘法，使其非常适合现代GPU加速器。我们提出了PRISM（用于矩阵函数计算的多项式拟合与随机迭代草图法），这是一个用于加速计算矩阵函数的迭代算法的通用框架。PRISM将自适应多项式逼近与随机草图法相结合：在每次迭代中，它通过一个草图最小二乘问题，用多项式代理拟合当前谱，以最小的开销自适应于当前实例。我们将PRISM应用于加速类牛顿-舒尔茨迭代的矩阵平方根和正交化计算，这些是机器学习中的核心原语。与现有方法不同，PRISM不需要显式的谱界或奇异值估计；它能自动适应不断演化的谱。经验表明，当PRISM集成到Shampoo和Muon优化器中时，能有效加速训练。