Why is Normalization Preferred? A Worst-Case Complexity Theory for Stochastically Preconditioned SGD under Heavy-Tailed Noise

1. Problem Statement

The paper addresses the convergence behavior of Stochastically Preconditioned Stochastic Gradient Descent (SPSGD) and its accelerated variants (e.g., Adam, RMSProp, Shampoo, K-FAC) in the presence of heavy-tailed gradient noise.

• Context: In modern large-scale machine learning (e.g., LLMs, image classification), gradient noise often violates the classical bounded-variance assumption. Instead, it exhibits heavy tails, meaning only a finite $p$-th moment exists for some $p \in (1, 2]$.
• The Conflict: Two primary techniques are used to stabilize training under heavy-tailed noise: Gradient Clipping (truncating large gradients) and Gradient Normalization (rescaling gradients to unit norm).
  • In standard (vanilla) SGD, both methods have been shown to achieve optimal convergence rates theoretically.
  • However, in practice, normalization is often preferred over clipping in large-scale adaptive methods.
• The Gap: There is a lack of theoretical understanding explaining why normalization outperforms clipping when combined with stochastic preconditioning (where the preconditioner $D_k$ depends on past gradients). The paper aims to close this gap by establishing a worst-case complexity theory.

2. Methodology and Framework

A. Unified Algorithmic Model

The authors propose a unified update rule for SPSGD:
$$x_{k+1} = x_k - \eta D_k m_k$$
$$m_k = \theta m_{k-1} + (1-\theta)\bar{g}_k$$
Where:

• $D_k$ is a symmetric positive definite stochastic preconditioning matrix (e.g., diagonal in Adam, Kronecker-factored in Shampoo).
• $\bar{g}_k$ is the stochastic gradient estimate.
• $m_k$ is the momentum buffer.

This framework encompasses a wide range of optimizers, including AdaGrad, Adam, RMSProp, Shampoo, and stochastic Newton methods.

B. Assumptions

1. Heavy-Tailed Noise: The gradient noise $\bar{g}_k - \nabla f(x_k)$ has a bounded $p$-th moment ($p \in (1, 2]$), i.e., $\mathbb{E}[\|\bar{g}_k - \nabla f(x_k)\|^p] \leq \sigma^p$.
2. Smoothness: The objective function $f$ is $L$-smooth.
3. Preconditioner Bounds: The condition number of the stochastic preconditioner $D_k$ is deterministically bounded ($\kappa(D_k) \leq \kappa_D$), though the eigenvalues themselves can be stochastic.

C. Analytical Tools

To handle the vector-valued nature of the momentum and preconditioning, the authors develop a novel vector-valued Burkholder-type inequality. This generalizes previous scalar results and allows for tighter bounds on the accumulation of martingale differences in the presence of heavy tails ($p \in (1, 2]$).

3. Key Contributions and Results

The paper establishes a sharp worst-case separation between normalization and clipping in the SPSGD setting.

A. Normalization (NSPSGD) is Robust

The authors prove that Step Normalization ensures convergence to a first-order stationary point even under heavy-tailed noise and stochastic preconditioning.

• Convergence Rates:
  • Known Parameters: $O(T^{-\frac{p-1}{3p-2}})$
  • Unknown Parameters: $O(T^{-\frac{p-1}{2p}})$
• Significance: These rates match the optimal rates for Normalized SGD (NSGD) in the vanilla setting. The analysis shows that normalization effectively decouples the update magnitude from the scale of the preconditioner $D_k$, preventing the "magnitude noise" of $D_k$ from destabilizing the optimization trajectory. The performance degradation is bounded only by a factor of $\sqrt{\kappa_D}$.

B. Clipping (ClipSPSGD) May Fail

In contrast, the paper proves that Step Clipping may fail to converge in the worst case when applied to SPSGD.

• The Mechanism of Failure: The failure arises from the statistical dependence between the stochastic preconditioner $D_k$ and the gradient estimate $\bar{g}_k$.
  • In vanilla SGD ($D_k = I$), the clipping error term vanishes because the noise is independent of the step direction.
  • In SPSGD, $D_k$ is constructed from past gradients, creating a correlation with the current gradient noise.
• The Error Term: The analysis reveals a persistent bias term: $-(\eta - L\eta^2 M_D) \nabla f(x_k)^T \mathbb{E}[D_k \epsilon_k^u]$, where $\epsilon_k^u$ is the unbiased noise component.
  • This term represents the covariance between the preconditioner and the clipped noise.
  • Unlike in vanilla SGD where this term is $O(\eta^2)$, in SPSGD it remains $O(\eta)$.
  • As the iteration count $T \to \infty$, the required clipping threshold $\tau$ must grow. This growth causes the error term to dominate, potentially leading to divergence or a failure to achieve the optimal convergence rate.

C. Geometric Interpretation

• Normalization: Projects the update strictly onto a sphere. This constrains the trajectory to a stable manifold, effectively removing magnitude information and breaking the dependence between the step size and the preconditioner's scale.
• Clipping: Allows the update to lie anywhere within a ball of radius $\tau$. Because the threshold $\tau$ must grow with time to ensure convergence, the "interior" of the ball expands. The statistical dependence between $D_k$ and $\bar{g}_k$ introduces a bias within this expanding region that clipping cannot eliminate.

4. Significance and Implications

1. Theoretical Justification for Empirical Practice: The paper provides the first rigorous theoretical explanation for why practitioners prefer normalization (e.g., in LARS, LAMB, and variants of Adam) over clipping for large-scale models. It explains that while clipping works for vanilla SGD, it is fundamentally fragile when combined with adaptive preconditioning under heavy tails.
2. Unified Complexity Theory: It establishes the first worst-case complexity bounds for SPSGD under heavy-tailed noise, covering a broad class of optimizers (Adam, Shampoo, etc.) under minimal structural assumptions.
3. New Mathematical Tool: The derived vector-valued Burkholder-type inequality is a significant contribution to probability theory, offering sharper constants for analyzing stochastic optimization algorithms with heavy-tailed noise.
4. Design Guidelines: The results suggest that for large-scale, adaptive optimization under heavy-tailed noise, algorithms should prioritize normalization mechanisms over clipping to ensure robust convergence guarantees.

Conclusion

The paper demonstrates that the "paradox" of normalization being preferred over clipping in practice is not merely a tuning issue but a fundamental theoretical property. Under stochastic preconditioning and heavy-tailed noise, clipping introduces an irreducible bias due to the dependence between the preconditioner and the gradient, whereas normalization inherently neutralizes this dependence, guaranteeing optimal convergence rates.