Reproduce “Why is Normalization Preferred? A Worst-Case Complexity Theory for Stochastically Preconditioned SGD under Heavy-Tailed Noise” on OpenResearch

About this paper

Yuchen Fang, James Demmel, Javad Lavaei

Researchers at UC Berkeley developed a worst-case complexity theory for Stochastically Preconditioned SGD (SPSGD) under heavy-tailed noise, demonstrating that step normalization guarantees robust convergence with optimal rates, while step clipping may fail due to statistical dependencies between the preconditioner and gradient noise. This work provides a theoretical basis for the empirical preference of normalization in adaptive optimizers.

Key insights

• Normalization effectively decouples the update direction from its magnitude, preventing adverse interactions with stochastic preconditioners and ensuring stability under heavy-tailed noise.
• Clipping's potential failure in SPSGD stems from the statistical dependence between the stochastic preconditioner and the clipped gradient noise, which generates a persistent, non-vanishing adverse bias term.
• The robustness of normalization is maintained even with suboptimally designed preconditioners, with performance degradation limited to a controlled multiplicative factor of the preconditioner's condition number.

Summary by alphaXiv.