SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\tilde{\mathcal{O}}$  denotes  $\mathcal{O}$  with logarithmic terms ignored. Consequently, the dependence of Eq. (2) over  $\beta_{2},\eta$  and  $T$  becomes
$$
\begin{array}{l} \mathbb {E} \sum_ {t = 1} ^ {T} \| \nabla f (\pmb {w} _ {t}) \| = \tilde {\mathcal {O}} \left(\frac {1}{\sqrt {1 - \beta_ {2}}} + \frac {\eta^ {2}}{\sqrt {(1 - \beta_ {2}) ^ {3}}} + \frac {1}{\eta} + T \sqrt {1 - \beta_ {2}} + \frac {\eta^ {2}}{\sqrt {1 - \beta_ {2}}} T\right) \\ + \tilde {\mathcal {O}} \left(\frac {\sqrt {T}}{\sqrt [ 4 ]{1 - \beta_ {2}}} + \frac {\eta \sqrt {T}}{\sqrt [ 4 ]{(1 - \beta_ {2}) ^ {3}}} + \frac {\sqrt {T}}{\sqrt {\eta}} + T \sqrt [ 4 ]{1 - \beta_ {2}} + \frac {\eta}{\sqrt [ 4 ]{1 - \beta_ {2}}} T\right). \\ \end{array}
$$
Here we consider two cases: (i).  $\beta_{2}$  and  $\eta$  are independent over  $T$ , and (ii).  $\beta_{2}$  and  $\eta$  are dependent over  $T$ . For case (i), based on the above equation, one can easily observe that the averaged gradient norm  $\frac{1}{T}\mathbb{E}\sum_{t = 1}^{T}\| \nabla f(\pmb {w}_t)\|$  will converge to the threshold  $\mathcal{O}\left(\frac{\eta^2}{\sqrt{1 - \beta_2}} +\sqrt[4]{1 - \beta_2} +\frac{\eta}{\sqrt[4]{1 - \beta_2}}\right)$  with rate  $\mathcal{O}(1 / \sqrt{T})$ . This aligns with the observation in [27, 34] that Adam will not converge to the stationary point with constant  $\beta_{2}$ .
For case (ii), in order to ensure convergence, i.e.,  $\min_{t\in [T]}\mathbb{E}\| \pmb {G}_t\| _1\to 0$  as  $T\rightarrow \infty$ , a sufficient condition is that the right-hand-side of the above equation is  $\pmb {o}(T)$ . Specifically, by choosing  $\eta = \Theta (T^{-a})$  and  $1 - \beta_{2} = \Theta (T^{-b})$ , we obtain that
$$
\begin{array}{l} \frac {1}{T} \mathbb {E} \sum_ {t = 1} ^ {T} \| \nabla f (\boldsymbol {w} _ {t}) \| = \tilde {\mathcal {O}} \left(T ^ {\frac {b}{2} - 1} + T ^ {- 2 a + \frac {3 b}{2} - 1} + T ^ {a - 1} + T ^ {- \frac {b}{2}} + T ^ {- 2 a + \frac {b}{2}}\right) \\ + \tilde {\mathcal {O}} \left(T ^ {- \frac {1}{2} + \frac {b}{4}} + T ^ {- \frac {1}{2} - a + \frac {3 b}{4}} + T ^ {- \frac {1}{2} + \frac {a}{2}} + T ^ {- \frac {b}{4}} + T ^ {- a + \frac {b}{4}}\right). \\ \end{array}
$$
By simple calculation, we obtain that the right-hand side of the above inequality is  $o(1)$  as  $T \to \infty$  if and only if  $b > 0$ ,  $1 > a > 0$  and  $b - a < 1$ . Moreover, the minimum of the right-hand side of the above inequality is  $\tilde{\mathcal{O}}\left(1 / T^{\frac{1}{4}}\right)$ , which is achieved at  $a = \frac{1}{2}$  and  $b = 1$ . Such a minimum implies an upper bound of the iteration complexity which at most differs from the lower bound by logarithmic factors as solving  $\tilde{\mathcal{O}}\left(1 / T^{\frac{1}{4}}\right) = \varepsilon$  gives  $T = \tilde{\mathcal{O}}\left(\frac{1}{\varepsilon^4}\right)$ . In Theorem 2, we will further remove the logarithmic factor by giving a refined proof when  $a = \frac{1}{2}$  and  $b = 1$  and close the gap between the upper and lower bounds.
Dependence over  $\lambda$ . Our analysis allows  $\lambda = 0$  in the adaptive learning rate  $\eta \frac{1}{\sqrt{\nu_t} + \lambda\mathbb{1}_d}$ . In contrast, some existing works [16, 21] require non-zero  $\lambda$  and their iteration complexity has polynomial dependence over  $\frac{1}{\lambda}$ , which is less desired as  $\lambda$  can be as small as  $10^{-8}$  in practice (e.g., in PyTorch's default setting). Furthermore, compared to their setting, our setting is more challenging as non-zero  $\lambda$  immediately provides an upper bound of the adaptive learning rate.
4.2 Proof Sketch of Theorem 1
In this section, we demonstrate the proof idea of Theorem 1. Generally speaking, our proof is inspired by (i). the construction of the Lyapunov function for SGDM [22] and (ii) the construction of auxiliary function and the conversion from regret bound to gradient bound for AdaGrad [31], but the adaptation of these techniques to Adam is highly non-trivial, as SGDM does not hold an adaptive learning rate, and the adaptive learning rate of AdaGrad is monotonously decreasing. Below we sketch the proof by identifying three key challenges in the proof and provide our solutions respectively.
Challenge I: Disentangle the stochasticity in stochastic gradient and adaptive learning rate. For simplicity, let us first consider the case where  $\beta_{1} = 0$ , i.e., where the momentum  $m_{t}$  degenerates to
the stochastic gradient  $g_{t}$ . According to the standard descent lemma, we have that
$$
\begin{array}{l} \mathbb {E} f (\boldsymbol {w} _ {t + 1}) \leq f (\boldsymbol {w} _ {t}) + \mathbb {E} \left[ \langle \boldsymbol {G} _ {t}, \boldsymbol {w} _ {t + 1} - \boldsymbol {w} _ {t} \rangle + \frac {L}{2} \| \boldsymbol {w} _ {t + 1} - \boldsymbol {w} _ {t} \| ^ {2} \right] \\ \leq \mathbb {E} f (\boldsymbol {w} _ {t}) + \underbrace {\mathbb {E} \left[ \left\langle \boldsymbol {G} _ {t} , - \eta \frac {1}{\sqrt {\boldsymbol {\nu} _ {t}}} \odot \boldsymbol {g} _ {t} \right\rangle \right]} _ {\text {F i r s t O r d e r}} + \underbrace {\frac {L}{2} \eta^ {2} \mathbb {E} \left\| \frac {1}{\sqrt {\boldsymbol {\nu} _ {t}}} \odot \boldsymbol {m} _ {t} \right\| ^ {2}} _ {\text {S e c o n d O r d e r}} \tag {3} \\ \end{array}
$$
The first challenge arises from bounding the "First Order" term above. To facilitate the understanding of the difficulty, we compare the "First Order" term of Adam to the corresponding "First Order" term of SGD, i.e.,  $-\eta \mathbb{E}\langle \pmb {G}_t,\pmb {g}_t\rangle$  . By directly applying  $\mathbb{E}^{\mid \mathcal{F}_t}g_t = G_t$  , we obtain that the "First-Order" term of SGD equals to  $-\eta \mathbb{E}\| \pmb {G}_t\| ^2$  . However, as for Adam, we do not even know what  $\mathbb{E}^{\mid \mathcal{F}_t}\frac{1}{\sqrt{\nu_t}}\odot \pmb {g}_t$  is given that the stochasticity in  $g_{t}$  and  $\nu_{t}$  entangles. A common practice is to use a surrogate adaptive learning rate  $\widetilde{\nu}_{t}$  measurable with respect to  $\mathcal{F}_t$  , to approximate the real adaptive learning rate  $\nu_{t}$  This leads to the following equation:
$$
\underbrace {\mathbb {E} \left[ \left\langle \boldsymbol {G} _ {t} , - \eta \frac {1}{\sqrt {\boldsymbol {\nu} _ {t}}} \odot \boldsymbol {g} _ {t} \right\rangle \right]} _ {\text {F i r s t O r d e r}} = \underbrace {\mathbb {E} \left[ \left\langle \boldsymbol {G} _ {t} , - \eta \frac {1}{\sqrt {\widetilde {\boldsymbol {\nu}} _ {t}}} \odot \boldsymbol {g} _ {t} \right\rangle \right]} _ {\text {F i r s t O r d e r M a i n}} + \underbrace {\mathbb {E} \left[ \left\langle \boldsymbol {G} _ {t} , - \eta \left(\frac {1}{\sqrt {\boldsymbol {\nu} _ {t}}} - \frac {1}{\sqrt {\widetilde {\boldsymbol {\nu}} _ {t}}}\right) \odot \boldsymbol {g} _ {t} \right\rangle \right]} _ {\text {E r r o r}}.
$$
One can immediately see that "First Order Main" terms equal to  $\mathbb{E}[\langle G_t, -\eta \frac{1}{\sqrt{\tilde{\nu}_t}} \odot G_t \rangle] < 0$ , but now we need to handle the "Error" term. In existing literature, such a term is mostly bypassed by applying the bounded gradient assumption [10, 36], which, however, we do not assume.
Solution to Challenge I. Inspired by recent advance in the analysis of AdaGrad [31], we consider the auxiliary function  $\xi_{t} = \mathbb{E}[\eta \langle \pmb{G}_{t}, - \frac{1}{\sqrt{\tilde{\nu}_{t + 1}}}\odot \pmb{G}_{t}\rangle ]$  , where we choose  $\widetilde{\pmb{\nu}}_t = \beta_2\pmb{\nu}_{t - 1} + (1 - \beta_2)\sigma_0^2\mathbb{1}_d.$  In the following lemma, we show that the error term can be controlled using  $\xi_{t}$  , parallel to (Lemma 4. [31]).
Lemma 1 (Informal version of Lemma 7 with  $\beta_{1} = 0$ ). Let all conditions in Theorem 1 hold. Then,
$$
\operatorname {E r r o r} \leq \frac {5}{8} \mathbb {E} \left[ \eta \left\langle \boldsymbol {G} _ {t}, - \frac {1}{\sqrt {\tilde {\nu} _ {t}}} \odot \boldsymbol {G} _ {t} \right\rangle \right] + \mathcal {O} \left(\frac {1}{\sqrt {\beta_ {2}}} \xi_ {t - 1} - \xi_ {t}\right) + S m a l l E r r o r. \tag {4}
$$
In the right-hand-side of inequality (4), one can easily observe that the first term can be controlled by "First Order Main" term, and the third term is as small as the "Second Order" term. However, the second term seems annoying – in the analysis of AdaGrad [31], there is no  $1 / \sqrt{\beta_2}$  factor, making the corresponding term a telescoping, but this is no longer true due to the existence of the  $1 / \sqrt{\beta_2}$  factor. We resolve this difficulty by looking at the sum of  $\frac{1}{\sqrt{\beta_2}} \xi_{t-1} - \xi_t$  over  $t$  from 1 to  $T$ , which gives  $\mathcal{O}((1 - \beta_2) \sum_{t=1}^{T-1} \xi_t)$ . By further noticing that  $\widetilde{\nu}_{t+1} \geq \beta_2 \widetilde{\nu}_t$ , we have
$$
\sum_ {t = 1} ^ {T} \left(\frac {1}{\sqrt {\beta_ {2}}} \xi_ {t - 1} - \xi_ {t}\right) \leq \mathcal {O} \left((1 - \beta_ {2}) \sum_ {t = 1} ^ {T - 1} \mathbb {E} \left[ \eta \left\langle \boldsymbol {G} _ {t}, - \frac {1}{\sqrt {\tilde {\nu} _ {t}}} \odot \boldsymbol {G} _ {t} \right\rangle \right]\right).
$$
The right-hand-side term can thus be controlled by the "First Order Main" term when  $\beta_{2}$  is close to 1.
Remark 1. Compared to the analysis of AdaGrad in [31], our proof technique has two-fold novelties. First, our auxiliary function has an additional  $(1 - \beta_{2})\sigma_{0}^{2}\mathbb{1}_{d}$  term, which is necessary for the analysis of Adam as it makes  $\widetilde{\nu}_t$  lower bounded from 0 (AdaGrad does not need this, as  $\pmb{\nu}_{t-1}$  of AdaGrad itself is lower bounded). Secondly, as discussed above, the "AdaGrad version" of second term in the right-hand-side of inequality (4) is a telescoping, the sum of which can be bounded straightforwardly.
Challenge II: Handle the mismatch between stochastic gradient and momentum. In the analysis above, we assume  $\beta_{1} = 0$ . Additional challenges arise when we move to the case where  $\beta_{1} \neq 0$ . Specifically, following the same routine, the "First Order Main" term now becomes  $\mathbb{E}\left[\langle G_t, -\eta \frac{1}{\sqrt{\tilde{\nu}_t}} \odot m_t\right]$ . It is hard to even estimate whether such a term is negative or not, given that  $m_t$  and  $\tilde{\nu}_t$  still has entangled stochasticity, and the conditional expectation of  $m_t$  also differs from  $G_t$ , both due to the existence of historical gradient.
Solution to Challenge II. Inspired by the state-of-art analysis of SGDM [22], which leverage the potential function  $f(v_{t})$  with  $v_{t} = \frac{\boldsymbol{w}_{t} - \beta\boldsymbol{w}_{t-1}}{1 - \beta}$ , we propose to use the potential function  $f(\boldsymbol{u}_{t})$  with  $\boldsymbol{u}_{t} = \frac{\boldsymbol{w}_{t} - \frac{\beta_{1}}{\sqrt{\beta_{2}}} \boldsymbol{w}_{t-1}}{1 - \frac{\beta_{1}}{\sqrt{\beta_{2}}}}$ . Applying descent lemma to  $f(\boldsymbol{u}_{t})$ , we obtain that
$$
\mathbb {E} [ f (\boldsymbol {u} _ {t + 1}) ] \leq \mathbb {E} f (\boldsymbol {u} _ {t}) + \underbrace {\mathbb {E} \left[ \langle \nabla f (\boldsymbol {u} _ {t}) , \boldsymbol {u} _ {t + 1} - \boldsymbol {u} _ {t} \rangle \right]} _ {\text {F i r s t O r d e r}} + \underbrace {\frac {L}{2} \mathbb {E} \| \boldsymbol {u} _ {t + 1} - \boldsymbol {u} _ {t} \| ^ {2}} _ {\text {S e c o n d O r d e r}}. \tag {5}
$$
We again focus on the "First Order" term, which can be written as
$$
\begin{array}{l} \mathbb {E} \left[ \left\langle \nabla f (\boldsymbol {u} _ {t}), \boldsymbol {u} _ {t + 1} - \boldsymbol {u} _ {t} \right\rangle \right] = \mathbb {E} \left[ \left\langle \nabla f (\boldsymbol {u} _ {t}), \frac {\boldsymbol {w} _ {t + 1} - \boldsymbol {w} _ {t}}{1 - \frac {\beta_ {1}}{\sqrt {\beta_ {2}}}} - \frac {\beta_ {1}}{\sqrt {\beta_ {2}}} \frac {\boldsymbol {w} _ {t} - \boldsymbol {w} _ {t - 1}}{1 - \frac {\beta_ {1}}{\sqrt {\beta_ {2}}}} \right\rangle \right] \\ \stackrel {(*)} {\approx} \mathbb {E} \left[ \left\langle \nabla f (\boldsymbol {w} _ {t}), - \frac {\eta}{1 - \frac {\beta_ {1}}{\sqrt {\beta_ {2}}}} \frac {1}{\sqrt {\boldsymbol {\nu} _ {t}}} \odot \boldsymbol {m} _ {t} + \frac {\eta}{1 - \frac {\beta_ {1}}{\sqrt {\beta_ {2}}}} \frac {\beta_ {1}}{\sqrt {\beta_ {2} \boldsymbol {\nu} _ {t - 1}}} \odot \boldsymbol {m} _ {t - 1} \right\rangle \right] \\ \stackrel {(\circ)} {\approx} \mathbb {E} \left[ \left\langle \nabla f (\boldsymbol {w} _ {t}), - \frac {\eta}{1 - \frac {\beta_ {1}}{\sqrt {\beta_ {2}}}} \frac {1}{\sqrt {\widetilde {\boldsymbol {\nu}} _ {t}}} \odot \boldsymbol {m} _ {t} + \frac {\eta}{1 - \frac {\beta_ {1}}{\sqrt {\beta_ {2}}}} \frac {\beta_ {1}}{\sqrt {\widetilde {\boldsymbol {\nu}} _ {t}}} \odot \boldsymbol {m} _ {t - 1} \right\rangle \right] \\ = \mathbb {E} \left[ \left\langle \boldsymbol {G} _ {t}, - \frac {\eta (1 - \beta_ {1})}{1 - \frac {\beta_ {1}}{\sqrt {\beta_ {2}}}} \frac {1}{\sqrt {\widetilde {\boldsymbol {\nu}} _ {t}}} \odot \boldsymbol {g} _ {t} \right\rangle \right] = \mathbb {E} \left[ \left\langle \boldsymbol {G} _ {t}, - \frac {\eta (1 - \beta_ {1})}{1 - \frac {\beta_ {1}}{\sqrt {\beta_ {2}}}} \frac {1}{\sqrt {\widetilde {\boldsymbol {\nu}} _ {t}}} \odot \boldsymbol {G} _ {t} \right\rangle \right]. \\ \end{array}
$$
Here approximate equation  $(\ast)$  is due to Assumption 1 and that  $\pmb{w}_t$  is close to  $\pmb{u}_t$ , and approximate equation  $(\circ)$  is due to Lemma 1 and  $\widetilde{\pmb{\nu}}_t = \beta_2\pmb{\nu}_{t-1} + (1 - \beta_2)\sigma_0^2 \approx \beta_2\pmb{\nu}_{t-1}$  (of course, these are informal statements. Please refer to Appendix C for the detailed proof). With the above methodology, we arrive at the following lemma.
Lemma 2 (Informal Version of Lemma 8). Let all conditions in Theorem 1 hold. Then,
$$
\mathbb {E} f (\boldsymbol {u} _ {t + 1}) \leq \mathbb {E} f (\boldsymbol {u} _ {t}) - \Omega \left(\mathbb {E} \left[ \eta \left\langle \boldsymbol {G} _ {t}, - \frac {1}{\sqrt {\tilde {\nu} _ {t}}} \odot \boldsymbol {G} _ {t} \right\rangle \right]\right) + \mathcal {O} \left(\frac {1}{\sqrt {\beta_ {2}}} \xi_ {t - 1} - \xi_ {t}\right) + S m a l l E r r o r.
$$
Summing the above lemma over  $t$  from 1 to  $T$ , we obtain
$$
\sum_ {t = 1} ^ {T} \mathbb {E} \left[ \left\| \frac {1}{\sqrt [ 4 ]{\widetilde {\boldsymbol {\nu}} _ {t}}} \odot \boldsymbol {G} _ {t} \right\| ^ {2} \right] \leq \mathcal {O} (1) + \sum_ {l = 1} ^ {d} \mathcal {O} \left(\mathbb {E} \ln \left(\frac {\boldsymbol {\nu} _ {t , i}}{\boldsymbol {\nu} _ {0 , l}}\right) - T \ln \beta_ {2}\right). \tag {6}
$$
We then encounter the second challenge.
Challenge III: Convert Eq. (6) to a bound of gradient norm. Although we have derived a regret bound, i.e., a bound of  $\sum_{t=1}^{T} \mathbb{E}\left[\left\|\frac{1}{\sqrt[4]{\tilde{\nu}_t}} \odot G_t\right\|^2\right]$ , we need to convert it into a bound of  $\mathbb{E}\left[\left\|G_t\right\|^2\right]$ . In existing works [36, 10, 16] which assumes bounded gradient, such a conversion is straightforward because (their version of)  $\widetilde{\nu}_t$  is upper bounded. However, we do not assume bounded gradient and  $\widetilde{\nu}_t$  can be arbitrarily large, making  $\mathbb{E}\left[\left\|\frac{1}{\sqrt[4]{\tilde{\nu}_t}} \odot G_t\right\|^2\right]$  arbitrarily small than  $\mathbb{E}\left[\left\|G_t\right\|^2\right]$ .
Solution to Challenge III. As this part involves coordinate-wise analysis, we define  $g_{t,i}, G_{t,i}, \nu_{t,i}$  and  $\widetilde{\nu}_{t,i}^{1}$  respectively as the  $l$ -th coordinate of  $g_{t}, G_{t}, \nu_{t}$ , and  $\widetilde{\nu}_{t}^{1}$ . To begin with, note that due to Cauchy's inequality and Hölder's inequality,
$$
\left(\mathbb {E} \sum_ {t = 1} ^ {T} \| \boldsymbol {G} _ {t} \|\right) ^ {2} \leq \left(\sum_ {t = 1} ^ {T} \mathbb {E} \left[ \left\| \frac {1}{\sqrt [ 4 ]{\widetilde {\boldsymbol {\nu}} _ {t}}} \odot \boldsymbol {G} _ {t} \right\| ^ {2} \right]\right) \left(\sum_ {t = 1} ^ {T} \mathbb {E} \left[ \left\| \sqrt [ 4 ]{\widetilde {\boldsymbol {\nu}} _ {t}} \right\| ^ {2} \right]\right). \tag {7}
$$
Therefore, we only need to derive an upper bound of  $\sum_{t=1}^{T} \mathbb{E}[\|\sqrt[4]{\widetilde{\boldsymbol{\nu}}_t}\|^2]$ , which is achieved by the following divide-and-conque methodology. Firstly, when  $|\pmb{G}_{t,i}| \geq \frac{\sigma_0}{\sigma_1}$ , we can show  $2\mathbb{E}^{|F_t|}|\pmb{g}_{t,i}|^2 \geq 2|\pmb{G}_{t,i}|^2 \geq \mathbb{E}^{|F_t|}|\pmb{g}_{t,i}|^2$ . Then, through a direct calculation, we obtain that
$$
\mathbb {E} \left[ \frac {\left| \boldsymbol {G} _ {t , i} \right| ^ {2}}{\sqrt {\tilde {\boldsymbol {\nu}} _ {t , i}}} \mathbf {1} _ {| G _ {t, i} | \geq \frac {\sigma_ {0}}{\sigma_ {1}}} \right] \geq \frac {\sqrt {\beta_ {2}}}{3 (1 - \beta_ {2}) \sigma_ {1} ^ {2}} \mathbb {E} \left[ \left(\sqrt {\tilde {\boldsymbol {\nu}} _ {t + 1 , i}} - \sqrt {\beta_ {2} \tilde {\boldsymbol {\nu}} _ {t , i}}\right) \mathbf {1} _ {| G _ {t, i} | \geq \frac {\sigma_ {0}}{\sigma_ {1}}} \right],
$$
$$
\sum_ {t = 1} ^ {T} \mathbb {E} \left[ \frac {| \boldsymbol {G} _ {t , i} | ^ {2}}{\sqrt {\widetilde {\boldsymbol {\nu}} _ {t , i}}} \right] \geq \frac {\sqrt {\beta_ {2}}}{3 (1 - \beta_ {2}) \sigma_ {1} ^ {2}} \sum_ {t = 1} ^ {T} \mathbb {E} \left[ \left(\sqrt {\widetilde {\boldsymbol {\nu}} _ {t + 1 , i}} - \sqrt {\beta_ {2} \widetilde {\boldsymbol {\nu}} _ {t , i}}\right) \mathbf {1} _ {| G _ {t, i} | \geq \frac {\sigma_ {0}}{\sigma_ {1}}} \right].
$$
Secondly, when  $|G_{t,i}| < \frac{\sigma_0}{\sigma_1}$ , define  $\{\bar{\nu}_{t,i}\}_{t=0}^{\infty}$  as  $\bar{\nu}_{0,l} = \nu_{0,l}, \bar{\nu}_{t,i} = \bar{\nu}_{t-1,i} + |g_{t,i}|^2 \mathbf{1}_{|G_{t,i}| < \frac{\sigma_0}{\sigma_1}}$ . One can easily observe that  $\bar{\nu}_{t,i} \leq \nu_{t,i}$ , and thus
$$
\begin{array}{l} \sum_ {t = 1} ^ {T} \mathbb {E} \left[ \left(\sqrt {\tilde {\boldsymbol {\nu}} _ {t + 1 , i}} - \sqrt {\beta_ {2} \tilde {\boldsymbol {\nu}} _ {t , i}}\right) \mathbf {1} _ {| \boldsymbol {G} _ {t, i} | <   \frac {\sigma_ {0} ^ {2}}{\sigma_ {1} ^ {2}}} \right] \\ \leq \sum_ {t = 1} ^ {T} \mathbb {E} \left(\sqrt {\beta_ {2} \bar {\nu} _ {t , i} + (1 - \beta_ {2}) \sigma_ {0} ^ {2}} - \sqrt {\beta_ {2} \left(\beta_ {2} \bar {\nu} _ {t - 1 , i} + (1 - \beta_ {2}) \sigma_ {0} ^ {2}\right)}\right) \\ = \mathbb {E} \sqrt {\beta_ {2} \bar {\nu} _ {t , i} + (1 - \beta_ {2}) \sigma_ {0} ^ {2}} + (1 - \sqrt {\beta_ {2}}) \sum_ {t = 1} ^ {T - 1} \mathbb {E} \sqrt {\beta_ {2} \bar {\nu} _ {t , i} + (1 - \beta_ {2}) \sigma_ {0} ^ {2}} - \mathbb {E} \sqrt {\beta_ {2} (\beta_ {2} \bar {\nu} _ {0 , i} + (1 - \beta_ {2}) \sigma_ {0} ^ {2})}. \\ \end{array}
$$
Putting the above two estimations together, we derive that
$$
(1 - \sqrt {\beta_ {2}}) \sum_ {t = 1} ^ {T + 1} \mathbb {E} \sqrt {\tilde {\boldsymbol {\nu}} _ {t , i}} \leq \frac {3 (1 - \beta_ {2}) \sigma_ {1} ^ {2}}{\sqrt {\beta_ {2}}} \sum_ {t = 2} ^ {T} \mathbb {E} \left[ \frac {| \boldsymbol {G} _ {t , i} | ^ {2}}{\sqrt {\tilde {\boldsymbol {\nu}} _ {t , i}}} \right] + (1 - \sqrt {\beta_ {2}}) (T + 1) \sqrt {\sigma_ {0} ^ {2} + \boldsymbol {\nu} _ {0 , i}}.
$$
The above methodology can be summarized as the following lemma.
Lemma 3. Let all conditions in Theorem 1 hold. Then,
$$
\sum_ {t = 1} ^ {T + 1} \sum_ {i = 1} ^ {d} \mathbb {E} \sqrt {\tilde {\boldsymbol {\nu}} _ {t , i}} \leq 2 (T + 1) \sum_ {i = 1} ^ {d} \sqrt {\boldsymbol {\nu} _ {0 , i} + \sigma_ {0} ^ {2}} + 2 4 d \frac {\sigma_ {1} ^ {2} C _ {1}}{\sqrt {\beta_ {2}}} \ln d \frac {\sigma_ {1} ^ {2} C _ {1}}{\sqrt {\beta_ {2}}} + C _ {2}.
$$
Based on Lemma 3, we can derive the estimation of  $\sum_{t=1}^{T} \mathbb{E}[\| \sqrt[4]{\widetilde{\nu}_t}^2]$  since  $\widetilde{\nu}_t$  is close to  $\pmb{\nu}_t$ . The proof is then completed by combining the estimation of  $\sum_{t=1}^{T} \mathbb{E}[\| \sqrt[4]{\widetilde{\nu}_t}^2]$  (Eq. (6)) and Eq. (7).
5 Gap-closing upper bound on the iteration complexity of Adam
In this section, based on a refined proof of Stage II of Theorem 1 (see Appendix C) under the specific case  $\eta = \Theta(1/\sqrt{T})$  and  $\beta_2 = 1 - \Theta(1/T)$ , we show that the logarithmic factor in Theorem 1 can be removed and the lower bound can be achieved. Specifically, we have the following theorem.
Theorem 2. Let Assumption 1 and Assumption 2 hold. Then, select the hyperparameters of Adam as  $\eta = \frac{a}{\sqrt{T}}$ ,  $\beta_{2} = 1 - \frac{b}{T}$  and  $\beta_{1} = c\sqrt{\beta_{2}}$ , where  $a, b > 0$  and  $0 \leq c < 1$  are independent of  $T$ . Then, let  $\mathbf{w}_{\tau}$  be the output of Adam in Algorithm 1, and we have
$$
\begin{array}{l} \mathbb {E} \| \nabla f (\boldsymbol {w} _ {r}) \| \leq \sqrt {2 \sum_ {i = 1} ^ {d} \sqrt {\nu_ {0 , i} + 3 b \sigma_ {0} ^ {2}} + \frac {4 D _ {2} \sigma_ {1} ^ {2} b}{\sqrt {T}} + \frac {2 5 6 \sigma_ {1} ^ {2} b}{(1 - c) ^ {2} T} \sum_ {i = 1} ^ {d} \mathbb {E} \frac {\boldsymbol {G} _ {1 , i} ^ {2}}{\sqrt {\tilde {\nu} _ {1 , i}}} + \frac {1 6 D _ {1} \sigma_ {1} ^ {2} b}{\sqrt {T}} \ln \left(e + \frac {4 \tilde {D} \sigma_ {1} ^ {2} b}{\sqrt {T}}\right)} \\ \times \sqrt {\frac {2 D _ {1}}{\sqrt {T}} \sum_ {i = 1} ^ {d} \ln \left(2 \sum_ {i = 1} ^ {d} \sqrt {\nu_ {0 , i} + 3 b \sigma_ {0} ^ {2}} + \frac {4 D _ {2} \sigma_ {1} ^ {2} b}{\sqrt {T}} + \frac {2 5 6 \sigma_ {1} ^ {2} b}{(1 - c) ^ {2} T} \sum_ {i = 1} ^ {d} \mathbb {E} \frac {\boldsymbol {G} _ {1 , i} ^ {2}}{\sqrt {\tilde {\nu} _ {1 , i}}} + \frac {1 6 D _ {1} \sigma_ {1} ^ {2} b}{\sqrt {T}} \ln \left(e + \frac {4 \tilde {D} \sigma_ {1} ^ {2} b}{\sqrt {T}}\right)\right)} + \frac {6 4}{(1 - c) ^ {2} T} \sum_ {i = 1} ^ {d} \mathbb {E} \frac {\boldsymbol {G} _ {1 , i} ^ {2}}{\sqrt {\tilde {\nu} _ {1 , i}}} + \frac {D _ {2}}{\sqrt {T}}} \\ \end{array}
$$
$$
D _ {1} \triangleq \frac {3 2 L a}{b} \frac {(1 + c) ^ {3}}{(1 - c) ^ {3}} + \frac {3 2 \sigma_ {0}}{\sqrt {b} (1 - c) ^ {3}} + \frac {(1 + \sigma_ {1} ^ {2}) \sigma_ {1} ^ {2} L ^ {2} d a ^ {2}}{(1 - c) ^ {4} \sigma_ {0} \sqrt {b ^ {3}}}, D _ {2} \triangleq \frac {8}{a} f (\boldsymbol {u} _ {1}) + D _ {1} \left(b d - \sum_ {i = 1} ^ {d} \ln \nu_ {0, i}\right).
$$
As a result, let  $\mathbf{A}$  be Adam in Algorithm 1, we have  $\mathcal{C}_{\varepsilon}(\mathbf{A}, \Delta, L, \sigma^2) = \mathcal{O}\left(\frac{1}{\varepsilon^4}\right)$ .
The proof of Theorem 2 is based on a refined solution of Challenge II in the proof of Theorem 1 under the specific hyperparameter settings, and we defer the concrete proof to Appendix D. Below we discuss on Theorem 2, comparing it with practice, with Theorem 1 and existing convergence rate of Adam, and with the convergence rate of AdaGrad.
Alignment with the practical hyperparameter choice. The hyperparameter setting in Theorem 2 indicates that to achieve the lower bound of iteration complexity, we need to select small  $\eta$  and close-to-1  $\beta_{2}$ , with less requirement over  $\beta_{1}$ . This agrees with the hyperparameter setting in deep learning libraries, for example,  $\eta = 10^{-3}$ ,  $\beta_{2} = 0.999$ , and  $\beta_{1} = 0.9$  in PyTorch.
Comparison with Theorem 1 and existing works. To our best knowledge, Theorem 2 is the first to derive the iteration complexity  $\mathcal{O}\left(\frac{1}{\varepsilon^4}\right)$ . Previously, the state-of-art iteration complexity is  $\mathcal{O}\left(\frac{\log 1 / \varepsilon}{\varepsilon^4}\right)$  [10] where they additionally assume bounded gradient. Theorem 2 is also tighter than Theorem 1 (while Theorem 1 holds for more general hyperparameter settings). As discussed in Section 4.1, if applying the hyperparameter setting in Theorem 2 (i.e.,  $\eta = \frac{a}{\sqrt{T}}$ ,  $\beta_{2} = 1 - \frac{b}{T}$  and  $\beta_{1} = c\sqrt{\beta_{2}}$ ) to Theorem 1, we will obtain that  $\mathbb{E}\| \nabla f(\boldsymbol{w}_{\tau})\| \leq \mathcal{O}(\mathrm{poly}(\log T) / \sqrt[4]{T})$  and  $\mathcal{C}_{\varepsilon}(\pmb{A},\Delta,L,\sigma^{2}) = \mathcal{O}\left(\frac{\log 1 / \varepsilon}{\varepsilon^4}\right)$ , worse than the upper bound in Theorem 2 and the lower bound in Proposition 1 by a logarithmic factor.