#  Adam with model exponential moving average is effective for nonconvex
optimization

Kwangjun Ahn
MIT/Microsoft Research
Cambridge, MA
<EMAIL_ADDRESS>Ashok Cutkosky
Boston University
Boston, MA
<EMAIL_ADDRESS>
###### Abstract

In this work, we offer a theoretical analysis of two modern optimization
techniques for training large and complex models: (i) adaptive optimization
algorithms, such as Adam, and (ii) the model exponential moving average (EMA).
Specifically, we demonstrate that a clipped version of Adam with model EMA
achieves the optimal convergence rates in various nonconvex optimization
settings, both smooth and nonsmooth. Moreover, when the scale varies
significantly across different coordinates, we demonstrate that the
coordinate-wise adaptivity of Adam is provably advantageous. Notably, unlike
previous analyses of Adam, our analysis crucially relies on its core
elements—momentum and discounting factors—as well as model EMA, motivating
their wide applications in practice.

##  1 Introduction

In neural network training, the training loss $F:\mathbb{R}^{d}\to\mathbb{R}$
is often optimized using an iterative optimization algorithm which starts with
the initial iterate ${\mathbf{x}}_{0}$ and then updates during each iteration
$t=1,2,\dots$ as follows:

$\displaystyle{\mathbf{x}}_{t}={\mathbf{x}}_{t-1}+{\mathbf{z}}_{t}\,,$ (1)

where ${\mathbf{z}}_{t}$ denotes the increment chosen by the algorithm during
the $t$-th iteration. One of the most popular optimization algorithms is Adam
(Kingma and Ba, 2014). Adam has gained significant attention due to its
effectiveness in training Transformer-based language models (Zhang et al.,
2020; Kunstner et al., 2023; Jiang et al., 2023; Pan and Li, 2023; Ahn et al.,
2024a; Kunstner et al., 2024; Zhang et al., 2024a).

The model exponential moving average (EMA) (Polyak and Juditsky, 1992;
Ruppert, 1988) is an optimization technique that has gained popularity in
conjunction with Adam for various recent applications. EMA maintains an
exponential moving average of the model iterates, ${\mathbf{x}}_{t}$, which
contributes to the stabilization of these iterates. There has been a
resurgence of interest in this technique due to its effectiveness in training
high-quality generative models (Yaz et al., 2018; Karras et al., 2019; Song et
al., 2021b; Dhariwal and Nichol, 2021; Nichol and Dhariwal, 2021; Song et al.,
2021a; Balaji et al., 2022; Karras et al., 2022; Rombach et al., 2022; Kang et
al., 2023; Karras et al., 2023). Moreover, a recent work by Block et al.
(2024) demonstrates the effectiveness of EMA for both language modeling and
imitation learning applications.

In this work, we theoretically study the effectiveness of these two modern
optimization techniques. Our main results can be informally summarized as
follows.

######  Theorem 1 (Informal).

clipped-Adam with the EMA on its iterates achieves the optimal convergence
rate for nonconvex optimization both for smooth and nonsmooth settings
(Section 5). The coordinate-wise adaptivity of Adam is particularly effective
when the scale varies across different coordinates (Section 6).

Our main results are based on the online-to-nonconvex conversion framework of
Cutkosky et al. (2023), which chooses the increment ${\mathbf{z}}_{t}$ based
on an online learner of choice. In particular, our approach is quite different
than the previous analyses of Adam (see Section 1.1 below). Notably, our
analysis relies on the key components of Adam (momentum and adaptive learning
rate) as well as EMA of the iterates, offering new, theoretical insight into
their success. See Section 7 for a more detailed discussion.

At a high level, our analysis combines the main insights from the two recent
works: Zhang and Cutkosky (2024) and Ahn et al. (2024b). We first carefully
modify the discounted-to-nonconvex conversion framework (7) of Zhang and
Cutkosky (2024) which converts an online learner that achieves a good
discounted regret (6) into a good noncovex optimizer. We then combine it with
the main insight of Ahn et al. (2024b) that an effective discounted online
learner can be designed based on scale-free Follow-the-Regularized-Leader
(FTRL) (Orabona and Pál, 2018). In particular, the way we arrive at Adam is
similar to Ahn et al. (2024b): choosing a discounted version of FTRL in the
discounted-to-nonconvex conversion leads to Adam.

###  1.1 Related work

Even though Adam is widely used in deep learning, our theoretical
understanding of its inner workings, especially the importance of its core
components—momentum and discounting factors—remains incomplete, as pointed out
by Ahn et al. (2024b). Most theoretical work on Adam and its variations focus
on characterizing the convergence rate for convex or smooth nonconvex
functions, where methods like SGD already achieve the minimax optimal
convergence rate. (Reddi et al., 2018; Zhou et al., 2019; Chen et al., 2019;
Zou et al., 2019; Alacaoglu et al., 2020; Guo et al., 2021; Défossez et al.,
2022; Zhang et al., 2022; Li et al., 2023; Wang et al., 2023). In fact, even
the most recent results (Li et al., 2023; Wang et al., 2023) are not
reflective of practice, in the sense that Adam’s convergence rate worsens with
momentum (Wang et al., 2023, §6) or is no better than that of SGD (Li et al.,
2023, §7). A notable exception is Crawshaw et al. (2022), which demonstrates
the advantages of momentum under the generalized smoothness conditions of
Zhang et al. (2020). However, the algorithm they analyze is signSGD, which
differs significantly from the original Adam. In contrast, as we will explore
in the subsequent sections, momentum and discounting factors are important in
our analysis. See Section 7 for more details.

We also highlight that our analysis relies on model EMA, a technique widely
used in practice as mentioned above (also see a recent work by Block et al.
(2024)). It is worth noting that EMA (or model averaging in general) has shown
to have generalization benefits in practice (Tarvainen and Valpola, 2017;
Izmailov et al., 2018). In this paper, we study EMA from an optimization
perspective, and show that the use of EMA leads to optimal guarantees for
nonconvex optimization. Interestingly, EMA naturally derives from the
discounted-to-online conversion (see Algorithm 1), which, we believe, provides
new theoretical insights into this practical method.

The use of EMA also represents a significant departure from most non-convex
optimization analyses. While EMA is a classical technique in the _convex_
setting, theoretical analyses in the non-convex setting typically randomly
select an iterate as the “final output” of the optimizer, rather than using
EMA. This random selection is intuitively extremely impractical (indeed, on
average it actually wastes half of the computation), and is never performed in
real implementations.

##  2 Setting for nonconvex and nonsmooth optimization

Throughout this paper, unless specified otherwise, $\left\|\cdot\right\|$
denotes the $L_{2}$ norm. Following (Cutkosky et al., 2023), we consider
optimizing a loss function $F$ that satisfies the following conditions,
accessing information about $F$ through a _stochastic gradient oracle_
$\textsc{StoGrad}:\mathbb{R}^{d}\times\mathcal{Z}\to\mathbb{R}^{d}$, for the
set of randomness $\mathcal{Z}$.

######  Assumption 2.

Let $F:\mathbb{R}^{d}\to\mathbb{R}$ be a differentiable function with the
following properties:

  * •

Let $\Delta\coloneqq F({\mathbf{x}}_{0})-\inf_{{\mathbf{x}}}F({\mathbf{x}})$.

  * •

For any two points ${\mathbf{x}}$ and $\mathbf{y}$,
$F({\mathbf{x}})-F(\mathbf{y})=\int_{0}^{1}\left\langle\nabla
F({\mathbf{x}}+t(\mathbf{y}-{\mathbf{x}})),\mathbf{y}-{\mathbf{x}}\right\rangle\mathrm{d}t$.

  * •

Lipschitzness. $F$ is $G$-Lipshitz, _i.e._ , for any point ${\mathbf{x}}$,
$\left\|\nabla F({\mathbf{x}})\right\|\leq G$.

  * •

Stochastic gradient variance. For any point ${\mathbf{x}}$, the stochastic
gradient $\mathbf{g}\leftarrow\textsc{StoGrad}({\mathbf{x}},r)$ for randomness
$r\in\mathcal{Z}$ satisfies $\mathop{\mathbb{E}}[\mathbf{g}]=\nabla
F({\mathbf{x}})$ and $\mathop{\mathbb{E}}\left\|\mathbf{g}-\nabla
F({\mathbf{x}})\right\|^{2}\leq\sigma^{2}$.

Here, the second condition, called _well-behavedess_ in (Cutkosky et al.,
2023, Definition 1), is a mild regularity condition. For any locally Lipschitz
function $F$, applying an arbitrarily small perturbation to the function is
sufficient to ensure this condition (Cutkosky et al., 2023, Proposition 2).

For the notion of optimality, we follow Zhang and Cutkosky (2024) and consider
the following notion of stationarity for nonconvex and nonsmooth functions.
This notion is a slight relaxation of the notion of a Goldstein stationarity
point (Goldstein, 1977), which was further studied by recent works (Zhang et
al., 2020; Davis et al., 2022; Tian et al., 2022; Jordan et al., 2023).

######  Definition 3 ($(\lambda,\varepsilon)$-stationary point).

Suppose $F:\mathbb{R}^{d}\to\mathbb{R}$ is differentiable. We say
${\mathbf{x}}$ is a $(\lambda,\varepsilon)$-stationary point of $F$ if
$\left\|\nabla F({\mathbf{x}})\right\|^{[\lambda]}\leq\varepsilon$, where

$\displaystyle\left\|\nabla
F({\mathbf{x}})\right\|^{[\lambda]}\coloneqq\inf_{\begin{subarray}{c}p\in\mathcal{P}(\mathbb{R}^{d}),\\\
\mathop{\mathbb{E}}_{\mathbf{y}\sim
p}[\mathbf{y}]={\mathbf{x}}\end{subarray}}\left\\{\left\|\mathop{\mathbb{E}}[\nabla
F(\mathbf{y})]\right\|+\lambda\cdot\mathop{\mathbb{E}}\left\|\mathbf{y}-{\mathbf{x}}\right\|^{2}\right\\}\,.$
(2)

To further motivate this definition, we remark that
$(\lambda,\epsilon)$-stationary points retain the desirable properties of
Goldstein stationary points. Specifically, the following result (Zhang and
Cutkosky, 2024, Lemma 2.3) demonstrates that, akin to Goldstein stationary
points, $(\lambda,\epsilon)$-stationary points can be reduced to first-order
stationary points with appropriate choices of $\lambda$ when the objective
function is smooth or second-order smooth.

######  Lemma 4.

If $F$ is $L$-smooth, then an $(L^{2}\varepsilon^{-1},\varepsilon)$-stationary
point ${\mathbf{x}}$ of $F$ satisfies $\left\|\nabla
F({\mathbf{x}})\right\|\leq 2\varepsilon$. Moreover, if $F$ is $H$-second-
order-smooth, then an $(H/2,\varepsilon)$-stationary point ${\mathbf{x}}$ of
$F$ satisfies $\left\|\nabla F({\mathbf{x}})\right\|\leq 2\varepsilon$.

Moreover, as shown by (Zhang and Cutkosky, 2024, Lemma 2.4),
$(\lambda,\varepsilon)$-stationary points can also be reduced to Goldstein
stationary points when $F$ is Lipschitz.

######  Lemma 5.

Suppose $F$ is $G$-Lipschitz. For any $\lambda,\varepsilon,\delta>0$, a
$(\lambda,\varepsilon)$-stationary point is a
$(\delta,\varepsilon^{\prime})$-Goldstein stationary point, where
$\varepsilon^{\prime}=(1+\frac{2G}{\lambda\delta^{2}})\cdot\varepsilon$.

Now we design algorithms that find $(\lambda,\varepsilon)$-stationary points
efficiently.

##  3 Discounted-to-nonconvex conversion: online learning of increments

Our main results are built on the online-to-nonconvex conversion framework of
Cutkosky et al. (2023). At its core, this framework involves selecting the
increment ${\mathbf{z}}_{t}$ using an online learner, as discussed by Ahn et
al. (2024b). Specifically, we follow a variant developed by Zhang and Cutkosky
(2024), which carefully incorporates the discounting factor in the conversion
process. Note that we make slight modifications to the version proposed by
Zhang and Cutkosky (2024) as follows. Here $\text{Exp}(1)$ denotes the
exponential random variable with mean $1$.

Algorithm 1 Discounted-to-nonconvex conversion (choosing increments via online
learning)

Input:  Initial point ${\mathbf{x}}_{0}$, $T\in\mathbb{N}$, online learning
algorithm $\mathcal{A}$, and discounting factor $\beta\in(0,1)$

for $t=1,2\dots,T$ do

Receive ${\mathbf{z}}_{t}$ from $\mathcal{A}{}$ // choose the increment using
an online learner

Update
${\mathbf{x}}_{t}\leftarrow{\mathbf{x}}_{t-1}+\alpha_{t}{\mathbf{z}}_{t}$,
where $\alpha_{t}{\sim}\text{Exp}(1)$ _i.i.d._

Compute $\mathbf{g}_{t}\leftarrow\textsc{StoGrad}({\mathbf{x}}_{t},r_{t})$
with freshly sampled randomness $r_{t}$

Send
$\ell^{[\beta]}_{t}({\mathbf{z}})\coloneqq\left\langle\beta^{-t}\mathbf{g}_{t},{\mathbf{z}}\right\rangle$
to $\mathcal{A}$

// Maintain exponential moving average (for output only):

Update
$\overline{\mathbf{x}}_{t}\leftarrow\frac{\beta-\beta^{t}}{1-\beta^{t}}\overline{\mathbf{x}}_{t-1}+\frac{1-\beta}{1-\beta^{t}}{\mathbf{x}}_{t}$
(Equivalently,
$\overline{\mathbf{x}}_{t}\leftarrow\frac{1-\beta}{1-\beta^{t}}\sum_{s=1}^{t}\beta^{t-s}{\mathbf{x}}_{s}$)

end for

Given Algorithm 1, it turns out we need to design an online learner that
minimizes the discounted regret, formally defined below. It is worth noting
that discounted regret has been recently studied with the goal of better
adapting online learners to dynamic environments (Ahn et al., 2024b; Zhang et
al., 2024b; Jacobsen and Cutkosky, 2024).

######  Definition 6 (Discounted regret).

For a comparator $\mathbf{u}$, the $\beta$-discounted regret is defined as

$\displaystyle\textup{Regret}^{[\beta]}_{t}(\mathbf{u})\coloneqq\beta^{t}\cdot\sum_{s=1}^{t}(\ell^{[\beta]}_{s}({\mathbf{z}}_{s})-\ell^{[\beta]}_{s}(\mathbf{u}))=\sum_{s=1}^{t}\left\langle\beta^{-t}{\mathbf{v}}_{s},{\mathbf{z}}_{s}-\mathbf{u}\right\rangle\,.$
(3)

The discounted regret of an online learner $\mathcal{A}$ can be used to upper
bound the norm of averaged gradients, as shown in the following result.

######  Lemma 7 (Discounted-to-nonconvex conversion).

Suppose that $F$ satisfies 2. Then for the comparator sequence chosen as
$\mathbf{u}_{t}\coloneqq-D\frac{\sum_{s=1}^{t}\beta^{-s}\nabla
F({\mathbf{x}}_{s})}{\left\|\sum_{s=1}^{t}\beta^{-s}\nabla
F({\mathbf{x}}_{s})\right\|}$, Algorithm 1 gives

$\displaystyle\mathop{\mathbb{E}}_{t\sim[T]}\mathop{\mathbb{E}}\left\|\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\nabla
F(\mathbf{y}_{t})\right\|$
$\displaystyle\leq\frac{\Delta}{DT}+\frac{2G+\sigma}{(1-\beta)T}+\sigma\sqrt{1-\beta}$
(4)
$\displaystyle\quad+\frac{1}{DT}\left[\beta\cdot\mathop{\mathbb{E}}\left[\textup{Regret}^{[\beta]}_{T}(\mathbf{u}_{T})\right]+(1-\beta)\cdot\sum_{t=1}^{T}\mathop{\mathbb{E}}\left[\textup{Regret}^{[\beta]}_{t}(\mathbf{u}_{t})\right]\right]\,,$
(5)

where $\mathbf{y}_{t}$ is distributed over $\\{{\mathbf{x}}_{s}\\}_{s=1}^{t}$
as
$\mathbb{P}(\mathbf{y}_{t}={\mathbf{x}}_{s})=\beta^{t-s}\cdot\frac{1-\beta}{1-\beta^{t}}$
for $s=1,2,\dots,t$.

The proof combines the techniques of (Cutkosky et al., 2023, Theorem 7) and
(Zhang and Cutkosky, 2024, Theorem 3.3). See Appendix A for details.

Given this, our goal is to design an online learner that achieves a low
discounted regret.

##  4 Scale-free Follow-the-Regularized-Leader (FTRL)

In this section, we introduce an algorithmic component, called the Followed-
The-Regularized-Leader (FTRL), a powerful online learning technique with
various applications (Gordon, 1999; Kalai and Vempala, 2005; Shalev-Shwartz
and Singer, 2006; Abernethy et al., 2008; Nesterov, 2009; Hazan and Kale,
2010).

For the setting, consider the online linear optimization (OLO) setting, where
during each round $t=1,\dots,T$, and online learner chooses
${\mathbf{z}}_{t}$, and then the linear loss
$\ell_{t}(\cdot)=\left\langle{\mathbf{v}}_{t},\cdot\right\rangle$ is revealed
by the environment. Here the goal of the online learner is to minimize the
regret defined as
$\sum_{t}\left\langle{\mathbf{v}}_{t},{\mathbf{z}}_{t}-\mathbf{u}\right\rangle$,
where $\mathbf{u}$ is the comparator in hindsight. For this setting, FTRL is
presented in Algorithm 2.

Algorithm 2 Follow-the-Regularized-Leader (FTRL)

0: Regularizers $\\{\psi_{t}\\}:\mathbb{R}^{d}\to\mathbb{R}$, the domain
$\mathcal{D}\subseteq\mathbb{R}^{d}$

1: for $t=1,2,\dots,T$ do

2: Update
${\mathbf{z}}_{t}\leftarrow\mathop{\text{argmin}}_{{\mathbf{z}}\in\mathcal{D}}\
\left[\psi_{t}({\mathbf{z}})+\sum_{s=1}^{t-1}\ell_{s}({\mathbf{z}})\right]$

3: Receive the next loss
$\ell_{t}(\cdot)=\left\langle{\mathbf{v}}_{t},\cdot\right\rangle$

4: end for

The key insight of Ahn et al. (2024b) and Zhang et al. (2024b) is that in
order to design an online learner for discounted regret, it is important that
the online learner is _scale-free_ as described below. In particular,
following Ahn et al. (2024b), we consider a gradient adaptive scale-free FTRL
algorithm called _scale-free FTRL_ (Orabona and Pál, 2018).

We will focus on the case where $\mathcal{D}=\mathbb{B}_{D}$, the
$d$-dimensional $L_{2}$-ball of radius $D>0$. Scale-free FTRL is given by
Algorithm 2 with the regularizer
$\psi_{t}(\cdot)=\frac{1}{\eta_{t}}\left\|\cdot\right\|^{2}$ step size
$\eta_{t}=\nicefrac{{D}}{{\sqrt{\sum_{s=1}^{t-1}\left\|{\mathbf{v}}_{s}\right\|^{2}}}}$.
Then using the clipping operator
$\mathrm{clip}_{D}({\mathbf{x}}):={\mathbf{x}}\min(\nicefrac{{D}}{{\left\|{\mathbf{x}}\right\|}},1)$,
we can write down the update rule more explicitly as follows:

$\displaystyle{\mathbf{z}}_{t}\leftarrow\mathop{\text{argmin}}_{{\mathbf{z}}\in\mathbb{B}_{D}}\left[\frac{1}{\eta_{t}}\left\|{\mathbf{z}}\right\|^{2}+\sum_{s=1}^{t-1}\left\langle{\mathbf{v}}_{s},{\mathbf{z}}\right\rangle\right]=-\mathrm{clip}_{D}\left(\eta\frac{\sum_{s=1}^{t-1}{\mathbf{v}}_{s}}{\sqrt{\sum_{s=1}^{t-1}\left\|{\mathbf{v}}_{s}\right\|^{2}}}\right)\,.$
(scale-free FTRL)

Here, if the denominator is zero, _i.e._ ,
${\mathbf{v}}_{1}=\cdots={\mathbf{v}}_{t-1}=\mathbf{0}$, then we set
${\mathbf{z}}_{t}\leftarrow 0$. Note that this algorithm is scale-free in the
sense that when the loss sequence is scaled by a scalar $c>0$, the updates
remain the same.

Let us now present the regret bound of scale-free FTRL.

######  Lemma 8 (Gradient-adaptive regret bound).

For any $T>0$, loss sequence ${\mathbf{v}}_{1:T}$ and comparator
$\mathbf{u}\in\mathbb{R}^{d}$ s.t. $\left\|\mathbf{u}\right\|\leq D$, scale-
free FTRL guarantees the following regret bound:

$\displaystyle\sum_{t=1}^{T}\left\langle{\mathbf{v}}_{t},{\mathbf{z}}_{t}-\mathbf{u}\right\rangle\leq
4D\sqrt{\sum_{t=1}^{T}\left\|{\mathbf{v}}_{t}\right\|^{2}}\,.$ (6)

We note that lemma 8 follows (with a slightly worse constant) from (Orabona
and Pál, 2018, Theorem 1), and the version we invoke is here due to (Ahn et
al., 2024b, Theorem A.1).

Recall from 7 that an online learner for the discounted-to-nonconvex
conversion (Algorithm 1) needs to have a low discounted regret. To achieve
this, following Ahn et al. (2024b) and Zhang et al. (2024b), we simply
substitute ${\mathbf{v}}_{t}\leftarrow\beta^{-t}\mathbf{g}_{t}$ into scale-
free FTRL, resulting in the update

$\displaystyle{\mathbf{z}}_{t}\leftarrow-\mathrm{clip}_{D}\left(D\frac{\sum_{s=1}^{t-1}\beta^{-s}\mathbf{g}_{s}}{\sqrt{\sum_{s=1}^{t-1}\beta^{-2s}\left\|\mathbf{g}_{s}\right\|^{2}}}\right)\,.$
($\beta$-FTRL)

Here again, if the denominator is zero, _i.e._ ,
$\mathbf{g}_{1}=\cdots=\mathbf{g}_{t-1}=\mathbf{0}$, then we set
${\mathbf{z}}_{t}\leftarrow 0$. Then, the following result characterizes the
discounted regret guarantee of $\beta$-FTRL.

######  Theorem 9 (Discounted regret bound).

Let $\beta\in(0,1]$. For any $T>0$, loss sequence $\mathbf{g}_{1:T}$ and
comparator $\mathbf{u}\in\mathbb{R}^{d}$ s.t. $\left\|\mathbf{u}\right\|\leq
D$, $\beta$-FTRL guarantees the following static regret bound

$\displaystyle\textup{Regret}^{[\beta]}_{T}(\mathbf{u})\leq
4D\sqrt{\sum_{t=1}^{T}\beta^{2(T-t)}\left\|\mathbf{g}_{t}\right\|^{2}}\,.$ (7)

We next use this result to design an algorithm for nonconvex optimization.

##  5  Discounted-FTRL leads to adaptive nonconvex optimization

In this section, as a warm-up, let us see the implications of choosing
$\mathcal{A}=\ref{dftrl}$ in Algorithm 1. First, let us obtain a bound on the
expected discounted regret. By Theorem 9 together with Jensen’s inequality, we
have the following regret bound for any $t=1,2,\dots,T$:

$\displaystyle\mathop{\mathbb{E}}\left[\textup{Regret}^{[\beta]}_{t}(\mathbf{u}_{t})\right]\leq
4D\mathop{\mathbb{E}}\sqrt{\sum_{t=1}^{T}\beta^{2(T-t)}\left\|\mathbf{g}_{t}\right\|^{2}}\leq
4D\sqrt{\sum_{t=1}^{T}\beta^{2(T-t)}\mathop{\mathbb{E}}\left\|\mathbf{g}_{t}\right\|^{2}}\,.$
(8)

Since $\mathop{\mathbb{E}}\left\|\mathbf{g}_{t}\right\|^{2}\leq
G^{2}+\sigma^{2}$ and
$\frac{1}{\sqrt{1-\beta^{2}}}\leq\frac{1}{\sqrt{1-\beta}}$, it follows that

$\displaystyle\mathop{\mathbb{E}}\left[\textup{Regret}^{[\beta]}_{t}(\mathbf{u}_{t})\right]\leq\frac{4D(G+\sigma)}{\sqrt{1-\beta^{2}}}\leq\frac{4D(G+\sigma)}{\sqrt{1-\beta}}\,.$
(9)

###  5.1 From gradient-adaptive regret to nonconvex optimization

In order to obtain nonconvex optimization guarantees in terms of the
$(\lambda,\varepsilon)$-stationarity (3), we need to handle the variance term.
Following (Zhang and Cutkosky, 2024, Lemma 3.2), the variance term can be
bounded as follows.

######  Lemma 10 (Variance bound).

Using the notations of 7, for any $t=1,2,\dots,T$, $\beta$-FTRL satisfies

$\displaystyle\mathop{\mathbb{E}}_{t\sim[T]}\mathop{\mathbb{E}}\left\|\mathbf{y}_{t}-\overline{\mathbf{x}}_{t}\right\|^{2}\leq
12\frac{D^{2}}{(1-\beta)^{2}}\,.$ (10)

###### Proof.

From (Zhang and Cutkosky, 2024, Lemma 3.2), it follows that
$\mathop{\mathbb{E}}\sum_{t=1}^{T}\left\|\mathbf{y}_{t}-\overline{\mathbf{x}}_{t}\right\|^{2}\leq\frac{12}{(1-\beta)^{2}}\mathop{\mathbb{E}}\sum_{t=1}^{T}\left\|{\mathbf{z}}_{t}\right\|^{2}$.
Now since $\left\|{\mathbf{z}}_{t}\right\|\leq D$ for all $t=1,2,\dots,T$,
after dividing each side by $T$, we get the desired inequality. ∎

Plugging the regret bound (9) into 7 and combining it with lemma 10, we arrive
at the following optimization guarantee in terms of the
$(\lambda,\varepsilon)$-stationarity. See Section B.1 for a proof.

######  Theorem 11.

Suppose that $F$ satisfies 2 and consider any $\lambda>0$. For $\alpha>0$,
choose $\mathcal{A}=\ref{dftrl}$ in Algorithm 1 with the following parameters:

$\displaystyle\beta=1-\left(\frac{\varepsilon}{10\alpha}\right)^{2},~{}~{}D=\frac{(1-\beta)\varepsilon^{1/2}}{4\lambda^{1/2}},~{}~{}\text{and}~{}~{}T=\frac{1}{1-\beta}\cdot\max\left\\{\frac{4\Delta\lambda^{1/2}}{\varepsilon^{3/2}},~{}\frac{12\alpha}{\varepsilon}\right\\}\,.$
(11)

Then we have $\mathop{\mathbb{E}}_{t\sim[T]}\mathop{\mathbb{E}}\left\|\nabla
F(\overline{\mathbf{x}}_{t})\right\|^{[\lambda]}\leq(1+\frac{G+\sigma}{\alpha})\varepsilon$.
In other words, a randomly chosen model EMA $\overline{\mathbf{x}}_{t}$ is a
$(\lambda,(1+\frac{G+\sigma}{\alpha})\varepsilon)$-stationary point, in
expectation.

###  5.2 Optimality and gradient adaptivity

Here, we discuss several notable aspects of the guarantee provided in 11.

####  5.2.1 Optimality

As shown in (Zhang and Cutkosky, 2024, Corollary 5.1), the lower bound on the
iteration complexity for finding a $(\lambda,\varepsilon)$-stationary point is
$\Omega((G+\sigma)^{2}\Delta\lambda^{1/2}\varepsilon^{-7/2})$, provided that
$\lambda\leq\frac{G^{4}}{\Delta^{2}}\varepsilon^{-1}$. 11 implies that setting
$\alpha=G+\sigma$ achieves this optimal iteration complexity.

######  Corollary 12.

In 11, choosing $\alpha=G+\sigma$ leads to the following iteration complexity
for finding a $(\lambda,\varepsilon)$-stationary point:

$\displaystyle
O\left(\max\left\\{\frac{(G+\sigma)^{2}\Delta\lambda^{1/2}}{\varepsilon^{7/2}},~{}\frac{(G+\sigma)^{3}}{\varepsilon^{3}}\right\\}\right)\,.$
(12)

In particular, treating $G$, $\sigma$, and $\Delta$ as constants, as long as
$\lambda\gtrsim\varepsilon$, this leads to the optimal complexity of
$O((G+\sigma)^{2}\Delta\lambda^{1/2}\varepsilon^{-7/2})$.

In light of lemma 4, the above optimal complexity can be converted into the
optimal complexities for smooth settings.

######  Corollary 13 (Smooth settings).

12 implies the following optimal iteration complexity for smooth settings.
Choosing $\lambda=O(\varepsilon^{-1})$, it implies the optimal complexity of
$O(\varepsilon^{-4})$ for smooth loss functions (Arjevani et al., 2023).
Similarly, with $\lambda=O(1)$, it achieves the optimal iteration complexitiy
of $O(\varepsilon^{-7/2})$ for second-order smooth loss functions (Arjevani et
al., 2020).

We next discuss the benefits of using the gradient-adaptive regret bound
(Theorem 9) by considering the case where we do not have knowledge of
$G,\sigma$.

####  5.2.2 Gradient adaptivity

A remarkable consequence of 11 is that, due to the gradient-adaptive regret
bound of Theorem 9, the final convergence guarantee has a better dependence on
$G,\sigma$ in the case when we do not have knowledge of them. For
concreteness, in the following discussion, we treat $G,\sigma,\Delta$ as
constants, and focus on the regime $\lambda\gtrsim\varepsilon$.

First, our 11 with $\alpha=1$ (since we do not know $G,\sigma$) leads to the
following iteration complexity for finding a
$(\lambda,\varepsilon)$-stationary point:

$\displaystyle
O\left((G+\sigma)^{7/2}\Delta\lambda^{1/2}\varepsilon^{-7/2}\right)$ (13)

The price we pay for not knowing $G,\sigma$ relative to the lower bound is a
multiplicative factor of $(G+\sigma)^{3/2}$. To see the benefit of this
adaptive regret approach, let us consider the guarantees given by Zhang and
Cutkosky (2024). Their approach is based on choosing online gradient descent
for $\mathcal{A}$ in Algorithm 1, when the learning rate is not properly tuned
with the knowledge of $G$ and $\sigma$, it would lead to the following
(suboptimal) discounted regret bound:

$\displaystyle\mathop{\mathbb{E}}\left[\textup{Regret}^{[\beta]}_{t}(\mathbf{u}_{t})\right]\leq
O\left(\frac{D(G+\sigma)^{2}}{\sqrt{1-\beta}}\right)\,.$ (14)

Then, the resulting iteration complexity becomes
$O(\Delta\lambda^{1/2}(\frac{\varepsilon}{(G+\sigma)^{2}})^{-7/2})$, which is
equal to $O\left((G+\sigma)^{7}\Delta\lambda^{1/2}\varepsilon^{-7/2}\right)$.
This is larger than the complexity due to our adaptive approach by a
multiplicative factor of $(G+\sigma)^{7/2}$.

Next, we build on the results from this section and consider a better approach
to design an adaptive nonconvex optimizer.

##  6 Coordinate-wise adaptivity via (clipped-)Adam

In this section, we consider the setting where the Lipschitzness constants
vary across different coordinates, which is empirically observed to be
reflective of practical neural network training (see, _e.g._ (Crawshaw et al.,
2022; Zhuang et al., 2022)). Formally, we consider the following setting.

######  Assumption 14.

Under the same setting as 2, we replace the last two conditions with the
following coordinate-wise version:

  * •

For each coordinate $i=1,2,\dots,d$, there is a Lipschitzness constant
$G_{i}>0$ and a variance constant $\sigma_{i}>0$ such that
$\forall{\mathbf{x}}$, $\left|\partial_{i}F({\mathbf{x}})\right|\leq G_{i}$
and the stochastic gradient
$\mathbf{g}\leftarrow\textsc{StoGrad}({\mathbf{x}},r)$ satisfies
$\mathop{\mathbb{E}}[\mathbf{g}[i]]=\partial_{i}F({\mathbf{x}}_{i})$ and
$\mathop{\mathbb{E}}\left|{\mathbf{g}[i]-\partial_{i}F({\mathbf{x}})}\right|^{2}\leq\sigma_{i}^{2}$.
(Here, $\partial_{i}F$ denotes the partial derivative of $F$ w.r.t. the $i$-th
coordinate.)

Let $\boldsymbol{G}\coloneqq(G_{1},\dots,G_{d})$ and
$\boldsymbol{\sigma}\coloneqq(\sigma_{1},\dots,\sigma_{d})$. Then, the above
condition implies the last two conditions in 2 with
$G=\left\|\boldsymbol{G}\right\|_{2}$ and
$\sigma=\left\|\boldsymbol{\sigma}\right\|_{2}$.

As we mentioned before, the previous approaches (Cutkosky et al., 2023; Zhang
and Cutkosky, 2024) choose the online learner $\mathcal{A}$ to be online
gradient descent, and hence choosing the learning rate requires the knowledge
of $G_{i},\sigma_{i}$ for all $i$. However, for neural network training, $d$
is equal to the number of parameters in the network, so tuning them
individually is computationally infeasible. We instead consider running
$\beta$-FTRL coordinate-wise in Algorithm 1, which will automatically adapt to
each coordinate. We begin with an important observation that such an approach
in fact leads to a popular optimizer widely used in practice.

###  6.1 Coordinate-wise discounted FTRL corresponds to (clipped-)Adam

For notational simplicity, fix a coordinate among $i=1,2,\dots,d$, and let us
denote the iterate by $x_{t}$, the stochastic gradient by $g_{t}$, and the
update by $z_{t}$. Then the resulting optimizer becomes:

$\displaystyle\boxed{x_{t+1}\leftarrow x_{t}-z_{t+1}\quad\text{where}\quad
z_{t+1}=-\mathrm{clip}_{D}\left(D\frac{\sum_{s=1}^{t}\beta_{1}^{t-s}g_{s}}{\sqrt{\sum_{s=1}^{t}\beta_{2}^{t-s}g_{s}^{2}}}\right)\,,}$
(clipped-Adam)

where $\beta_{1}=\beta$ and $\beta_{2}=\beta^{2}$. Here, again if the
denominator is zero, _i.e._ , if $g_{1}=\dots=g_{t}=0$, then we set the update
to be zero, _i.e._ , $z_{t}=0$. Note that this is almost exactly the Adam
optimizer (Kingma and Ba, 2014), except that now we add clipping to control
the variance of the iterates relative to their EMA. Notably, clipped-Adam
retains one of the most important properties of Adam: it is _scale-invariant_.
The scale invariance causes the optimizer to make updates of the same
magnitude on each coordinate even when the scale differs across different
coordinates.

We also remark that clipped-Adam does not consider the “bias correction” terms
in the original updates of Adam (Kingma and Ba, 2014). However, note that the
bias correction terms are coordinate-independent, and they can be merged into
the scalar $D$.

###  6.2 Nonconvex optimization guarantees of clipped-Adam

We next discuss the theoretical guarantees of clipped-Adam for nonconvex and
nonsmooth optimizaton. Inspired by (Duchi et al., 2010; McMahan and Streeter,
2010), where the coordinate-wise online learners lead to regret bounds with
respect to the $L_{1}$ norms of stochastic gradients, we consider the
following variant of 3, in the same vein as (Cutkosky et al., 2023, Section
4).

######  Definition 15 ($(\lambda,\varepsilon)$-$L_{1}$-stationary point).

Suppose $F:\mathbb{R}^{d}\to\mathbb{R}$ is differentiable. We say
${\mathbf{x}}$ is a $(\lambda,\varepsilon)$-$L_{1}$-stationary point of $F$ if
$\left\|\nabla F({\mathbf{x}})\right\|^{[\lambda]}_{1}\leq\varepsilon$, where

$\displaystyle\left\|\nabla
F({\mathbf{x}})\right\|^{[\lambda]}_{1}\coloneqq\inf_{\begin{subarray}{c}p\in\mathcal{P}(\mathbb{R}^{d}),\\\
\mathop{\mathbb{E}}_{\mathbf{y}\sim
p}[\mathbf{y}]={\mathbf{x}}\end{subarray}}\left\\{\left\|\mathop{\mathbb{E}}[\nabla
F(\mathbf{y})]\right\|_{1}+\lambda\cdot\mathop{\mathbb{E}}\left\|\mathbf{y}-{\mathbf{x}}\right\|_{2}^{2}\right\\}\,.$
(15)

Using the fact $\left\|\cdot\right\|_{1}\leq\sqrt{d}\left\|\cdot\right\|_{2}$,
one can connect the two notions of $(\lambda,\varepsilon)$-stationary points.

######  Lemma 16.

A $({\lambda}/{\sqrt{d}},{\varepsilon}/{\sqrt{d}})$-stationary point is a
$(\lambda,\varepsilon)$-$L_{1}$-stationary point.

In order to obtain the guarantee in terms of $L_{1}$-norm, we consider the
coordinate-wise version of discounted-to-online conversion, in the same vein
as (Cutkosky et al., 2023, Appendix G). See Section A.1 for details.

######  Lemma 17 ($L_{1}$-variant of 7).

Suppose that $F$ satisfies 14. Consider the comparator sequence chosen as
$\mathbf{u}_{t}$ defined as
$\mathbf{u}_{t}[i]\coloneqq-D\frac{\sum_{s=1}^{t}\beta^{-s}\partial_{i}F({\mathbf{x}}_{s})}{\left|\sum_{s=1}^{t}\beta^{-s}\partial_{i}F({\mathbf{x}}_{s})\right|}$
for $i=1,2,\dots,d$. Then, Algorithm 1 gives

$\displaystyle\mathop{\mathbb{E}}_{t\sim[T]}\mathop{\mathbb{E}}\left\|\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\nabla
F(\mathbf{y}_{t})\right\|_{1}$
$\displaystyle\leq\frac{\Delta}{DT}+\frac{2\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{1}}{(1-\beta)T}+\left\|\boldsymbol{\sigma}\right\|_{1}\sqrt{1-\beta}$
(16)
$\displaystyle\quad+\frac{1}{DT}\left[\beta\cdot\mathop{\mathbb{E}}\left[\textup{Regret}^{[\beta]}_{T}(\mathbf{u}_{T})\right]+(1-\beta)\cdot\sum_{t=1}^{T}\mathop{\mathbb{E}}\left[\textup{Regret}^{[\beta]}_{t}(\mathbf{u}_{t})\right]\right]\,,$
(17)

where $\mathbf{y}_{t}$ is distributed over $\\{{\mathbf{x}}_{s}\\}_{s=1}^{t}$
as
$\mathbb{P}(\mathbf{y}_{t}={\mathbf{x}}_{s})=\beta^{t-s}\cdot\frac{1-\beta}{1-\beta^{t}}$
for $s=1,2,\dots,t$.

Next, let us consider the (expected) regret bound. Fix a coordinate
$i=1,\dots,d$. Then, by the one-dimensional version of Theorem 9 together with
Jensen’s inequality, we have the following regret bound for any
$t=1,2,\dots,T$:

$\displaystyle\mathop{\mathbb{E}}\left[\textup{Regret}^{[\beta]}_{t}(\mathbf{u}_{t}[i])\right]\leq
4D\sqrt{\sum_{t=1}^{T}\beta^{2(T-t)}\mathop{\mathbb{E}}\left|{\mathbf{g}_{t}[i]}\right|^{2}}\leq\frac{4D(G_{i}+\sigma_{i})}{\sqrt{1-\beta}}\,.$
(18)

Hence, taking the sum over all coordinates $i=1,\dots,d$, we obtain

$\displaystyle\mathop{\mathbb{E}}\left[\textup{Regret}^{[\beta]}_{t}(\mathbf{u}_{t})\right]\leq\frac{4D\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{1}}{\sqrt{1-\beta}}\,.$
(19)

Combining these together, one get the following guarantee in terms of the
$L_{1}$ norm. See Section B.2 for a proof.

######  Theorem 18.

Suppose that $F$ satisfies 14 and consider any $\lambda>0$. For $\alpha>0$,
choose the coordinate-wise optimizer clipped-Adam in Algorithm 1 with the
following parameters:

$\displaystyle\beta=1-\left(\frac{\varepsilon}{10\alpha}\right)^{2},~{}~{}D=\frac{(1-\beta)\varepsilon^{1/2}}{4d^{1/2}\lambda^{1/2}},~{}~{}\text{and}~{}~{}T=\frac{1}{1-\beta}\cdot\max\left\\{\frac{4\Delta
d^{1/2}\lambda^{1/2}}{\varepsilon^{3/2}},~{}\frac{12\alpha}{\varepsilon}\right\\}\,.$
(20)

Then we have $\mathop{\mathbb{E}}_{t\sim[T]}\mathop{\mathbb{E}}\left\|\nabla
F(\overline{\mathbf{x}}_{t})\right\|^{[\lambda]}_{1}\leq(1+\frac{\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{1}}{\alpha})\varepsilon$.
In other words, a randomly chosen model EMA $\overline{\mathbf{x}}_{t}$ is a
$(\lambda,(1+\frac{\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{1}}{\alpha})\varepsilon)$-$L_{1}$-stationary
point, in expectation.

###  6.3 Benefits of coordinate-wise adaptivity of clipped-Adam

In this section, we discuss the benefits of coordinate-wise adaptivity by
examining the guarantee from 18 and compare it with that of 11. We begin with
the $(\lambda,\varepsilon)$-$L_{1}$-stationary point guarantee due to 18. We
consider the scenario where $\beta$ is carefully tuned by making the optimal
choice of $\alpha$.

######  Corollary 19.

In 18, choosing $\alpha=\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{1}$
leads to the following iteration complexity for finding a
$(\lambda,\varepsilon)$-$L_{1}$-stationary point:

$\displaystyle
O\left(\max\left\\{\frac{\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{1}^{2}\Delta
d^{1/2}\lambda^{1/2}}{\varepsilon^{7/2}},~{}\frac{\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{1}^{3}}{\varepsilon^{3}}\right\\}\right)\,.$
(21)

In order to better appreciate the benefits of coordinate-wise adaptivity, let
us compare the above iteration complexity with that of 11.

For concreteness, we treat $G=\left\|\boldsymbol{G}\right\|_{2}$ and
$\sigma=\left\|\boldsymbol{\sigma}\right\|_{2}$ as constants throughout, and
more importantly, we assume that the coordinates are heterogeneous in the
sense that

$\displaystyle\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{1}\approx\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{2}\,.$
(22)

The assumption (22) roughly says that a few coordinates of
$\boldsymbol{G}+\boldsymbol{\sigma}$ take much larger values than the rest; if
all the coordinates of $\boldsymbol{G}+\boldsymbol{\sigma}$ have similar
magnitudes, then
$\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{1}\approx\sqrt{d}\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{2}$.
In the case $\lambda\gtrsim\varepsilon$, 19 implies that the iteration
complexity is

$\displaystyle
O(\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{1}^{2}\Delta
d^{1/2}\lambda^{1/2}\varepsilon^{-7/2})\,.$ (23)

Next, let us consider the counterpart that does not adapt to each coordinate
separately. In this case, we apply lemma 16, which tells us that to find a
$(\lambda,\varepsilon)$-$L_{1}$ stationary point it suffices to find a
$({\lambda}/{\sqrt{d}},{\varepsilon}/{\sqrt{d}})$-stationary point. Then, from
12, the iteration complexity is
$O(\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{2}^{2}\Delta(\lambda/\sqrt{d})^{1/2}(\varepsilon/\sqrt{d})^{-7/2})$,
_i.e._ ,

$\displaystyle
O(\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{2}^{2}\Delta
d^{3/2}\lambda^{1/2}\varepsilon^{-7/2})\,.$ (24)

Hence, when (22) holds, (23) can be lower than (24) by a multiplicative factor
of $d$, showing the benefits of coordinate-wise adaptivity.

##  7 Discussion

Our analyses of Adam based on the discounted-to-online conversion is quite
different than the previous ones. As discussed in Section 1.1, the previous
analyses often result in guarantees that are not quite reflective of practice—
_e.g._ , the rates get better without momentum and the rates are no better
than that of non-adaptive methods. In contrast, our analyses and results
highlight the role of the practical components as highlighted below.

  * •

Momentum. In order to obtain a low discounted regret, any sensible online
learners should combine the past history of stochastic gradients
$\mathbf{g}_{1:t}$ to make the decision ${\mathbf{z}}_{t+1}$. Such online
learners under the discounted-to-online conversion lead to momentum methods
that combine $\mathbf{g}_{1:t}$ to obtain the next increment
${\mathbf{z}}_{t+1}$. In particular, non-momentum methods would correspond to
aggressive online learners that only use the last gradient $\mathbf{g}_{t}$ to
make the decision ${\mathbf{z}}_{t+1}$. This perspective provides new insights
into understanding the role of momentum, as echoed by Ahn et al. (2024b).

  * •

Adaptive learning rates. The adaptive learning rate due to $\beta$-FTRL leads
to a gradient-adaptive regret bound (Theorem 9), which is important to obtain
a better Lipshitzness dependence (Section 5.2) as well as the coordinate-wise
adaptivity for high-dimension settings (Section 6.3). Our analysis offers
theoretical benefits of adaptive learning rate from a discounted regret
perspective.

  * •

Model EMA. Lastly, the discounted-to-nonconvex conversion (Algorithm 1)
naturally leads to guarantees in terms of the model EMA,
$\overline{\mathbf{x}}_{t}$. At a high level (see Appendix A precise details),
this is because for a dynamic environment, it is important to discount the
losses such that online learners adapt to changing environments. The
appearance of model EMA in the discounted-to-nonconvex conversion provides a
new perspective on its role.

Our analyses and results have several limitations and raise several
interesting questions. Firstly, clipped-Adam does not precisely match the
original Adam algorithm, warranting further investigation into the original
Adam update. Specifically, our analysis suggests choosing $\beta_{1}=\beta$
and $\beta_{2}=\beta^{2}$, which does not align with the commonly used
practical choices. Understanding the exact roles of these practical choices
for $\beta_{1}$ and $\beta_{2}$ would be valuable.

In Section 5.2, we observed that our iteration complexity for finding a
$(\lambda,\varepsilon)$-stationary point is
$O(\Delta(G+\sigma)^{7/2}\lambda^{1/2}\varepsilon^{-7/2})$ when $G$ and
$\sigma$ are unknown. Investigating whether this complexity is optimal
presents another intriguing direction for future research.

Lastly, from a practical standpoint, developing a more advanced online learner
for discounted regret and designing an algorithm that surpasses Adam in
practicality would have significant practical implications.

## References

  * Abernethy et al. (2008) Jacob D. Abernethy, Elad Hazan, and Alexander Rakhlin.  Competing in the dark: An efficient algorithm for bandit linear optimization.  In _Conference on Learning Theory (COLT)_ , pages 263–274. Omnipress, 2008.
  * Ahn et al. (2024a) Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra.  Linear attention is (maybe) all you need (to understand transformer optimization).  In _The Twelfth International Conference on Learning Representations_ , 2024a.  URL https://openreview.net/forum?id=0uI5415ry7.
  * Ahn et al. (2024b) Kwangjun Ahn, Zhiyu Zhang, Yunbum Kook, and Yan Dai.  Understanding Adam optimizer via online learning of updates: Adam is FTRL in disguise.  In _International Conference on Machine Learning_. PMLR, 2024b.
  * Alacaoglu et al. (2020) Ahmet Alacaoglu, Yura Malitsky, Panayotis Mertikopoulos, and Volkan Cevher.  A new regret analysis for Adam-type algorithms.  In _International conference on machine learning (ICML)_ , pages 202–210. PMLR, 2020.
  * Arjevani et al. (2020) Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Ayush Sekhari, and Karthik Sridharan.  Second-order information in non-convex stochastic optimization: Power and limitations.  In _Conference on Learning Theory_ , pages 242–299, 2020.
  * Arjevani et al. (2023) Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Nathan Srebro, and Blake Woodworth.  Lower bounds for non-convex stochastic optimization.  _Mathematical Programming_ , 199(1):165–214, 2023.
  * Balaji et al. (2022) Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Qinsheng Zhang, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, et al.  eDiff-I: Text-to-image diffusion models with an ensemble of expert denoisers.  _arXiv preprint arXiv:2211.01324_ , 2022.
  * Block et al. (2024) Adam Block, Dylan J Foster, Akshay Krishnamurthy, Max Simchowitz, and Cyril Zhang.  Butterfly effects of SGD noise: Error amplification in behavior cloning and autoregression.  In _The Twelfth International Conference on Learning Representations_ , 2024.  URL https://openreview.net/forum?id=CgPs04l9TO.
  * Chen et al. (2019) Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong.  On the convergence of A class of Adam-type algorithms for non-convex optimization.  In _International Conference on Learning Representations (ICLR)_ , 2019.
  * Crawshaw et al. (2022) Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang.  Robustness to unbounded smoothness of generalized signsgd.  _Advances in neural information processing systems_ , 35:9955–9968, 2022.
  * Cutkosky et al. (2023) Ashok Cutkosky, Harsh Mehta, and Francesco Orabona.  Optimal stochastic non-smooth non-convex optimization through online-to-non-convex conversion.  In _International Conference on Machine Learning (ICML)_ , volume 202, pages 6643–6670. PMLR, 2023.
  * Davis et al. (2022) Damek Davis, Dmitriy Drusvyatskiy, Yin Tat Lee, Swati Padmanabhan, and Guanghao Ye.  A gradient sampling method with complexity guarantees for lipschitz functions in high and low dimensions.  In _Advances in Neural Information Processing Systems_ , volume 35, pages 6692–6703. Curran Associates, Inc., 2022.
  * Défossez et al. (2022) Alexandre Défossez, Leon Bottou, Francis Bach, and Nicolas Usunier.  A simple convergence proof of Adam and AdaGrad.  _Transactions on Machine Learning Research (TMLR)_ , 2022.
  * Dhariwal and Nichol (2021) Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  _Advances in neural information processing systems_ , 34:8780–8794, 2021.
  * Duchi et al. (2010) J. Duchi, E. Hazan, and Y. Singer.  Adaptive subgradient methods for online learning and stochastic optimization.  In _Conference on Learning Theory (COLT)_ , pages 257–269, 2010.
  * Goldstein (1977) AA Goldstein.  Optimization of lipschitz continuous functions.  _Mathematical Programming_ , 13:14–22, 1977.
  * Gordon (1999) Geoffrey J. Gordon.  Regret bounds for prediction problems.  In _Conference on Learning Theory (COLT)_ , pages 29–40. ACM, 1999.
  * Guo et al. (2021) Zhishuai Guo, Yi Xu, Wotao Yin, Rong Jin, and Tianbao Yang.  A novel convergence analysis for algorithms of the Adam family.  _arXiv preprint arXiv:2112.03459_ , 2021.
  * Hazan and Kale (2010) Elad Hazan and Satyen Kale.  Extracting certainty from uncertainty: Regret bounded by variation in costs.  _Machine Learning_ , 80:165–188, 2010.
  * Izmailov et al. (2018) Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson.  Averaging weights leads to wider optima and better generalization.  In _34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018_ , pages 876–885. Association For Uncertainty in Artificial Intelligence (AUAI), 2018.
  * Jacobsen and Cutkosky (2024) Andrew Jacobsen and Ashok Cutkosky.  Online linear regression in dynamic environments via discounting.  In _International Conference on Machine Learning_. PMLR, 2024.
  * Jiang et al. (2023) Kaiqi Jiang, Dhruv Malik, and Yuanzhi Li.  How does adaptive optimization impact local neural network geometry?  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2023.  URL https://openreview.net/forum?id=gIG8LvTLuc.
  * Jordan et al. (2023) Michael Jordan, Guy Kornowski, Tianyi Lin, Ohad Shamir, and Manolis Zampetakis.  Deterministic nonsmooth nonconvex optimization.  In _The Thirty Sixth Annual Conference on Learning Theory_ , pages 4570–4597. PMLR, 2023.
  * Kalai and Vempala (2005) Adam Kalai and Santosh Vempala.  Efficient algorithms for online decision problems.  _Journal of Computer and System Sciences (JCSS)_ , 71(3):291–307, 2005.
  * Kang et al. (2023) Minguk Kang, Jun-Yan Zhu, Richard Zhang, Jaesik Park, Eli Shechtman, Sylvain Paris, and Taesung Park.  Scaling up gans for text-to-image synthesis.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 10124–10134, 2023.
  * Karras et al. (2019) Tero Karras, Samuli Laine, and Timo Aila.  A style-based generator architecture for generative adversarial networks.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 4401–4410, 2019.
  * Karras et al. (2022) Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.  Elucidating the design space of diffusion-based generative models.  _Advances in Neural Information Processing Systems_ , 35:26565–26577, 2022.
  * Karras et al. (2023) Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine.  Analyzing and improving the training dynamics of diffusion models.  _arXiv preprint arXiv:2312.02696_ , 2023.
  * Kingma and Ba (2014) Diederik Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In _International Conference on Learning Representations (ICLR)_ , 2014.
  * Kunstner et al. (2023) Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt.  Noise is not the main factor behind the gap between SGD and Adam on Transformers, but sign descent might be.  In _International Conference on Learning Representations (ICLR)_ , 2023.
  * Kunstner et al. (2024) Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti.  Heavy-tailed class imbalance and why adam outperforms gradient descent on language models.  _arXiv preprint arXiv:2402.19449_ , 2024.
  * Li et al. (2023) Haochuan Li, Alexander Rakhlin, and Ali Jadbabaie.  Convergence of Adam under relaxed assumptions.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2023.  URL https://openreview.net/forum?id=yEewbkBNzi.
  * McMahan and Streeter (2010) H. Brendan McMahan and Matthew Streeter.  Adaptive bound optimization for online convex optimization.  In _Proceedings of the 23rd Annual Conference on Learning Theory (COLT)_ , pages 244–256, 2010.
  * Nesterov (2009) Yurii Nesterov.  Primal-dual subgradient methods for convex problems.  _Mathematical programming_ , 120(1):221–259, 2009.
  * Nichol and Dhariwal (2021) Alexander Quinn Nichol and Prafulla Dhariwal.  Improved denoising diffusion probabilistic models.  In _International conference on machine learning_ , pages 8162–8171. PMLR, 2021.
  * Orabona and Pál (2018) Francesco Orabona and Dávid Pál.  Scale-free online learning.  _Theoretical Computer Science_ , 716:50–69, 2018.
  * Pan and Li (2023) Yan Pan and Yuanzhi Li.  Toward understanding why Adam converges faster than SGD for Transformers.  _arXiv preprint arXiv:2306.00204_ , 2023.
  * Polyak and Juditsky (1992) Boris T Polyak and Anatoli B Juditsky.  Acceleration of stochastic approximation by averaging.  _SIAM journal on control and optimization_ , 30(4):838–855, 1992.
  * Reddi et al. (2018) Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar.  On the convergence of adam and beyond.  In _International Conference on Learning Representations_ , 2018.
  * Rombach et al. (2022) Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 10684–10695, 2022.
  * Ruppert (1988) David Ruppert.  Efficient estimations from a slowly convergent robbins-monro process.  Technical report, Cornell University Operations Research and Industrial Engineering, 1988.
  * Shalev-Shwartz and Singer (2006) Shai Shalev-Shwartz and Yoram Singer.  Online learning meets optimization in the dual.  In _Conference on Learning Theory (COLT)_ , volume 4005, pages 423–437. Springer, 2006.
  * Song et al. (2021a) Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  In _International Conference on Learning Representations_ , 2021a.  URL https://openreview.net/forum?id=St1giarCHLP.
  * Song et al. (2021b) Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  In _International Conference on Learning Representations_ , 2021b.  URL https://openreview.net/forum?id=PxTIG12RRHS.
  * Tarvainen and Valpola (2017) Antti Tarvainen and Harri Valpola.  Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results.  _Advances in neural information processing systems_ , 30, 2017.
  * Tian et al. (2022) Lai Tian, Kaiwen Zhou, and Anthony Man-Cho So.  On the finite-time complexity and practical computation of approximate stationarity concepts of lipschitz functions.  In _International Conference on Machine Learning_ , pages 21360–21379. PMLR, 2022.
  * Wang et al. (2023) Bohan Wang, Jingwen Fu, Huishuai Zhang, Nanning Zheng, and Wei Chen.  Closing the gap between the upper bound and lower bound of Adam’s iteration complexity.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2023.
  * Yaz et al. (2018) Yasin Yaz, Chuan-Sheng Foo, Stefan Winkler, Kim-Hui Yap, Georgios Piliouras, Vijay Chandrasekhar, et al.  The unusual effectiveness of averaging in GAN training.  In _International Conference on Learning Representations_ , 2018.
  * Zhang et al. (2020) Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie.  Why gradient clipping accelerates training: A theoretical justification for adaptivity.  In _International Conference on Learning Representations (ICLR)_ , 2020.
  * Zhang and Cutkosky (2024) Qinzi Zhang and Ashok Cutkosky.  Random scaling and momentum for non-smooth non-convex optimization.  In _International Conference on Machine Learning_. PMLR, 2024.
  * Zhang et al. (2022) Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo.  Adam can converge without any modification on update rules.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , volume 35, pages 28386–28399, 2022.
  * Zhang et al. (2024a) Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo.  Why transformers need adam: A hessian perspective.  _arXiv preprint arXiv:2402.16788_ , 2024a.
  * Zhang et al. (2024b) Zhiyu Zhang, David Bombara, and Heng Yang.  Discounted adaptive online prediction.  In _International Conference on Machine Learning_. PMLR, 2024b.
  * Zhou et al. (2019) Zhiming Zhou, Qingru Zhang, Guansong Lu, Hongwei Wang, Weinan Zhang, and Yong Yu.  Adashift: Decorrelation and convergence of adaptive learning rate methods.  In _International Conference on Learning Representations (ICLR)_ , 2019.
  * Zhuang et al. (2022) Zhenxun Zhuang, Mingrui Liu, Ashok Cutkosky, and Francesco Orabona.  Understanding AdamW through proximal methods and scale-freeness.  _Transactions on Machine Learning Research_ , 2022.
  * Zou et al. (2019) Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu.  A sufficient condition for convergences of Adam and RMSProp.  In _The IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 11127–11135, 2019.

[section] [section]l1

##  Appendix A Proof of discounted-to-nonconvex conversion (7)

Note first that via a change of summation, we get

$\displaystyle\sum_{n=1}^{T}\sum_{t=1}^{n}\beta^{n-t}(1-\beta)(F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1}))$
$\displaystyle=\sum_{t=1}^{T}\sum_{n=t}^{T}\beta^{n-t}(1-\beta)(F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1}))$
(25)
$\displaystyle=\sum_{t=1}^{T}(1-\beta^{T-t+1})(F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1}))$
(26)
$\displaystyle=F({\mathbf{x}}_{T})-F({\mathbf{x}}_{0})-\sum_{t=1}^{T}\beta^{T-t+1}(F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1}))\,.$
(27)

Rearranging the above together with the fact
$F({\mathbf{x}}_{0})-F({\mathbf{x}}_{T})\leq
F({\mathbf{x}}_{0})-\inf_{{\mathbf{x}}}F({\mathbf{x}})=:\Delta$, we get

$\displaystyle-\Delta\leq\underbrace{\mathop{\mathbb{E}}\left[\sum_{n=1}^{T}\sum_{t=1}^{n}\beta^{n-t}(1-\beta)(F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1}))\right]}_{\leavevmode\hbox
to15.2pt{\vbox to15.2pt{\pgfpicture\makeatletter\hbox{\hskip
7.5977pt\lower-7.5977pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{7.3977pt}{0.0pt}\pgfsys@curveto{7.3977pt}{4.08568pt}{4.08568pt}{7.3977pt}{0.0pt}{7.3977pt}\pgfsys@curveto{-4.08568pt}{7.3977pt}{-7.3977pt}{4.08568pt}{-7.3977pt}{0.0pt}\pgfsys@curveto{-7.3977pt}{-4.08568pt}{-4.08568pt}{-7.3977pt}{0.0pt}{-7.3977pt}\pgfsys@curveto{4.08568pt}{-7.3977pt}{7.3977pt}{-4.08568pt}{7.3977pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
} }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-3.375pt}{-3.07498pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{\small A}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}+\underbrace{\mathop{\mathbb{E}}\left[\sum_{t=1}^{T}\beta^{T-t+1}(F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1}))\right]}_{\leavevmode\hbox
to14.91pt{\vbox to14.91pt{\pgfpicture\makeatletter\hbox{\hskip
7.45659pt\lower-7.45659pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{7.25659pt}{0.0pt}\pgfsys@curveto{7.25659pt}{4.00775pt}{4.00775pt}{7.25659pt}{0.0pt}{7.25659pt}\pgfsys@curveto{-4.00775pt}{7.25659pt}{-7.25659pt}{4.00775pt}{-7.25659pt}{0.0pt}\pgfsys@curveto{-7.25659pt}{-4.00775pt}{-4.00775pt}{-7.25659pt}{0.0pt}{-7.25659pt}\pgfsys@curveto{4.00775pt}{-7.25659pt}{7.25659pt}{-4.00775pt}{7.25659pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
} }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-3.1875pt}{-3.07498pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{\small B}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}\,.$ (28)

We also recall the following fact about exponential random variable due to
[Zhang and Cutkosky, 2024, Lemma 3.1].

######  Lemma 20.

Let $\alpha\sim\text{Exp}(\lambda)$ for some $\lambda>0$, then

$\displaystyle\mathop{\mathbb{E}}_{\alpha}[F({\mathbf{x}}+\alpha{\mathbf{z}})-F({\mathbf{x}})]=\mathop{\mathbb{E}}_{\alpha}[\left\langle\nabla
F({\mathbf{x}}+\alpha{\mathbf{z}}),{\mathbf{z}}\right\rangle]/\lambda\,.$ (29)

Now with this lemma, we will handle each term separately.

Analysis of A

Note that for each $t\leq n$, since
${\mathbf{x}}_{t}\leftarrow{\mathbf{x}}_{t-1}+\alpha_{t}{\mathbf{z}}_{t}$ for
$\alpha_{t}{\sim}\text{Exp}(1)$, lemma 20 yields

$\displaystyle\mathop{\mathbb{E}}[F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1})]$
$\displaystyle=\mathop{\mathbb{E}}\left\langle\nabla
F({\mathbf{x}}_{t}),{\mathbf{z}}_{t}\right\rangle=\mathop{\mathbb{E}}\left\langle\nabla
F({\mathbf{x}}_{t}),\mathbf{u}_{n}\right\rangle+\mathop{\mathbb{E}}\left\langle\nabla
F({\mathbf{x}}_{t}),{\mathbf{z}}_{t}-\mathbf{u}_{n}\right\rangle$ (30)
$\displaystyle=\mathop{\mathbb{E}}\left\langle\nabla
F({\mathbf{x}}_{t}),\mathbf{u}_{n}\right\rangle+\mathop{\mathbb{E}}\left\langle\nabla
F({\mathbf{x}}_{t})-\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{n}\right\rangle+\mathop{\mathbb{E}}\left\langle\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{n}\right\rangle$
(31) $\displaystyle=\underbrace{\mathop{\mathbb{E}}\left\langle\nabla
F({\mathbf{x}}_{t}),\mathbf{u}_{n}\right\rangle}_{\leavevmode\hbox
to13.33pt{\vbox to13.33pt{\pgfpicture\makeatletter\hbox{\hskip
6.66582pt\lower-6.66582pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{6.46582pt}{0.0pt}\pgfsys@curveto{6.46582pt}{3.57101pt}{3.57101pt}{6.46582pt}{0.0pt}{6.46582pt}\pgfsys@curveto{-3.57101pt}{6.46582pt}{-6.46582pt}{3.57101pt}{-6.46582pt}{0.0pt}\pgfsys@curveto{-6.46582pt}{-3.57101pt}{-3.57101pt}{-6.46582pt}{0.0pt}{-6.46582pt}\pgfsys@curveto{3.57101pt}{-6.46582pt}{6.46582pt}{-3.57101pt}{6.46582pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
} }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.25pt}{-2.9pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{\small 1}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}+\underbrace{\mathop{\mathbb{E}}\left\langle\nabla
F({\mathbf{x}}_{t})-\mathbf{g}_{t},-\mathbf{u}_{n}\right\rangle}_{\leavevmode\hbox
to13.33pt{\vbox to13.33pt{\pgfpicture\makeatletter\hbox{\hskip
6.66582pt\lower-6.66582pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{6.46582pt}{0.0pt}\pgfsys@curveto{6.46582pt}{3.57101pt}{3.57101pt}{6.46582pt}{0.0pt}{6.46582pt}\pgfsys@curveto{-3.57101pt}{6.46582pt}{-6.46582pt}{3.57101pt}{-6.46582pt}{0.0pt}\pgfsys@curveto{-6.46582pt}{-3.57101pt}{-3.57101pt}{-6.46582pt}{0.0pt}{-6.46582pt}\pgfsys@curveto{3.57101pt}{-6.46582pt}{6.46582pt}{-3.57101pt}{6.46582pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
} }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.25pt}{-2.9pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{\small 2}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}+\underbrace{\mathop{\mathbb{E}}\left\langle\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{n}\right\rangle}_{\leavevmode\hbox
to13.33pt{\vbox to13.33pt{\pgfpicture\makeatletter\hbox{\hskip
6.66582pt\lower-6.66582pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{6.46582pt}{0.0pt}\pgfsys@curveto{6.46582pt}{3.57101pt}{3.57101pt}{6.46582pt}{0.0pt}{6.46582pt}\pgfsys@curveto{-3.57101pt}{6.46582pt}{-6.46582pt}{3.57101pt}{-6.46582pt}{0.0pt}\pgfsys@curveto{-6.46582pt}{-3.57101pt}{-3.57101pt}{-6.46582pt}{0.0pt}{-6.46582pt}\pgfsys@curveto{3.57101pt}{-6.46582pt}{6.46582pt}{-3.57101pt}{6.46582pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
} }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.25pt}{-2.9pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{\small 3}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}},$ (32)

where the last line follows from the fact
$\mathop{\mathbb{E}}[\left\langle\nabla
F({\mathbf{x}}_{t})-\mathbf{g}_{t},{\mathbf{z}}_{t}\right\rangle]=0$.

Now let us handle each term.

  * 1:

Note that using the definition of $\mathbf{y}_{n}$, we have

$\displaystyle\mathop{\mathbb{E}}\sum_{t=1}^{n}\beta^{n-t}(1-\beta)\left\langle\nabla
F({\mathbf{x}}_{t}),\mathbf{u}_{n}\right\rangle=(1-\beta)\mathop{\mathbb{E}}\left\langle\sum_{t=1}^{n}\beta^{n-t}\nabla
F({\mathbf{x}}_{t}),-D\frac{\sum_{t=1}^{n}\beta^{n-t}\nabla
F({\mathbf{x}}_{t})}{\left\|\sum_{t=1}^{n}\beta^{n-t}\nabla
F({\mathbf{x}}_{t})\right\|}\right\rangle$ (33)
$\displaystyle\quad=(1-\beta^{n})\mathop{\mathbb{E}}\left\langle\sum_{t=1}^{n}\frac{1-\beta}{1-\beta^{n}}\beta^{n-t}\nabla
F({\mathbf{x}}_{t}),-D\frac{\sum_{t=1}^{n}\frac{1-\beta}{1-\beta^{n}}\beta^{n-t}\nabla
F({\mathbf{x}}_{t})}{\left\|\sum_{t=1}^{n}\frac{1-\beta}{1-\beta^{n}}\beta^{n-t}\nabla
F({\mathbf{x}}_{t})\right\|}\right\rangle$ (34)
$\displaystyle\quad=-D(1-\beta^{n})\mathop{\mathbb{E}}\left\|\mathop{\mathbb{E}}_{\mathbf{y}_{n}}\nabla
F(\mathbf{y}_{n})\right\|\leq-D\mathop{\mathbb{E}}\left\|\mathop{\mathbb{E}}_{\mathbf{y}_{n}}\nabla
F(\mathbf{y}_{n})\right\|+DG\beta^{n}\,.$ (35)

Therefore, summing over $n=1,\dots,T$, we obtain:

$\displaystyle\mathop{\mathbb{E}}\sum_{n=1}^{T}\sum_{t=1}^{n}\beta^{n-t}(1-\beta)\left\langle\nabla
F({\mathbf{x}}_{t}),\mathbf{u}_{n}\right\rangle\leq-D\mathop{\mathbb{E}}\sum_{t=1}^{T}\left\|\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\nabla
F(\mathbf{y}_{t})\right\|+\frac{DG}{1-\beta}\,.$ (36)

  * 2:

For the second term, using Cauchy-Schwartz inequality, we have

$\displaystyle\mathop{\mathbb{E}}\sum_{t=1}^{n}\beta^{n-t}\left\langle\nabla
F({\mathbf{x}}_{t})-\mathbf{g}_{t},-\mathbf{u}_{n}\right\rangle\leq\sqrt{\mathop{\mathbb{E}}\left\|\sum_{t=1}^{n}\beta^{n-t}(\nabla
F({\mathbf{x}}_{t})-\mathbf{g}_{t})\right\|^{2}\mathop{\mathbb{E}}\left\|\mathbf{u}_{n}\right\|^{2}}\,.$
(37)

Using the bounded variance assumption on the stochastic gradient oracle, we
have

$\displaystyle\mathop{\mathbb{E}}\left\|\sum_{t=1}^{n}\beta^{n-t}(\nabla
F({\mathbf{x}}_{t})-\mathbf{g}_{t})\right\|^{2}=\mathop{\mathbb{E}}\sum_{t=1}^{n}\beta^{2(n-t)}\left\|\nabla
F({\mathbf{x}}_{t})-\mathbf{g}_{t}\right\|^{2}\leq\frac{\sigma^{2}}{1-\beta^{2}}\,.$
(38)

Therefore, summing over $n=1,\dots,T$, and using the fact that
$\frac{1}{1-\beta^{2}}\leq\frac{1}{1-\beta}$, we get the following bound on
the second term:

$\displaystyle\mathop{\mathbb{E}}\sum_{n=1}^{T}\sum_{t=1}^{n}\beta^{n-t}(1-\beta)\left\langle\nabla
F({\mathbf{x}}_{t})-\mathbf{g}_{t},-\mathbf{u}_{n}\right\rangle\leq\sum_{n=1}^{T}(1-\beta)\cdot\frac{\sigma
D}{\sqrt{1-\beta^{2}}}\leq\sigma DT\sqrt{1-\beta}\,.$ (39)

  * 3:

Lastly, for the third term, we have

$\displaystyle\mathop{\mathbb{E}}\sum_{n=1}^{T}\sum_{t=1}^{n}\beta^{n-t}(1-\beta)\left\langle\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{n}\right\rangle$
$\displaystyle=(1-\beta)\mathop{\mathbb{E}}\sum_{n=1}^{T}\left[\sum_{t=1}^{n}\mathop{\mathbb{E}}\left\langle\beta^{n-t}\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{n}\right\rangle\right]$
(40)
$\displaystyle=(1-\beta)\mathop{\mathbb{E}}\sum_{t=1}^{T}\textup{Regret}^{[\beta]}_{t}(\mathbf{u}_{t})\,.$
(41)

Analysis of B

Note that for each $t$, since
${\mathbf{x}}_{t}\leftarrow{\mathbf{x}}_{t-1}+\alpha_{t}{\mathbf{z}}_{t}$ for
$\alpha_{t}{\sim}\text{Exp}(1)$, lemma 20 yields

$\displaystyle\mathop{\mathbb{E}}[F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1})]$
$\displaystyle=\mathop{\mathbb{E}}\left\langle\nabla
F({\mathbf{x}}_{t}),{\mathbf{z}}_{t}\right\rangle=\mathop{\mathbb{E}}\left\langle\mathbf{g}_{t},{\mathbf{z}}_{t}\right\rangle$
(42)
$\displaystyle=\mathop{\mathbb{E}}\left\langle\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{T}\right\rangle+\mathop{\mathbb{E}}\left\langle\mathbf{g}_{t},\mathbf{u}_{T}\right\rangle\leq\mathop{\mathbb{E}}\left\langle\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{T}\right\rangle+D(G+\sigma)\,.$
(43)

Thus,

$\displaystyle\mathop{\mathbb{E}}\left[\sum_{t=1}^{T}\beta^{T-t+1}(F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1}))\right]$
$\displaystyle=\beta\mathop{\mathbb{E}}\sum_{t=1}^{T}\left[\left\langle\beta^{T-t}\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{T}\right\rangle+\beta^{T-t}D(G+\sigma)\right]$
(44)
$\displaystyle\leq\beta\mathop{\mathbb{E}}[\textup{Regret}^{[\beta]}_{T}(\mathbf{u}_{T})]+\frac{D(G+\sigma)}{1-\beta}\,.$
(45)

Combining A and B

Combining the above analyses and rearranging, it follows that

$\displaystyle
D\mathop{\mathbb{E}}\sum_{t=1}^{T}\left\|\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\nabla
F(\mathbf{y}_{t})\right\|$ $\displaystyle\leq\Delta+\frac{DG}{1-\beta}+\sigma
DT\sqrt{1-\beta}+(1-\beta)\mathop{\mathbb{E}}\sum_{t=1}^{T}\left[\textup{Regret}^{[\beta]}_{t}(\mathbf{u}_{t})\right]$
(46)
$\displaystyle\quad+\beta\mathop{\mathbb{E}}[\textup{Regret}^{[\beta]}_{T}(\mathbf{u}_{T})]+\frac{D(G+\sigma)}{1-\beta}\,.$
(47)

Dividing both sides by $DT$, we get the desired result.

###  A.1 Proof of the coordinate-wise version (lemma 17)

The proof closely follows that of 7. In particular, with $\Delta\coloneqq
F({\mathbf{x}}_{0})-\inf_{{\mathbf{x}}}F({\mathbf{x}})$, we have

$\displaystyle-\Delta\leq\underbrace{\mathop{\mathbb{E}}\left[\sum_{n=1}^{T}\sum_{t=1}^{n}\beta^{n-t}(1-\beta)(F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1}))\right]}_{\leavevmode\hbox
to15.2pt{\vbox to15.2pt{\pgfpicture\makeatletter\hbox{\hskip
7.5977pt\lower-7.5977pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{7.3977pt}{0.0pt}\pgfsys@curveto{7.3977pt}{4.08568pt}{4.08568pt}{7.3977pt}{0.0pt}{7.3977pt}\pgfsys@curveto{-4.08568pt}{7.3977pt}{-7.3977pt}{4.08568pt}{-7.3977pt}{0.0pt}\pgfsys@curveto{-7.3977pt}{-4.08568pt}{-4.08568pt}{-7.3977pt}{0.0pt}{-7.3977pt}\pgfsys@curveto{4.08568pt}{-7.3977pt}{7.3977pt}{-4.08568pt}{7.3977pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
} }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-3.375pt}{-3.07498pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{\small A}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}+\underbrace{\mathop{\mathbb{E}}\left[\sum_{t=1}^{T}\beta^{T-t+1}(F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1}))\right]}_{\leavevmode\hbox
to14.91pt{\vbox to14.91pt{\pgfpicture\makeatletter\hbox{\hskip
7.45659pt\lower-7.45659pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{7.25659pt}{0.0pt}\pgfsys@curveto{7.25659pt}{4.00775pt}{4.00775pt}{7.25659pt}{0.0pt}{7.25659pt}\pgfsys@curveto{-4.00775pt}{7.25659pt}{-7.25659pt}{4.00775pt}{-7.25659pt}{0.0pt}\pgfsys@curveto{-7.25659pt}{-4.00775pt}{-4.00775pt}{-7.25659pt}{0.0pt}{-7.25659pt}\pgfsys@curveto{4.00775pt}{-7.25659pt}{7.25659pt}{-4.00775pt}{7.25659pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
} }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-3.1875pt}{-3.07498pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{\small B}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}\,.$ (48)

We begin with the term  B. Using the same decomposition as before, we have

$\displaystyle\mathop{\mathbb{E}}[F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1})]$
$\displaystyle=\mathop{\mathbb{E}}\left\langle\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{T}\right\rangle+\mathop{\mathbb{E}}\left\langle\mathbf{g}_{t},\mathbf{u}_{T}\right\rangle=\mathop{\mathbb{E}}\left\langle\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{T}\right\rangle+\sum_{i=1}^{d}\mathop{\mathbb{E}}\mathbf{g}_{t}[i]\mathbf{u}_{T}[i]$
(49)
$\displaystyle\leq\mathop{\mathbb{E}}\left\langle\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{T}\right\rangle+D\sum_{i=1}^{d}(G_{i}+\sigma_{i})\,.$
(50)

Thus,

$\displaystyle\mathop{\mathbb{E}}\left[\sum_{t=1}^{T}\beta^{T-t+1}(F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1}))\right]$
$\displaystyle=\beta\mathop{\mathbb{E}}\sum_{t=1}^{T}\left[\left\langle\beta^{T-t}\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{T}\right\rangle+\beta^{T-t}\sum_{i=1}^{d}D(G_{i}+\sigma_{i})\right]$
(51)
$\displaystyle\leq\beta\mathop{\mathbb{E}}[\textup{Regret}^{[\beta]}_{T}(\mathbf{u}_{T})]+\frac{D\sum_{i=1}^{d}(G_{i}+\sigma_{i})}{1-\beta}\,.$
(52)

Moving onto the term  A, we again use the same decomposition:

$\displaystyle\mathop{\mathbb{E}}[F({\mathbf{x}}_{t})-F({\mathbf{x}}_{t-1})]$
$\displaystyle=\underbrace{\mathop{\mathbb{E}}\left\langle\nabla
F({\mathbf{x}}_{t}),\mathbf{u}_{n}\right\rangle}_{\leavevmode\hbox
to13.33pt{\vbox to13.33pt{\pgfpicture\makeatletter\hbox{\hskip
6.66582pt\lower-6.66582pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{6.46582pt}{0.0pt}\pgfsys@curveto{6.46582pt}{3.57101pt}{3.57101pt}{6.46582pt}{0.0pt}{6.46582pt}\pgfsys@curveto{-3.57101pt}{6.46582pt}{-6.46582pt}{3.57101pt}{-6.46582pt}{0.0pt}\pgfsys@curveto{-6.46582pt}{-3.57101pt}{-3.57101pt}{-6.46582pt}{0.0pt}{-6.46582pt}\pgfsys@curveto{3.57101pt}{-6.46582pt}{6.46582pt}{-3.57101pt}{6.46582pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
} }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.25pt}{-2.9pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{\small 1}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}+\underbrace{\mathop{\mathbb{E}}\left\langle\nabla
F({\mathbf{x}}_{t})-\mathbf{g}_{t},-\mathbf{u}_{n}\right\rangle}_{\leavevmode\hbox
to13.33pt{\vbox to13.33pt{\pgfpicture\makeatletter\hbox{\hskip
6.66582pt\lower-6.66582pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{6.46582pt}{0.0pt}\pgfsys@curveto{6.46582pt}{3.57101pt}{3.57101pt}{6.46582pt}{0.0pt}{6.46582pt}\pgfsys@curveto{-3.57101pt}{6.46582pt}{-6.46582pt}{3.57101pt}{-6.46582pt}{0.0pt}\pgfsys@curveto{-6.46582pt}{-3.57101pt}{-3.57101pt}{-6.46582pt}{0.0pt}{-6.46582pt}\pgfsys@curveto{3.57101pt}{-6.46582pt}{6.46582pt}{-3.57101pt}{6.46582pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
} }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.25pt}{-2.9pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{\small 2}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}+\underbrace{\mathop{\mathbb{E}}\left\langle\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{n}\right\rangle}_{\leavevmode\hbox
to13.33pt{\vbox to13.33pt{\pgfpicture\makeatletter\hbox{\hskip
6.66582pt\lower-6.66582pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}{{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{6.46582pt}{0.0pt}\pgfsys@curveto{6.46582pt}{3.57101pt}{3.57101pt}{6.46582pt}{0.0pt}{6.46582pt}\pgfsys@curveto{-3.57101pt}{6.46582pt}{-6.46582pt}{3.57101pt}{-6.46582pt}{0.0pt}\pgfsys@curveto{-6.46582pt}{-3.57101pt}{-3.57101pt}{-6.46582pt}{0.0pt}{-6.46582pt}\pgfsys@curveto{3.57101pt}{-6.46582pt}{6.46582pt}{-3.57101pt}{6.46582pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
} }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.25pt}{-2.9pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{\small 3}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}\,.$ (53)

As before, let us handle each term one by one separately.

  * 1:

Note that using the definition of $\mathbf{y}_{n}$, for each coordinate
$i=1,\dots,d$, we have

$\displaystyle\mathop{\mathbb{E}}\sum_{t=1}^{n}\beta^{n-t}(1-\beta)\partial_{i}F({\mathbf{x}}_{t})\mathbf{u}_{n}[i]=(1-\beta)\mathop{\mathbb{E}}\left[\left(\sum_{t=1}^{n}\beta^{n-t}\nabla
F({\mathbf{x}}_{t})\right)\left(-D\frac{\sum_{t=1}^{n}\beta^{n-t}\partial_{i}F({\mathbf{x}}_{t})}{\left|{\sum_{t=1}^{n}\beta^{s-t}\partial_{i}F({\mathbf{x}}_{t})}\right|}\right)\right]$
(54)
$\displaystyle\quad=(1-\beta^{n})\mathop{\mathbb{E}}\left[\left(\sum_{t=1}^{n}\frac{1-\beta}{1-\beta^{n}}\beta^{n-t}\nabla
F({\mathbf{x}}_{t})\right)\left(-D\frac{\sum_{t=1}^{n}\frac{1-\beta}{1-\beta^{n}}\beta^{n-t}\nabla
F({\mathbf{x}}_{t})}{\left|{\sum_{t=1}^{n}\frac{1-\beta}{1-\beta^{n}}\beta^{n-t}\nabla
F({\mathbf{x}}_{t})}\right|}\right)\right]$ (55)
$\displaystyle\quad=-D(1-\beta^{n})\mathop{\mathbb{E}}\left|{\mathop{\mathbb{E}}_{\mathbf{y}_{n}}\partial_{i}F(\mathbf{y}_{n})}\right|\leq-D\mathop{\mathbb{E}}\left|{\mathop{\mathbb{E}}_{\mathbf{y}_{n}}\partial_{i}F(\mathbf{y}_{n})}\right|+DG_{i}\beta^{n}\,.$
(56)

Therefore, summing over $i=1,\dots,d$ and then $n=1,\dots,T$, we obtain:

$\displaystyle\mathop{\mathbb{E}}\sum_{n=1}^{T}\sum_{t=1}^{n}\beta^{n-t}(1-\beta)\left\langle\nabla
F({\mathbf{x}}_{t}),\mathbf{u}_{n}\right\rangle\leq-D\mathop{\mathbb{E}}\sum_{t=1}^{T}\left\|\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\nabla
F(\mathbf{y}_{t})\right\|_{1}+\frac{D\sum_{i=1}G_{i}}{1-\beta}\,.$ (57)

  * 2:

For each coordinate $i=1,\dots,d$, we have

$\displaystyle\mathop{\mathbb{E}}\sum_{t=1}^{n}\beta^{n-t}(\partial_{i}F({\mathbf{x}}_{t})-\mathbf{g}_{t}[i])(-\mathbf{u}_{n}[i])\leq\sqrt{\mathop{\mathbb{E}}\left|{\sum_{t=1}^{n}\beta^{n-t}(\partial_{i}F({\mathbf{x}}_{t})-\mathbf{g}_{t}[i])}\right|^{2}\mathop{\mathbb{E}}\left|{\mathbf{u}_{n}[i]}\right|^{2}}\,.$
(58)

Using the coordinate-wise bounded variance assumption on the stochastic
gradient oracle,

$\displaystyle\mathop{\mathbb{E}}\left|{\sum_{t=1}^{n}\beta^{n-t}(\partial_{i}F({\mathbf{x}}_{t})-\mathbf{g}_{t}[i])}\right|^{2}=\mathop{\mathbb{E}}\sum_{t=1}^{n}\beta^{2(n-t)}\left|{\partial_{i}F({\mathbf{x}}_{t})-\mathbf{g}_{t}[i]}\right|^{2}\leq\frac{\sigma_{i}^{2}}{1-\beta^{2}}\,.$
(59)

Therefore, summing over $n=1,\dots,T$, and using the fact that
$\frac{1}{1-\beta^{2}}\leq\frac{1}{1-\beta}$, we get the following bound on
the second term:

$\displaystyle\mathop{\mathbb{E}}\sum_{n=1}^{T}\sum_{t=1}^{n}\beta^{n-t}(1-\beta)\left\langle\nabla
F({\mathbf{x}}_{t})-\mathbf{g}_{t},-\mathbf{u}_{n}\right\rangle$
$\displaystyle\leq\sum_{n=1}^{T}(1-\beta)\cdot\frac{D\sum_{i=1}^{d}\sigma_{i}}{\sqrt{1-\beta^{2}}}$
(60) $\displaystyle\leq
DT\left(\sum_{i=1}^{d}\sigma_{i}\right)\sqrt{1-\beta}\,.$ (61)

  * 3:

We use the same manipulation as before:

$\displaystyle\mathop{\mathbb{E}}\sum_{n=1}^{T}\sum_{t=1}^{n}\beta^{n-t}(1-\beta)\left\langle\mathbf{g}_{t},{\mathbf{z}}_{t}-\mathbf{u}_{n}\right\rangle$
$\displaystyle=(1-\beta)\mathop{\mathbb{E}}\sum_{t=1}^{T}\textup{Regret}^{[\beta]}_{t}(\mathbf{u}_{t})\,.$
(62)

Combining the above, we get the desired result in lemma 17.

##  Appendix B Proof of main theorems

###  B.1 Proof of 11

By 3, since $\mathop{\mathbb{E}}[\mathbf{y}_{t}]=\overline{\mathbf{x}}_{t}$,
it holds that

$\displaystyle\mathop{\mathbb{E}}_{t\sim[T]}\left\|\nabla
F(\overline{\mathbf{x}}_{t})\right\|^{[\lambda]}\leq\mathop{\mathbb{E}}_{t\sim[T]}\left[\left\|\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\nabla
F(\mathbf{y}_{t})\right\|+\lambda\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\left\|\mathbf{y}_{t}-\overline{\mathbf{x}}_{t}\right\|^{2}\right]\,.$
(63)

We begin with the second term (the variance term). By lemma 10, we have

$\displaystyle\lambda\mathop{\mathbb{E}}_{t\sim[T]}\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\left\|\mathbf{y}_{t}-\overline{\mathbf{x}}_{t}\right\|^{2}$
$\displaystyle\leq 12\frac{\lambda D^{2}}{(1-\beta)^{2}}\,.$ (64)

Hence, by choosing $D=\frac{(1-\beta)\varepsilon^{1/2}}{4\lambda^{1/2}}$, it
follows that
$\lambda\cdot\mathop{\mathbb{E}}_{t\sim[T]}\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\left\|\mathbf{y}_{t}-\overline{\mathbf{x}}_{t}\right\|^{2}\leq\varepsilon$.

Next consider the first term (the norm of the averaged gradients). Plugging
the regret bound (9) into 7, we get

$\displaystyle\mathop{\mathbb{E}}_{t\sim[T]}\left\|\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\nabla
F(\mathbf{y}_{t})\right\|$
$\displaystyle\leq\frac{\Delta}{DT}+\frac{2G+\sigma}{(1-\beta)T}+\sigma\sqrt{1-\beta}+\frac{4(G+\sigma)}{T\sqrt{1-\beta}}+4(G+\sigma)\sqrt{1-\beta}$
(65)
$\displaystyle\leq\frac{4\Delta\lambda^{1/2}}{(1-\beta)\varepsilon^{1/2}T}+\frac{6G+5\sigma}{(1-\beta)T}+(4G+5\sigma)\sqrt{1-\beta}\,,$
(66)

where the last line follows since
$\frac{1}{\sqrt{1-\beta}}\leq\frac{1}{1-\beta}$ and
$D=\frac{(1-\beta)\varepsilon^{1/2}}{4\lambda^{1/2}}$. Choosing
$\beta=1-(\frac{\varepsilon}{10\alpha})^{2}$, the last term is bounded by
$\frac{G+\sigma}{2\alpha}\varepsilon$. Moreover, choosing
$T=(1-\beta)^{-1}\cdot\max\left\\{4\Delta\lambda^{1/2}\varepsilon^{-3/2},~{}12\alpha\varepsilon^{-1}\right\\}$,
the first and second terms are bounded by $\varepsilon$ and
$\frac{G+\sigma}{2\alpha}\varepsilon$, respectively. This concludes the proof.

###  B.2 Proof of 18

By 15, since $\mathop{\mathbb{E}}[\mathbf{y}_{t}]=\overline{\mathbf{x}}_{t}$,
it holds that

$\displaystyle\mathop{\mathbb{E}}_{t\sim[T]}\left\|\nabla
F(\overline{\mathbf{x}}_{t})\right\|^{[\lambda]}_{1}\leq\mathop{\mathbb{E}}_{t\sim[T]}\left[\left\|\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\nabla
F(\mathbf{y}_{t})\right\|_{1}+\lambda\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\left\|\mathbf{y}_{t}-\overline{\mathbf{x}}_{t}\right\|_{2}^{2}\right]\,.$
(67)

We begin with the second term (the variance term). This time, given that now
each coordinate of update ${\mathbf{z}}_{t}$ is bounded by $D$, _i.e._ ,
$\left|{{\mathbf{z}}_{t}[i]}\right|\leq D$, applying the variance bound due to
lemma 10 coordinate-wise implies:

$\displaystyle\lambda\mathop{\mathbb{E}}_{t\sim[T]}\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\left\|\mathbf{y}_{t}-\overline{\mathbf{x}}_{t}\right\|_{2}^{2}\leq
12\frac{\lambda dD^{2}}{(1-\beta)^{2}}\,.$ (68)

Hence, by choosing
$D=\frac{(1-\beta)\varepsilon^{1/2}}{4d^{1/2}\lambda^{1/2}}$, it follows that
$\lambda\mathop{\mathbb{E}}_{t\sim[T]}\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\left\|\mathbf{y}_{t}-\overline{\mathbf{x}}_{t}\right\|^{2}\leq\varepsilon$.

Next consider the first term (the $L_{1}$-norm of the averaged gradients).
Plugging the regret bound (19) into lemma 17, and doing similar manipulations
as the proof of 11, we get

$\displaystyle\mathop{\mathbb{E}}_{t\sim[T]}\left\|\mathop{\mathbb{E}}_{\mathbf{y}_{t}}\nabla
F(\mathbf{y}_{t})\right\|_{1}$
$\displaystyle\leq\frac{\Delta}{DT}+\frac{\left\|6\boldsymbol{G}+5\boldsymbol{\sigma}\right\|_{1}}{(1-\beta)T}+\left\|4\boldsymbol{G}+5\boldsymbol{\sigma}\right\|_{1}\sqrt{1-\beta}$
(69) $\displaystyle=\frac{4\Delta
d^{1/2}\lambda^{1/2}}{(1-\beta)\varepsilon^{1/2}T}+\frac{\left\|6\boldsymbol{G}+5\boldsymbol{\sigma}\right\|_{1}}{(1-\beta)T}+\left\|4\boldsymbol{G}+5\boldsymbol{\sigma}\right\|_{1}\sqrt{1-\beta}\,,$
(70)

where the last line follows since
$D=\frac{(1-\beta)\varepsilon^{1/2}}{4d^{1/2}\lambda^{1/2}}$. Choosing
$\beta=1-(\frac{\varepsilon}{10\alpha})^{2}$, the last term is bounded by
$\frac{\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{1}}{2\alpha}\cdot\varepsilon$.
Moreover, choosing $T=(1-\beta)^{-1}\cdot\max\left\\{4\Delta
d^{1/2}\lambda^{1/2}\varepsilon^{-3/2},~{}12\alpha\varepsilon^{-1}\right\\}$,
the first and second terms are bounded by $\varepsilon$ and
$\frac{\left\|\boldsymbol{G}+\boldsymbol{\sigma}\right\|_{1}}{2\alpha}\cdot\varepsilon$,
respectively. This concludes the proof.