The Power of Adaptivity in SGD: Self-Tuning Step Sizes with Unbounded Gradients and Affine Variance
Matthew Faw, Isidoros Tziotis, Constantine Caramanis, Aryan Mokhtari, Sanjay Shakkottai, Rachel Ward
Introduction
Due to its simplicity, an enormous amount of literature, starting by [RM51], has sought to understand convergence guarantees for variants of stochastic gradient descent (SGD):
for minimizing a function using stochastic gradients and a step size schedule . When the (non-convex) objective function is smooth (i.e., has -Lipschitz-continuous gradients) and the stochastic gradients are unbiased and have affine varianceWhile the proof of convergence under affine variance is not given explicitly in [GL13], by slightly modifying the step size choice, the analysis given in this work continues to hold with no additional modifications. Indeed, this observation is made explicitly by [BCN18, Theorem 4.8 ]., i.e.,
then it is well-known that SGD with a properly-tuned step size (depending on and ) converges to a first-order stationary point with error after iterations [GL13, BCN18]. Moreover, [ACDFSW19] showed this rate is tight under these assumptions.
Given these results, it is natural to ask if knowledge of and is necessary to obtain this optimal rate of convergence. Indeed, this has been the motivation for adaptive step size algorithms such as AdaGrad-Norm, where for any parameters , the step size, , is given by
[WWB19] showed that AdaGrad-Norm enjoys a convergence rate even when neither nor is used to tune the step size-schedule. However, their analysis only holds when and the gradients are uniformly upper-bounded – an assumption which is violated even by strongly convex functions such as . In fact, [LO19, Section 4] suggests that, due to the correlation between and in the standard AdaGrad-Norm, the assumption that the gradients are uniformly-bounded might be necessary to prove their convergence guarantee. Although some works on similar adaptive SGD algorithms do not require the gradients to be uniformly upper-bounded [LO19, LO20], their analysis only holds when the step-size is (conditionally) independent of the current stochastic gradient , and require subgaussian noise (a condition which forces ). However, disentangling from is detrimental to the normalization scheme, rendering these methods crucially dependent on the knowledge of the Lipschitz constant for determining their step size.
Extending these results from the bounded variance setting () to the affine variance setting is important. Indeed, results that hold only for the case of bounded variance effectively require that one has noiseless access to gradients when their magnitudes are large (see Remark 1 for more discussion). As opposed to the non-adaptive SGD setting where this extension is immediate (discussed above), in AdaGrad-Norm (and more generally, in adaptive methods), the bias introduced by the correlation between and causes this additional variance to be significantly more problematic.
Our analysis must overcome two main challenges: (i) possibly unbounded gradients, and (ii) an additional bias term introduced by affine variance. Prior work avoided or circumvented these challenges via additional assumptions. Our work requires several new insights that we believe may be of independent interest. Furthermore, as we state in Remark 14, these insights are broadly applicable to related adaptive algorithms such as coordinate-wise AdaGrad. We outline these below.
Main Challenge 2: Additional bias from affine variance.
In the affine variance setting, the expected difference in function value between consecutive time steps is bounded as:
Related Work. [GL13] were the first to study the convergence of SGD for opimizing a non-convex, smooth objective function. They proved that a properly-tuned SGD converges to a first-order stationary point at rate , if the step sizes are chosen as for a constant . Further, [ACDFSW19] proved that the rate is unimprovable for any algorithm with only first-order oracle access, assuming the function is non-convex, smooth, and the stochastic gradients are unbiased with bounded variance.
The original AdaGrad algorithm was proposed simultaneously by [DHS11, MS10] whereas [SM10] were the first to consider a variant of AdaGrad referred to as AdaGrad-Norm. [WWB19] analyzed AdaGrad-Norm for minimizing a smooth, non-convex function with uniformly-bounded gradients. They showed that AdaGrad-Norm converges at essentially the same rate as SGD, but without the need to know the smoothness constant (albeit under the restrictive assumption that the gradients are uniformly upper-bounded). In a simultaneous work, [LO19] studied a variant of AdaGrad-Norm where step size is conditionally independent of the current stochastic gradient , unlike in the standard AdaGrad setting. They provided a similar convergence guarantee without needing a uniform upper-bound on the stochastic gradients, but requiring that the noise have bounded support and additionally requiring knowledge of the smoothness parameter to tune their step sizes. In a followup work [LO20], the same authors proved high-probability convergence of a class of adaptive algorithms (including their variant of AdaGrad-Norm, as well as coordinate-wise AdaGrad with momentum) under the assumption of subgaussian noise. Note that, like the earlier result, their step sizes needed to be tuned with knowledge of the smoothness parameter, and further needed to be conditionally independent of the current gradient. [KLC22] established high probability results for AdaGrad without knowledge of the smoothness parameter in the bounded variance regime, assuming that the norm of the gradients are uniformly upper-bounded (i.e., the objective function is Lipschitz). They were further able to remove the Lipschitz assumption, but only when in addition to bounded variance, the noise of the stochastic gradients is subgaussian. [GG20] studied the asymptotic convergence of AdaGrad (as well as and RMSProp), where their analysis requires uniform gradient bounds as well as uniform bounds on the nd and th moments of the gradient noise. Very recently, [JXH22] established asymptotic almost-sure convergence of the AdaGrad-Norm iterates to first-order stationary points. Unlike our work, they do not provide rates of convergence, and their focus on asymptotics makes their analysis and results significantly different. [ZSJSL18] studied a weighted version of coordinate-wise AdaGrad with momentum, where they assumed the gradients were uniformly bounded. [DBBU20] later improved upon these results with respect to the dependence on the momentum parameter.
Several recent works have studied the convergence of other adaptive algorithms, all of which are based on the assumption of uniformly-bounded stochastic gradients. For instance, [KLBC19] developed an adaptive, accelerated algorithm that achieves optimal rates in the constrained, convex (smooth and non-smooth) regime, without knowledge of the smoothness or noise parameters. [CLSH18] studied the convergence of a class of Adam-like algorithms (originally introduced by [KB15]). Later, building on the results of [WWB19], [DBBU20] improved on this analysis of Adam with respect to the dependence on the momentum parameter and range of valid hyperparameters. [GXYJY21] provide an alternate analysis of a class of Adam-like algorithms for different momentum parameter scaling. [SMBM21] studied “delayed” versions of Adam (as well as a new algorithm they called AvaGrad), which makes the step sizes conditionally independent of the current stochastic gradient, .
Preliminaries
(Motivation for Affine Variance) This scaling is important for machine learning applications with feature noise (including missing features) [Ful09, KL20], in robust linear regression [XCM08], and generally whenever the model parameters are multiplicatively perturbed by noise (e.g., a multilayer network, where noise from a previous layer multiplies the parameters in subsequent layers). More broadly, restricting to bounded variance (i.e., assuming ) is equivalent to assuming “noiseless” access to the gradient when the magnitude of the gradient grows (e.g., a strongly convex function); this is because the stochastic gradient is an arbitrarily small perturbation of the true gradient in this regime. Finally, as discussed earlier, the analysis for non adaptive SGD is essentially unaffected by affine variance [BCN18].
Further, we will assume that the function is -smooth:
A key property of AdaGrad-Norm is that the step-size sequence is tightly controlled:
In fact, variations of this observation have been noted for a number of AdaGrad variants [WWB19, DBBU20]. While simple, it is crucially important to our analysis, since, taken together with Assumption 3, it implies that the gradient at time scales at most polynomially in .
Consider any times during a run of algorithm (AG-Norm). Then, deterministically,
Moreover, with probability at least , the following bound also holds
As a consequence of Lemma 2, we derive deterministically, and an analogous bound of with probability . Of course, Lemma 2 only gives a much weaker control over than a uniform bound, and has not (to the best of our knowledge) been previously exploited. However loose, this bound nonetheless is one of the key steps to removing the uniform gradient bound, and may be of independent interest (e.g., useful for refining the convergence rates for strongly convex problems).
As mentioned earlier, a key difficulty in analyzing adaptive algorithms is the bias introduced by the correlation between the step size and the stochastic gradient at each time . To analyze the convergence of such algorithms, it is useful to introduce the following “decorrelated” step size.
Motivating the Proof
We have discussed the two main challenges in Section 1.1: unbounded gradients and affine variance. Now that we have the required mathematical definitions from Section 2, we discuss these challenges in more detail. Adaptive stochastic gradient methods exhibit two difficulties not present in the non-adaptive regime: (i) Since the step size depends on the trajectory of stochastic gradients, one must argue about the scaling of these stochastic gradients, and (ii) the step size is correlated with the current gradient, , as well as the past gradients. These manifest themselves as follows: by -smoothness (Assumption 3) and the AdaGrad-Norm algorithm (AG-Norm), we have that
When and are conditionally independent, then the inner product term above is mean-zero. As a consequence, as long as the step size , (5) immediately implies that
Although in the non-adaptive setting, we could simply choose , in the adaptive regime it is no longer obvious that such a condition holds. One may observe, however, that by Jensen’s inequality and Definition 3
Affine Variance: Upper-bounding the bias.
The bias term in (3) presents another difficulty in analyzing the rate of convergence in the adaptive setting. Specifically, in the affine variance setting
Main Results
In this section, we sketch out the key ideas that go into deriving a bound on the convergence rate of AdaGrad-Norm to a first order stationary point. Our main result is the following:
With probability at least , the iterates of (AG-Norm) satisfy:
where .We use the notation to mean for some absolute constant independent of all problem parameters. Moreover, when , then with probability at least ,
where , B\propto(1+\sigma_{1}^{\nicefrac{{3}}{{2}}})(b_{0}+\sigma_{0}+\left\lVert\nabla F(\mathbf{w}_{1})\right\rVert+\eta L{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}+\nicefrac{{F(\mathbf{w}_{1})-F^{*}}}{{\eta}})^{2}}, and C^{\prime}\propto{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}(1+\sigma_{1}^{2})(b_{0}+\sigma_{0}+\eta L+\nicefrac{{F(\mathbf{w}_{1})-F^{*}}}{{\eta}})^{2}}.
4 demonstrates two interesting regimes for our guarantee. Namely, (11) shows a convergence rate for any choices of , thus establishing our parameter-free guarantee. However, this bound does not recover the convergence rate in the “small-noise” regime. Through a minor modification to the proof technique used to obtain (11), we are able to derive (12), which demonstrates that (AG-Norm) recovers an rate of convergence when – the rate obtainable by a well-tuned gradient descent in the noiseless regime up to logarithmic factors. We emphasize that (AG-Norm) does not require a priori knowledge of the smoothness parameter or the variance parameters to obtain either of the convergence rates in (11) or (12). Indeed, (AG-Norm) adapts automatically to obtain the faster rate in the “small-noise” regime.
As displayed in (9), the affine variance scaling introduces additional bias that our analysis must handle. Indeed, this bound taken together with (3) implies the following lemma.
Recall the step size proxy of Definition 3 and the notation in Definition 6. With , we obtain
Key Idea: Compensating for the “bad” times.
Lemma 7 shows that, even when we focus on the good times, we still must argue about the deviations at bad times to obtain a convergence guarantee. In order to address this problem, we begin by rewriting Lemma 7 by: (i) upper bounding the “bad” times using the (potentially quite large) bound obtained from Lemma 5, and (ii) subtracting some of the “good” deviation terms from both sides to compensate for the bad terms.
In the same setting as Lemma 7, we have that
2 Bounding the Expected Sum of Gradients via Recursive Improvement
Suppose that, for some parameters , , a sufficiently large polynomial function of , and sufficiently large constant ,
In particular, as a consequence of Lemma 2,
3 Wrapping up
With these bounds from Lemmas 12 and 13 in place, obtaining the convergence result for (AG-Norm) in 4 is immediate. Indeed, we note that Lemma 12 gives us essentially the same bound as the one obtainable in the uniformly-bounded variance case (10) (modulo the summation over the set instead of all times ). Therefore, we may apply (essentially) the same Hölder’s inequality argument as in [WWB19], replacing their application of the uniform gradient bound with our bound on the expected sum of gradients from Lemma 13, and taking extra care that our summation from Lemma 12 is over a random set . We give the full proof of this theorem in Appendix F.
While we focus in this paper on the convergence rate of one particular adaptive SGD method, our methods are not overly specialized to AdaGrad-Norm. Indeed, using nearly identical arguments per coordinate, we can obtain similar convergence rates under similar assumptions for coordinate-wise AdaGrad, albeit with an additional polynomial dependence on .
Conclusion
Acknowledgements
This research is supported in part by NSF Grants 1952735, 1934932, 2019844, 2127697, and 2112471, ARO Grant W911NF2110226, AFOSR MURI FA9550-19-1-0005, the Machine Learning Lab (MLL) at UT Austin, and the Wireless Networking and Communications Group (WNCG) Industrial Affiliates Program.
References
Appendix A Preliminaries
Here, we provide proofs for claims from Section 2, as well as some auxiliary results and notation. We additionally state some definitions that will be useful for proving our results.
For any sequence such that and for all ,
The base case of holds with equality. Let us now assume that the claim holds at . Then, we have that
where the first inequality holds by the induction hypothesis, and the second because of the fact (where denotes the natural logarithm). ∎
Our analysis will focus on adaptive gradient algorithms with a particularly convenient structure, which we refer to as the Bounded Step-Size Property
We say that an optimization algorithm has -Bounded Step-Sizes if, for any pair of adjacent iterates generated by the algorithm, the following inequality holds deterministically:
Another convenient property of the algorithms we study is what we call the Decay Property:
We say that an optimization algorithm satisfies the -Decay Property if the iterate sequence satisfies the following inequality deterministically:
We observe that these property is satisfied by a number of interesting adaptive gradient algorithms.
AdaGrad-Norm has -Bounded Step-Sizes and -Decay. The first follows since for any time ,
The second is an immediate consequence of Lemma 15, taking and for .
Coordinate-wise AdaGrad (with coordinate-dependent step sizes
has -Bounded Step-Sizes and -Decay. The first follows since since for every coordinate . The second follows by applying Lemma 15 to the sum of for each coordinate.
We note here that all of the remaining results in this section could be stated in more generality by using Definitions 16 and 17. To showcase our ideas in the simplest manner, we will state everything in the context of the AdaGrad-Norm algorithm (AG-Norm).
By Assumption 3 and 18, we also have the following simple, but quite useful, facts, which give us crude but, crucially, polynomial (in ) bound on :
Consider the AdaGrad-Norm algorithm (AG-Norm) running on an -smooth objective function . Then, for any times ,
The proof follows by first applying the triangle inequality and using a telescoping sum to bound
then noting that, for each , by Assumptions 3 and 18,
The above bound on is quite useful, since it guarantees a polynomial (in ) bound for . However, note that this bound is much more crude than the bound assumed by [WWB19, DBBU20] (where they assumed for every ). It turns out that, on “nice” sample paths, a significantly tighter bound can be derived. Intuitively, these sample paths are those for which the quantity is bounded by a polynomial in
For any time s\in{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\{0\}\cup}[T] and failure probability , we define the following “nice event”:
and take for .
As we will soon see, bounding the quantity will be crucial in many parts of our analysis. Under the “nice” events from Definition 22, this quantity can be easily controlled:
For any choice of , and any sample path, (AG-Norm) satisfies
Further, assuming that the “nice event” (20) () is true at time , and taking as in (21),
In particular, since is (trivially) always true, the above implies that
Additionally, when (the nice event at time ) is true,
We already established (22) in 18. For the remaining inequalities, we may assume without loss of generality that . Indeed, whenever , then by Definition 22, and thus is never true, so all of the claims follow trivially.
To show (23), we note that, on any sample path, by (22) and Jensen’s inequality,
To bound this term above, first observe that, as noted in (3), Assumptions 1 and 2 imply that
Further, when (20) () is true at time , we have that, by Lemma 21,
Combining the above bounds, we conclude that
as claimed. Finally, observe that (24) and (25) follow immediately from (23), taking (noting that is true deterministically) and , respectively. ∎
With the above construction in place, we are ready to give a slightly stronger bound for , improving upon Lemma 21 (with high probability) in many interesting regimes.
Consider any time during a run of (AG-Norm) initialized at a starting point , and is currently at iterate . Then,
and additionally, assuming that from Definition 22 is true, and taking as in (21), then
The proof follows effectively from the same arguments used to prove Lemma 21, only using the improved bound from Lemma 23 in place of Lemma 21. Indeed, using the same decomposition, and applying Cauchy-Schwarz, we have that
where the first inequality follows from the decomposition used in the proof of Lemma 21, the second follows by Cauchy-Schwarz, and the third from Lemma 15.
The second claim follows immediately from the above, combined with Lemma 23. ∎
Appendix B Deriving the Starting Point
Here, we provide the proof for the starting point of our analysis, Lemma 5, from Section 4.
We will begin by using our assumption of -smoothness, along with the definition of the algorithm, to get the bound:
Hence, by taking expectations of our first inequality and adding this mean-zero quantity to the resulting expression, we have that
We will now focus on bounding the second term. Observe that, denoting and ,
Plugging this bound into the above, and taking expectation with respect to the filtration at , we have shown that
We will now show that the second and third terms above have the same upper bound. Focus on the second term above, we apply Hölder’s inequality and the affine variance assumption to conclude that
Now, focusing on the third term, by Jensen’s inequality to the concave function , we know that
which show that the second and third terms have exactly the same upper bound. Combining these expressions and rearranging, we find
Appendix C Most Times are (Typically) Good
Here, we provide proofs regarding properties and consequences of the “good” times (Definition 6) from Section 4.
Recalling the step size proxy of Definition 3 and the notation in Definition 6, we obtain
where , and is the function defined in (21).
Observe that an equivalent condition for a time to be “good” in the sense of Definition 6 is:
Now, summing the above expression over all times , and applying Lemma 23, we find that
Now, by (23) in Lemma 23, we know that, whenever is true, then
By (27), we additionally know that, for each time ,
Appendix D Compensating for “Bad” Time-Steps
Here, we provide proofs for the compensation arguments presented in Section 4
Let us begin by proving that, for any times ,
The claim is trivial when so we focus on the case when . Let us denote and . Then, observe that
Therefore, we can observe that the step sizes are sufficiently close, since
where the last line follows by Lemma 21. We will now use this observation in order to prove the claimed inequality. We will proceed by considering two cases.
In the first case, if , then by Lemma 21, . This implies that
In the alternative case, when , then
Now, we will use Lemma 10 to bound the first term above. We will use the trivial bound for the second term: by Definition 3 and Lemma 21, we may bound each term inside of the sum of the second expression above as:
where
The result follows immediately by combining Lemmas 29 and 27. Note that this result, up to logarithmic factors, takes essentially the same form as in the uniformly-bounded setting (10). ∎
Appendix E Bounding the Expected Sum of Gradients via Recursive Improvement
Here, we provide a proof for the recursive improvement argument presented in Section 4.
Suppose that, for some constants and , the following inequality is true:
Then, in fact, the following tighter bound also holds:
In particular, as a consequence of Lemmas 24 and 23,
The proof will proceed in three steps, in which we will invoke the auxiliary Lemmas 32, 33 and 34. It is straightforward to verify that the constant specified in this lemma, as well as the choice of , satisfy the constraints from those lemmas. Thus, we are free to use these results to prove our desired result.
In order to “remove” the indicator from the expectation above, we will need to show that, when is false, cannot be too large. Recall that we have two main tools to upper bound the size of this sum: Lemma 21, which gives a deterministic upper bound of , and Lemma 24, which gives a high-probability upper bound of . These insights allow us to prove Lemma 33, which tells us that
Step 2: Bounding the “good” terms.
With the indicator removed from the above expression, we are now ready to use Lemma 30 together with (40) and (41) to obtain a bound on the expected size of the gradients at the good times:
where the second inequality follows by upper bounding Hence, by choosing and Note that these choices of satisfy the requirements of Lemmas 32, 33 and 34. Indeed, by construction. Further, since , we have that and, whenever , , and when , . Finally, since ., we conclude that
Step 3: Bounding the “bad” terms.
E.2 Technical Lemmas
We divide the proof in two cases: (1) , and (2) . In the first case, the claimed result (45) holds trivially, since by definition (see Definition 22), and thus,
Let us assume that (the “nice” event from Definition 22) is true. Then, we have that
where the first inequality follows by definition of and by the assumed bound (44). The second inequality follows since , and the third since and . The final inequality follows by plugging in our choice of , and using the fact that .
Now, since , the above inequality implies that
where the first inequality follows since and by Lemma 21, and the second since and .
Now, in order to “remove” the indicator from the expectation, we will need to show that, when is false, cannot be too large. Recall that we have two main tools to upper bound the size of this sum: Lemma 21, which gives a deterministic upper bound of , and Lemma 24, which gives a high-probability upper bound of . To exploit this “lighter” regime of Lemma 24, it will be useful to introduce the following event:
where is the same choice as in Lemma 32. By definition, , so
where in the last inequality, we use the following facts: (chosen in Lemma 32) and which hold since , , and by the initial conditions on . Now, since by assumption (which implies that ), we may simplify the above to conclude that
By our assumption that , the claimed bound is immediate. ∎
Therefore, choosing , and assuming that , we conclude that
Appendix F Obtaining the Convergence Rate for AdaGrad-Norm
Here, we provide a proof for the main result of this paper, a proof of convergence for the AdaGrad-Norm algorithm.
With probability at least , the AdaGrad-Norm algorithm (AG-Norm) for any choice of parameters satisfies:
Furthermore, whenever , then with probability at least , (AG-Norm) also satisfies:
where .
We note that the second bound in 35, (35), is particularly interesting in the regime when . Indeed, in this setting, our bound yields a convergence rate which one should expect in the noiseless regime.
Focusing now on lower bounding the numerator of (49),
where the lower bound above follows since the average is always larger than the minimum. If it were the case that , then, at this point, we would essentially be done with our proof. However, since is a random set, we must take some additional care. Because is in expectation by Lemma 26, this is only a minor technicality. Indeed,
Therefore, collecting the results we have derived so far into a lower bound on the right-hand side of (49), and applying the result of Lemma 30 to upper bound the left-hand side of (49), we have obtained the following upper bound:
where the inequality follows since . The above failure probability can be easily upper bounded via Markov’s inequality:
Hence, by a final application of Markov’s inequality, we obtain, for any ,
Solving this quadratic inequality, we conclude that
In particular, this implies by Markov’s inequality that, with probability at least ,
This shows that, in the setting when , then we recover a convergence rate, as in the noiseless setting. ∎