改进了Shampoo类优化器，支持半精度存储并降低计算开销。
Reparametrizing Shampoo and SOAP for Subspace Basis Updates and BFloat16 Storage
- 重构预条件矩阵更新机制，仅在子空间中进行部分QR分解
- 在BFloat16下保持性能，使KL-SOAP达到甚至超过KL-Shampoo水平
- 适用于多种基于QR的优化器，提升训练效率与内存利用率
基于Shampoo的方法（如KL-Shampoo和SOAP）在训练神经网络时表现优异，依赖于QR分解。然而现有QR实现需单精度（FP32）运算，计算成本高，导致大规模预条件矩阵时时间与内存消耗大。使用BFloat16存储虽可减小内存占用，但会损害这类方法的性能。本文提出一种预条件器重参数化方法，支持BFloat16存储，并通过结合更新与未更新的基向量形成完整基。通过仅在子空间内对部分基向量执行QR分解，显著降低计算开销，缓解BFloat16带来的性能下降。该方法广泛适用于依赖QR分解的Shampoo类方法，包括KL-Shampoo、SOAP和KL-SOAP。尤其在BFloat16下，显著提升SOAP与KL-SOAP性能，使KL-SOAP达到或超越KL-Shampoo。整体上，使基于Shampoo的方法更高效、更节省内存。
原文摘要 · Abstract (English)
Shampoo-based methods, such as KL-Shampoo and SOAP, have demonstrated strong performance in training neural networks and rely on QR decomposition. Because existing QR implementations require single-precision (FP32) arithmetic and remain computationally expensive, these methods become time- and memory-intensive when their preconditioning matrices are large. Moreover, using BFloat16 (BFP16) storage to reduce memory usage can degrade the performance of Shampoo-based methods. We propose a reparametrization of the preconditioner that supports BFP16 storage and forms a complete basis by combining updated basis vectors with unchanged ones. By updating only part of the basis through QR decomposition in a subspace, our approach reduces computational overhead while mitigating the performance degradation caused by BFP16 storage. Our approach applies broadly to Shampoo-based methods that employ QR decomposition, including KL-Shampoo, SOAP, and KL-SOAP. In particular, it improves the performance of SOAP and KL-SOAP under BFP16 storage, enabling KL-SOAP to match or exceed KL-Shampoo. Overall, our approach makes Shampoo-based methods more memory- and time-efficient.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。