提出新型优化器PowerStep，训练大模型时内存减半且不降速。
PowerStep: Memory-Efficient Adaptive Optimization via $\ell_p$-Norm Steepest Descent
- 通过非线性变换动量缓冲区实现自适应，无需存储二阶矩统计量。
- 在124M到235B参数的Transformer上，收敛速度媲美Adam，内存减半。
- 支持int8量化，内存再降8倍，适合资源受限的大模型训练。
自适应优化器（如Adam）已成为训练大规模神经网络（如Transformer）的标准，但需维护梯度一阶与二阶矩的运行估计，导致显著内存开销。本文提出PowerStep，一种内存高效的优化器，在不存储二阶矩统计量的前提下实现逐坐标自适应。受ℓ_p-范数几何下最陡下降的启发，我们证明对动量缓冲区施加非线性变换即可获得自适应性。理论证明PowerStep在非凸随机优化中以最优的O(1/√T)速率收敛。大量实验表明，在参数量从124M到235B的Transformer模型上，PowerStep的收敛速度与Adam相当，同时将优化器内存减半。结合激进的int8量化后，PowerStep仍保持数值稳定，相较全精度Adam进一步降低约8倍优化器内存。因此，PowerStep为大规模训练提供了一种原理清晰、可扩展且资源高效的选择。代码已开源。
原文摘要 · Abstract (English)
Adaptive optimizers, most notably Adam, have become the default standard for training large-scale neural networks such as Transformers. These methods maintain running estimates of gradient first and second moments, incurring substantial memory overhead. We introduce PowerStep, a memory-efficient optimizer that achieves coordinate-wise adaptivity without storing second-moment statistics. Motivated by steepest descent under an $\ell_p$-norm geometry, we show that applying a nonlinear transform directly to a momentum buffer yields coordinate-wise adaptivity. We prove that PowerStep converges at the optimal $O(1/\sqrt{T})$ rate for non-convex stochastic optimization. Extensive experiments on Transformer models ranging from 124M to 235B parameters demonstrate that PowerStep matches Adam's convergence speed while halving optimizer memory. Furthermore, when combined with aggressive \texttt{int8} quantization, PowerStep remains numerically stable and reduces optimizer memory by $\sim\!8\times$ compared to full-precision Adam. PowerStep thus provides a principled, scalable and resource-efficient alternative for large-scale training. Code is available at https://github.com/yaolubrain/PowerStep.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。