用于高效神经网络训练的稳定白化优化器
机器学习
2025-10-27 v3
摘要
在本工作中，我们对神经网络优化进行实验性研究。基于 Shampoo 系列算法，我们识别并消除了三个关键问题，提出了 SPlus 方法。首先，我们发现，当矩阵逆矩阵被长期缓存时，原始 Shampoo 方法容易发散。我们引入一种替代的有界更新方法，结合历史特征基和即时归一化，实现了普遍的稳定性和显著降低的计算要求。其次，我们采用面积感知的缩放，使学习率能够跨越网络宽度进行迁移。最后，我们发现高学习率会导致大参数噪声，提出一种简单的迭代平均方案以加快学习速度。为充分验证这些发现，我们引入了一个针对性的 Transformer 训练基准，考虑三个目标（语言建模、图像分类和扩散建模），并在训练的不同阶段进行评估。在平均情况下，SPlus 能够在梯度步数中达到 Adam 的验证性能的 44-58%，在墙钟时间中达到 62-83%。
关键词
引用
@article{arxiv.2506.07254,
title = {A Stable Whitening Optimizer for Efficient Neural Network Training},
author = {Kevin Frans and Sergey Levine and Pieter Abbeel},
journal= {arXiv preprint arXiv:2506.07254},
year = {2025}
}