General Tail Bounds for Non-Smooth Stochastic Mirror Descent
This work addresses robustness in optimization for machine learning practitioners dealing with non-smooth and heavy-tailed data, though it is incremental as it builds on prior tail bound analyses.
The paper tackles the problem of deriving tail bounds for Stochastic Mirror Descent under heavy-tailed noise, extending existing results from light-tailed to heavier-tailed regimes, and shows that the average of iterates often outperforms the last iterate in such settings.
In this paper, we provide novel tail bounds on the optimization error of Stochastic Mirror Descent for convex and Lipschitz objectives. Our analysis extends the existing tail bounds from the classical light-tailed Sub-Gaussian noise case to heavier-tailed noise regimes. We study the optimization error of the last iterate as well as the average of the iterates. We instantiate our results in two important cases: a class of noise with exponential tails and one with polynomial tails. A remarkable feature of our results is that they do not require an upper bound on the diameter of the domain. Finally, we support our theory with illustrative experiments that compare the behavior of the average of the iterates with that of the last iterate in heavy-tailed noise regimes.