Convergence of Adam Under Relaxed Assumptions

Introduction

In this paper, we study the non-convex unconstrained stochastic optimization problem

The Adaptive Moment Estimation (Adam) algorithm (Kingma and Ba, 2014) has become one of the most popular optimizers for solving (1) when $f$ is the loss for training deep neural networks. Owing to its efficiency and robustness to hyper-parameters, it is widely applied or even sometimes the default choice in many machine learning application domains such as natural language processing (Vaswani et al., 2017; Brown et al., 2020; Devlin et al., 2019), generative adversarial networks (Radford et al., 2015; Isola et al., 2016; Zhu et al., 2017), computer vision (Dosovitskiy et al., 2020), and reinforcement learning (Lillicrap et al., 2015; Mnih et al., 2016; Schulman et al., 2017). It is also well known that Adam significantly outperforms stochastic gradient descent (SGD) for certain models like transformer (Zhang et al., 2020b; Kunstner et al., 2023; Ahn et al., 2023).

Despite its success in practice, theoretical analyses of Adam are still limited. The original proof of convergence in (Kingma and Ba, 2014) was later shown by Reddi et al. (2018) to contain gaps. The authors in (Reddi et al., 2018) also showed that for a range of momentum parameters chosen independently with the problem instance, Adam does not necessarily converge even for convex objectives. However, in deep learning practice, the hyper-parameters are in fact problem-dependent as they are usually tuned after given the problem and weight initialization. Recently, there have been many works proving the convergence of Adam for non-convex functions with various assumptions and problem-dependent hyper-parameter choices. However, these results leave significant room for improvement. For example, (D’efossez et al., 2020; Guo et al., 2021) prove the convergence to stationary points assuming the gradients are bounded by a constant, either explicitly or implicitly. On the other hand, (Zhang et al., 2022; Wang et al., 2022) consider weak assumptions, but their convergence results are still limited. See Section 2 for more detailed discussions of related works.

To address the above-mentioned gap between theory and practice, we provide a new convergence analysis of Adam without assuming bounded gradients, or equivalently, Lipschitzness of the objective function. In addition, we also relax the standard global smoothness assumption, i.e., the Lipschitzness of the gradient function, as it is far from being satisfied in deep neural network training. Instead, we consider a more general, relaxed, and non-uniform smoothness condition according to which the local smoothness (i.e., Hessian norm when it exists) around $x$ is bounded by a sub-quadratic function of the gradient norm $\left\|\nabla f(x)\right\|$ (see Definition 3.2 and Assumption 2 for the details). This generalizes the $(L_{0},L_{1})$ smoothness condition proposed by (Zhang et al., 2019) based on language model experiments. Even though our assumptions are much weaker and more realistic, we can still obtain the same $\mathcal{O}(\epsilon^{-4})$ gradient complexity for convergence to an $\epsilon$-stationary point.

The key to our analysis is a new technique to obtain a high probability, constant upper bound on the gradients along the optimization trajectory of Adam, without assuming Lipschitzness of the objective function. In other words, it essentially turns the bounded gradient assumption into a result that can be directly proven. Bounded gradients imply bounded stepsize at each step, with which the analysis of Adam essentially reduces to the simpler analysis of AdaBound (Luo et al., 2019). Furthermore, once the gradient boundedness is achieved, the analysis under the generalized non-uniform smoothness assumption is not much harder than that under the standard smoothness condition. We will introduce the technique in more details in Section 5. We note that the idea of bounding gradient norm along the trajectory of the optimization algorithm can be use in other problems as well. For more details, we refer the reader to our concurrent work (Li et al., 2023) in which we present a set of new techiniques and methods for bounding gradient norm for other optimization algorithms under a generalized smoothness condition.

Another contribution of this paper is to show that the gradient complexity of Adam can be further improved with variance reduction methods. To this end, we propose a variance-reduced version of Adam by modifying its momentum update rule, inspired by the idea of the STORM algorithm (Cutkosky and Orabona, 2019). Under additional generalized smoothness assumption of the component function $f(\cdot,\xi)$ for each $\xi$, we show that this provably accelerates the convergence with a gradient complexity of $\mathcal{O}(\epsilon^{-3})$. This rate improves upon the existing result of (Wang and Klabjan, 2022) where the authors obtain an asymptotic convergence of their approach to variance reduction for Adam in the non-convex setting, under the bounded gradient assumption.

In light of the above background, we summarize our main contributions as follows.

We develop a new analysis to show that Adam converges to stationary points under relaxed assumptions. In particular, we do not assume bounded gradients or Lipschitzness of the objective function. Furthermore, we also consider a generalized non-uniform smoothness condition where the local smoothness or Hessian norm is bounded by a sub-quadratic function of the gradient norm. Under these more realistic assumptions, we obtain a dimension free gradient complexity of $\mathcal{O}(\epsilon^{-4})$ if the gradient noise is centered and bounded.

We generalize our analysis to the setting where the gradient noise is centered and has sub-Gaussian norm, and show the convergence of Adam with a gradient complexity of $\mathcal{O}(\epsilon^{-4}\log^{3.25}(1/\epsilon))$.

We propose a variance-reduced version of Adam (VRAdam) with provable convergence guarantees. In particular, we obtain the accelerated $\mathcal{O}(\epsilon^{-3})$ gradient complexity.

Related work

In this section, we discuss the relevant literature related to different aspects of our work.

Adam was first proposed by Kingma and Ba (2014) with a theoretical convergence guarantee for convex functions. However, Reddi et al. (2018) found a gap in the proof of this convergence analysis, and also constructed counter-examples for a range of hyper-parameters on which Adam does not converge. That being said, the counter-examples depend on the hyper-parameters of Adam, i.e., they are constructed after picking the hyper-parameters. Therefore, it does not rule out the possibility of obtaining convergence guarantees for problem-dependent hyper-parameters, as also pointed out by (Shi et al., 2021; Zhang et al., 2022).

Many recent works have developed convergence analyses of Adam with various assumptions and hyper-parameter choices. Zhou et al. (2018b) show Adam with certain hyper-parameters can work on the counter-examples of (Reddi et al., 2018). De et al. (2018) prove convergence for general non-convex functions assuming gradients are bounded and the signs of stochastic gradients are the same along the trajectory. The analysis in (D’efossez et al., 2020) also relies on the bounded gradient assumption. Guo et al. (2021) assume the adaptive stepsize is upper and lower bounded by two constants, which is not necessarily satisfied unless assuming bounded gradients or considering the AdaBound variant (Luo et al., 2019). (Zhang et al., 2022; Wang et al., 2022) consider very weak assumptions. However, they show either 1) “convergence” only to some neighborhood of stationary points with a constant radius, unless assuming the strong growth condition; or 2) convergence to stationary points but with a slower rate.

After Reddi et al. (2018) pointed out the non-convergence issue with Adam, various variants of Adam that can be proved to converge were proposed (Zou et al., 2018; Gadat and Gavra, 2022; Chen et al., 2018, 2022; Luo et al., 2019; Zhou et al., 2018b). For example, AMSGrad (Reddi et al., 2018) and AdaFom (Chen et al., 2018) modify the second order momentum so that it is non-decreasing. AdaBound (Luo et al., 2019) explicitly imposes upper and lower bounds on the second order momentum so that the stepsize is also bounded. AdaShift (Zhou et al., 2018b) uses a new estimate of the second order momentum to correct the bias. There are also some works (Zhou et al., 2018a; Gadat and Gavra, 2022; Iiduka, 2023) that provide convergence guarantees of these variants. One closely related work to ours is (Wang and Klabjan, 2022), which considers a variance-reduced version of Adam by combining Adam and SVRG (Johnson and Zhang, 2013). However, they assume bounded gradients and can only get an asymptotic convergence in the non-convex setting.

Generalizing the standard smoothness condition in a variety of settings has been a focus of many recent papers. Recently, (Zhang et al., 2019) proposed a generalized smoothness condition called $(L_{0},L_{1})$ smoothness, which assumes the local smoothness or Hessian norm is bounded by an affine function of the gradient norm. The assumption was well-validated by extensive experiments conducted on language models. Various analyses of different algorithms under this condition were later developed (Zhang et al., 2020a; Qian et al., 2021; Zhao et al., 2021; Faw et al., 2023; Reisizadeh et al., 2023; Crawshaw et al., 2022). One recent closely-related work is (Wang et al., 2022) which studies converges of Adam under the $(L_{0},L_{1})$ smoothness condition. However, their results are still limited, as we have mentioned above. In this paper, we consider an even more general smoothness condition where the local smoothness is bounded by a sub-quadratic function of the gradient norm, and prove the convergence of Adam under this condition. In our concurrent work (Li et al., 2023), we further analyze various other algorithms in both convex and non-convex settings under similar generalized smoothness conditions following the same key idea of bounding gradients along the trajectory.

The technique of variance reduction was introduced to accelerate convex optimization in the finite-sum setting (Roux et al., 2012; Johnson and Zhang, 2013; Shalev-Shwartz and Zhang, 2013; Mairal, 2013; Defazio et al., 2014). Later, many works studied variance-reduced methods in the non-convex setting and obtained improved convergence rates for standard smooth functions. For example, SVRG and SCSG improve the $\mathcal{O}(\epsilon^{-4})$ gradient complexity of stochastic gradient descent (SGD) to $\mathcal{O}(\epsilon^{-10/3})$ (Allen-Zhu and Hazan, 2016; Reddi et al., 2016; Lei et al., 2017). Many new variance reduction methods (Fang et al., 2018; Tran-Dinh et al., 2019; Liu et al., 2020; Li et al., 2021; Cutkosky and Orabona, 2019; Liu et al., 2023) were later proposed to further improve the complexity to $\mathcal{O}(\epsilon^{-3})$, which is optimal and matches the lower bound in (Arjevani et al., 2023). Recently, (Reisizadeh et al., 2023; Chen et al., 2023) obtained the $\mathcal{O}(\epsilon^{-3})$ complexity for the more general $(L_{0},L_{1})$ smooth functions. Our variance-reduced Adam is motivated by the STORM algorithm proposed by (Cutkosky and Orabona, 2019), where an additional term is added in the momentum update to correct the bias and reduce the variance.

Preliminaries

The formal definition of Adam proposed in (Kingma and Ba, 2014) is shown in Algorithm 1, where Lines 5–9 describe the update rule of iterates $\{x_{t}\}_{1\leq t\leq T}$. Lines 5–6 are the updates for the first and second order momentum, $m_{t}$ and $v_{t}$, respectively. In Lines 7–8, they are re-scaled to $\hat{m}_{t}$ and $\hat{v}_{t}$ in order to correct the initialization bias due to setting $m_{0}=v_{0}=0$. Then the iterate is updated by $x_{t+1}=x_{t}-h_{t}\odot\hat{m}_{t}$ where $h_{t}=\eta/(\sqrt{\hat{v}_{t}}+\lambda)$ is the adaptive stepsize vector for some parameters $\eta$ and $\lambda$.

2 Assumptions

In what follows below, we will state our main assumptions for analysis of Adam.

We start with a standard assumption in optimization on the objective function $f$ whose domain lies in a Euclidean space with dimension $d$.

Besides Assumption 1, the only additional assumption we make regarding $f$ is that its local smoothness is bounded by a sub-quadratic function of the gradient norm. More formally, we consider the following $(\rho,L_{0},L_{\rho})$ smoothness condition with $0\leq\rho<2$.

A differentiable real-valued function $f$ is $(\rho,L_{0},L_{\rho})$ smooth for some constants $\rho,L_{0},L_{\rho}\geq 0$ if the following inequality holds almost everywhere in $\operatorname*{dom}(f)$

When $\rho=1$ , Definition 3.2 reduces to the $(L_{0},L_{1})$ smoothness condition in (Zhang et al., 2019). When $\rho=0$ or $L_{\rho}=0$, it reduces to the standard smoothness condition.

The objective function $f$ is $(\rho,L_{0},L_{\rho})$ smooth with $0\leq\rho<2$.

The standard smooth function class is very restrictive as it only contains functions that are upper and lower bounded by quadratic functions. The $(L_{0},L_{1})$ smooth function class is more general since it also contains, e.g., univariate polynomials and exponential functions. Assumption 2 is even more general and contains univariate rational functions, double exponential functions, etc. See Appendix B.1 for the formal propositions and proofs. We also refer the reader to our concurrent work (Li et al., 2023) for more detailed discussions of examples of $(\rho,L_{0},L_{\rho})$ smooth functions for different $\rho$s.

It turns out that bounded Hessian norm at a point $x$ implies local Lipschitzness of the gradient in the neighborhood around $x$. In particular, we have the following lemma.

Lemma 3.4 can be actually used as the definition of $(\rho,L_{0},L_{\rho})$ smooth functions in place of Assumption 2. Besides the local gradient Lipschitz condition, it also suggests that, as long as the update at each step is small enough, the iterates will not go outside of the domain.

For the special case of $\rho=1$, choosing $a=\max\{\left\|\nabla f(x)\right\|,L_{0}/L_{1}\}$, one can verify that the required locality size in Lemma 3.4 satisfies $\frac{a}{L_{0}+L_{1}(\left\|\nabla f(x)\right\|+a)}\geq\frac{1}{3L_{1}}$. In this case, Lemma 3.4 states that $\left\|x-y\right\|\leq 1/(3L_{1})$ implies $\left\|\nabla f(y)-\nabla f(x)\right\|\leq 2(L_{0}+L_{1}\left\|\nabla f(x)\right\|)\left\|y-x\right\|.$ Therefore, it reduces to the local gradient Lipschitz condition for $(L_{0},L_{1})$ smooth functions in (Zhang et al., 2019, 2020a) up to numerical constant factors. For $\rho\neq 1$, the proof is more involved because Grönwall’s inequality used in (Zhang et al., 2019, 2020a) no longer applies. Therefore we defer the detailed proof of Lemma 3.4 to Appendix B.2.

2.2 Stochastic gradient

We consider one of the following two assumptions on the stochastic gradient $\nabla f(x_{t},\xi_{t})$ in our analysis of Adam.

The gradient noise is centered and almost surely bounded. In particular, for some $\sigma\geq 0$ and all $t\geq 1$,

The gradient noise is centered with sub-Gaussian norm. In particular, for some $R\geq 0$ and all $t\geq 1$,

Results

In the section, we provide our convergence results for Adam under Assumptions 1, 2, and 3 or 4. To keep the statements of the theorems concise, we first define several problem-dependent constants. First, we let $\Delta_{1}:=f(x_{1})-f^{*}<\infty$ be the initial sub-optimality gap. Next, given a large enough constant $G>0$, we define

where $L$ can be viewed as the effective smoothness constant along the trajectory if one can show $\left\|\nabla f(x_{t})\right\|\leq G$ and $\left\|x_{t+1}-x_{t}\right\|\leq r$ at each step (see Section 5 for more detailed discussions). We will also use $c_{1},c_{2}$ to denote some small enough numerical constants and $C_{1},C_{2}$ to denote some large enough ones. The formal convergence results under Assumptions 1, 2, and 3 are presented in the following theorem, whose proof is deferred in Appendix C.

Suppose Assumptions 1, 2, and 3 hold. Denote $\iota:=\log(1/\delta)$ for any $0<\delta<1$. Let $G$ be a constant satisfying $G\geq\max\left\{2\lambda,2\sigma,\sqrt{C_{1}\Delta_{1}L_{0}},(C_{1}\Delta_{1}L_{\rho})^{\frac{1}{2-\rho}}\right\}$. Choose

Let $T=\max\left\{\frac{1}{\beta^{2}},\;\frac{C_{2}\Delta_{1}G}{\eta\epsilon^{2}}\right\}$. Then with probability at least $1-\delta$, we have $\left\|\nabla f(x_{t})\right\|\leq G$ for every $1\leq t\leq T$, and $\frac{1}{T}\sum_{t=1}^{T}\left\|\nabla f(x_{t})\right\|^{2}\leq\epsilon^{2}$.

Note that $G$, the upper bound of gradients along the trajectory, is a constant that depends on $\lambda,\sigma,L_{0},L_{\rho}$, and the initial sub-optimality gap $\Delta_{1}$, but not on $\epsilon$. There is no requirement on the second order momentum parameter $\beta_{\operatorname{sq}}$, although many existing works like (D’efossez et al., 2020; Zhang et al., 2022; Wang et al., 2022) need certain restrictions on it. We choose very small $\beta$ and $\eta$, both of which are $\mathcal{O}(\epsilon^{2})$. Therefore, from the choice of $T$, it is clear that we obtain a gradient complexity of $\mathcal{O}(\epsilon^{-4})$, where we only consider the leading term. We are not clear whether the dependence on $\epsilon$ is optimal or not, as the $\Omega(\epsilon^{-4})$ lower bound in (Arjevani et al., 2023) assumes the weaker bounded variance assumption than our Assumpion 3. However, it matches the state-of-the-art complexity among existing analyses of Adam.

One limitation of the dependence of our complexity on $\lambda$ is $\mathcal{O}(\lambda^{-2})$, which might be large since $\lambda$ is usually small in practice, e.g., the default choice is $\lambda=10^{-8}$ in the PyTorch implementation. There are some existing analyses on Adam (D’efossez et al., 2020; Zhang et al., 2022; Wang et al., 2022) whose rates do not depend explicitly on $\lambda$ or only depend on $\log(1/\lambda)$. However, all of them depend on $\operatorname*{poly}(d)$, whereas our rate is dimension free. The dimension $d$ is also very large, especially when training transformers, for which Adam is widely used. We believe that independence on $d$ is better than that on $\lambda$, because $d$ is fixed given the architecture of the neural network but $\lambda$ is a hyper-parameter which we have the freedom to tune. In fact, based on our preliminary experimental results on CIFAR-10 shown in Figure 1, the performance of Adam is not very sensitive to the choice of $\lambda$. Although the default choice of $\lambda$ is $10^{-8}$, increasing it up to $0.01$ only makes minor differences.

As discussed in Section 3.2.2, we can generalize the bounded gradient noise condition in Assumption 3 to the weaker sub-Gaussian noise condition in Assumption 4. The following theorem formally shows the convergence result under Assumptions 1, 2, and 4, whose proof is deferred in Appendix C.6.

Suppose Assumptions 1, 2, and 4 hold. Denote $\iota:=\log(2/\delta)$ and $\sigma:=R\sqrt{2\log(4T/\delta)}$ for any $0<\delta<1$. Let $G$ be a constant satisfying $G\geq\max\left\{2\lambda,2\sigma,\sqrt{C_{1}\Delta_{1}L_{0}},(C_{1}\Delta_{1}L_{\rho})^{\frac{1}{2-\rho}}\right\}$. Choose

Let $T=\max\left\{\frac{1}{\beta^{2}},\;\frac{C_{2}\Delta_{1}G}{\eta\epsilon^{2}}\right\}$. Then with probability at least $1-\delta$, we have $\left\|\nabla f(x_{t})\right\|\leq G$ for every $1\leq t\leq T$, and $\frac{1}{T}\sum_{t=1}^{T}\left\|\nabla f(x_{t})\right\|^{2}\leq\epsilon^{2}$.

Note that the main difference of Theorem 4.2 from Theorem 4.1 is that $\sigma$ is now $\mathcal{O}(\sqrt{\log T})$ instead of a constant. With some standard calculations, one can show that the gradient complexity in Theorem 4.2 is bounded by $\mathcal{O}(\epsilon^{-4}\log^{p}(1/\epsilon))$, where $p=\max\left\{3,\frac{9+2\rho}{4}\right\}<3.25$.

Analysis

We want to bound the gradients along the optimization trajectory mainly for two reasons. First, as discussed in Section 2, many existing analyses of Adam rely on the assumption of bounded gradients, because unbounded gradient norm leads to unbounded second order momentum $\hat{v}_{t}$ which implies very small stepsize, and slow convergence. On the other hand, once the gradients are bounded, it is straightforward to control $\hat{v}_{t}$ as well as the stepsize, and therefore the analysis essentially reduces to the easier one for AdaBound. Second, informally speakingThe statement is informal because here we can only show bounded gradients and Hessians at the iterate points, which only implies local smoothness near the neighborhood of each iterate point (see Section 5.2). However, the standard smoothness condition is a stronger global condition which assumes bounded Hessian at every point within a convex set., under Assumption 2, bounded gradients also imply bounded Hessians, which essentially reduces the $(\rho,L_{0},L_{\rho})$ smoothness to the standard smoothness. See Section 5.2 for more formal discussions.

In this paper, instead of imposing the strong assumption of globally bounded gradients, we develop a new analysis to show that with high probability, the gradients are always bounded along the trajectory of Adam until convergence. The essential idea can be informally illustrated by the following “circular" reasoning that we will make precise later. On the one hand, if $\left\|\nabla f(x_{t})\right\|\leq G$ for every $t\geq 1$, it is not hard to show the gradient converges to zero based on our discussions above. On the other hand, we know that a converging sequence must be upper bounded. Therefore there exists some $G^{\prime}$ such that $\left\|\nabla f(x_{t})\right\|\leq G^{\prime}$ for every $t\geq 1$. In other words, the bounded gradient condition implies the convergence result and the convergence result also implies the boundedness condition, forming a circular argument.

This circular argument is of course flawed. However, we can break the circularity of reasoning and rigorously prove both the bounded gradient condition and the convergence result using a contradiction argument. Before introducing the contradiction argument, we first need to provide the following useful lemma, which is the reverse direction of a generalized Polyak-Lojasiewicz (PL) inequality, whose proof is deferred in Appendix B.3.

Under Assumptions 1 and 2, we have $\left\|\nabla f(x)\right\|^{2}\!\leq 3(3L_{0}+4L_{\rho}\left\|\nabla f(x)\right\|^{\rho})(f(x)-f^{*})$.

Define the function $\zeta(u):=\frac{u^{2}}{3(3L_{0}+4L_{\rho}u^{\rho})}$ over $u\geq 0$. It is easy to verify that if $\rho<2$, $\zeta$ is increasing and its range is $[0,\infty)$. Therefore, $\zeta$ is invertible and $\zeta^{-1}$ is also increasing. Then, for any constant $G>0$, denoting $F=\zeta(G)$, Lemma 5.1 suggests that if $f(x)-f^{*}\leq F$, we have

In other words, if $\rho<2$, the gradient is bounded within any sub-level set, even though the sub-level set could be unbounded. Then, let $\tau$ be the first time the sub-optimality gap is strictly greater than $F$, truncated at $T+1$, or formally,

Then at least when $t<\tau$, we have $f(x_{t})-f^{*}\leq F$ and thus $\left\|\nabla f(x_{t})\right\|\leq G$. Based on our discussions above, it is not hard to analyze the updates before time $\tau$, and one can contruct some Lyapunov function to obtain an upper bound on $f(x_{\tau})-f^{*}$. On the other hand, if $\tau\leq T$, we immediately obtain a lower bound on $f(x_{\tau})$, that is $f(x_{\tau})-f^{*}>F$, by the definition of $\tau$ in (3). If the lower bound is greater than the upper bound, it leads to a contradiction, which shows $\tau=T+1$, i.e., the sub-optimality gap and the gradient norm are always bounded by $F$ and $G$ respectively before the algorithm terminates. We will illustrate the technique in more details in the simple deterministic setting in Section 5.3, but first, in Section 5.2, we introduce several prerequisite lemmas on the $(\rho,L_{0},L_{\rho})$ smoothness.

2 Local smoothness

In Section 5.1, we informally mentioned that $(\rho,L_{0},L_{\rho})$ smoothness essentially reduces to the standard smoothness if the gradient is bounded. In this section, we will make the statement more precise. First, note that Lemma 3.4 implies the following useful corollary.

The proof of Corollary 5.2 is deferred in Appendix B.4. Although the inequalities in Corollary 5.2 look very similar to the standard global smoothness condition with constant $L$, it is still a local condition as it requires $\left\|x-y\right\|\leq r$. Fortunately, at least before $\tau$, such a requirement is easy to satisfy for small enough $\eta$, according to the following lemma whose proof is deferred in Appendix C.5.

Under Assumption 3, if $t<\tau$ and choosing $G\geq\sigma$, we have $\left\|x_{t+1}-x_{t}\right\|\leq\eta D$ where $D:=2G/\lambda$.

Then as long as $\eta\leq r/D$, we have $\left\|x_{t+1}-x_{t}\right\|\leq r$ which satisfies the requirement in Corollary 5.2. Then we can apply the inequalities in it in the same way as the standard smoothness condition. In other words, most classical inequalities derived for standard smooth functions also apply to $(\rho,L_{0},L_{\rho})$ smooth functions.

3 Warm-up: analysis in the deterministic setting

In this section, we consider the simpler deterministic setting where the stochastic gradient $\nabla f(x_{t},\xi_{t})$ in Algorithm 1 or (18) is replaced with the exact gradient $\nabla f(x_{t})$. As discussed in Section 5.1, the key in our contradiction argument is to obtain both upper and lower bounds on $f(x_{\tau})-f^{*}$. In the following derivations, we focus on illustrating the main idea of our analysis technique and ignore minor proof details. In addition, all of them are under Assumptions 1, 2, and 3.

In order to obtain the upper bound, we need the following two lemmas. First, denoting $\epsilon_{t}:=\hat{m}_{t}-\nabla f(x_{t})$, we can obtain the following informal descent lemma for deterministic Adam.

For any $t<\tau$, choosing $G\geq\lambda$ and a small enough $\eta$,

where “$\lessapprox$” omits less important terms.

By the definition of $\tau$, for all $t<\tau$, we have $f(x_{t})-f^{*}\leq F$ which implies $\left\|\nabla f(x_{t})\right\|\leq G$. Then from the update rule (18) in Proposition C.1 provided later in Appendix C, it is easy to verify $\hat{v}_{t}\preceq G^{2}$ since $\hat{v}_{t}$ is a convex combination of $\{(\nabla f(x_{s}))^{2}\}_{s\leq t}$. Let $h_{t}:=\eta/(\sqrt{\hat{v}_{t}}+\lambda)$ be the stepsize vector and denote $H_{t}:=\textnormal{diag}(h_{t})$. We know

As discussed in Section 5.2, when $\eta$ is small enough, we can apply Corollary 5.2 to obtain

where in the first (approximate) inequality we ignore the second order term $\frac{1}{2}L\left\|x_{t+1}-x_{t}\right\|^{2}\propto\eta^{2}$ in Corollary 5.2 for small enough $\eta$; the equality applies the update rule $x_{t+1}-x_{t}=-H_{t}\hat{m}_{t}=-H_{t}(\nabla f(x_{t})+\epsilon_{t})$; in the second inequality we use $2a^{\top}Ab\leq\left\|a\right\|_{A}^{2}+\left\|b\right\|_{A}^{2}$ for any PSD matrix $A$ and vectors $a$ and $b$; and the last inequality is due to (5). ∎

Compared with the standard descent lemma for gradient descent, there is an additional term of $\left\|\epsilon_{t}\right\|^{2}$ in Lemma 5.4. In the next lemma, we bound this term recursively.

Choosing $\beta=\Theta(\eta G^{\rho+1/2})$, if $t<\tau$, we have

By the update rule (18) in Proposition C.1, we have

For small enough $\eta$, we can apply Corollary 5.2 to get

where the second inequality is due to $L=\mathcal{O}(G^{\rho})$ and $\left\|x_{t+1}-x_{t}\right\|=\mathcal{O}(\eta)\left\|\hat{m}_{t}\right\|$; and the last inequality uses $\hat{m}_{t}=\nabla f(x_{t})+\epsilon_{t}$ and Young’s inequality $\left\|a+b\right\|^{2}\leq 2\left\|a\right\|^{2}+2\left\|b\right\|^{2}$. Therefore,

where the first inequality uses (7) and Young’s inequality $\left\|a+b\right\|^{2}\leq(1+u)\left\|a\right\|^{2}+(1+1/u)\left\|b\right\|^{2}$ for any $u>0$; the second inequality uses $(1-\alpha_{t+1})(1+\alpha_{t+1}/2)\leq 1-\alpha_{t+1}/2$ and (8); and in the last inequality we use $\beta\leq\alpha_{t+1}$ and choose $\beta=\Theta(\eta G^{\rho+1/2})$ which implies $\mathcal{O}(\eta^{2}G^{2\rho}/\alpha_{t+1})\leq\frac{\lambda\beta}{16G}\leq\beta/4$. ∎

Now we combine them to get the upper bound on $f(x_{\tau})-f^{*}$. Define the function $\Phi_{t}:=f(x_{t})-f^{*}+\frac{2\eta}{\lambda\beta}\left\|\epsilon_{t}\right\|^{2}$. Note that for any $t<\tau$, (4)$+\frac{2\eta}{\lambda\beta}\times$(6) gives

The above inequality shows $\Phi_{t}$ is non-increasing and thus a Lyapunov function. Therefore, we have

where in the last inequality we use $\Phi_{1}=f(x_{1})-f^{*}=\Delta_{1}$ since $\epsilon_{1}=\hat{m}_{1}-\nabla f(x_{1})=0$ in the deterministic setting.

As discussed in Section 5.1, if $\tau\leq T$, we have $F<f(x_{\tau})-f^{*}\leq\Delta_{1}$. Note that we are able to choose a large enough constant $G$ so that $F=\frac{G^{2}}{3(3L_{0}+4L_{\rho}G^{\rho})}$ is greater than $\Delta_{1}$, which leads to a contradiction and shows $\tau=T+1$. Therefore, (9) holds for all $1\leq t\leq T$. Taking a summation over $1\leq t\leq T$ and re-arranging terms, we get

if choosing $T\geq\frac{8G\Delta_{1}}{\eta\epsilon^{2}}$, i.e., it shows convergence with a gradient complexity of $\mathcal{O}(\epsilon^{-2})$ since both $G$ and $\eta$ are constants independent of $\epsilon$ in the deterministic setting.

4 Extension to the stochastic setting

Fortunately, $\tau$ is in fact a stopping time with nice properties. If the noise is almost surely bounded as in Assumption 3, by a more careful analysis, we can obtain a high probability upper bound on $f(x_{\tau})-f^{*}$ using concentration inequalities. Then we can still obtain a contradiction and convergence under this high probability event. If the noise has sub-Gaussian norm as in Assumption 4, one can change the definition of $\tau$ to

for appropriately chosen $F$ and $\sigma$. Then at least when $t<\tau$, the noise is bounded by $\sigma$. Hence we can get the same upper bound on $f(x_{\tau})-f^{*}$ as if Assumption 3 still holds. However, when $t\leq T$, the lower bound $f(x_{\tau})-f^{*}>F$ does not necessarily holds, which requires some more careful analyses. The details of the proofs are involved and we defer them in Appendix C.

Variance-reduced Adam

Aside from the adaptive stepsize, one major difference between Algorithm 2 and STORM is that our hyper-parameters $\eta$ and $\beta$ are fixed constants whereas theirs are decreasing as a function of $t$. Choosing constant hyper-parameters requires a more accurate estimate at the initialization. That is why we use a mega-batch $\mathcal{S}_{1}$ to evaluate the gradient at the initial point to initialize $m_{1}$ and $v_{1}$ (Lines 2–3). In practice, one can also do a full-batch gradient evaluation at initialization. Note that there is no initialization bias for the momentum, so we do not re-scale $m_{t}$ and only re-scale $v_{t}$. We also want to point out that although the initial mega-batch gradient evaluation makes the algorithm a bit harder to implement, constant hyper-parameters are usually easier to tune and more common in training deep neural networks. It should be not hard to extend our analysis to time-decreasing $\eta$ and $\beta$ and we leave it as an interesting future work.

In addition to Assumption 1, we need to impose the following assumptions which can be viewed as stronger versions of Assumptions 2 and 3, respectively.

The objective function $f$ and the component function $f(\cdot,\xi)$ for each fixed $\xi$ are $(\rho,L_{0},L_{\rho})$ smooth with $0\leq\rho<2$.

The random variables $\{\xi_{t}\}_{1\leq t\leq T}$ are sampled i.i.d. from some distribution $\mathcal{P}$ such that for any $x\in\operatorname*{dom}(f)$,

Assumption 6 is stronger than Assumption 3. Assumption 3 applies only to the iterates generated by the algorithm, while Assumption 6 is a pointwise assumption over all $x\in\operatorname*{dom}(f)$ and further assumes an i.i.d. nature of the random variables $\{\xi_{t}\}_{1\leq t\leq T}$. Also note that, similar to Adam, it is straightforward to generalize the assumption to noise with sub-Gaussian norm as in Assumption 4.

In this part, we briefly discuss challenges in the analysis of VRAdam. The detailed analysis is deferred in Appendix D. Note that Corollary 5.2 requires bounded update $\left\|x_{t+1}-x_{t}\right\|\leq r$ at each step. For Adam, it is easy to satisfy for a small enough $\eta$ according to Lemma 5.3. However, for VRAdam, obtaining a good enough almost sure bound on the update is challenging even though the gradient noise is bounded. To bypass this difficulty, we directly impose a bound on $\left\|\nabla f(x_{t})-m_{t}\right\|$ by changing the definition of the stopping time $\tau$, similar to how we deal with the sub-Gaussian noise condition for Adam. In particular, we define

Then by definition, both $\left\|\nabla f(x_{t})\right\|$ and $\left\|\nabla f(x_{t})-m_{t}\right\|$ are bounded by $G$ before time $\tau$, which directly implies bounded update $\left\|x_{t+1}-x_{t}\right\|$. Of course, the new definition brings new challenges to lower bounding $f(x_{\tau})-f^{*}$, which requires more careful analyses specific to the VRAdam algorithm. Please see Appendix D for the details.

2 Convergence guarantees for VRAdam

In the section, we provide our main results for convergence of VRAdam under Assumptions 1, 5, and 6. We consider the same definitions of problem-dependent constants $\Delta_{1},r,L$ as those in Section 4 to make the statements of theorems concise. Let $c$ be a small enough numerical constant and $C$ be a large enough numerical constant. The formal convergence result is shown in the following theorem.

Suppose Assumptions 1, 5, and 6 hold. For any $0<\delta<1$, let $G>0$ be a constant satisfying $G\geq\max\left\{2\lambda,2\sigma,\sqrt{C\Delta_{1}L_{0}/\delta},(C\Delta_{1}L_{\rho}/\delta)^{\frac{1}{2-\rho}}\right\}.$ Choose $0\leq\beta_{\operatorname{sq}}\leq 1$ and $\beta=a^{2}\eta^{2}$, where $a=40L\sqrt{G}\lambda^{-3/2}$. Choose

Then with probability at least $1-\delta$, we have $\left\|\nabla f(x_{t})\right\|\leq G$ for every $1\leq t\leq T$, and $\frac{1}{T}\sum_{t=1}^{T}\left\|\nabla f(x_{t})\right\|^{2}\leq\epsilon^{2}.$

Note that the choice of $G$, the upper bound of gradients along the trajectory of VRAdam, is very similar to that in Theorem 4.1 for Adam. The only difference is that now it also depends on the failure probability $\delta$. Similar to Theorem 4.1, there is no requirement on $\beta_{\operatorname{sq}}$ and we choose a very small $\beta=\mathcal{O}(\epsilon^{2})$. However, the variance reduction technique allows us to take a larger stepsize $\eta=\mathcal{O}(\epsilon)$ (compared with $\mathcal{O}(\epsilon^{2})$ for Adam) and obtain an accelerated gradient complexity of $\mathcal{O}(\epsilon^{-3})$, where we only consider the leading term. We are not sure whether it is optimal as the $\Omega(\epsilon^{-3})$ lower bound in (Arjevani et al., 2023) assumes the weaker bounded variance condition. However, our result significantly improves upon (Wang and Klabjan, 2022), which considers a variance-reduced version of Adam by combining Adam and SVRG (Johnson and Zhang, 2013) and only obtains asymptotic convergence in the non-convex setting. Similar to Adam, our gradient complexity for VRAdam is dimension free but its dependence on $\lambda$ is $\mathcal{O}(\lambda^{-2})$. Another limitation is that, the dependence on the failure probability $\delta$ is polynomial, worse than the poly-log dependence in Theorem 4.1 for Adam.

Conclusion and future works

In this paper, we proved the convergence of Adam and its variance-reduced version under less restrictive assumptions compared to those in the existing literature. We considered a generalized non-uniform smoothness condition, according to which the Hessian norm is bounded by a sub-quadratic function of the gradient norm almost everywhere. Instead of assuming the Lipschitzness of the objective function as in existing analyses of Adam, we use a new contradiction argument to prove that gradients are bounded by a constant along the optimization trajectory. There are several interesting future directions that one could pursue following this work.

Our analysis relies on the assumption of bounded noise or noise with sub-Gaussian norm. However, the existing lower bounds in (Arjevani et al., 2023) consider the weaker bounded variance assumption. Hence, it is not clear whether the $\mathcal{O}(\epsilon^{-4})$ complexity we obtain for Adam is tight in this setting. It will be interesting to see whether one can relax the assumption to the bounded variance setting. One may gain some insights from recent papers such as (Faw et al., 2022; Wang et al., 2023) that analyze AdaGrad under weak noise conditions. An alternative way to show the tightness of the $\mathcal{O}(\epsilon^{-4})$ complexity is to prove a lower bound under the bounded noise assumption.

Another interesting future direction is to see if the techniques developed in this work for bounding gradients (including those in the the concurrent work (Li et al., 2023)) can be generalized to improve the convergence results for other optimization problems and algorithms. We believe it is possible so long as the function class is well behaved and the algorithm is efficient enough so that $f(x_{\tau})-f^{*}$ can be well bounded for some appropriately defined stopping time $\tau$.

We want to note that our results can not explain why Adam is better than SGD for training transformers, because (Li et al., 2023) shows that non-adaptive SGD converges with the same $\mathcal{O}(\epsilon^{-4})$ gradient complexity under even weaker conditions. It would be interesting and impactful if one can find a reasonable setting (function class, gradient oracle, etc) under which Adam or other adaptive methods provably outperform SGD.

Acknowledgments

This work was supported, in part, by the MIT-IBM Watson AI Lab and ONR Grants N00014-20-1-2394 and N00014-23-1-2299. We also acknowledge support from DOE under grant DE-SC0022199, and NSF through awards DMS-2031883 and DMS-1953181.

References

Appendix A Probabilistic lemmas

In this section, we state several well-known and useful probabilistic lemmas without proof.

Let $\{Z_{t}\}_{t\geq 1}$ be a martingale with respect to a filtration $\{\mathcal{F}_{t}\}_{t\geq 0}$. Assume that $\left|Z_{t}-Z_{t-1}\right|\leq c_{t}$ almost surely for all $t\geq 0$. Then for any fixed $T$, with probability at least $1-\delta$,

In this section, we provide proofs related to $(\rho,L_{0},L_{\rho})$ smoothness. In what follows, we first provide a formal proposition in Appendix B.1 showing that univariate rational functions and double exponential functions are $(\rho,L_{0},L_{\rho})$ smooth with $\rho<2$, as we claimed in Section 3.2.1, and then provide the proofs of Lemma 3.4, Lemma 5.1, and Corollary 5.2 in Appendix B.2, B.3 and B.4 respectively.

Any univariate rational function $P(x)/Q(x)$, where $P,Q$ are two polynomials, and any double exponential function $a^{(b^{x})}$, where $a,b>1$, are $(\rho,L_{0},L_{\rho})$ smooth with $1<\rho<2$. However, they are not necessarily $(L_{0},L_{1})$ smooth.

We prove the proposition in the following four parts:

1. Univariate rational functions are $(\rho,L_{0},L_{\rho})$ smooth with $1<\rho<2$. Let $f(x)=P(x)/Q(x)$ where $P$ and $Q$ are two polynomials. Then the partial fractional decomposition of $f(x)$ is given by

Then we have $f(x)=w(x)+\sum_{i=1}^{m}\sum_{r=1}^{j_{i}}p_{ir}(x)+\sum_{i=1}^{n}\sum_{r=1}^{k_{i}}q_{ir}(x)$. For any $3/2<\rho<2$, we know that $\rho>\frac{r+2}{r+1}$ for any $r\geq 1$. Then we can show that

where the first equality is because one can easily verify that the first and second order derivatives of $p_{ij_{i}}$ dominate those of all other terms when $x$ goes to $a_{i}$, and the second equality is because $\left|p^{\prime}_{ij_{i}}(x)\right|^{\rho}=\mathcal{O}\left((x-a_{i})^{-\rho(j_{i}+1)}\right)$, $\left|p^{\prime\prime}_{ij_{i}}(x)\right|=\mathcal{O}\left((x-a_{i})^{-(j_{i}+2)}\right)$, and $\rho(j_{i}+1)>j_{i}+2$ (here we assume $j_{i}\geq 1$ since otherwise there is no need to prove (10) for $i$). Note that (10) implies that, for any $L_{\rho}>0$, there exists $\delta_{i}>0$ such that

Similarly, one can show $\lim_{x\to\infty}\frac{\left|f^{\prime}(x)\right|^{\rho}}{\left|f^{\prime\prime}(x)\right|}=\infty$, which implies there exists $M>0$ such that

We know $\mathcal{B}$ is a compact set and therefore the continuous function $f^{\prime\prime}$ is bounded within $\mathcal{B}$, i.e., there exists some constant $L_{0}>0$ such that

Combining (11), (12), and (13), we have shown

which completes the proof of the first part.

2. Rational functions are not necessarily $(L_{0},L_{1})$ smooth. Consider the ration function $f(x)=1/x$. Then we know that $f^{\prime}(x)=-1/x^{2}$ and $f^{\prime\prime}(x)=2/x^{3}$. Note that for any $0<x\leq\min\{(L_{0}+1)^{-1/3},(L_{1}+1)^{-1}\}$, we have

which shows $f$ is not $(L_{0},L_{1})$ smooth for any $L_{0},L_{1}\geq 0$.

3. Double exponential functions are $(\rho,L_{0},L_{\rho})$ smooth with $1<\rho<2$. Let $f(x)=a^{(b^{x})}$, where $a,b>1$, be a double exponential function. Then we know that

where the first equality is a direct calculation; the second equality uses change of variable $y=b^{x}$; and the last equality is because exponential function grows faster than linear function. Then we complete the proof following a similar argument to that in Part 1.

4. Double exponential functions are not necessarily $(L_{0},L_{1})$ smooth. Consider the double exponential function $f(x)=e^{(e^{x})}$. Then we have

For any $x\geq\max\left\{\log(L_{0}+1),\log(L_{1}+1)\right\}$, we can show that

which shows $f$ is not $(L_{0},L_{1})$ smooth for any $L_{0},L_{1}\geq 0$. ∎

B.2 Proof of Lemma 3.4

Before proving Lemma 3.4, we need the following lemma that generalizes (a special case of) Grönwall’s inequality.

It is straightforward to verify that the solution to (14) satisfies

Then it suffices to show $\phi(u(t))\leq\phi(v(t))\;\forall t\in[a,b]$. Note that

With Lemma B.2, one can bound the gradient norm within a small enough neighborhood of a given point as in the following lemma.

Suppose $f$ is $(\rho,L_{0},L_{\rho})$ smooth for some $\rho,\rho,L_{0},L_{\rho}\geq 0$. For any $a>0$ and points $x,y\in\operatorname*{dom}(f)$ satisfying $\left\|y-x\right\|\leq\frac{a}{L_{0}+L_{\rho}(\left\|\nabla f(x)\right\|+a)^{\rho}}$, we have

Denote functions $z(t):=(1-t)x+ty$ and $u(t):=\left\|\nabla f(z(t))\right\|$ for $0\leq t\leq 1$. Note that for any $0\leq t\leq s\leq 1$, by triangle inequality,

We know that $u(t)=\left\|\nabla f(z(t))\right\|$ is differentiable since $f$ is second order differentiableHere we assume $u(t)>0$ for $0<t<1$. Otherwise, we can define $t_{m}=\sup\{0<t<1\mid u(t)=0\}$ and consider the interval $[t_{m},1]$ instead.. Then we have

Let $\phi(w):=\int_{0}^{w}\frac{1}{(L_{0}+L_{\rho}v^{\rho})\left\|y-x\right\|}dv$. By Lemma B.2, we know that

Denote $\psi(w):=\int_{0}^{w}\frac{1}{(L_{0}+L_{\rho}v^{\rho})}dv=\phi(w)\cdot\left\|y-x\right\|$. We have

Since $\psi$ is increasing, we have $\left\|\nabla f(y)\right\|\leq\left\|\nabla f(x)\right\|+a$. ∎

With Lemma B.3, we are ready to prove Lemma 3.4.

On the other hand, by the definition of $t_{\text{b}}$, we know $z(t)\in\mathcal{X}$ for every $0\leq t<t_{\text{b}}$. Then by Lemma B.3, for all $0\leq t<t_{\text{b}}$, we have $\left\|\nabla f(z(t))\right\|\leq\left\|\nabla f(x)\right\|+a$. Therefore for all $0\leq t<t_{\text{b}}$

which contradicts with (15). Therefore we have shown $y\in\operatorname*{dom}(f)$. We have

where the last inequality is due to Lemma B.3. ∎

B.3 Proof of Lemma 5.1

Denote $G:=\left\|\nabla f(x)\right\|$ and $L:=3L_{0}+4L_{\rho}G^{\rho}$. Let $y:=x-\frac{1}{2L}\nabla f(x)$. Then we have

where the inequality can be easily verified considering both cases of $G\leq(L_{0}/L_{\rho})^{1/\rho}$ and $G\geq(L_{0}/L_{\rho})^{1/\rho}$. Then based on Corollary 5.2, we have $y\in\operatorname*{dom}(f)$ and

B.4 Proof of Corollary 5.2

First, Lemma 3.4 states that for any $a>0$,

If $\left\|\nabla f(x)\right\|\leq G$, we choose $a=\max\{G,(L_{0}/L_{\rho})^{1/\rho}\}$. Then it is straightforward to verify that

Therefore we have shown for any $x,y$ satisfying $\left\|y-x\right\|\leq r$,

Next, let $z(t):=(1-t)x+ty$ for $0\leq t\leq 1$. We know

Appendix C Convergence analysis of Adam

In this section, we provide detailed convergence analysis of Adam. We will focus on proving Theorem 4.1 under the bounded noise assumption (Assumption 3) in most parts of this section except Appendix C.6 where we will show how to generalize the results to noise with sub-Gaussian norm (Assumption 4) and provide the proof of Theorem 4.2.

For completeness, we repeat some important technical definitions here. First, we define

as the deviation of the re-scaled momentum from the actual gradient. Given a large enough constant $G$ defined in Theorem 4.1, denoting $F=\frac{G^{2}}{3(3L_{0}+4L_{\rho}G^{\rho})}$, we formally define the stopping time $\tau$ as

i.e., $\tau$ is the first time when the sub-optimality gap is strictly greater than $F$, truncated at $T+1$ to make sure it is bounded in order to apply Lemma A.2. Based on Lemma 5.1 and the discussions below it, we know that if $t<\tau$, we have both $f(x_{t})-f^{*}\leq F$ and $\left\|\nabla f(x_{t})\right\|\leq G$. It is clear to see that $\tau$ is a stopping timeIndeed, $\tau-1$ is also a stopping time because $\nabla f(x_{t})$ only depends on $\{\xi_{s}\}_{s<t}$, but that is unnecessary for our analysis. with respect to $\{\xi_{t}\}_{t\geq 1}$ because the event $\{\tau\geq t\}$ is a function of $\{\xi_{s}\}_{s<t}$ and independent of $\{\xi_{s}\}_{s\geq t}$. Next, let

be the stepsize vector and $H_{t}:=\textnormal{diag}(h_{t})$ be the diagonal stepsize matrix. Then the update rule can be written as

Finally, as in Corollary 5.2 and Lemma 5.3, we define the following constants.

The bias correction steps in Lines 7–8 make Algorithm 1 a bit complicated. In the following proposition, we provide an equivalent yet simpler update rule of Adam.

Denote $\alpha_{t}=\frac{\beta}{1-(1-\beta)^{t}}$ and $\alpha^{\operatorname{sq}}_{t}=\frac{\beta_{\operatorname{sq}}}{1-(1-\beta_{\operatorname{sq}})^{t}}$. Then the update rule in Algorithm 1 is equivalent to

where initially we set $\hat{m}_{1}=\nabla f(x_{1},\xi_{1})$ and $\hat{v}_{1}=(\nabla f(x_{1},\xi_{1}))^{2}$. Note that since $1-\alpha_{1}=1-\alpha^{\operatorname{sq}}_{1}=0$, there is no need to define $\hat{m}_{0}$ and $\hat{v}_{0}$.

Denote $Z_{t}=1-(1-\beta)^{t}$. Then we know $\alpha_{t}=\beta/Z_{t}$ and $m_{t}=Z_{t}\hat{m}_{t}$. By the momentum update rule in Algorithm 1, we have

Note that $Z_{t}$ satisfies the following property

Next, we verify the initial condition. By Algorithm 1, since we set $m_{0}=0$, we have $m_{1}=\beta\nabla f(x_{1},\xi_{1})$. Therefore we have $\hat{m}_{1}=m_{1}/Z_{1}=\nabla f(x_{1},\xi_{1})$ since $Z_{1}=\beta$. Then the proof is completed by applying the same analysis on $v_{t}$ and $\hat{v}_{t}$. ∎

C.2 Useful lemmas for Adam

In this section, we list several useful lemmas for the convergence analysis. Their proofs are all deferred in Appendix C.5.

First note that when $t<\tau$, all the quantities in the algorithm are well bounded. In particular, we have the following lemma.

Next, we provide a useful lemma regarding the time-dependent re-scaled momentum parameters in (18).

Let $\alpha_{t}=\frac{\beta}{1-(1-\beta)^{t}}$, then for all $T\geq 2$, we have $\sum_{t=2}^{T}\alpha_{t}^{2}\leq 3(1+\beta^{2}T)$.

In the next lemma, we provide an almost sure bound on $\epsilon_{t}$ in order to apply Azuma-Hoeffding inequality (Lemma A.1).

Denote $\gamma_{t-1}=(1-\alpha_{t})(\epsilon_{t-1}+\nabla f(x_{t-1})-\nabla f(x_{t}))$. Choosing $\eta\leq\min\left\{\frac{r}{D},\frac{\sigma\beta}{DL}\right\}$, if $t\leq\tau$, we have $\left\|\epsilon_{t}\right\|\leq 2\sigma$ and $\left\|\gamma_{t-1}\right\|\leq 2\sigma$.

Finally, the following lemma hides messy calculations and will be useful in the contradiction argument.

Under the parameter choices in either Theorem 4.1 or Theorem 4.2, we have $I_{1}\leq I_{2}$ and $I_{1}/T\leq\epsilon^{2}$.

C.3 Proof of Theorem 4.1

Before proving the main theorems, several important lemmas are needed. First, we provide a descent lemma for Adam.

If $t<\tau$, choosing $G\geq\sigma+\lambda$ and $\eta\leq\min\left\{\frac{r}{D},\frac{\lambda}{6L}\right\}$, we have

Since we choose $\eta\leq\frac{r}{D}$, by Lemma 5.3, we have $\left\|x_{t+1}-x_{t}\right\|\leq r$ if $t<\tau$. Then we can apply Corollary 5.2 to show that for any $t<\tau$,

where the second inequality uses (17) and (19); the third inequality is due to Young’s inequality $a^{\top}Ab\leq\frac{1}{3}\left\|a\right\|_{A}^{2}+\frac{3}{4}\left\|b\right\|_{A}^{2}$ and $\left\|a+b\right\|_{A}^{2}\leq 2\left\|a\right\|_{A}^{2}+2\left\|b\right\|_{A}$ for any PSD matrix $A$; the second last inequality uses $\eta\leq\frac{\lambda}{6L}$; and the last inequality is due to (19). ∎

The following lemma bounds the sum of the error term $\left\|\epsilon_{t}\right\|^{2}$ before the stopping time $\tau$. Since its proof is complicated, we defer it in Appendix C.4.

If $G\geq 2\sigma$ and $\eta\leq\min\left\{\frac{r}{D},\frac{\lambda^{3/2}\beta}{6L\sqrt{G}},\frac{\sigma\beta}{DL}\right\}$, with probability $1-\delta$,

Combining Lemma C.6 and Lemma C.7, we obtain the following useful lemma, which simultaneously bounds $f(x_{t})-f^{*}$ and $\sum_{t=1}^{\tau-1}\left\|\nabla f(x_{t})\right\|^{2}$.

If $G\geq 2\max\{\lambda,\sigma\}$ and $\eta\leq\min\left\{\frac{r}{D},\frac{\lambda^{3/2}\beta}{6L\sqrt{G}},\frac{\sigma\beta}{DL}\right\}$, then with probability at least $1-\delta$,

With Lemma C.8, we are ready to complete the contradiction argument and the convergence analysis. Below we provide the proof of Theorem 4.1.

By the definition of $\tau$, if $\tau\leq T$, we have

Based on Lemma C.5, we have $I_{1}\leq I_{2}$, which leads to a contradiction. Therefore, we must have $\tau=T+1$ conditioned on $\mathcal{E}$. Then, Lemma C.8 also implies that under $\mathcal{E}$,

where the last inequality is due to Lemma C.5. ∎

C.4 Proof of Lemma C.7

In order to prove Lemma C.7, we need the following several lemmas.

Denote $\gamma_{t-1}=(1-\alpha_{t})(\epsilon_{t-1}+\nabla f(x_{t-1})-\nabla f(x_{t}))$. If $G\geq 2\sigma$ and $\eta\leq\min\left\{\frac{r}{D},\frac{\lambda^{3/2}\beta}{6L\sqrt{G}}\right\}$, we have for every $2\leq t\leq\tau$,

According to the update rule (18), we have

Since we choose $\eta\leq\frac{r}{D}$, by Lemma 5.3, we have $\left\|x_{t}-x_{t-1}\right\|\leq r$ if $t\leq\tau$. Therefore by Corollary 5.2, for any $2\leq t\leq\tau$,

where the first inequality uses Young’s inequality $\left\|a+b\right\|^{2}\leq(1+u)\left\|a\right\|^{2}+(1+1/u)\left\|b\right\|^{2}$ for any $u>0$; the second inequality is due to

the third inequality uses (24) and Young’s inequality; and in the last inequality we choose $\eta\leq\frac{\lambda^{3/2}\beta}{6L\sqrt{G}}$, which implies $\frac{2\eta^{2}L^{2}}{\lambda^{2}\beta}\leq\frac{\lambda\beta}{16G}\leq\frac{\beta}{2}\leq\frac{\alpha_{t}}{2}$. Then by (23), we have

Denote $\gamma_{t-1}=(1-\alpha_{t})(\epsilon_{t-1}+\nabla f(x_{t-1})-\nabla f(x_{t}))$. If $G\geq 2\sigma$ and $\eta\leq\min\left\{\frac{r}{D},\frac{\sigma\beta}{DL}\right\}$, with probability $1-\delta$,

Since $\tau$ is a stopping time, we know that $1_{\tau\geq t}$ is a function of $\{\xi_{s}\}_{s<t}$. Also, by definition, we know $\gamma_{t-1}$ is a function of $\{\xi_{s}\}_{s<t}$. Then, denoting

Then by the Azuma-Hoeffding inequality (Lemma A.1), we have with probability at least $1-\delta$,

where in the last inequality we use Lemma C.3. ∎

By Lemma C.9, we have for every $2\leq t\leq\tau$,

Taking a summation over $t$ from $2$ to $\tau$, we have

where the first inequality uses Lemma C.10; and the second inequality uses Lemma C.3 and $\left\|\epsilon_{1}\right\|^{2}=\left\|\nabla f(x_{1},\xi_{1})-\nabla f(x_{1})\right\|^{2}\leq\sigma^{2}$. Then we complete the proof by multiplying both sides by $2/\beta$. ∎

C.5 Omitted proofs for Adam

In this section, we provide all the omitted proofs for Adam including those of Lemma 5.3 and all the lemmas in Appendix C.2.

By definition of $\tau$, we have $\left\|\nabla f(x_{t})\right\|\leq G$ if $t<\tau$. Then Assumption 3 directly implies $\left\|\nabla f(x_{t},\xi_{t})\right\|\leq G+\sigma$. $\left\|\hat{m}_{t}\right\|$ can be bounded by a standard induction argument as follows. First note that $\left\|\hat{m}_{1}\right\|=\left\|\nabla f(x_{1},\xi_{1})\right\|\leq G+\sigma$. Supposing $\left\|\hat{m}_{k-1}\right\|\leq G+\sigma$ for some $k<\tau$, then we have

Then we can show $\hat{v}_{t}\preceq(G+\sigma)^{2}$ in a similar way noting that $(\nabla f(x_{t},\xi_{t}))^{2}\preceq\left\|\nabla f(x_{t},\xi_{t})\right\|^{2}\leq(G+\sigma)^{2}$. Given the bound on $\hat{v}_{t}$, it is straight forward to bound the stepsize $h_{t}$. ∎

First, when $t\geq 1/\beta$, we have $(1-\beta)^{t}\leq 1/e$. Therefore,

Next, note that when $t<1/\beta$, we have $(1-\beta)^{t}\leq 1-\frac{1}{2}\beta t$. Then we have

Therefore we have $\sum_{t=2}^{T}\alpha_{t}^{2}\leq 3(1+\beta^{2}T).$ ∎