Simplicity Suffices for Parameter Noise Injection in Stochastic Gradient Descent

Noise injection is a well-known technique in stochastic optimization. We report its surprising effectiveness with an interleaved (on-off-on-off...) rather than the usual monotonic decay schedule. We present a theoretical analysis of noise injection, which confirms that corruption by impulse noise approximates a Jacobian regularization, whereas Gaussian noise acts as a curvature penalty. This regularization behavior has been invoked to explain why noise injection increases model robustness. But the interleaved nature of our proposed schedule produces superior results even for the optimization objective: mixing phases of noisy data permits the optimizer to escape local minima and increase exploration without the risk of catastrophically forgetting the important features from the clean data. To stabilize this training scheme against the rapid changes of the loss when switching between clean and noisy data, we introduce a gradient-norm stabilization technique that scales noisy updates based on clean gradient magnitudes. We compare this method with other common augmentation methods and find substantial improvements in corruption tolerance and robustness to real-world distribution shifts on CIFAR-100-C, ImageNet-C, and ImageNet-R for ResNet and ViT architectures, with the best results being achieved by stacking our method on top of other augmentations. Through saliency and attention maps we show that the effect of interleaved noise injection stems from penalizing the failure modes encouraged by the inductive bias of the models: impulse noise works against the locality bias of convolutional (ResNet) architectures, and Gaussian noise reduces the tendency of attention-based models to pick up large-scale spurious features. Interleaved noise injection is therefore an effective tool to improve the test performance on clean, noisy, and out-of-distribution data at essentially zero computational cost.

Accelerating stochastic gradient methods with classical momentum schemes, such as Polyak's heavy ball, has proven highly successful in training large-scale machine learning models, particularly when combined with the hardware acceleration of large mini-batch computations. Yet, the effect of classical momentum on stochastic mini-batch optimization has been poorly understood theoretically, with prior works requiring strong noise assumptions and extremely large mini-batches. In this work, we develop a general theory of stochastic momentum acceleration for optimizing over quadratics in the interpolation regime, a popular abstraction for studying deep learning dynamics which also includes classical methods such as randomized Kaczmarz and coordinate descent. Our framework encompasses both heavy ball and Nesterov-style momentum, allows for arbitrary mini-batch sizes, and makes minimal assumptions on the stochastic noise. In particular, we show that acceleration from classical momentum is directly proportional to the gradient mini-batch size (up to a natural saturation point), thereby enabling perfect parallelization of mini-batch computations. Our theory also provides a simple choice for the momentum parameter, which is shown to be effective empirically.

A growing lesson from neural network optimization is that optimizer design should respect how the model is parametrized. The layerwise input-output structure of neural networks motivates scale-invariant optimizers, such as Muon and Scion, whose updates also support hyperparameter transfer. At the same time, stochastic gradient noise in deep learning is often far from sub-Gaussian and may exhibit heavy tails. These observations have shaped recent algorithmic principles for training neural networks, yet their joint theoretical consequences are underexplored. In particular, it remains unclear what dimension dependence is unavoidable for gradient-based methods given the problem class is defined by input-output norm and under heavy-tailed noise, and whether higher-order smoothness can accelerate training. We study these questions through nonconvex smooth stochastic optimization over $\mathbb R^{m\times n}$ equipped with general norms and under $p^\mathrm{th}$-moment heavy-tailed noise, where the goal is to achieve an $ε$-stationary point in the dual norm. Our first contribution is a dimension-dependent lower bound: when $\frac{\max\{m,n\}}{(\min\{m,n\})^2}$ is large enough, any gradient-based method requires $Ω(\min\{m, n\}ε^{-\frac{3p-2}{p-1}})$ oracles for the problem class defined by the spectral norm, which is a common input-output norm. We prove that a scale-invariant Scion method with the spectral norm can achieve the matching upper bound of $O(\min\{m, n\}ε^{-\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\}ε^{-\frac{5p-3}{2p-2}})$ when the Hessian is Lipschitz. Finally, we incorporate heuristics into our transported method and evaluate it across multiple architectures and model sizes, demonstrating its flexibility and compatibility with neural network training.