重新参数化 Shampoo 和 SOAP 以进行子空间基更新和 BFloat16 存储
机器学习
2026-05-27 v1
摘要
Shampoo 基于的方法，如 KL-Shampoo 和 SOAP，已在训练神经网络方面显示出强大的性能，并依赖于 QR 分解。由于现有 QR 实现需要单精度 (FP32) 算术且仍昂贵，这些方法在其 preconditioning 矩阵较大时变得时间和内存密集。此外，使用 BFloat16 (BFP16) 存储以减少内存使用会降低 Shampoo 基于方法的性能。我们提出一种重新参数化的 preconditioner，支持 BFP16 存储并通过结合更新的基向量与不变的基向量来形成完整基。通过在子空间中仅对部分基进行 QR 分解来更新，我们的方法减少了计算开销，同时缓解了 BFP16 存储导致的性能降低。我们的方法广泛适用于采用 QR 分解的各种 Shampoo 基于方法，包括 KL-Shampoo、SOAP 和 KL-SOAP。特别是，它在 BFP16 存储下提高了 SOAP 和 KL-SOAP 的性能，使 KL-SOAP 能够匹配或超过 KL-Shampoo。总体而言，我们的方法使 Shampoo 基于方法更加在内存和时间上高效。
引用
@article{arxiv.2605.26327,
title = {Reparametrizing Shampoo and SOAP for Subspace Basis Updates and BFloat16 Storage},
author = {Alan Milligan and Zikun Xu and Simon Lacoste-Julien and Felix Dangel and Wu Lin},
journal= {arXiv preprint arXiv:2605.26327},
year = {2026}
}
备注
Preprint, working in progress