重新参数化Shampoo和SOAP以支持子空间基更新和BFloat16存储

arXiv cs.LG 论文
shampoo soap preconditioning subspace-basis bfloat16 qr-decomposition memory-efficiency

摘要

本文提出了一种对基于Shampoo的优化方法（如KL-Shampoo和SOAP）中预条件器的重新参数化，以支持BFloat16存储，并通过在子空间中仅通过QR分解更新部分基来减少计算开销，从而使这些方法在内存和时间上更高效。

arXiv:2605.26327v1 公告类型：新 摘要：基于Shampoo的方法（如KL-Shampoo和SOAP）在训练神经网络中表现出色，且依赖于QR分解。由于现有的QR实现需要单精度（FP32）算术且计算成本高昂，当预条件矩阵较大时，这些方法会变得耗时且占用大量内存。此外，使用BFloat16（BFP16）存储来减少内存占用会降低基于Shampoo的方法的性能。我们提出了一种预条件器的重新参数化，该参数化支持BFP16存储，并通过将更新的基向量与未更改的基向量结合形成完整的基。通过在子空间中仅通过QR分解更新部分基，我们的方法减少了计算开销，同时缓解了BFP16存储带来的性能下降。我们的方法广泛适用于使用QR分解的基于Shampoo的方法，包括KL-Shampoo、SOAP和KL-SOAP。特别地，它提升了SOAP和KL-SOAP在BFP16存储下的性能，使KL-SOAP能够匹配或超越KL-Shampoo。总体而言，我们的方法使基于Shampoo的方法在内存和时间上更高效。
查看原文
查看缓存全文

缓存时间: 2026/05/27 09:08 

# 重参数化Shampoo与SOAP以实现子空间基更新与BFloat16存储

来源：https://arxiv.org/abs/2605.26327  
查看PDF（https://arxiv.org/pdf/2605.26327）

> **摘要：**基于Shampoo的方法（如KL-Shampoo和SOAP）在训练神经网络中表现出色，且依赖于QR分解。由于现有的QR实现需要单精度（FP32）算术且计算成本依然高昂，当预处理矩阵较大时，这些方法在时间和内存上变得非常密集。此外，使用BFloat16（BFP16）存储来减少内存占用会降低基于Shampoo的方法的性能。我们提出了一种预处理器的重参数化方法，该方法支持BFP16存储，并通过将更新后的基向量与未改动的基向量相结合来形成完整的基。通过子空间中的QR分解仅更新部分基，我们的方法减少了计算开销，同时缓解了BFP16存储带来的性能下降。我们的方法广泛适用于使用QR分解的基于Shampoo的方法，包括KL-Shampoo、SOAP和KL-SOAP。特别是，它提升了SOAP和KL-SOAP在BFP16存储下的性能，使KL-SOAP能够达到或超越KL-Shampoo。总体而言，我们的方法使基于Shampoo的方法在内存和时间上更加高效。

## 提交历史

来自：Wu Lin [查看邮件（https://arxiv.org/show-email/8842e08e/2605.26327）]  
**[v1]** 2026年5月25日星期一 21:03:03 UTC（545 KB）

相似文章

SOAP、Muon 及更多：推动LLM预训练规模扩展

arXiv cs.LG

本文改进了高阶优化器 SOAP 和 Muon，用于大规模 LLM 预训练，解决了大批量训练中的不稳定性问题，并引入了一种与 Megatron-LM 兼容的层式分布式优化器。实验表明，在数十亿参数规模下，它们始终优于 AdamW。

稀疏设计（5分钟阅读）

TLDR AI

Moonshot的Kimi K3是一个2.8万亿参数的开源权重模型，拥有896个专家（每个token激活16个），体现了在保持活跃计算不变的情况下扩大总参数的趋势，并使用注意力压缩来减少KV缓存大小，使得前沿推理更容易实现，但存储成本较高。