提出SPlus优化器，训练更稳定高效，加速收敛并减少计算开销。
A Stable Whitening Optimizer for Efficient Neural Network Training
- 用历史特征基结合瞬时归一化，解决缓存矩阵逆导致的发散问题。
- 在不同网络宽度间实现学习率迁移，提升训练泛化性。
- 通过参数迭代平均，缓解高学习率带来的噪声，支持更快训练。
本文基于实验研究神经网络优化。在Shampoo算法基础上，识别并解决三大关键问题，提出SPlus方法。首先，发现原始Shampoo在长期缓存矩阵逆时易发散，引入结合历史特征基与瞬时归一化的有界更新，显著提升稳定性并降低计算量。其次，采用形状感知缩放，实现学习率在不同网络宽度间的迁移。第三，发现高学习率引发大参数噪声，提出简单迭代平均方案，突破快速训练瓶颈。为验证结论，构建针对性Transformer训练基准，涵盖语言建模、图像分类与扩散建模三种任务。结果表明，SPlus平均仅需Adam 44-58%的梯度步数和62-83%的运行时间，即可达到相当的验证性能。
原文摘要 · Abstract (English)
In this work, we take an experimentally grounded look at neural network optimization. Building on the Shampoo family of algorithms, we identify and alleviate three key issues, resulting in the proposed SPlus method. First, we find that naive Shampoo is prone to divergence when matrix-inverses are cached for long periods. We introduce an alternate bounded update combining a historical eigenbasis with instantaneous normalization, resulting in across-the-board stability and significantly lower computational requirements. Second, we adapt a shape-aware scaling to enable learning rate transfer across network width. Third, we find that high learning rates result in large parameter noise, and propose a simple iterate-averaging scheme which unblocks faster learning. To properly confirm these findings, we introduce a pointed Transformer training benchmark, considering three objectives (language modelling, image classification, and diffusion modelling) across different stages of training. On average, SPlus is able to reach the validation performance of Adam within 44-58% of the gradient steps and 62-83% of the wallclock time.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。