arXiv:2509.03378stat.MLcs.LG2025-09被引 13

基于KL散度重设计优化器,提升神经网络训练效率与性能。

Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimization

  • 用KL散度最小化重构二阶动量估计,揭示原有方法理论缺陷
  • 新算法在预训练中超越原版Shampoo/SOAP,且无需额外内存开销
  • 适合追求高效、低内存优化器的模型训练研究者

Shampoo及其高效变体SOAP采用结构化二阶矩估计,在训练神经网络方面表现优异。然而实践中,Shampoo通常需结合Adam进行步长调整才能具备竞争力,而SOAP通过在Shampoo的特征基下应用Adam缓解此问题——但代价是引入额外内存开销。以往分析多依赖Frobenius范数解释这些方案,本文改从KL散度最小化角度重新建模其估计过程,揭示此前被忽视的理论局限,并推动更优设计。基于此视角,我们提出$ extbf{KL-Shampoo}$和$ extbf{KL-SOAP}$,在神经网络预训练中达到或超过原版性能,同时实现与SOAP相当的每迭代运行时间。值得注意的是,KL-Shampoo无需依赖Adam即可取得良好效果,消除了由Adam引入的内存负担。实验表明,KL-Shampoo持续优于SOAP、Shampoo乃至KL-SOAP,证明基于KL的方法为结构化优化器设计提供了有前景的新路径。代码已开源:https://github.com/yorkerlin/KL-Methods

原文摘要 · Abstract (English)

Shampoo and its efficient variant, SOAP, employ structured second-moment estimations and have shown strong performance for training neural networks (NNs). In practice, however, Shampoo typically requires step-size grafting with Adam to be competitive, and SOAP mitigates this by applying Adam in Shampoo's eigenbasis -- at the cost of additional memory overhead from Adam in both methods. Prior analyses have largely relied on the Frobenius norm to motivate these estimation schemes. We instead recast their estimation procedures as covariance estimation under Kullback-Leibler (KL) divergence minimization, revealing a previously overlooked theoretical limitation and motivating principled redesigns. Building on this perspective, we develop $\textbf{KL-Shampoo}$ and $\textbf{KL-SOAP}$, practical schemes that match or exceed the performance of Shampoo and SOAP in NN pre-training while achieving SOAP-level per-iteration runtime. Notably, KL-Shampoo does not rely on Adam to attain competitive performance, eliminating the memory overhead introduced by Adam. Across our experiments, KL-Shampoo consistently outperforms SOAP, Shampoo, and even KL-SOAP, establishing the KL-based approach as a promising foundation for designing structured methods in NN optimization. An implementation of KL-Shampoo/KL-SOAP is available at https://github.com/yorkerlin/KL-Methods

优化器神经网络结构化方法KL散度

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。