Large Deviation Upper Bounds and Improved MSE Rates of Nonlinear SGD: Heavy-tailed Noise and Power of Symmetry
Aleksandar Armacki, Shuhua Yu, Dragana Bajovic, Dusan Jakovetic, Soummya Kar
TL;DR
The paper develops a unified online nonlinear SGD framework that accommodates a broad class of nonlinearities (including sign, clipping, normalization, and quantization) under symmetric heavy-tailed noise without moment bounds. It establishes a sharp large deviation upper bound for the minimum gradient norm with a decay rate of sqrt(t)/log t and provides an explicit rate function that captures dependence on noise, nonlinearity, and problem parameters. For non-convex costs, it proves the optimal tilde O(t^{-1/2}) MSE rate, while for strongly convex costs the last-iterate MSE is arbitrarily close to the optimal O(t^{-1}); it also proves almost-sure convergence with rates approaching o(t^{-1/4}) for non-convex problems. The work improves upon prior results by avoiding moment assumptions, treating nonlinearities in a black-box way, and delivering sharp, parameter-free tail and convergence guarantees applicable to online optimization tasks in modern learning systems.
Abstract
We study large deviation upper bounds and mean-squared error (MSE) guarantees of a general framework of nonlinear stochastic gradient methods in the online setting, in the presence of heavy-tailed noise. Unlike existing works that rely on the closed form of a nonlinearity (typically clipping), our framework treats the nonlinearity in a black-box manner, allowing us to provide unified guarantees for a broad class of bounded nonlinearities, including many popular ones, like sign, quantization, normalization, as well as component-wise and joint clipping. We provide several strong results for a broad range of step-sizes in the presence of heavy-tailed noise with symmetric probability density function, positive in a neighbourhood of zero and potentially unbounded moments. In particular, for non-convex costs we provide a large deviation upper bound for the minimum norm-squared of gradients, showing an asymptotic tail decay on an exponential scale, at a rate $\sqrt{t} / \log(t)$. We establish the accompanying rate function, showing an explicit dependence on the choice of step-size, nonlinearity, noise and problem parameters. Next, for non-convex costs and the minimum norm-squared of gradients, we derive the optimal MSE rate $\widetilde{\mathcal{O}}(t^{-1/2})$. Moreover, for strongly convex costs and the last iterate, we provide an MSE rate that can be made arbitrarily close to the optimal rate $\mathcal{O}(t^{-1})$, improving on the state-of-the-art results in the presence of heavy-tailed noise. Finally, we establish almost sure convergence of the minimum norm-squared of gradients, providing an explicit rate, which can be made arbitrarily close to $o(t^{-1/4})$.