PowerStep：通过 $\ell_p$-Norm 最速下降进行内存高效自适应优化

Adam 等自适应优化器是训练 Transformer 的标准，但存储梯度一阶矩和二阶矩会产生大量内存开销。我们引入了 PowerStep，一种内存高效的优化器，无需存储二阶统计数据即可实现坐标自适应。受 $\ell_p$-norm 最速下降的推动，PowerStep 将有符号幂变换直接应用于一个动量缓冲区。我们使用 $O(1/\sqrt{T})$ 项和依赖于噪声的残差，为精确的、非正则化的更新建立了有限范围的平稳性界限。 Transformer 从 124M 到 235B 参数的实验显示了具有竞争力的验证质量，同时相对于 AdamW 将 $\texttt{fp32}$ 优化器状态内存减半。与统一的 $\texttt{int8}$ 量化相结合，PowerStep 保持数值稳定，并且与 $\texttt{fp32}$ AdamW 相比，将优化器状态内存减少了 $\sim8\times$。因此，PowerStep 为大规模训练提供了一种简单、内存高效的替代方案。