High Probability Convergence Bounds for Non-convex Stochastic Gradient Descent with Sub-Weibull Noise

Stochastic gradient descent (SGD) and its variants are some of the most commonly used algorithms in machine learning. In particular, they are used for training neural network and transformer models, models that have achieved considerable success on image classification and language processing tasks in recent years. The training in this case is non-convex and smooth for many activation/attention functions, such as sigmoid, GELU, and softmax. Even for ReLU, which is not differentiable, the training is smooth over most of the parameter space and avoidance of the non-smooth part is dealt with separately. Thus, the smooth non-convex convergence analysis of SGD has far-reaching influence on the field of machine learning.

There is a large literature on the almost sure and mean convergence of SGD assuming strong smoothness. Bertsekas and Tsitsiklis (2000) prove that SGD almost surely converges to a first-order stationary point assuming strong smoothness and the relaxed growth noise condition; more recently, Patel (2021) showed the same result with a different proof technique. Ghadimi and Lan (2013) prove that the mean of the squared gradient norm of SGD converges to zero at a rate of $O(1/\sqrt{T})$ assuming strong smoothness and the bounded variance noise condition, where $T$ is the number of SGD iterations. Khaled and Richtárik (2020) prove the same but with the expected smoothness noise condition. Sebbouh et al. (2021) strengthen this to an almost sure convergence rate. Assuming strong smoothness and convexity, the tight mean convergence rate of the squared gradient norm of SGD is $O(1/\sqrt{T})$ (Nemirovsky and Yudin, 1983, Thm. 5.3.1). In fact, this is the optimal rate for all stochastic first-order methods assuming only strong smoothness (Arjevani et al., 2019). Thus, the $O(1/\sqrt{T})$ mean convergence rate of Ghadimi and Lan (2013) is tight and shows that SGD is optimal in the smooth non-convex setting.

However, mean convergence is not the end of the story. A convergence guarantee is generally required with some arbitrary probability, $1-\delta$. For a single run of SGD, using Markov’s inequality gives a $1/\delta$ scaling to the bound. If one can re-run SGD many times (say, $r$ times), then by choosing the best run, $\delta$ can be relatively large since the overall success is at least by $1-\delta^{r}$. On the other hand, if just a single run of SGD is allowed, then the $1/\delta$ factor leads to nearly vacuous results, and hence a direct high probability convergence analysis is necessary to understand the behavior of unique runs of SGD.

Li and Orabona (2020) prove a $O(\log(T)\log(T/\delta)/\sqrt{T})$ high probability convergence rate for a weighted average of the squared gradient norms of SGD assuming strong smoothness and norm sub-Gaussian noise. However, a series of recent papers (Gürbüzbalaban et al., 2021; Şimşekli et al., 2019; Panigrahi et al., 2019) suggest that norm sub-Gaussian noise is often not satisfied. While the central limit theorem can be used to heuristically justify norm sub-Gaussian noise for mini-batch SGD with large batch-sizes, it cannot for small batch-sizes. In particular, Panigrahi et al. (2019) provide examples where the noise is Gaussian for a batch-size of 4096, is not Gaussian for a batch-size of 32, and starts out Gaussian then becomes non-Gaussian for a batch-size of 256. Gürbüzbalaban et al. (2021) and Şimşekli et al. (2019) suggest instead assuming that the $p$th moment of the norm of the noise is bounded, where $p\in(1,2)$. Scaman and Malherbe (2020) prove a $O(1/T^{(p-1)/p}/\delta^{(2+p)/(2p)})$ high probability convergence rate for SGD assuming strong smoothness and bounded $p$th moment noise for $p\in(1,2)$. By allowing noise with possibly infinite variance, SGD converges at a slower rate in terms of $T$. Moreover, the dependence on $1/\delta$ is polynomial rather than logarithmic. Cutkosky and Mehta (2021) prove a $O(\log(T/\delta)/T^{(p-1)/(3p-2)})$ high probability convergence rate for clipped SGD (which uses clipping, momentum, and normalization) assuming strong smoothness and bounded $p$th moment noise for $p\in(1,2)$, thus improving the dependence on $1/\delta$ to logarithmic, but clipped SGD will still have a slower convergence rate in terms of $T$ even when the noise is norm sub-Gaussian. So, knowledge about the type of noise should actually affect whether we use SGD or clipped SGD. This motivates the question: How heavy can the noise be for SGD to have a high probability convergence rate that depends logarithmically on $1/\delta$?

Towards this end, we consider norm sub-Weibull noise. While it is much lighter tailed than bounded $p$th moment noise for $p\in(1,2)$, it is a natural extension of norm sub-Gaussian noise and it turns out that it admits a convergence rate for SGD with logarithmic dependence on $1/\delta$. This leads us to the first contribution of our paper. Assuming strong smoothness and norm sub-Weibull noise, we prove a $O((\log(T)\log(1/\delta)^{2\theta}+\log(T/\delta)^{\min\{0,\theta-1\}}\log(1/\delta))/\sqrt{T})$ high probability convergence rate for a weighted average of the squared gradient norms of SGD, where $\theta$ is the sub-Weibull tail weight. This is given in Theorem 15. In the special case of norm sub-Gaussian noise, which is $\theta=1/2$, this becomes $O(\log(T)\log(1/\delta)/\sqrt{T})$, which slightly improves the result of Li and Orabona (2020). For more general $\theta>1/2$, we make the additional assumption that the objective function is Lipschitz continuous. At its most basic, the Lipschitz continuity assumption says that the norm of the true gradient is bounded for all of the iterates. Li and Liu (2022), building off of our pre-print, relax this to only assuming that the step-size times the true gradient is bounded for all of the iterates, which is a weaker assumption since the step-size goes to zero.

However, hidden within these convergence rates is a subtle issue that requires a post-processing algorithm. All of these high probability convergence rates are for a weighted average of the squared gradient norms of the iterates rather than being for the squared gradient norm of a single iterate. While this implies a high probability convergence rate for the iterate with the smallest squared gradient norm, to actually determine which iterate has the smallest squared gradient norm, we would need to estimate the true gradient for each iterate! This leads us to another question: Is there a post-processing strategy for a single run of SGD that is as efficient as 2-RSG?

2-RSG is the algorithm of Ghadimi and Lan (2013) that takes multiple runs of SGD and picks the best one to get a high probability convergence rate. First, it goes from a mean convergence rate for the iterate with the smallest squared gradient norm to a mean convergence rate for a particular iterate by randomly choosing the iterate. It does this for $\Theta(\log(1/\delta))$ runs of SGD. Then, it uses $\Theta(\log(1/\delta)\sigma^{2}/\delta\epsilon)$ samples to estimate the $\Theta(\log(1/\delta))$ gradients and pick the run with the smallest squared estimated gradient norm, where $\sigma^{2}$ is the variance of the noise and $\epsilon$ is the convergence tolerance. To compete with this, we introduce a post-processing algorithm that randomly chooses $\Theta(\log(1/\delta))$ iterates of a single run of SGD and then picks the one with the smallest squared estimated gradient norm using $\Theta(\log(1/\delta)\sigma^{2}/\delta\epsilon)$ samples and we prove that our high probability convergence rate applies to the squared gradient norm of this iterate. The algorithm and the bound on the squared gradient norm of its output are in Theorem 21. The result can be combined with any high probability convergence rate on a weighted average of the squared gradient norm of the iterates of SGD, such as the results of Li and Orabona (2020), Scaman and Malherbe (2020), and Cutkosky and Mehta (2021). Moreover, the first part of it, going from a high probability bound on a weighted average of random variables to a high probability bound on the smallest of a small subset of those random variables, applies to more general stochastic sequences as shown in Theorem 19, which is a probability result of independent interest.

Underlying our convergence analysis are concentration inequalities. Vladimirova et al. (2020), Wong et al. (2020), and Bakhshizadeh et al. (2023) prove concentration inequalities for sub-Weibull random variables with deterministic scale parameters. However, we need a concentration inequality for a sub-Weibull martingale difference sequence (MDS) where the scale parameters are themselves random variables making up an auxiliary sequence. MDS concentration inequalities upper bound the partial sum of the main sequence and lower bound the partial sum of the auxiliary sequence simultaneously. When the lower bound depends on the partial sum of the main sequence, we call the concentration inequality self-normalized. Freedman (1975) proves a sub-Gaussian MDS concentration inequality. Fan et al. (2015) proves a sub-exponential MDS concentration inequality. Harvey et al. (2019) proves a sub-Gaussian MDS self-normalized concentration inequality. The proofs for these results all rely on the moment generating function (MGF), which does not exist for sub-Weibull random variables with $\theta>1$, i.e., heavier than sub-exponential. Nevertheless, we are able to prove a sub-Weibull martingale difference sequence self-normalized concentration inequality. This is given in Theorem 11. The proof uses the MGF truncation technique of Bakhshizadeh et al. (2023) but applied to an MDS rather than i.i.d. random variables. The truncation level is determined by the tail decay rate and shows up as the $\log(T/\delta)^{\theta-1}$ in the SGD convergence rate.

We also consider convergence under the Polyak-Łojasiewicz (PŁ) inequality. Some examples of problems with PŁ objectives are logistic regression over compact sets (Karimi et al., 2016) and matrix factorization (Sun and Luo, 2016). Deep linear neural networks satisfy the PŁ inequality in large regions of the parameter space (Charles and Papailiopoulos, 2018, Thm. 4.5), as do two-layer neural networks with an extra identity mapping (Li and Yuan, 2017). Furthermore, sufficiently wide neural networks satisfy the PŁ inequality locally around random initialization (Allen-Zhu et al., 2019a, b; Du et al., 2018, 2019; Liu et al., 2022). While strong convexity implies the PŁ inequality, a PŁ function need not even be convex, hence the PŁ condition is considerably more applicable than strong convexity in the context of neural networks. Moreover, as pointed out by Karimi et al. (2016), the convergence proof for gradient descent assuming strong smoothness and strong convexity is actually simplified if we use the PŁ inequality instead of strong convexity. Thus, we would like to find a simple, elegant proof for the high probability convergence of SGD assuming strong smoothness and the PŁ inequality.

Assuming strong smoothness and strong convexity, the tight mean convergence rate of SGD is $O(1/T)$ (Nemirovski et al., 2009). The same mean convergence rate can be shown assuming strong smoothness and the PŁ inequality (Karimi et al., 2016; Orvieto and Lucchi, 2019). But neither of these papers address high probability convergence. To fill this gap, we prove a $O(\log(1/\delta)/T)$ high probability convergence rate for the objective value of SGD assuming strong smoothness, the PŁ inequality, and norm sub-Gaussian noise. This is given in Theorem 8. The proof relies on a novel probability result, Theorem 9, that essentially shows that adding two kinds of noise to a contracting sequence—sub-Gaussian noise with variance depending on the sequence itself and sub-exponential noise—results in a sub-exponential sequence. Unfortunately, the proof does not extend to the case of norm sub-Weibull noise. It is unclear if this is just an artifact of the proof or if the faster convergence of SGD under the PŁ inequality really cannot be maintained under norm sub-Weibull noise.

Applicability of assumptions. Given the motivation from neural networks, it is necessary to say a few words regarding the applicability of the Lipschitz continuity and strong smoothness assumptions. First, we show in Lemmas 1 and 2 that these assumptions are satisfied by simple neural network models that are, nevertheless, powerful enough to interpolate generic data sets using only twice the minimal number of parameters required for any model to do so. Second, they are satisfied if the iterates remain in a bounded set, which is precisely what happens in the NTK setting. Third, while it is true that Lipschitz continuity and strong smoothness may not be satisfied in practice, our paper is guided by the principle of “how far can we relax assumptions while still being able to prove high probability convergence rates?” An alternative principle is “what can we prove given assumptions that are completely justified in practice?” This is the principle guiding Patel et al. (2022). They provide a neural network counterexample that does not even satisfy $(L_{0},L_{1})$-smoothness (Zhang et al., 2020), an alternative assumption to strong smoothness. Then, they prove—assuming only (1) a lower bound on the objective function, (2) local $\alpha$-Hölder continuity of the gradient, (3) equality of the true and mean gradients, and (4) that the $p$th moment of the norm of the stochastic gradient, for some $p\in(1,2]$, is bounded by an upper semi-continuous function—that, with probability 1, either (1) the norm of the iterates of SGD go to infinity, or (2) SGD almost surely converges to a first order stationary point. These assumptions are much weaker than ours, but so is the conclusion. We leave it as a future research direction to further relax our assumptions.

Organization. Section 2 establishes the optimization framework and reviews optimization, probability, and sub-Weibull terminology and results. Section 3 proves the PŁ convergence rate. Section 4 proves the MDS concentration inequality. Section 5 proves the non-convex convergence rate. Section 6 establishes the post-processing algorithm. Section 7 trains a two layer neural network model on synthetic data with Weibull noise injected into the gradient, showing the dependence of the tail of the convergence error on the tail of the noise.

Notation. We use $x$, $y$, $z$ for vectors and $X$, $Y$, $Z$, $\xi$, $\psi$ for random variables. We use $e_{t}$ for the error vector and distinguish it from Euler’s number with the subscript (the subscript will denote the iteration index). We use $\sigma^{2}$ for the variance and so spell out sigma for sigma algebra. We use $O$ and $\Theta$ for big-$O$ notation. When comparing two sequences of real numbers, $(a_{t})$ and $(b_{t})$: $a_{t}=o(b_{t})$ if $\lim a_{t}/b_{t}=0$, $a_{t}=O(b_{t})$ if $\limsup|a_{t}|/b_{t}<\infty$, and $a_{t}=\Theta(b_{t})$ if $a_{t}=O(b_{t})$ and $b_{t}=O(a_{t})$. We use $[n]$ to denote the set $\{1,\ldots,n\}$ and $\Gamma$ to denote the gamma function.

We are interested in the unconstrained optimization problem

where $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ is a differentiable function, and the SGD iteration

where $g_{t}\in\mathbb{R}^{d}$ is an estimate of $\nabla f(x_{t})$, $\eta_{t}$ is the step-size, and $x_{0}\in\mathbb{R}^{d}$ is the initial point. We restrict our attention to the setting where $x_{0}$ is deterministic, but the results easily extend to the setting where $x_{0}$ is a random vector. We define the noise as the vector $e_{t}=\nabla f(x_{t})-g_{t}$ and assume that, conditioned on the previous iterates, it is unbiased and its norm is sub-Weibull (which, as we will see in Lemma 6, implies that the variance of the noise is bounded).

One of the main examples of Eq. (1) is the stochastic approximation problem: $f=\mathbb{E}\left[F(\cdot,\xi)\right]$ where $(\Omega,\mathcal{F},P)$ is a probability space and $F:\mathbb{R}^{d}\times\Omega\rightarrow\mathbb{R}$ (Nemirovski et al., 2009; Ghadimi and Lan, 2013; Bottou and Bousquet, 2008). In this case, we independently sample $\xi_{0},\xi_{1},\ldots$ and set $g_{t}=\nabla F(x_{t},\xi_{t})$.

Another example is the sample average approximation problem which is the special case of the stochastic approximation problem where there is a finite set $\{\xi_{1},\ldots,\xi_{n}\}$ such that $\xi=\xi_{i}$ with probability $1/n$. In this setting, we define $F_{i}=F(\cdot,\xi_{i})$ so that

In this section, we prove the following convergence bound, proving Theorem 9, a probability result, along the way. The convergence bound matches the optimal convergence rate in mean of $O(1/T)$ and has $\log(1/\delta)$ dependence on $\delta$.

Note that

and, since $\eta_{t}$ is decreasing in $t$ for $t\geq 0$,

so we can apply Eq. (4).

Define

Then multiplying both sides of Eq. (4) by $(t+2\kappa)^{2}$ and noting $1-\mu\eta_{t}=\frac{(t+2\kappa-1)^{2}}{(t+2\kappa)^{2}}$ we get the recursion

Recall that we defined $\mathcal{F}_{t}$ as the sigma algebra generated by $e_{0},\ldots,e_{t}$ for all $t\geq 0$ and $\mathcal{F}_{-1}=\{\emptyset,\Omega\}$. So, $X_{t}$ is $\mathcal{F}_{t-1}$-measurable, $Y_{t}$ is $\mathcal{F}_{t}$-measurable, and $Z_{t}$ is $\mathcal{F}_{t}$-measurable.

For all $t\geq 0$, assume, conditioned on $\mathcal{F}_{t-1}$, that $e_{t}$ is centered and $\|e_{t}\|$ is $\sigma$-sub-Gaussian. Then, $Y_{t}$ is centered conditioned on $\mathcal{F}_{t-1}$. Note that we can bound $Y_{t}$ using Cauchy-Schwarz and Eq. (2) in the following way:

But, if we use this to get a new recursion, then we get a sub-optimal convergence rate. Instead, we need to keep the same recursion but use the bound on $Y_{t}$ in its MGF:

Thus, $Y_{t}$ is $\frac{18L}{\mu^{2}}X_{t}\sigma^{2}$-sub-Gaussian conditioned on $\mathcal{F}_{t-1}$. Similarly,

and so $Z_{t}$ is $\frac{2L}{\mu^{2}}\sigma^{2}$-sub-exponential conditioned on $\mathcal{F}_{t-1}$.

So, we have a recursion for a stochastic sequence with two types of additive noise—sub-Gaussian noise, with variance depending on the sequence itself, and sub-exponential noise. We prove that this makes the main sequence sub-exponential in the following theorem.

The proof is similar to the proof of Theorem 4.1 of Harvey et al. (2019) but with some key differences. There, the recursion is $X_{i+1}\leq\alpha_{i}X_{i}+Y_{i}\sqrt{X_{i}}+\gamma_{i}$ where $\mathbb{E}\left[\exp(\lambda Y_{i})\mid\mathcal{F}_{i-1}\right]\leq\exp\left(\frac{\lambda^{2}}{2}\beta_{i}^{2}\right)$ and $\gamma_{i}$ is deterministic. We had to move the implicit dependence of the sub-Gaussian term inside of the MGF. We also had to allow $\gamma_{i}$ to be a sub-exponential random variable $Z_{i}$, and so applied Cauchy-Schwarz, which contributed the 2 in the recursion for $(K_{i})$.

Note that if $\alpha_{i}=1,\gamma_{i}=\gamma,\beta_{i}=\beta~{}\forall i$, then the assumptions on $(K_{i})$ are satisfied by, for example, $K_{0}=X_{0}$ and $K_{i+1}=K_{i}+2\gamma+\beta^{2}~{}\forall i$. Returning to the proof of Theorem 8, for absolute constants $a$ and $b$, by Proposition 2.5.2 and Proposition 2.7.1 of Vershynin (2018), we can set $\beta_{t}^{2}=\frac{18aL\sigma^{2}}{\mu^{2}}$ and $\gamma_{t}=\frac{2bL\sigma^{2}}{\mu^{2}}$ and apply Theorem 9 with

Dividing by $(T+2\kappa)^{2}$ completes the proof.  

A big take-away from the PŁ analysis is its simplicity, matching the simplicity of gradient descent’s convergence analysis in the same setting. It also serves as a good warm-up for the non-convex analysis that follows. Induction on the convergence recursion does not work for the non-convex analysis. Instead it requires an MDS self-normalized concentration inequality, which we state and prove in the next section.

In this section, we prove the following MDS concentration inequality. For $\theta=1/2$ without $\alpha$ (i.e., $\alpha=0$), Eq. (6), we recover the classical Freedman’s inequality (Freedman, 1975). For $\theta\in(1/2,1]$ without $\alpha$, Eq. (8), we recover Theorem 2.6 of Fan et al. (2015). For $\theta=1/2$ with $\alpha$, Eq. (5), we recover Theorem 3.3 of Harvey et al. (2019), called the “Generalized Freedman” inequality.