Technical Summary: PowerStep: Memory-Efficient Adaptive Optimization via ℓp-Norm Steepest Descent

Problem Statement

Adaptive optimizers, particularly Adam and AdamW, have become the standard for training large-scale neural networks (e.g., Transformers) by preconditioning gradient updates with accumulated first and second-moment statistics. While effective, these methods incur substantial memory overhead because they maintain a second-moment accumulator ($v_t$) with the same dimensionality as the model parameters, typically requiring full 32-bit floating-point (fp32) precision for stability. This "memory wall" doubles the optimizer state footprint relative to Stochastic Gradient Descent (SGD) with momentum, limiting scalability for models with hundreds of billions of parameters.

Existing attempts to mitigate this include factorizing the second-moment matrix, blockwise sharing, or applying low-precision quantization. However, quantization often leads to training instability or divergence in adaptive methods due to the high sensitivity of the second-moment estimator (specifically the reciprocal square root operation) to quantization errors. Alternatively, some methods discard the second moment entirely but often fail to match the convergence speed of Adam on large-scale tasks.

Methodology

The authors introduce PowerStep, a memory-efficient optimizer that achieves coordinate-wise adaptivity without storing second-moment statistics. The method is derived from the principles of $\ell_p$-norm steepest descent.

Theoretical Foundation

In general normed spaces, the steepest descent direction $v^*$ that maximally decreases an objective $f(\theta)$ under an $\ell_p$-norm constraint is given by:
$$v^*_i \propto -\text{sign}(g_i) \cdot |g_i|^{\frac{1}{p}-1}$$
where $g_i$ is the gradient component. This leads to a "signed power transform" $\Phi_\beta(x) = \text{sign}(x) \odot |x|^\beta$, where $\beta = 1/(p-1) \in [0, 1]$.

Algorithm Design

Unlike the Powerball method, which applies this transform to raw gradients, PowerStep applies the transform to a heavy-ball momentum buffer ($m_t$). The update rule is:

1. Momentum Update: $m_t = \gamma m_{t-1} + g_t$
2. Transformed Update: $\theta_t = \theta_{t-1} - \eta_t \Phi_\beta(m_t)$

This design yields a diagonal preconditioner $D_t = \text{diag}(|m_t|^{\beta-1})$. For $\beta \in (0, 1)$, the exponent $\beta-1$ is negative, meaning the step size is inversely related to the magnitude of the accumulated momentum. Small momentum magnitudes (flat directions) are amplified, while large magnitudes (steep directions) are dampened. This provides coordinate-wise adaptivity similar to Adam but computed instantaneously from the momentum buffer, eliminating the need for a second-moment estimator.

Low-Precision Implementation

PowerStep's single-buffer design is naturally amenable to aggressive quantization. The authors propose storing only the momentum buffer $m_t$ in int8 precision while performing the signed power transform and parameter updates in full precision. This contrasts with Adam, where quantizing the second-moment buffer $v_t$ causes catastrophic error accumulation due to the nonlinearity of the square root and division operations.

Key Contributions

1. Algorithm Design: PowerStep is derived from $\ell_p$-norm steepest descent principles, demonstrating that applying a signed power transform to a heavy-ball momentum buffer yields coordinate-wise adaptivity without second-moment statistics.
2. Theoretical Analysis: The paper proves that PowerStep achieves the optimal $O(1/\sqrt{T})$ convergence rate for non-convex stochastic optimization under standard regularity assumptions (L-smoothness, bounded gradients, etc.).
3. Empirical Validation: Extensive experiments on Transformer models ranging from 124M to 235B parameters (including dense and Mixture-of-Experts architectures) show that PowerStep matches the convergence speed of AdamW while using only half the optimizer state memory.
4. Quantization Robustness: When combined with aggressive int8 quantization, PowerStep remains numerically stable and reduces optimizer memory by approximately 8× compared to full-precision AdamW. In contrast, AdamW diverges under the same quantization conditions.

Experimental Results

• Convergence Speed: Across models from 124M to 235B parameters, PowerStep matches the training loss trajectories of AdamW. Related memory-efficient methods (pbSGDM, SignSGD, AdamS) exhibited slower convergence or instability on larger models.
• Memory Efficiency: PowerStep reduces the optimizer state memory footprint by 50% compared to AdamW (fp32) because it stores only one buffer instead of two.
• Quantization Performance: Under int8 quantization of the optimizer state:
  • AdamW: Training collapses immediately due to quantization error accumulation in the second-moment estimator.
  • PowerStep: Remains stable and matches the convergence of its full-precision counterpart.
  • Memory Reduction: PowerStep (int8) achieves an ~8× reduction in optimizer memory compared to AdamW (fp32). For example, on a 235B parameter model, memory per NPU drops from 6768 MB (AdamW) to 872.4 MB (PowerStep int8).
• Hyperparameter Sensitivity: The method is robust to the momentum coefficient $\gamma$ in the range $[0.85, 0.95]$. A power exponent $\beta = 0.1$ provides the best trade-off between rapid early progress and long-run stability; $\beta=0$ (equivalent to SignSGD) leads to collapse, while $\beta=0.2$ converges too slowly.

Significance and Claims

The paper claims that PowerStep provides a principled, scalable, and resource-efficient alternative for large-scale training. By establishing a direct connection between adaptive optimization and $\ell_p$-norm geometry, PowerStep achieves the theoretical optimality of adaptive methods without the memory cost of second-moment estimation.

The authors emphasize that the elimination of the second-moment buffer not only halves memory usage but also fundamentally improves robustness to low-precision quantization. This allows for an additional order-of-magnitude memory reduction (total ~8× vs. AdamW) without sacrificing convergence speed or final model quality. The work suggests that heavy-ball momentum combined with nonlinear transforms can replace the complex second-moment machinery of Adam, offering a simpler and more scalable path for training future large-scale models.

Future work directions identified by the authors include integrating PowerStep with structured matrix-valued updates (e.g., the Muon optimizer) and exploring even more aggressive quantization (e.g., int4).