arXiv:2409.11321v2 宣告类型: replace-cross 摘要：文献中有越来越多的证据表明，在深度学习优化任务中，Shampoo，一种高阶预条件方法，比Adam更有效。然而，与仅仅更新一阶和二阶矩的运行平均值的Adam相比，Shampoo 的缺点在于其额外的超参数和计算开销。本文建立了Shampoo（实现时使用1/2次方）与Adafactor之间的正式联系——Adafactor是Adam的一个内存高效的近似方法，表明Shampoo等价于在Shampoo预条件器的特征基中运行Adafactor。这一洞见导致设计了一个更简单且计算效率更高的算法：ShampoO with Adam in the Preconditioner's eigenbasis（SOAP）。 关于提高Shampoo的计算效率，最直接的方法似乎是更少地计算Shampoo的特征分解。不幸的是，如我们的实验证实所示，这会导致性能下降，并且这种下降随着计算频率的增加而恶化。SOAP通过持续更新二阶矩的运行平均值（就像Adam所做的那样），但由于是在当前（缓慢变化）的坐标基中进行的，从而减轻了这种下降。此外，由于SOAP等价于在旋转空间中运行Adam，它比Adam多引入了一个额外的超参数（预条件频率）。我们对360M和660M的规模的语言模型预训练进行了实证评估。在大批次处理的情况下，SOAP与AdamW相比，迭代次数减少了超过40%，墙钟时间减少了超过35%，且这两项指标均比Shampoo分别提高了约20%。SOAP的实现可以在这里找到：https://github.com/nikhilvyas/SOAP。