Scale-Invariant Neural Network Optimization: Norm Geometry and Heavy-Tailed Noise

TL;DR

This paper investigates the theoretical limits and algorithmic strategies for scale-invariant neural network optimization under heavy-tailed noise, emphasizing the role of norm geometry and higher-order smoothness.

Contribution

It provides dimension-dependent lower bounds and matching upper bounds for scale-invariant methods, introduces a transported Scion method leveraging higher-order smoothness, and demonstrates practical effectiveness across architectures.

Findings

Dimension dependence is unavoidable for certain scale-invariant methods with general norms.

A batched Scion method achieves optimal bounds under spectral norm.

Transported Scion improves convergence when higher-order smoothness is exploited.

Abstract

A growing lesson from neural network optimization is that optimizer design should respect how the model is parametrized. Scale-invariant methods become important because their normalized layerwise updates can not only support hyperparameter transfer across model sizes but exploit input-output matrix norm geometry. At the same time, stochastic gradient noises in deep learning are often far from sub-Gaussian and may exhibit heavy tails. These crucial observations have shaped recent algorithmic principles for training neural networks, yet their joint theoretical consequences remain underexplored. In particular, it is unclear what dimension dependence is unavoidable for scale-invariant methods with general input-output matrix norm, and whether higher-order smoothness can accelerate training under heavy-tailed noise. We study these questions through nonconvex smooth stochastic optimization over $\mathbb{R}^{m\times n}$ with general norms, where the goal is to achieve an $\epsilon$-stationary point under $p^{\mathrm{th}}$-moment heavy-tailed noise. Our first contribution is a dimension-dependent lower bound: when $\frac{\max\{m,n\}}{(\min\{m,n\})^2}$ is large enough, any scale-invariant first-order method with spectral norm requires $\Omega(\min\{m, n\}\epsilon^{-\frac{3p-2}{p-1}})$ oracle calls. We prove that a batched Scion method with spectral norm achieves the matching upper bound of $O(\min\{m, n\}\epsilon^{-\frac{3p-2}{p-1}})$. To exploit higher-order smoothness, we propose a transported Scion method and improve the bound to $O(\min\{m, n\}\epsilon^{-\frac{5p-3}{2p-2}})$ when the norm is spectral and the Hessian is Lipschitz. Finally, we incorporate practical heuristics into our transported method and evaluate it across multiple architectures and model sizes, demonstrating its flexibility and compatibility in training neural networks.