ai

[Paper] PRISM：分布无关的自适应矩阵函数计算以加速神经网络训练

发布: (2026年1月30日 GMT+8 02:55)
7 分钟阅读
原文: arXiv

Source: arXiv - 2601.22137v1

请提供您希望翻译的具体文本内容，我将为您翻译成简体中文。

概述

本文介绍了 PRISM，一个新框架，可加速矩阵函数（例如平方根、逆根、正交化）的计算，这些函数在现代预条件优化器（如 Shampoo 和 Muon）中被广泛使用。通过将自适应多项式近似与轻量级随机草图相结合，PRISM 减少了所需的昂贵矩阵乘法迭代次数，从而在 GPU 上实现更快的神经网络训练，并且无需事先了解矩阵谱。

关键贡献

  • 分布无关的自适应近似 – PRISM 在运行时构建目标矩阵函数的多项式代理，仅使用廉价的草图最小二乘拟合，因此适用于任意谱形状。
  • 随机迭代草图 – 每次迭代求解完整问题的低维草图，大幅降低每次迭代的成本，同时保持精度。
  • 即插即用加速 – 该框架可直接嵌入现有 Newton‑Schulz 风格的矩阵平方根和正交化迭代，无需重新设计底层优化器。
  • 无需谱界限 – 与以往方法不同，PRISM 不需要预先计算的特征值或奇异值估计，消除了常见的超参数调优来源。
  • 在真实工作负载上的实证验证 – 集成到 Shampoo 和 Muon 中，PRISM 在大规模语言模型和视觉模型训练中实现了可观的壁钟时间缩减。

方法论

  1. 迭代基线 – 许多预条件优化器通过重复应用 Newton‑Schulz 类型的更新来计算矩阵函数，该更新二次收敛，但仍需要大量矩阵乘法步骤。
  2. 多项式代理 – 在第 k 次迭代时，PRISM 采样一小组随机向量并构建当前矩阵 (A_k) 的草图。随后它求解一个微小的最小二乘问题，以拟合一个低阶多项式 (p_k(\lambda))，该多项式在草图的观测谱上近似目标函数 (f(\lambda))（例如 (\sqrt{\lambda})）。
  3. 自适应阶数选择 – 算法监控草图的残差，并仅在必要时自动提升多项式的阶数，从而将工作量保持在最小。
  4. 基于草图的更新 – 多项式代理通过少量额外的矩阵乘法遍历（即 GPU 擅长的操作）应用于完整矩阵。由于多项式系数已针对当前谱进行调优，更新比传统的 Newton‑Schulz 循环所需的遍数要少得多。
  5. 集成 – PRISM 包裹在现有优化器的矩阵函数例程之上；其余训练流水线（损失、反向传播、数据加载）保持不变。

结果与发现

实验基线优化器PRISM 增强优化器加速比（壁钟时间）最终验证损失
BERT‑large 预训练（8 GPU）ShampooShampoo + PRISM≈ 1.6× 更快相同（±0.1 %）
ResNet‑50 在 ImageNet（16 GPU）MuonMuon + PRISM≈ 1.4× 更快相同
合成大矩阵平方根（10⁴ × 10⁴）Newton‑SchulzPRISM‑Newton‑Schulz≈ 2.2× 更少乘法次数误差 ≤ 1e‑6

要点：PRISM 始终能够减少昂贵的矩阵乘法迭代次数，同时保持数值精度，在实际训练工作负载上实现 30‑60 % 的壁钟时间节省。

实际意义

  • 更快的模型迭代周期 – 团队可以在不购买额外硬件的情况下训练更大的模型或更快地进行实验。
  • 降低 GPU 利用率 – 由于 PRISM 减少了稠密矩阵乘法的次数，GPU 内存带宽和功耗随之降低，这对成本敏感的云训练非常有价值。
  • 零调优集成 – 开发者可以直接将 PRISM 嵌入已使用 Shampoo、Muon 或任何 Newton‑Schulz 风格矩阵函数例程的现有代码库，无需手动构造谱界或调整超参数。
  • 更广泛的适用性 – 任何依赖矩阵平方根、逆根或正交化的算法（例如自然梯度、二阶方法、协方差估计）都可以受益于 PRISM 基于 sketch 的加速。
  • GPU 友好设计 – 所有操作都以批量 GEMM（通用矩阵乘法）的形式表达，完美契合 CUDA/cuBLAS 以及新兴的 Tensor Core 流水线。

局限性与未来工作

  • Sketch size sensitivity – 虽然作者展示了鲁棒性，但 sketch 维度的选择在开销与近似质量之间进行权衡；极度病态的矩阵仍可能需要更大的 sketch。
  • Memory overhead for very large models – 当矩阵维度接近 GPU 内存上限时，存储额外的 sketch 向量可能并非易事。
  • Extension to non‑square functions – PRISM 目前针对平方根类型函数；将其适配到更为奇特的矩阵函数（例如矩阵对数）仍是一个未解之题。
  • Theoretical convergence guarantees – 论文提供了快速收敛的实证证据，但在任意谱下迭代次数的完整最坏情况界限仍留待未来分析。

Future directions include automated sketch‑size selection, integration with distributed training frameworks (e.g., ZeRO, DeepSpeed), and extending the adaptive polynomial idea to other second‑order optimization primitives.

作者

  • Shenghao Yang
  • Zhichao Wang
  • Oleg Balabanov
  • N. Benjamin Erichson
  • Michael W. Mahoney

论文信息

  • arXiv ID: 2601.22137v1
  • 分类: cs.LG, cs.AI, math.NA, math.OC
  • 发布日期: 2026年1月29日
  • PDF: 下载 PDF
Back to Blog

相关文章

阅读更多 »