🔢 mathematics

PowerStep: Memory-Efficient Adaptive Optimization via $\ell_p$ -Norm Steepest Descent

PowerStep 是一种内存高效的自适应优化器，它通过 $\ell_p$ -范数最速下降法实现坐标级自适应而无需存储二阶矩统计量，在匹配 Adam 收敛速度的同时，显著降低了大规模 Transformer 训练的内存开销。

原作者： Yao Lu, Dengdong Fan, Shixun Zhang, Yonghong Tian

发布于 2026-05-12

📖 1 分钟阅读🧠 深度阅读

原作者： Yao Lu, Dengdong Fan, Shixun Zhang, Yonghong Tian

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在尝试教一个巨大且复杂的机器人（即神经网络）说一门新语言。为此，你在机器人尝试每句话后都会给予反馈。这种反馈就是“梯度”，它告诉机器人应向哪个方向移动才能变得更好。

多年来，训练这些机器人的标准方法是一种名为Adam的算法。Adam 就像一位非常谨慎、高度有条理的图书管理员。每当机器人犯了一个错误，Adam 不仅会查看当前的错误，还会维护一本庞大而详细的账本（即“二阶矩缓冲区”），记录机器人曾经犯下的每一个错误的历史。

问题所在：随着机器人变得更大（拥有数十亿个参数），这本账本变得极其庞大。它占用的内存如此之多，以至于拖慢了整个系统，甚至导致系统崩溃。这就像试图背着一装满百科全书的沉重背包跑马拉松。

登场：PowerStep——“智能直觉”方法

这篇论文的作者提出了一种名为PowerStep的新优化器。PowerStep 不再背负那装满历史数据的沉重背包，而是基于几何学和直觉运用了一个巧妙的技巧。

以下是 PowerStep 的工作原理，使用简单的类比来说明：

1. “重球”与“账本”

Adam（账本）：“我记得你昨天在单词'A'上犯了一个大错，在单词'B'上犯了一个小错。我会根据这些过去错误的平方来调整你的路径。”这需要存储大量数据。
PowerStep（重球）：PowerStep 使用一种称为“动量”的概念。想象一个重球从山坡上滚下。如果球滚得很快（强信号），它就会继续滚动；如果它滚得很慢（弱信号），就需要一点推力。
- PowerStep 不需要记住球速的历史。它只需观察球此刻滚动的速度。
- 它对当前速度应用一种特殊的“魔法过滤器”（带符号的幂变换）。如果球滚得太快，过滤器会温和地让它慢下来；如果它滚得太慢，过滤器就会给它一个助推。

2. “音量旋钮”类比

将机器人的学习过程想象成一个拥有成千上万个音量旋钮（对应机器人的每一个部分）的音响系统。

Adam 会聆听音乐的整个历史来决定如何调节每个旋钮。它很准确，但需要一台庞大的计算机来处理所有这些历史数据。
PowerStep 则聆听声音的当前音量。
- 如果某个旋钮已经调得很高（机器人很自信或梯度很大），PowerStep 会将音量调低一点，以防止它变得太响（过冲）。
- 如果某个旋钮几乎没开（机器人不确定或梯度很小），PowerStep 会将音量调高，以帮助它听得更清楚。
- 魔法之处：它无需写下历史书就能瞬间完成这一操作。它只是对当下做出反应。

为什么这很重要？

该论文宣称 PowerStep 取得了三大主要胜利：

体积减半：由于 PowerStep 不需要存储那本庞大的“二阶矩账本”，它比 Adam 节省50% 的内存。这就像将那装满百科全书的沉重背包换成了一本轻便的笔记本。
速度相同：尽管更轻量，PowerStep 的学习速度与 Adam 完全相同。它能在相同的时间内到达终点。
经受住“压缩”考验：研究人员尝试将数据压缩成极小、低质量的格式（称为“int8 量化”），这通常会破坏 Adam，因为它丢失了太多细节。
- Adam：被压缩后，Adam 的“账本”因缺失细节而变得极其混乱，导致机器人开始原地打转（发散）。
- PowerStep：因为它依赖于当前的动量，而不是脆弱的历史账本，即使数据被压缩，它也能保持稳定。这使得与标准方法相比，内存使用量减少了8 倍，而不会破坏机器人。

总结

该论文证明，你无需背负厚重的历史书就能有效训练巨型 AI 模型。通过使用一种基于几何学的智能方法，对当下的“动量”做出反应，PowerStep 实现了与行业标准（Adam）相同的结果，但内存成本仅为一半。而且，当将其与数据压缩结合时，它就成为了训练未来巨型 AI 模型的高效工具。

注：该论文在从小型（1.24 亿参数）到超大型（2350 亿参数）的模型上验证了这一点，证明其在任何规模下都有效。

技术摘要：PowerStep：基于 ℓp-范数最速下降的内存高效自适应优化

问题陈述

自适应优化器，特别是 Adam 和 AdamW，已成为训练大规模神经网络（例如 Transformer）的标准方法，它们通过利用累积的一阶和二阶矩统计量对梯度更新进行预条件处理。尽管这些方法行之有效，但它们会产生巨大的内存开销，因为它们维护一个与模型参数维度相同的二阶矩累加器（ $v_t$ ），通常为了稳定性需要完整的 32 位浮点数（fp32）精度。这种“内存墙”使得优化器状态的占用空间相对于带动量的随机梯度下降（SGD）增加了一倍，限制了拥有数千亿参数模型的扩展性。

现有的缓解尝试包括对二阶矩矩阵进行分解、分块共享或应用低精度量化。然而，由于二阶矩估计器（特别是倒数平方根运算）对量化误差高度敏感，量化往往会导致自适应方法的训练不稳定或发散。或者，一些方法完全丢弃二阶矩，但在大规模任务上通常无法达到 Adam 的收敛速度。

方法论

作者引入了 PowerStep，这是一种内存高效的优化器，它在不存储二阶矩统计量的情况下实现了坐标自适应。该方法源于 $\ell_p$ -范数最速下降 的原理。

理论基础

在一般的赋范空间中，在 $\ell_p$ -范数约束下使目标函数 $f(\theta)$ 下降幅度最大的最速下降方向 $v^*$ 由下式给出：
$v^*_i \propto -\text{sign}(g_i) \cdot |g_i|^{\frac{1}{p}-1}$
其中 $g_i$ 是梯度分量。这导出了一个“有符号幂变换” $\Phi_\beta(x) = \text{sign}(x) \odot |x|^\beta$ ，其中 $\beta = 1/(p-1) \in [0, 1]$ 。

算法设计

与将变换应用于原始梯度的 Powerball 方法不同，PowerStep 将该变换应用于 重球动量缓冲区（ $m_t$ ）。更新规则如下：

动量更新： $m_t = \gamma m_{t-1} + g_t$
变换更新： $\theta_t = \theta_{t-1} - \eta_t \Phi_\beta(m_t)$

该设计产生了一个对角预条件器 $D_t = \text{diag}(|m_t|^{\beta-1})$ 。对于 $\beta \in (0, 1)$ ，指数 $\beta-1$ 为负数，意味着步长与累积动量的大小成反比。小动量大小（平坦方向）被放大，而大动量大小（陡峭方向）被抑制。这提供了类似于 Adam 的坐标自适应，但它是从动量缓冲区即时计算得出的，从而消除了对二阶矩估计器的需求。

低精度实现

PowerStep 的单缓冲区设计天然适合激进的量化。作者提出仅以 int8 精度存储动量缓冲区 $m_t$ ，同时以全精度执行有符号幂变换和参数更新。这与 Adam 形成对比，在 Adam 中，由于平方根和除法运算的非线性，量化二阶矩缓冲区 $v_t$ 会导致灾难性的误差累积。

主要贡献

算法设计： PowerStep 源于 $\ell_p$ -范数最速下降原理，证明了将有符号幂变换应用于重球动量缓冲区可以在没有二阶矩统计量的情况下实现坐标自适应。
理论分析： 该论文证明了在标准正则性假设（L-平滑、有界梯度等）下，PowerStep 实现了非凸随机优化的最优 $O(1/\sqrt{T})$ 收敛速率。
实证验证： 在参数范围从 124M 到 235B 的 Transformer 模型（包括稠密架构和混合专家架构）上进行的广泛实验表明，PowerStep 在收敛速度上与 AdamW 相当，同时仅使用 一半的优化器状态内存。
量化鲁棒性： 当与激进的 int8 量化结合时，PowerStep 保持数值稳定，与全精度 AdamW 相比，优化器内存减少了约 8 倍。相比之下，AdamW 在相同的量化条件下会发散。

实验结果

收敛速度： 在从 124M 到 235B 参数的模型中，PowerStep 与 AdamW 的训练损失轨迹相匹配。相关的内存高效方法（pbSGDM、SignSGD、AdamS）在较大模型上表现出收敛较慢或不稳定的情况。
内存效率： 由于 PowerStep 仅存储一个缓冲区而不是两个，与 AdamW（fp32）相比，优化器状态内存占用减少了 50%。
量化性能： 在优化器状态进行 int8 量化时：
- AdamW： 由于二阶矩估计器中的量化误差累积，训练立即崩溃。
- PowerStep： 保持稳定，并与其全精度对应版本具有相同的收敛性。
- 内存减少： 与 AdamW（fp32）相比，PowerStep（int8）实现了优化器内存 约 8 倍 的减少。例如，在 235B 参数的模型上，每个 NPU 的内存从 6768 MB（AdamW）降至 872.4 MB（PowerStep int8）。
超参数敏感性： 该方法对动量系数 $\gamma$ 在 $[0.85, 0.95]$ 范围内具有鲁棒性。幂指数 $\beta = 0.1$ 在早期快速进展和长期稳定性之间提供了最佳权衡； $\beta=0$ （等同于 SignSGD）会导致崩溃，而 $\beta=0.2$ 收敛过慢。

意义与主张

该论文声称，PowerStep 为大规模训练提供了一种 有原则的、可扩展的且资源高效的替代方案。通过建立自适应优化与 $\ell_p$ -范数几何之间的直接联系，PowerStep 实现了自适应方法的理论最优性，同时避免了二阶矩估计的内存成本。

作者强调，消除二阶矩缓冲区不仅将内存使用量减半，而且从根本上提高了对低精度量化的鲁棒性。这使得在不牺牲收敛速度或最终模型质量的情况下，能够实现额外数量级的内存减少（与 AdamW 相比总计约 8 倍）。这项工作表明，结合非线性变换的重球动量可以替代 Adam 复杂的二阶矩机制，为训练未来的大规模模型提供更简单、更具可扩展性的路径。

作者指出的未来工作方向包括将 PowerStep 与结构化矩阵值更新（例如 Muon 优化器）集成，以及探索更激进的量化（例如 int4）。