PRISM无需谱信息，自适应加速神经网络训练中的矩阵函数计算。
PRISM: Distribution-free Adaptive Computation of Matrix Functions for Accelerating Neural Network Training
- 用随机采样拟合多项式逼近矩阵谱，动态适应变化的矩阵特征。
- 在Shampoo和Muon优化器中应用，训练速度提升显著。
- 无需预先设定奇异值范围，适合复杂场景的自适应优化。
矩阵函数如平方根、逆平方根和正交化在神经网络训练的预条件梯度方法中起核心作用。这推动了避免显式特征分解、主要依赖矩阵乘法的迭代算法发展，非常适合现代GPU加速。本文提出PRISM（Polynomial-fitting and Randomized Iterative Sketching for Matrix functions computation），一种通用框架，用于加速矩阵函数的迭代计算。PRISM结合自适应多项式逼近与随机采样：每轮迭代通过采样最小二乘问题拟合当前谱的多项式代理，以极低开销适应具体实例。我们将PRISM应用于加速Newton-Schulz类迭代计算矩阵平方根和正交化，这些是机器学习的核心基础操作。与先前方法不同，PRISM无需显式谱界或奇异值估计；能自动适应随时间演化的谱结构。实验表明，将其集成到Shampoo和Muon优化器中可显著加速训练。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。