SOAP：利用Adam改进并稳定Shampoo算法
机器学习
2025-02-03 v2 人工智能
摘要
越来越多的证据表明，在深度学习优化任务中，高阶预条件方法Shampoo比Adam更有效。然而，与Adam相比，Shampoo的缺点包括额外的超参数和计算开销，而Adam仅更新一阶矩和二阶矩量的滑动平均值。本工作建立了Shampoo（以1/2次幂实现）与Adafactor——一种内存高效的Adam近似——之间的形式化联系，表明Shampoo等价于在Shampoo预条件子的特征基中运行Adafactor。这一洞察促成了一个更简单且计算效率更高的算法设计：在预条件子特征基中运行Adam的Shampoo（SOAP）。关于提高Shampoo的计算效率，最直接的方法是简单地降低Shampoo特征分解的频率。不幸的是，正如我们的经验结果所示，这会导致性能下降，且下降程度随该频率的降低而加剧。SOAP通过在当前（缓慢变化的）坐标基中持续更新二阶矩的滑动平均值（正如Adam所做的那样）来缓解这种性能下降。此外，由于SOAP等价于在旋转空间中运行Adam，与Adam相比，它仅引入一个额外的超参数（预条件频率）。我们在360m和660m规模的语言模型预训练上对SOAP进行了经验评估。在大规模批次训练模式下，与AdamW相比，SOAP将迭代次数减少了40%以上，挂钟时间减少了35%以上；与Shampoo相比，这两项指标均改善了约20%。SOAP的实现可在https://github.com/nikhilvyas/SOAP获取。
引用
@article{arxiv.2409.11321,
title = {SOAP: Improving and Stabilizing Shampoo using Adam},
author = {Nikhil Vyas and Depen Morwani and Rosie Zhao and Mujin Kwun and Itai Shapira and David Brandfonbrener and Lucas Janson and Sham Kakade},
journal= {arXiv preprint arXiv:2409.11321},
year = {2025}
}