DASH：基于批量块预条件化与高效逆根求解器的加速Shampoo
机器学习
2026-02-03 v1
摘要
Shampoo是众多近似二阶优化器之一，其变体已获 MLCommons AlgoPerf 竞赛冠军，并被证明能够产生激活异常更少、更易压缩的模型。然而，当前应用Shampoo的计算成本较高，主要源于其内部操作的高昂开销。本文通过提出一种新的实现方法（即\textbf{D}istributed \textbf{A}ccelerated \textbf{SH}ampoo，简称DASH），显著提升了其速度。DASH基于两项主要新技术：其一，我们展示了预条件化块可堆叠为3D张量，以显著提高GPU利用率；其二，我们引入Newton-DB迭代法和切比雪夫多项式逼近，作为计算Shampoo所需逆矩阵根的新型高速方法。除算法贡献外，我们还对矩阵缩放对Shampoo收敛的关键影响进行了首次深入分析。就实际应用而言，我们的GPU感知实现在优化后的Distributed Shampoo基础上实现了最高的加速；而Newton-DB在所有测试方法中实现了最低的验证困惑度。我们的代码已公开于https://github.com/IST-DASLab/DASH。
引用
@article{arxiv.2602.02016,
title = {DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers},
author = {Ionut-Vlad Modoranu and Philip Zmushko and Erik Schultheis and Mher Safaryan and Dan Alistarh},
journal= {arXiv preprint arXiv:2602.02016},
year = {2026}
}