text:
\section{A General Reduction from sub-Gaussian to Bounded Noise}
\label{sec:sub_gaussian_reduction}

Given an unbiased gradient oracle one can consider different noise assumptions.
For establishing high probability convergence either sub-Gaussian noise or bounded noise are most commonly used.
Sub-Gaussian noise is less restrictive but can make the analysis cumbersome.
In this section we explain how obtaining guarantees with a bounded oracle can yield the same guarantees with a sub-Gaussian oracle (up to logarithmic factors) with high probability.

Let $X$ be a zero-mean random vector such that for some $\sigma>0$, for all $t > 0$,
\begin{align}\label{eq:sub_gaussian_tail}
    \Pr\brk*{\norm{X} \geq t}
    &\leq
    2 \exp\brk{-\ifrac{t^2}{\sigma^2}}
    .
\end{align}
This definition is equivalent to the sub-Gaussian definition used by \citet{li2020high} and is similar to the one used by \citet{kavis2021high}.

The following lemma is a reduction from $X$ to a new bounded random vector with zero-mean which is equal to $X$ with high probability.

\begin{lemma}\label{lem:truncation}
    For any $\delta \in (0,1)$ there exist a random variable $\bar{X}$ such that:
    \begin{enumerate}[nosep,label=(\roman*)]
        \item $\bar{X}$ is zero-mean: $\mathbb{E}\brk[s]{\bar{X}}=0;$
        \item $\bar{X}$ is equal to $X$ w.h.p.: $\Pr\brk*{\bar{X}=X} \geq 1-\delta;$
        \item $\bar{X}$ is bounded with probability 1: $\Pr\brk!{\norm{\bar{X}}<3 \sigma \sqrt{\ln(4/\delta)}}=1.$
    \end{enumerate}
\end{lemma}
The lemma implies that given a sequence of oracle queries $(w_t)_{t=1}^T$ to a $\sigma$-sub-Gaussian unbiased oracle $g(w)$ (where $w_t$ might depend on the previous queries) there exist a $O(\sigma \sqrt{\log(T/\delta')})$-bounded unbiased oracle $\tilde{g}(w)$ such that with probability at least $1-\delta'$ return the same observations.
This happens by the following procedure (Given a query $w$):
\begin{enumerate}
    \item Let $X=g(w)-\mathbb{E}[g(w)].$
    \item Let $\bar{X}$ be the truncated random vector of \cref{lem:truncation} for $\delta=\ifrac{\delta'}{T}$.
    \item Return $\mathbb{E}[g(w)]+\bar{X}$.
\end{enumerate}
Hence, up to logarithmic factors, one can assume a bounded unbiased oracle and obtain the same guarantee with high probability.
Note that this reduction holds even for an affine noise oracle $g$ where
\begin{align*}
    \Pr\brk!{\norm{g(w)-\mathbb{E}[g(w)]} \geq t} \leq 2 \exp\brk{-t^2/\brk{\sigma_0^2+\sigma_1^2 \norm{\mathbb{E}[g(w)]}^2}}
\end{align*}
by creating a new gradient oracle $\widetilde{g}(w)$ such that
\begin{align*}
    \norm{\widetilde{g}(w)-\mathbb{E}[\widetilde{g}(w)]}^2
    &\leq
    9 \sigma_0^2 \ln(4T/\delta')
    + 9\sigma_1^2 \ln(4T/\delta')\norm{\mathbb{E}[\widetilde{g}(w)]}^2
    .
\end{align*}
\begin{proof}[Proof of \cref{lem:truncation}]
Denote $r=\sigma \sqrt{\ln(4/\delta)}$; note that $\Pr(\norm{X} > r) \leq \delta/2$ (due to \cref{eq:sub_gaussian_tail}).
Define a random vector $Z$ by rejection sampling $X$ on the $r$-unit ball: if $\norm{X} \leq r$, let $Z=X$; otherwise, continue resampling from the distribution of $X$ and reject samples as long as their norm is larger than $r$; let $Z$ be the first accepted sample.
To correct the bias of $Z$, let
\begin{align*}
    \bar{X}
    &=
    \begin{cases}
        Z
        &
        \text{with prob.~$1-\delta/2;$}
        \\
        - \frac{2-\delta}{\delta} \mathbb{E}[Z]
        &
        \text{with prob.~$\delta/2.$}
    \end{cases}
\end{align*}
By construction, $\mathbb{E}[\bar{X}]=0$.
Also, $\Pr(\bar{X}=X) \geq 1-\delta$ due to a union bound of $\Pr(\bar{X}=Z)$ and $\Pr(X=Z)$, each happens with probability at least $1-\delta/2$.
We are left with bounding $\bar{X}$.
As $\norm{Z} \leq r$ it remains to bound the correction term.
For this, we first bound $\mathbb{E}[Z]=\mathbb{E}[X \mid \norm{X} \leq r]$.
By the law of total expectation,
\begin{align*}
    \mathbb{E}\brk[s]*{X \mathbf{1}\brk[c]{\norm{X} \leq r}}
    &=
    \Pr(\norm{X} \leq r)\mathbb{E}\brk[s]*{X \mathbf{1}\brk[c]{\norm{X} \leq r} \mid \norm{X} \leq r}
    + \Pr(\norm{X} > r)\mathbb{E}\brk[s]*{X \mathbf{1}\brk[c]{\norm{X} \leq r} \mid \norm{X} > r}
    \\
    &=
    \Pr(\norm{X} \leq r)\mathbb{E}\brk[s]*{X \mathbf{1}\brk[c]{\norm{X} \leq r} \mid \norm{X} \leq r}
    \\
    &=
    \Pr(\norm{X} \leq r)\mathbb{E}\brk[s]*{X \mid \norm{X} \leq r}
    .
\end{align*}
Hence, as $\mathbb{E}[X]=0$,
\begin{align*}
    \mathbb{E}\brk[s]*{X \mid \norm{X} \leq r}
    &=
    \frac{\mathbb{E}\brk[s]*{X \mathbf{1}\brk[c]{\norm{X} \leq r}}}{\Pr(\norm{X} \leq r)}
    =
    -\frac{\mathbb{E}\brk[s]*{X \mathbf{1}\brk[c]{\norm{X} > r}}}{\Pr(\norm{X} \leq r)}
    .
\end{align*}
Thus, using the tail bound (\cref{eq:sub_gaussian_tail}),
\begin{align*}
    \norm{\mathbb{E}[Z]}
    &\leq
    2 \mathbb{E}[\norm{X} \mathbf{1}\brk[c]{\norm{X} > r}]
    \tag{$\Pr(\norm{X} \leq r)>1/2$}
    \\
    &\leq
    2 \int_0^\infty \Pr(\norm{X} \mathbf{1}\brk[c]{\norm{X} > r} \geq x) dx
    \tag{tail formula for expectations}
    \\
    &\leq
    2 r \Pr(\norm{X} > r)
    + 2 \int_r^\infty \Pr(\norm{X} \mathbf{1}\brk[c]{\norm{X} > r} \geq x) dx
    \tag{split integral at $r$}
    \\
    &\leq
    \delta r
    + 2 \int_r^\infty \Pr(\norm{X} \geq x) dx
    \tag{$\Pr(\norm{X} > r) \leq \delta/2$}
    \\
    &\leq
    \delta r
    + 4 \int_r^\infty e^{-x^2/\sigma^2} dx
    \tag{\cref{eq:sub_gaussian_tail}}
    \\
    &\leq
    \delta r
    + \frac{2 \sigma^2}{r} \int_r^\infty \frac{2 x}{\sigma^2} e^{-x^2/\sigma^2} dx
    \tag{$x/r \geq 1$ for $x \geq r$}
    \\
    &=
    \delta r
    -\frac{2 \sigma^2}{r}e^{-x^2/\sigma^2} \big|_r^\infty
    =
    \delta r
    + \frac{2 \sigma^2}{r}e^{-r^2/\sigma^2}
    \\
    &\leq
    \delta r
    + \frac{\delta \sigma}{2}
    .
\end{align*}
We conclude that $\bar{X}$ is bounded in absolute value (with probability one) by
\begin{align*}
    \max\brk[c]*{r, (r + \sigma/2) (2-\delta)}
    &\leq
    3 \sigma \sqrt{\ln(4/\delta)}
    .
    \qedhere
\end{align*}
\end{proof}

\section{Convergence analysis in the convex case}
\label{sec:convex}
In this section we give a convergence analysis for \adasgd provided that the objective $f$ is a convex and $\beta$-smooth function over ${\mathbb{R}}^d$.
We prove the following:
\begin{theorem}\label{thm:convex_convergence_affine}
    Assume that $f$ is convex and $\beta$-smooth.
    Then for
    \adasgd, for any $\delta \in (0,1/4)$,
    it holds with probability at least $1-\delta$ that
    \begin{enumerate}[label=(\roman*)]
        \item%
        $\begin{aligned}[t]
            f\brk3{\frac{1}{T}\sum_{t=1}^T w_t}-f^\star
            &
            \leq
            O\bigg(
                \brk2{\tfrac{D^2}{\eta}+\eta}
                + D \sqrt{\log \tfrac{1}{\delta}}
            \bigg)
            \frac{\sigma_0}{\sqrt{T}}
            \\&
            +
            O\bigg(
                \tfrac{\gamma d_1^2}{\eta}
                + \beta (1+\sigma_1^2) \brk2{\tfrac{D^2}{\eta}+\eta}^2
                + \beta \sigma_1^2 D^2 \log \tfrac{1}{\delta}
            \bigg)
            \frac{1}{T}
            ;
        \end{aligned}$
        \item%
        and $\norm{w_t-w^\star} \leq D$ for all $1 \leq t \leq T+1$,
    \end{enumerate}
    where $d_1=\norm{w_1-w^\star}$, and
    \begin{align*}
        D^2
        &=
        \smash{\wt{O}}
        \Big(
            d_1^2
            +\eta^2
            \big(
                1+\sigma_1^2 + \gamma^{-2}\brk{\sigma_0^2+\beta \sigma_1^2 F}
                + \gamma^{-4} (\sigma_0^4+\beta^2 \sigma_1^4 F^2)
            \big)
        \Big)
        ,
        \\
        F
        &=
        \widetilde{O}\brk*{
        f(w_1)-f^\star
        + \eta \sigma_0
        + \eta^2 \beta \sigma_1^2
        + \eta^2 \beta
        }
        .
    \end{align*}
\end{theorem}
Few remarks are in order.
First, note that when $\sigma_0 \rightarrow 0$ we obtain a rate of order $1/T$, matching the optimal rate for non-stochastic smooth optimization (note that this holds even when $\sigma_1 \gg 0$).
Second, unlike the result in the non-convex case, the bound here contains a polynomial dependence on $1/\gamma$ (in the $D$ term).
This is not an artifact of our analysis: setting a overly small $\gamma$ can in fact lead \adasgd to gain a distance of $\Omega(T)$ to a minimizer; see \cref{sec:convex_lower_bound} for more details.
That said, the algorithm is fairly robust to the choice of $\gamma$; for example, consider the noisy regime with no affine noise (i.e., $\sigma_0 \gg 0$, $\sigma_1=0$): setting $\eta=\Theta(\norm{w_1-w^\star})$
and any $\gamma$ within the wide range
$\gamma \in \brk[s]{\sigma_0,\sigma_0 \sqrt{T}}$
leads to the optimal rate (up to logarithmic factors).

\subsection{Proof of~\texorpdfstring{\cref{thm:convex_convergence_affine}}{Theorem}}

In the following, we will require some additional notation.
First, for all $t$ we let
$
    d_t \triangleq \norm{w_t-w^\star} \text{ and } \bar{d}_t \triangleq \max_{s \leq t} \norm{w_s-w^\star}.
$
In the proof below, we use a slightly different ``decorrelated step size,'' given by
\begin{align*}
    \hat{\eta}_s
    \triangleq
    \frac{\eta}{\sqrt{G_{s-1}^2+\norm{\nabla f(w_s)}^2}}
    .
\end{align*}
The proof of the main result is based on a regret bound (\cref{lem:convex_regret_affine}) and a bound $d_t \leq D$ (\cref{lem:wbound_affine}) which holds with high probability, analogous to \cref{lem:non-convex-regret,lem:fbound_affine} of the previous section. See \cref{sec:convex_appendix} for their statements and proofs.
We now prove our main result in this section.
\begin{proof}[Proof of \cref{thm:convex_convergence_affine}]
With probability at least $1-3\delta$ we have from \cref{lem:wbound_affine} that $d_t \leq D$ for all $t \leq T+1$.
Using \cref{lem:convex_regret_affine} we have
\begin{align*}
    \sum_{t=1}^T g_t \boldsymbol{\cdot} (w_t-w^\star)
    &\leq
    \brk*{\frac{D^2}{\eta} + \eta} \sqrt{\sum_{t=1}^T \norm{g_t}^2}
    + \frac{\gamma d_1^2}{2 \eta}
    .
\end{align*}
Under the noise assumption, smoothness and convexity,
\begin{align*}
    \norm{g_t}^2
    &\leq
    2 \norm{\nabla f(w_t)}^2 + 2\norm{g_t - \nabla f(w_t)}^2
    \leq
    2 (1+\sigma_1^2) \norm{\nabla f(w_t)}^2 + 2 \sigma_0^2
    \\
    &\leq
    4 \beta (1+\sigma_1^2) \nabla f(w_t) \boldsymbol{\cdot} (w_t-w^\star) + 2 \sigma_0^2
    .
\end{align*}
We obtain
\begin{align*}
    \sum_{t=1}^T g_t \boldsymbol{\cdot} (w_t-w^\star)
    &
    \leq
    \brk*{\frac{D^2}{\eta}+\eta} \sigma_0 \sqrt{2 T}
    + \frac{\gamma d_1^2}{2 \eta}
    +2 \brk*{\frac{D^2}{\eta}+\eta} \sqrt{\beta (1+\sigma_1^2) \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} (w_t-w^\star)}
    .
\end{align*}
Using the inequality $2 a b \leq a^2+b^2$, we
then have
\begin{align}\label{eq:convex_regret_affine}
    \!\sum_{t=1}^T g_t \boldsymbol{\cdot} (w_t-w^\star)
    &
    \leq
    \brk*{\frac{D^2}{\eta}+\eta} \sigma_0 \sqrt{2 T}
    + \frac{\gamma d_1^2}{2 \eta}
    + \frac{1}{4} \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} (w_t-w^\star)
    + 4 \beta (1+\sigma_1^2) \brk*{\frac{D^2}{\eta}+\eta}^2
    .
\end{align}
By \cref{lem:sub_gaussian} with
$Z_t=(\nabla f(w_t)-g_t) \boldsymbol{\cdot} (w_t-w^\star)$,
$\abs{Z_t} \leq \sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_t)}^2} \norm{w_t-w^\star}$,
and
\begin{align*}
    \lambda
    &
    =
    \frac{1}{\sigma_0 D \sqrt{T/\log \tfrac{1}{\delta}}+6 \beta \sigma_1^2 D^2}
    ,
\end{align*}
with probability at least $1-4\delta$ (union bound with the event of \cref{lem:wbound_affine}),
\begin{align*}
    \sum_{t=1}^T (\nabla f(w_t)-g_t) \boldsymbol{\cdot} (w_t-w^\star)
    &\leq
    \frac{3 \lambda}{4} \sum_{t=1}^T \brk*{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_t)}^2} \norm{w_t-w^\star}^2
    + \frac{1}{\lambda} \log \tfrac{1}{\delta}
    \\
    &\leq
    \sigma_0 \sqrt{T \log \tfrac{1}{\delta}} \brk*{\frac{\bar{d}_T^2}{D}+ D}
    + \frac{\bar{d}_T^2 \sum_{t=1}^T \norm{\nabla f(w_t)}^2}{8 \beta D^2}
    + 6 \beta \sigma_1^2 D^2 \log \tfrac{1}{\delta}
    .
\end{align*}
Using the bound $\bar{d}_T \leq D$, with probability at least $1-4\delta$,
\begin{align*}
    \sum_{t=1}^T (\nabla f(w_t)-g_t) \boldsymbol{\cdot} (w_t-w^\star)
    &\leq
    2 \sigma_0 D \sqrt{T \log \tfrac{1}{\delta}}
    + \frac{1}{8 \beta} \sum_{t=1}^T \norm{\nabla f(w_t)}^2
    + 6 \beta \sigma_1^2 D^2 \log \tfrac{1}{\delta}
    \\
    &\leq
    2 \sigma_0 D \sqrt{T \log \tfrac{1}{\delta}}
    + \frac14 \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} (w_t-w^\star)
    + 6 \beta \sigma_1^2 D^2 \log \tfrac{1}{\delta}
    ,
\end{align*}
where the last inequality follows from \cref{lem:smooth_norm_bound} and convexity.
Summing with \cref{eq:convex_regret_affine} and rearranging, with probability at least $1-4 \delta$,
\begin{align*}
    \sum_{s=1}^t \nabla f(w_s) \boldsymbol{\cdot} (w_s-w^\star)
    &
    \leq
    8 \beta (1+\sigma_1^2) \brk*{\frac{D^2}{\eta}+\eta}^2
    + \brk*{
        \brk*{\frac{D^2}{\eta}+\eta} \sqrt{8}
        + 4 D \sqrt{\log \tfrac{1}{\delta}}
    } \sigma_0 \sqrt{T}
    + \frac{\gamma d_1^2}{\eta}
    \\
    &
    + 12 \beta \sigma_1^2 D^2 \log \tfrac{1}{\delta}
    .
\end{align*}
A standard application of Jensen's inequality concludes the proof (up to a constant factor in the failure probability, which we subsumed in the big-O notation).
\end{proof}

\section{Proofs of \texorpdfstring{\cref{sec:convex}}{Convex Section}}
\label{sec:convex_appendix}
Below are the statements of the two main lemmas used to prove \cref{thm:convex_convergence_affine} (their proofs follow).
The first is a regret analysis which depends on $\bar{d}_t$.
\begin{lemma}\label{lem:convex_regret_affine}
    Let $f$ be a convex and $\beta$-smooth function. Then
    \begin{align*}
        \sum_{t=1}^T g_t \boldsymbol{\cdot} (w_t-w^\star)
        \leq
        \brk*{\frac{\bar{d}_T^2}{\eta} + \eta} \sqrt{\sum_{t=1}^T \norm{g_t}^2}
        + \frac{\gamma d_1^2}{2 \eta}
        .
    \end{align*}
\end{lemma}
Similarly to the non-convex case, we need a high probability bound of $\norm{w_t-w^\star}$ in order to use the regret analysis of \cref{lem:convex_regret_affine}; this is given in our next lemma.
\begin{lemma}\label{lem:wbound_affine}
    With probability at least $1-3 \delta$, it holds that for all $t \leq T+1$,
    \begin{align*}
        d_t^2
        &\leq
        D^2
        \triangleq
        2 d_1^2
        + \eta^2 \brk*{
        \frac{1}{2}
        + 2 C_1
        + 8 C^2
        + A_T C
        + B_T \brk*{\frac{\sigma_0}{\gamma}+\sigma_1}^2
        }
        ,
    \end{align*}
    where $C=2 \log(1+\sigma^2 T / 2\gamma^2) + 7 \sigma^2 \log(T/\delta)/\gamma^2$, $C_1$ is defined at \cref{lem:sum_eta_g_t_2_affine},
    $A_T=512 \log\brk*{\frac{60 \log^2(6T)}{\delta}}$ and $B_T=512 \log^2\brk*{\frac{60 \log^2(6T)}{\delta}}$.
\end{lemma}
\subsection{Proof of \texorpdfstring{\cref{lem:convex_regret_affine}}{Lemma}}
\begin{proof}%
From the standard analysis of SGD, we have for all $t$ that
\begin{align*}
    g_t \boldsymbol{\cdot} (w_t-w^\star)
    &=
    \frac{1}{2\eta_t} \brk!{ \norm{w_t-w^\star}^2 - \norm{w_{t+1}-w^\star}^2 }
        + \frac{\eta_t}{2} \norm{g_t}^2
    .
\end{align*}
Summing over $t=1,\ldots,T$, we have
\begin{align*}
    \sum_{t=1}^T g_t \boldsymbol{\cdot} (w_t-w^\star)
    &\leq
    \frac{\gamma \norm{w_1-w^\star}^2}{2 \eta}
    + \frac12 \sum_{t=1}^T \brk*{\frac{1}{\eta_t}
    - \frac{1}{\eta_{t-1}}} \norm{w_t-w^\star}^2
    + \frac12 \sum_{t=1}^T \eta_t \norm{g_t}^2
    .
\end{align*}
Now, since $\eta_t \leq \eta_{t-1}$, we have for all $t$ that
\begin{align*}
    \frac{1}{\eta_t} - \frac{1}{\eta_{t-1}}
    \leq
    \eta_t \brk*{ \frac{1}{\eta_t^2} - \frac{1}{\eta_{t-1}^2} }
    =
    \frac{\eta_t}{\eta^2} \brk*{\gamma^2 + \sum_{s=1}^t \norm{g_s}^2 - \gamma^2 - \sum_{s=1}^{t-1} \norm{g_s}^2}
    =
    \frac{\eta_t}{\eta^2} \norm{g_t}^2
    ,
\end{align*}
thus
\begin{align*}
    \sum_{t=1}^T g_t \boldsymbol{\cdot} (w_t-w^\star)
    &\leq
    \frac{\gamma \norm{w_1-w^\star}^2}{2 \eta}
    + \frac12 \sum_{t=1}^T \eta_t \norm{g_t}^2 \brk3{1 + \frac{\norm{w_t-w^\star}^2}{\eta^2}}
    .
\end{align*}
Finally, bounding $\norm{w_t-w^\star} \leq \bar{d}_T$ and applying \cref{lem:sum_eta_g_t}
we obtain
\begin{align*}
    \sum_{t=1}^T g_t \boldsymbol{\cdot} (w_t-w^\star)
    &\leq
    \brk*{\frac{\bar{d}_T^2}{\eta} + \eta} \sqrt{\sum_{t=1}^T \norm{g_t}^2}
    + \frac{\gamma \norm{w_1-w^\star}^2}{2 \eta}
    .\qedhere
\end{align*}
\end{proof}
\subsection{Proof of \texorpdfstring{\cref{lem:wbound_affine}}{Lemma}}
In order to obtain the high probability bound of $\bar{d}_T$ we use of the following empirical Bernstein concentration bound.
\begin{lemma}[Corollary 1 of \citet{carmon2022making}]\label{lem:emp_bernstein}
    Let $X_t$ be adapted to $\mathcal{F}_t$ such that $\abs{X_t} \leq 1$ with probability $1$ for all $t$.
    Then, for every $\delta \in (0,1)$ and any $\hat{X}_t \in \mathcal{F}_{t-1}$ such that $\abs{\hat{X}_t}\leq 1$ with probability $1$,
    \begin{align*}
        \Pr\brk*{\exists t < \infty : \abs*{\sum_{s \leq t} (X_s-\mathbb{E}\brk[s]{X_s | \mathcal{F}_{s-1}})}
        \geq \sqrt{A_t(\delta) \sum_{s \leq t} (X_s-\hat{X}_s)^2 + B_t(\delta)}}
        &\leq
        \delta
        ,
    \end{align*}
    where $A_t(\delta)=16 \log\brk*{\frac{60 \log(6t)}{\delta}}$ and $B_t(\delta)=16 \log^2\brk*{\frac{60 \log(6t)}{\delta}}$.
\end{lemma}
We use the concentration bound to prove the following technical martingale bound.
\begin{lemma}\label{lem:convex_martingale_affine}
    Let $\bar{d'}_t = \max\brk[c]{\bar{d}_t,\eta}$. Then with probability at least $1-\delta$, it holds that for all $t \leq T$,
    \begin{align*}
        \sum_{s=1}^t & \hat{\eta}_{s} (\nabla f(w_s)-g_s) \boldsymbol{\cdot} (w_s-w^\star)
        \\
        &\leq
        2 \bar{d'}_t
        \sqrt{A_t(\ifrac{\delta}{\log(4T)}) \sum_{s \leq t} \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2 + \eta^2 \brk*{\frac{\sigma_0}{\gamma}+\sigma_1}^2 B_t(\ifrac{\delta}{\log(4T)})}
        ,
    \end{align*}
    where $A(\cdot)$ and $B(\cdot)$ are defined at \cref{lem:emp_bernstein}.
\end{lemma}
The next lemma bounds the $\sum_{s \leq t} \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2$ term which appear in \cref{lem:convex_martingale_affine}.
\begin{lemma}\label{lem:sum_eta_noise}
    With probability at least $1-2\delta$, for all $1 \leq t \leq T$,
    \begin{align*}
            \sum_{s=1}^t \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2
            \leq
            2 \eta^2 \log(1 + \sigma^2 T / 2 \gamma^2)
            + 7 \eta_0^2 \sigma^2 \log(T/\delta)
        .
    \end{align*}
\end{lemma}
Now we move to prove the lemma.
\begin{proof}[Proof of \cref{lem:wbound_affine}]
    Rolling a single step of SGD,
    \begin{align*}
        \norm{w_{s+1}-w^\star}^2
        &=
        \norm{w_{s}-w^\star}^2
        - 2 \eta_s g_s \boldsymbol{\cdot} (w_s-w^\star)
        + \eta_s^2 \norm{g_s}^2
        .
    \end{align*}
    Summing for $s=1 \ldots t$ and applying \cref{lem:sum_eta_g_t_2_affine},
    \begin{align}\label{eq:roll_convex_sgd}
        \norm{w_{t+1}-w^\star}^2
        &=
        \norm{w_{1}-w^\star}^2
        - 2 \sum_{s=1}^t \eta_s g_s \boldsymbol{\cdot} (w_s-w^\star)
        + \sum_{s=1}^t \eta_s^2 \norm{g_s}^2
        \nonumber
        \\
        &\leq
        \norm{w_{1}-w^\star}^2
        \underbrace{- 2 \sum_{s=1}^t \eta_s g_s \boldsymbol{\cdot} (w_s-w^\star)}_{(*)}
        + \eta^2 C_1
        .
    \end{align}
    Focusing on $(*)$, due to convexity,
    \begin{align*}
        - 2 \sum_{s=1}^t \eta_s g_s \boldsymbol{\cdot} (w_s-w^\star)
        &=
        - 2 \sum_{s=1}^t \eta_s \nabla f(w_s) \boldsymbol{\cdot} (w_s-w^\star)
        + 2 \sum_{s=1}^t \eta_s (\nabla f(w_s)-g_s) \boldsymbol{\cdot} (w_s-w^\star)
        \\
        &\leq
        2 \sum_{s=1}^t \eta_s (\nabla f(w_s)-g_s) \boldsymbol{\cdot} (w_s-w^\star)
        .
    \end{align*}
    In order to create a martingale we replace $\eta_s=\ifrac{\eta}{\sqrt{G_{s-1}^2+\norm{g_s}^2}}$ with $\hat{\eta}_s = \ifrac{\eta}{\sqrt{G_{s-1}^2+\norm{\nabla f(w_s)}^2}}$,
    \begin{align*}
        (*)
        &\leq
        2 \sum_{s=1}^t \hat{\eta}_{s} (\nabla f(w_s)-g_s) \boldsymbol{\cdot} (w_s-w^\star)
        + 2 \sum_{s=1}^t (\eta_s-\hat{\eta}_{s}) (\nabla f(w_s)-g_s) \boldsymbol{\cdot} (w_s-w^\star)
        .
    \end{align*}
    We observe that
    \begin{align*}
        \abs{\eta_s-\hat{\eta}_s}
        &=
        \eta \frac{\abs2{\sqrt{G_{s-1}^2+\norm{\nabla f(w_s)}^2}-\sqrt{G_{s-1}^2+\norm{g_s}^2}}}{\sqrt{G_{s-1}^2+\norm{g_s}^2}\sqrt{G_{s-1}^2+\norm{\nabla f(w_s)}^2}}
        \\
        &=
        \eta \frac{\abs!{\norm{\nabla f(w_s)}^2-\norm{g_s}^2}}{\sqrt{G_{s-1}^2+\norm{g_s}^2}\sqrt{G_{s-1}^2+\norm{\nabla f(w_s)}^2}(\sqrt{G_{s-1}^2+\norm{g_s}^2}+\sqrt{G_{s-1}^2+\norm{\nabla f(w_s)}^2})}
        \\
        &\leq
        \eta \frac{\norm{\nabla f(w_s)-g_s}(\norm{\nabla f(w_s)}+\norm{g_s})}{\sqrt{G_{s-1}^2+\norm{g_s}^2}\sqrt{G_{s-1}^2+\norm{\nabla f(w_s)}^2}(\sqrt{G_{s-1}^2+\norm{g_s}^2}+\sqrt{G_{s-1}^2+\norm{\nabla f(w_s)}^2})}
        \\
        &\leq
        \eta \frac{\norm{\nabla f(w_s)-g_s}}{\sqrt{G_{s-1}^2+\norm{g_s}^2}\sqrt{G_{s-1}^2+\norm{\nabla f(w_s)}^2}}
        .
    \end{align*}
    Thus
    \begin{align*}
        \sum_{s=1}^{t} (\eta_s-\hat{\eta}_s)(\nabla f(w_s)-g_s) \boldsymbol{\cdot} (w_s-w^\star)
        &\leq
        \sum_{s=1}^{t} \abs{\eta_s-\hat{\eta}_s}\norm{\nabla f(w_s)-g_s}d_s
        \\
        &\leq
        \bar{d}_t \sum_{s=1}^{t} \abs{\eta_s-\hat{\eta}_s}\norm{\nabla f(w_s)-g_s}
        \\
        &\leq
        \eta \bar{d}_t \sum_{s=1}^t \frac{\norm{\nabla f(w_s)-g_s}^2}{\sqrt{G_{s-1}^2+\norm{g_s}^2}\sqrt{G_{s-1}^2+\norm{\nabla f(w_s)}^2}}
        \\
        &\leq \frac{\bar{d}_t}{\eta} \sum_{s=1}^t \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2
        .
    \end{align*}
    Combining the inequality with \cref{lem:convex_martingale_affine}, with probability at least $1-\delta$,
    \begin{align*}
        (*)
        &\leq
        \frac{2 \bar{d}_t}{\eta} \sum_{s=1}^t \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2
        \\
        &+ 4 \bar{d'}_t
            \sqrt{A_t(\ifrac{\delta}{\log(4T)}) \sum_{s \leq t} \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2 + \eta^2 \brk*{\frac{\sigma_0}{\gamma}+\sigma_1}^2 B_t(\ifrac{\delta}{\log(4T)})}
        ,
    \end{align*}
    and using $ab \leq a^2/2 + b^2/2$,
    \begin{align*}
        (*)
        &\leq
        \frac{\bar{d}_t^2}{4}
        + \frac{4}{\eta^2} \brk*{\sum_{s=1}^t \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2}^2
        + \frac{\bar{d'}_t^2}{4}
        \\
        &+ 16
            \brk*{A_t(\ifrac{\delta}{\log(4T)}) \sum_{s \leq t} \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2 + \eta^2 \brk*{\frac{\sigma_0}{\gamma}+\sigma_1}^2 B_t(\ifrac{\delta}{\log(4T)})}
        \\
        &\leq
        \frac{\bar{d}_t^2}{2}
        + \frac{4}{\eta^2} \brk*{\sum_{s=1}^t \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2}^2
        + \frac{\eta^2}{4}
        \tag{$\bar{d'}_t=\max\brk[c]{\bar{d}_t,\eta}$}
        \\
        &+ 16
            \brk*{A_t(\ifrac{\delta}{\log(4T)}) \sum_{s \leq t} \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2 + \eta^2 \brk*{\frac{\sigma_0}{\gamma}+\sigma_1}^2 B_t(\ifrac{\delta}{\log(4T)})}
        .
    \end{align*}
    Applying \cref{lem:sum_eta_noise}, with probability at least $1-3\delta$ (via a union bound),
    \begin{align*}
        (*)
        &\leq
        \frac{\bar{d}_t^2}{2}
        + 4 \eta^2 C^2
        + \frac{\eta^2}{4}
        + 16 \eta^2
            \brk*{A_t(\ifrac{\delta}{\log(4T)}) C + \brk*{\frac{\sigma_0}{\gamma}+\sigma_1}^2 B_t(\ifrac{\delta}{\log(4T)})}
        ,
    \end{align*}
    where $C=2 \log(1+\sigma^2 T / 2\gamma^2) + 7 \sigma^2/\gamma^2 \log(T/\delta)$.
    Hence, returning to \cref{eq:roll_convex_sgd} and applying \cref{lem:sum_eta_g_t_2_affine},
    \begin{align*}
        d_{t+1}^2
        &\leq
        \frac{\bar{d}_t^2}{2}
        + d_1^2
        + \eta^2 \brk*{
            \frac{1}{4}
            + C_1
            + 4 C^2
            + 16
            \brk*{A_t(\ifrac{\delta}{\log(4T)}) C + \brk*{\frac{\sigma_0}{\gamma}+\sigma_1}^2 B_t(\ifrac{\delta}{\log(4T)})}
        }
        \\
        &\leq
        \frac{\bar{d}_t^2}{2}
        + \frac{D}{2}
        .
        \tag{def of $D$, $32 A_t(\ifrac{\delta}{\log(4T)}) \leq A_T$ and $32 B_t(\ifrac{\delta}{\log(4T)}) \leq B_T$}
    \end{align*}
    Hence, rolling to $t=1$, with probability at least $1-3\delta$, $d_t \leq D$ for all $1 \leq t \leq T+1$.
\end{proof}
\begin{proof}[Proof of \cref{lem:convex_martingale_affine}]
    In order to invoke \cref{lem:emp_bernstein} we will replace $w_s-w^\star$ with a version which is scaled and projected to the unit ball.
    We denote
    $
        a_k = 2^{k-1} \bar{d'}_1
        \text{ and }
        k_t=\ceil{\log (\bar{d'}_t/\bar{d'}_1)}+1
        .
    $
    Thus, $\bar{d}_t \leq \bar{d'}_t \leq a_{k_t} \leq 2\bar{d'}_t$.
    Since $\norm{w_{s+1}-w^\star} \leq \norm{w_{s}-w^\star} + \eta$ for all $s$, $\bar{d}_t \leq d_1 + \eta (t-1)$ and $1 \leq k_t \leq \ceil{\log (t)} + 1 \leq \log(4T)$.
    Defining the projection to the unit ball, $\proj{x}=x/\max\brk[c]*{1,\norm{x}}$,
    \begin{align*}
        \frac{w_s-w^\star}{a_{k_t}}
        &=
        \Pi_1 \brk*{\frac{w_s-w^\star}{a_{k_t}}}
        .
    \end{align*}
    To simplify notation, let $C = \eta(\ifrac{\sigma_0}{\gamma}+\sigma_1)$ and note that
    \begin{align*}
        \norm{\hat{\eta}_s (\nabla f(w_s)-g_s)}
        &\leq
        \frac{\eta \sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_s)}^2}}{\sqrt{\gamma^2 + \norm{\nabla f(w_s)}^2}}
        \leq
        \frac{\eta \sigma_0}{\gamma}
        + \eta \sigma_1
        =
        C
        .
    \end{align*}
    Thus,
    \begin{align}\label{eq:sum_martingales}
        \sum_{s=1}^t & \frac{\hat{\eta}_s (\nabla f(w_s)-g_s) \boldsymbol{\cdot} (w_s-w^\star)}{C a_{k_t}}
        \nonumber
        \\&=
        \sum_{s=1}^t \frac{\hat{\eta}_s (\nabla f(w_s)-g_s)}{C} \boldsymbol{\cdot} \Pi_1 \brk*{\frac{w_s-w^\star}{a_{k_t}}}
        \nonumber
        \\
        &\leq
        \abs*{\sum_{s=1}^t \frac{\hat{\eta}_s (\nabla f(w_s)-g_s)}{C} \boldsymbol{\cdot} \Pi_1 \brk*{\frac{w_s-w^\star}{a_{k_t}}}}
        \nonumber
        \\
        &\leq
        \max_{1 \leq k \leq \floor{\log(4T)}}\abs*{\sum_{s=1}^t \frac{\hat{\eta}_s (\nabla f(w_s)-g_s)}{C} \boldsymbol{\cdot} \Pi_1 \brk*{\frac{w_s-w^\star}{a_{k}}}}
        .
    \end{align}
    Let $X_s^{(k)}=\frac{\hat{\eta}_s (\nabla f(w_s)-g_s)}{C} \boldsymbol{\cdot} \Pi_1 \brk*{\frac{w_s-w^\star}{a_{k}}}$ for some $k$.
    $$
    \condE{g_s}{s-1}=\nabla f(w_s)
    \quad\implies\quad
    \condE{X_s^{(k)}}{s-1}=0
    .
    $$
    Also note that $X_s^{(k)}\leq 1$ with probability $1$.
    Using \cref{lem:emp_bernstein} with the $X_s^{(k)}$ we defined and $\hat{X_s}=0$, for any $k$ and $\delta' \in (0,1)$, with probability at least $1-\delta'$, for all $t \leq T$,
    \begin{align*}
        \abs*{\sum_{s \leq t} X_s^{(k)}}
        &\leq
        \sqrt{A_t(\delta') \sum_{s \leq t} (X_s^{(k)})^2 + B_t(\delta')}
        .
    \end{align*}
    We can upper bound $(X_s^{(k)})^2$,
    \begin{align*}
        (X_s^{(k)})^2
        &\leq
        \frac{\hat{\eta}_s^2 \norm{\nabla f(w_s)-g_s}^2}{C^2} \norm*{\Pi_1 \brk*{\frac{w_s-w^\star}{a_{k}}}}^2
        \tag{Cauchy-Schwarz inequality}
        \\
        &\leq
        \frac{\hat{\eta}_s^2 \norm{\nabla f(w_s)-g_s}^2}{C^2}
        \tag{$\norm{\Pi_1(x)} \leq 1$}
        \\
        &\leq
        \frac{\eta_{s-1}^2}{C^2} \norm{\nabla f(w_s)-g_s}^2
        \tag{$\eta_{s-1} \geq \hat{\eta}_s$}
        .
    \end{align*}
    Thus, returning to \cref{eq:sum_martingales} multiplied by $C a_{k_t}$, with probability at least $1-\delta' \log(4T)$ (union bound for all
    $
    1
    \leq k
    \leq
    \floor{\log(4T)})
    ,
    $
    \begin{align*}
        \sum_{s=1}^t
        \hat{\eta}_s (\nabla f(w_s)-g_s) \boldsymbol{\cdot} (w_s-w^\star)
        &\leq
        C a_{k_t}
        \sqrt{A_t(\delta') \sum_{s \leq t} \frac{\eta_{s-1}^2}{C^2} \norm{\nabla f(w_s)-g_s}^2 + B_t(\delta')}
        .
    \end{align*}
    As $a_{k_t} \leq 2\bar{d'}_t$, picking $\delta'=\ifrac{\delta}{\log(4T)}$, with probability at least $1-\delta$,
    \begin{align*}
        \sum_{s=1}^t
        \eta_{s-1} (\nabla f(w_s)-g_s) \boldsymbol{\cdot} (w_s-w^\star)
        &\leq
        2 \bar{d'}_t
        \sqrt{A_t(\ifrac{\delta}{\log(4T)}) \sum_{s \leq t} \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2 + C^2 B_t(\ifrac{\delta}{\log(4T)})}
        .
        \qedhere
    \end{align*}
\end{proof}
\begin{proof}[Proof of \cref{lem:sum_eta_noise}]
    Let $Z_s = \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s)-g_s)$
    and
    note that $\condE{Z_s}{s-1}=0$. Hence, $\brk{Z_s}_{s=1}^t$ is a martingale difference sequence.
    By \cref{lem:sub_gaussian} (with $\abs{Z_s} \leq \sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_s)}^2} \norm{\nabla f(w_s)}$ and $\lambda=2/3 \sigma^2$) and a union bound (over $1 \leq t \leq T$), with probability at least $1-\delta$,
    \begin{align*}
        \sum_{s=1}^t \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s)-g_s)
        &\leq
        \frac{3 \lambda}{4} \sum_{s=1}^t (\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_s)}^2) \norm{\nabla f(w_s)}^2
        + \frac{1}{\lambda} \log(T/\delta)
        \\
        &
        \leq
        \frac{\sigma_0^2 + 2 \beta \sigma_1^2 \bar{\Delta}_t}{2\sigma^2} \sum_{s=1}^t \norm{\nabla f(w_s)}^2
        + \frac{3}{2} \sigma^2 \log(T/\delta)
        \tag{\cref{lem:smooth_norm_bound}}
    \end{align*}
    for all $t \in [T]$, and under a union bound with \cref{lem:fbound_affine}, with probability at least $1-2\delta$, since $\sigma_0^2 + 2 \beta \sigma_1^2 \bar{\Delta}_t \leq \sigma^2$,
    \begin{align}\label{eq:convex_martingale_1_affine}
        \sum_{s=1}^t \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s)-g_s)
        &\leq
        \frac{1}{2} \sum_{s=1}^t \norm{\nabla f(w_s)}^2
        + \frac{3}{2} \sigma^2 \log(T/\delta)
        .
    \end{align}
    Let
    $$
    k = \min\set*{t \leq T : \sum_{s=1}^t \norm{\nabla f(w_s)-g_s}^2 > 6 \sigma^2 \log(T/\delta)}.
    $$
    For $t < k$,
    \begin{align*}
        \sum_{s=1}^t \eta_{s-1}^2 \norm{\nabla f(w_s)-g_t}^2
        &\leq
        \eta_0^2 \sum_{s=1}^t \norm{\nabla f(w_s)-g_t}^2
        \leq
        6 \eta_0^2 \sigma^2 \log(T/\delta)
        .
    \end{align*}
    For $t \geq k$,
    \begin{align*}
        G_t^2
        &=
        \gamma^2
        + \sum_{s=1}^t \norm{\nabla f(w_s)}^2
        + \sum_{s=1}^t \norm{\nabla f(w_s) - g_s}^2
        - 2 \sum_{s=1}^t \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s)-g_s)
        \\
        &\geq
        \gamma^2
        + \sum_{s=1}^t \norm{\nabla f(w_s) - g_s}^2
        - 3 \sigma^2 \log(T/\delta)
        \tag{\cref{eq:convex_martingale_1_affine}}
        \\
        &\geq
        \gamma^2
        + \frac{1}{2} \sum_{s=1}^t \norm{\nabla f(w_s) - g_s}^2
        \tag{$t \geq k$}
        .
    \end{align*}
    Hence,
    \begin{align*}
        \sum_{s=k}^t \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2
        &=
        \sum_{s=k}^t \eta_{s}^2 \norm{\nabla f(w_s)-g_s}^2
        +
        \sum_{s=k}^t (\eta_{s-1}^2-\eta_{s}^2) \norm{\nabla f(w_s)-g_s}^2
        \\
        &\leq
        \sum_{s=k}^t \eta_{s}^2 \norm{\nabla f(w_s)-g_s}^2
        +
        \eta_0^2 \sigma^2
        \\
        &\leq
        2 \eta^2 \sum_{s=k}^t \frac{\norm{\nabla f(w_s)-g_s}^2}{2 \gamma^2 +
        \sum_{k=1}^s \norm{\nabla f(w_k)-g_k}^2}
        + \eta_0^2 \sigma^2
        .
    \end{align*}
    Using \cref{lem:sum_eta_g_t} with $G_0=\sqrt{2 \gamma^2} \text{ and } g_t=\nabla f(w_s)-g_s,$
    \begin{align*}
        \sum_{s=k}^t \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2
        &\leq
        2 \eta^2 \log\brk*{\frac{2 \gamma^2 + \sum_{s=k}^t \norm{\nabla f(w_s)-g_s}^2}{2 \gamma^2}}
        + \eta_0^2 \sigma^2
        \\
        &\leq
        2 \eta^2 \log\brk*{1+\frac{\sigma_0^2 T + \sigma_1^2 \sum_{s=k}^t \norm{\nabla f(w_s)}^2}{2 \gamma^2}}
        + \eta_0^2 \sigma^2
        \tag{noise bound}
        \\
        &\leq
        2 \eta^2 \log(1 + \sigma^2 T / 2 \gamma^2)
        + \eta_0^2 \sigma^2
        \tag{\cref{lem:smooth_norm_bound,lem:fbound_affine}}
        .
    \end{align*}
    Hence, combining with the case of $t < k$, with probability at least $1-2\delta$, for all $t \leq T$,
    \begin{align*}
        \sum_{s=1}^t \eta_{s-1}^2 \norm{\nabla f(w_s)-g_s}^2
        &\leq
        2 \eta^2 \log(1 + \sigma^2 T / 2 \gamma^2)
        + 7 \eta_0^2 \sigma^2 \log(T/\delta)
        .
        \qedhere
    \end{align*}
\end{proof}
\section{Lower Bound for the Effective Diameter in the Convex Case}
\label{sec:convex_lower_bound}

In \cref{sec:convex} we establish a bound of $\max_{t \leq T+1} \norm{w_{t}-w^\star}^2$ which has polynomial dependence in $\ifrac{1}{\gamma}$.
In this section we construct a one-dimensional example which shows that such polynomial dependence is necessary in order to obtain a bound of $\max_{t \leq T+1} \norm{w_{t}-w^\star}^2$ which is $o(T)$ with high probability.

Let $f(w)=\tfrac{\beta}{2}w^2$ for some $0 < \beta < \frac{\sigma}{T^{3/2}}$ ($\beta$ is extremely small with respect to $\sigma$).
Let $g(w)$ be a stochastic oracle such that
\begin{align*}
    g(w)
    &=
    \begin{cases}
        \nabla f(w) + \sigma & \text{w.p. $\frac{1}{T}$;}
        \\
        \nabla f(w) - \frac{\sigma}{T-1} & \text{w.p. $1-\frac{1}{T}$,}
    \end{cases}
\end{align*}
where $T > 1$.
Hence, $E\brk[s]{g(w)-\nabla f(w)}=0$ and
$\norm{g(w)-\nabla f(w)}^2 \leq \sigma^2$ with probability $1$.
Let $w_1=w^\star=0$ and $\eta=1$.

With constant probability ($\approx e^{-1}$), $g_t=\nabla f(w_t) - \sigma/(T-1)$ for all $t \in [T]$.
As $\max_{t \leq T+1} \norm{w_{t}-w^\star}^2=o(T)$ with high probability, we assume that for large enough $T$, $\abs{w_t-w^\star} < \frac{\sqrt{T}}{2}$ for all $t \leq T+1$ (still with constant probability under a union bound).
Hence, with constant probability,
\begin{align*}
    w_{T+1}
    &= -\sum_{t=1}^{T} \eta_t g_t
    = -\sum_{t=1}^{T} \eta_t \brk*{\beta w_t-\frac{\sigma}{T-1}}
    \geq \frac{\sigma}{2(T-1)} \sum_{t=1}^{T} \eta_t
    ,
\end{align*}
where the last inequality is the assumption of small $\beta$.
Thus, as $\norm{g_s} = \norm{\beta w_t - \ifrac{\sigma}{(T-1)}} \leq \ifrac{\sigma}{(T-1)}$ (as $w_t \geq 0$ for all $t$ since the stochastic gradients are always negative),
\begin{align*}
    \sum_{t=1}^T \eta_t
    &=
    \sum_{t=1}^T \frac{1}{\sqrt{\gamma^2 + \sum_{s=1}^t \norm{g_s}^2}}
    \geq
    \sum_{t=1}^T \frac{1}{\sqrt{\gamma^2 + \frac{\sigma^2 t}{(T-1)^2}}}
    =
    \frac{T-1}{\sigma} \sum_{t=1}^T \frac{1}{\sqrt{\frac{\gamma^2 (T-1)^2}{\sigma^2} + t}}
    .
\end{align*}
Returning to $w_{T+1}$, bounding sum by integration,
\begin{align*}
    w_{T+1}
    &\geq
    \frac{1}{2} \sum_{t=1}^T \frac{1}{\sqrt{\frac{\gamma^2 (T-1)^2}{\sigma^2} + t}}
    \geq
        \sqrt{\frac{\gamma^2 (T-1)^2}{\sigma^2} + T+1}-\sqrt{\frac{\gamma^2 (T-1)^2}{\sigma^2} + 1}
    =
    \Omega\brk*{\ifrac{\sigma}{\gamma}}
    .
\end{align*}
Thus, $\norm{w_{T+1}-w^\star}^2 = \Omega\brk{\ifrac{\sigma^2}{\gamma^2}}$.
\section{Introduction}
Stochastic Gradient Descent (SGD) \citep{robbins1951stochastic} is the de facto standard optimization method in modern machine learning.
SGD updates take the following form; $$w_{t+1}=w_{t}-\eta_{t} g_t,$$ starting from an arbitrary point $w_{1}$, where $\eta_{t}$ is the step-size at time $t$ and $g_{t}$ is a stochastic gradient at $w_{t}$.
The sequence of stepsizes $\eta_1,\eta_2,\ldots$ is notoriously hard to tune in machine learning applications~\citep[e.g.,][]{bottou2012stochastic,schaul2013no}, especially when the problem parameters (e.g., smoothness or Lipschitz constant of the objective) or estimates thereof are unknown.
This has popularized the study of \emph{adaptive gradient methods}~\citep[e.g.,][]{duchi2011adaptive,kingma2014adam,reddi2018convergence,tran2019convergence,alacaoglu2020new}
where stepsizes are computed adaptively based on magnitudes of previous stochastic gradients.

Although their non-adaptive counterparts are well studied, optimization theory of adaptive methods is still lacking and attract much interest, particularly in the smooth and non-convex setting.
The different types of analyses performed by previous works in this setting revolve around the assumptions taken in order to establish convergence which includes the type of noise (bounded variance, sub-Gaussian, bounded support, affine variance, etc.), whether the true gradients or the stochastic gradients are uniformly bounded \citep{ward2019adagrad,kavis2021high}, and also whether to use a step-size that does not depends on the current stochastic gradient \citep{li2019convergence,li2020high} or to handle the bias introduced by such a dependency \citep{ward2019adagrad,Faw2022ThePO}.

We focus on perhaps the simplest and most well-studied adaptive optimization method:
SGD with AdaGrad-type stepsizes, which we refer to as~\adasgd (also known in the literature as AdaGrad-Norm; e.g.,
\citealt{ward2019adagrad,Faw2022ThePO}).
For parameters $\eta,\gamma > 0$, \adasgd performs SGD updates with stepsizes
\begin{align*}
    \eta_t = \frac{\eta}{\sqrt{\gamma^2 + \sum_{s=1}^t \norm{g_s}^2}},
\end{align*}
where $g_t$ is a stochastic gradient at $w_t$.

Although numerous convergence results for smooth \adasgd exist in the literature (mostly in the non-convex case, except for \citet{li2019convergence} discussed below), they all fall short of at least one important aspect.
First, many of the existing analyses require a strong global Lipschitz assumption on the objective \citep[][]{ward2019adagrad,zhou2018convergence,kavis2021high} which does not hold even for simple quadratic functions in one dimension, and one that is not required for the analysis of analogous non-adaptive methods.
Second,
several results assume partial knowledge of the problem parameters, such as the smoothness parameter~\citep[e.g.,][]{li2019convergence,li2020high}, and thus in a sense do not fully establish the desired adaptivity of \adasgd.

To our knowledge, only \citet{Faw2022ThePO} do not impose neither of the above limiting assumptions, but their bounds do not hold with high probability: namely, they give convergence bounds with a polynomial dependence on the probability margin (granted, under a somewhat weaker tail assumption on the gradient noise).
On the other hand, their results apply in a significantly more challenging setting with an ``affine variance'' noise model, where the variance of the gradient noise is not uniformly bounded and may scale with the squared norm of the true gradient.

\subsection{Our contributions}

We provide a comprehensive analysis of \adasgd that does not suffer from any of the former issues:
for a $\beta$-smooth objective and under a general affine variance noise model, we establish high probability convergence guarantees in both the non-convex and convex regimes, without any knowledge of the problem parameters and without imposing a global Lipschitz condition on the objective nor a bound on the distance to a minimizer.
Further, our analysis is considerably simpler than previous analyses of \adasgd in similar generality~\citep[e.g.,][]{Faw2022ThePO}

In more detail, we consider a stochastic first-order model with affine variance noise, similar to the one of \citet{Faw2022ThePO}; given a point $w \in {\mathbb{R}}^d$ we assume that the gradient oracle returns a random vector $g(w)$ such that
$
    \mathbb{E}[g(w) \mid w] = \nabla f(w)
$
and, with probability one, satisfies%
\footnote{While our analysis assumes for simplicity affine noise bounded with probability one, we provide a general reduction to allow for sub-Gaussian noise with bounded affine variance; see \cref{sec:sub_gaussian_reduction} for details.}
\begin{align*}
    \norm{g(w)-\nabla f(w)}^2 \leq \sigma_0^2 + \sigma_1^2 \norm{\nabla f(w)}^2
    ,
\end{align*}
where $\sigma_0,\sigma_1 \geq 0$ are noise parameters, unknown to the algorithm.
Our main results in this setup are as follows:
\begin{itemize}[leftmargin=*]
\item
Our first main result establishes a high probability convergence guarantee in the general $\beta$-smooth (not necessarily convex) setting.
Given a $\beta$-smooth objective $f$, we show that for $T$-steps \adasgd, with probability at least $1-\delta$,
\begin{align*}
\frac{1}{T} \! \sum_{t=1}^T \norm{\nabla f(w_t)}^2
&
\leq
C_1 \frac{\operatorname{poly}\brk{\log\tfrac{T}{\delta}}}{T}
+ C_2 \frac{\sigma_0 \operatorname{poly}\brk{\log\tfrac{T}{\delta}}}{\sqrt{T}}
,
\end{align*}
where $C_1$ and $C_2$ are quantities that depend polynomially on the problem parameters ($\beta,\sigma_0,\sigma_1,\eta,\gamma$; the dependence on $1/\gamma$ is in fact logarithmic).
The formal statement is given in \cref{thm:non_convex_convergence}.
\item
Our second main result is a convergence guarantee for \adasgd in the same unconstrained optimization setup but with an additional convexity assumption.
Given a $\beta$-smooth and \emph{convex} objective $f$, we show that for $T$-steps \adasgd, with probability at least $1-\delta$, it holds that
\begin{align*}
    f \brk3{\frac1T \! \sum_{t=1}^{T} w_t \!} \!-\! f^\star \!
    &
    \leq
    C_3 \frac{\operatorname{poly}\brk{\log\tfrac{T}{\delta}}}{T}
    + C_4 \frac{\sigma_0 \operatorname{poly}\brk{\log\tfrac{T}{\delta}}}{\sqrt{T}}
    ,
\end{align*}
where $f^\star = \min_{w \in {\mathbb{R}}^d} f(w)$ and $C_3$ and $C_4$ are again quantities that depend polynomially on the problem parameters ($\beta,\sigma_0,\sigma_1,\eta,\gamma$; here the dependence on $1/\gamma$ is polynomial).
The formal statement is provided in \cref{thm:convex_convergence_affine}.
\end{itemize}

These convergence bounds match the rates of the form  $O(1/T+\sigma_0/\sqrt{T})$ for \emph{non-adaptive, properly tuned SGD} in the convex~\citep{lan2012optimal} and non-convex cases~\citep{ghadimi2013stochastic} cases, up to problem dependent constants and logarithmic factors.
However, we emphasize again that our results hold true, with high probability, \emph{without} assuming uniformly bounded gradients and \emph{without} requiring any knowledge of the problem parameters ($\beta,\sigma_0,\sigma_1$).
To our knowledge, in the non-convex case our result is the first convergence guarantee to hold with high probability (with a proper poly-log dependence on the probability margin) without requiring a global Lipschitz assumption---and in fact, it holds even in a more general affine variance noise model.
Similarly, as far as we are aware, our result in the convex case constitutes the first (high probability) analysis of adaptive SGD updates that applies in an unconstrained optimization setup, i.e., without restricting the updates to a domain of bounded (and a-priori known) diameter using projections.

Further, our bounds maintain a remarkable property of the non-adaptive bounds for smooth SGD~\citep{lan2012optimal,ghadimi2013stochastic}:
as can be seen from the bounds above,
they exhibit optimal rates (in the context of non-accelerated gradient methods) in both the high-noise and low-noise regimes, in the sense that they improve from the traditional stochastic $1/\sqrt{T}$ rate to a fast $1/T$ rate, up to logarithmic factors and problem-dependent constants, when the additive noise parameter $\sigma_0$ is sufficiently small: namely, when $\sigma_0 = O(1/\sqrt{T})$.
In contrast to the non-adaptive bounds, however, this interpolation between different rates occurs automatically, without knowledge of the underlying parameters and their relevant regime.
Unlike prior work on the affine noise model, our rate-interpolation takes place even when the affine term is large, i.e., when $\sigma_1 \gg 0$; for comparison, \citet{Faw2022ThePO} required that $\sigma_1 = O(1/\sqrt{T})$ for drawing a similar conclusion.
For completeness, we also provide a high probability analysis of non-adaptive, smooth SGD (with known parameters) under the affine noise model, complementing the result in expectation of \citet[Theorem 4.8]{bottou2018optimization}; see details in \cref{sec:known-non-convex}.

Finally, it is worth highlighting another remarkable adaptivity property of \adasgd, demonstrated by our results: the algorithm is able to automatically adapt to the optimization problem at hand and the same step-size schedule simultaneously achieves the optimal rates of convergence of (non-accelerated) SGD in both the convex and non-convex smooth (even unconstrained) optimization settings.

\subsection{Analysis overview}

Let us highlight some of the key challenges in analyzing SGD with adaptive stepsizes and outline our approach for overcoming them.
We begin with the non-convex smooth case, which has been the focus of prior work.

\paragraph{Non-convex regret analysis.}

The starting point of our analysis is a ``non-convex regret bound'' for \adasgd, similar to~\citet{kavis2021high};
for a $\beta$-smooth objective $f$ it can be shown (see \cref{lem:non-convex-regret} for details and a proof) that
\begin{align*}
    \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} g_t
    &\leq
    \frac{\gamma \Delta_1}{\eta}
    + \brk*{\frac{2 \bar{\Delta}_T}{\eta}+\eta \beta} \sqrt{\sum_{t=1}^T \norm{g_t}^2}
    ,
\end{align*}
where $\Delta_1=f(w_1)-f^\star$ and $\bar{\Delta}_T = \max_{t \leq T} f(w_t)-f^\star$.

The challenging term to handle in this analysis turns out to be the one that involves $\bar{\Delta}_T$.
Indeed, assume for example a deterministic bound $\bar{\Delta}_T \leq F$ (e.g., this is the case whenever $f$ is upper bounded uniformly); taking  expectations of both sides, we obtain
\begin{align*}
    \sum_{t=1}^T & \mathbb{E}\norm{\nabla f(w_t)}^2
    \leq
    \frac{\gamma \Delta_1}{\eta}
    + \brk*{\frac{2 F}{\eta}+\eta \beta} \sqrt{\sum_{t=1}^T \mathbb{E}\norm{g_t}^2}
    .
\end{align*}
Then, using $\norm{g_t}^2 \leq 2 \norm{\nabla f(w_t)-g_t}^2 + 2 \norm{\nabla f(w_t)}^2$ and our assumption on the noise yields the bound
\begin{align*}
    \mathbb{E}\brk[s]3{ \frac1T \sum_{t=1}^T \norm{\nabla f(w_t)}^2}
    &
    \leq
    \frac{2 \gamma \Delta_1}{\eta T}
    + (1+\sigma_1^2) \brk*{\frac{2 F}{\eta}+\eta \beta}^2 \frac1T
    + \brk*{\frac{2 F}{\eta}+\eta \beta} \sigma_0 \sqrt{\frac8T}
    ,
\end{align*}
which concludes an in-expectation analysis.

\paragraph*{Bounding $\bm{\bar{\Delta}_T}$ w.h.p.}
The issue, however, is that $\bar{\Delta}_T$ is a random variable which is not bounded a-priori.
Due to the probabilistic dependence between the different factors in the regret bound, our approach proceeds by establishing a bound over $\bar{\Delta}_T$ that holds with high probability (rather than in expectation).
Due to smoothness, we can bound $\bar{\Delta}_T$ via
\begin{align*}
    f (w_{t+1})
    &\leq
    f(w_{1})
    - \sum_{s=1}^{t} \eta_s \nabla f(w_s) \boldsymbol{\cdot} g_s
    + \frac{\beta}{2} \sum_{s=1}^{t} \eta_s^2 \norm{g_s}^2
    .
\end{align*}
The second summation in the bound is straightforward to handle and can be bounded (deterministically) by a logarithmic term (see \cref{lem:sum_eta_g_t_2_affine} for details).
The first summation is more problematic, again due to (conditional) probabilistic dependencies involving the stepsize $\eta_s$, which is crucially a random variable that depends on stochastic gradient $g_s$.

\paragraph{Decorrelated stepsizes.}

To break the correlation between $\eta_{s}$ and $g_{s}$ in the second term above, \citet{kavis2021high} replace $\eta_s$ with $\eta_{s-1}$.
Then, for bounding the difference between $\eta_s$ and $\eta_{s-1}$ they assumed a \emph{global Lipschitz assumption} which limits the effect of each individual gradient on the stepsize $\eta_s$.
To avoid this limiting assumption, we use a different ``decorrelated'' proxy for the stepsize, defined by replacing $\norm{g_t}^2$ with $(1+\sigma_1^2)\norm{\nabla f(w_t)}^2 + \sigma_0^2$;
this was introduced first by \citet{ward2019adagrad} and then modified by \citet{Faw2022ThePO} to allow for affine-variance noise.
We can then decompose using the proxy,
\begin{align*}
    - \eta_s \nabla f(w_s) \boldsymbol{\cdot} g_s
    &
    =
    - \tilde{\eta}_s \norm{\nabla f(w_s)}^2
    + \tilde{\eta}_s \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s) - g_s)
    + (\tilde{\eta}_s-\eta_s) \nabla f(w_s) \boldsymbol{\cdot} g_s
    ,
\end{align*}
and bound the cumulative sum of each of the terms.

\paragraph*{Concentration with affine variance.}
The challenging term to bound is the martingale
\begin{align*}
    \sum_{s=1}^t \tilde{\eta}_s \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s) - g_s).
\end{align*}
Bounding the martingale is challenging because the magnitude $\norm{\nabla f(w_s)-g_s}$ depends on the previous steps with affine-variance noise.
The key observation is that since we wish to establish a bound for $f(w_{t+1})-f^\star$ for all $t$, and thanks to the well-known property of smooth functions that $\norm{\nabla f(w)}^2 \leq 2 \beta (f(w)-f^\star)$, it is suffice to perform an inductive argument and use the bound for the previous steps to bound the affine noise terms.
As the noise bound in a given step is in itself a random variable that depends on previous stochastic gradients, we appeal to a Bernstein-type martingale concentration inequality (specifically, we use one due to \citealt{li2020high}) and a careful analysis yields a bound on $- \sum_{s=1}^t \eta_s \nabla f(w_s) \boldsymbol{\cdot} g_s$ and establishes a bound on $\bar{\Delta}_T$ without the need for a global Lipschitz assumption.

\paragraph*{Additional challenges in the convex case.}

We follow a similar approach using a regret bound and a bound of $\norm{w_t-w^\star}^2$ which holds for all $t \leq T$ with high probability.
In the convex case we need to perform a similar decomposition to bound $- \sum_{s=1}^t \eta_s g_s \boldsymbol{\cdot} (w_s-w^\star)$ where $w^\star$ is a minimizer of $f$.
Using the proxy stepsize sequence $\tilde{\eta}_t$
does not work.
Instead, we employ a slightly different decorrelated stepsize sequence $(\hat{\eta}_s)_s$
that (A) suffices for bounding $\sum_{s=1}^t (\hat{\eta}_s-\eta_s) g_s \boldsymbol{\cdot} (w_s-w^\star)$ and (B) keeps the norm of the zero-mean vector $\hat{\eta}_s (\nabla f(w_s) - g_s)$ bounded with probability one, which is crucial for the martingale analysis.
The martingale lemma of \citet{li2020high} is too loose for bounding $\sum_{s=1}^t \hat{\eta}_s (\nabla f(w_s) - g_s) \boldsymbol{\cdot} (w_s-w^\star)$ and we rely instead on the martingale analysis technique of \citet{carmon2022making} which uses an empirical Bernstein inequality introduced by \citet{howard2021time}.
The empirical Bernstein is necessary in order to tightly bound the martingale
and establish the bound on $\norm{w_t-w^\star}^2$.

\section{Non-convex SGD with known parameters}
\label{sec:known-non-convex}

In this section we provide a high probability convergence analysis of Stochastic Gradient Descent (SGD) for a $\beta$-smooth objective function $f$ under an affine noise model with known parameters $(\beta,\sigma_0,\sigma_1)$.
The tuned step-size we use is the following,
\begin{align}\label{eq:known-eta}
    \eta
    &=
    \min\brk[c]*{\frac{1}{4 \beta \brk{1 + \sigma_1^2 \log \tfrac{T}{\delta}}},\frac{\alpha}{\sigma_0 \sqrt{T}}}
    ,
\end{align}
for some parameter $\alpha>0$.
Following is the convergence guarantee.
\begin{theorem}\label{thm:known-non-convex-convergence}
    Assume that $f$ is $\beta$-smooth.
    Then for SGD with step-size (\cref{eq:known-eta}), for any $\delta \in (0,\tfrac{1}{2})$,
    it holds with probability at least $1-2 \delta$ that
    \begin{enumerate}[label=(\roman*)]
        \item%
        $\begin{aligned}[t]
            \frac{1}{T} \! \sum_{t=1}^T \norm{\nabla f(w_t)}^2
            &
            \leq
            \brk2{\beta \alpha + \frac{\Delta_1}{\alpha}}\frac{2 \sigma_0 }{\sqrt{T}}
            + \brk*{8 \beta \Delta_1 (1+ 4 \sigma_1^2) \log \tfrac{T}{\delta}
            + 24 \sigma_1^2 \beta^2 \alpha^2 \log \tfrac{1}{\delta}
            + 15 \sigma_0^2 \log \tfrac{1}{\delta}}
            \frac{1}{T}
            ;
        \end{aligned}$
        \item%
        and $f(w_t)-f^\star \leq F$ for all $1 \leq t \leq T+1$,
    \end{enumerate}
    where $\Delta_1=f(w_1)-f^\star$ and
    \begin{align}
    \label{eq:known-f-bound}
    F
    &=
    \begin{aligned}[t]
        2 \Delta_1
        + 2 \beta \alpha^2
        + 3 \min\brk[c]*{\frac{\sigma_0^2}{4 \beta \brk{1 + \sigma_1^2 \log \tfrac{T}{\delta}}},\frac{\sigma_0 \alpha}{\sqrt{T}}} \log \tfrac{T}{\delta}
        .
    \end{aligned}
    \end{align}
\end{theorem}
Similarly to \cref{thm:non_convex_convergence}, \cref{thm:known-non-convex-convergence} establishes a high probability bound of $f(w_t)-f^\star$ for all $t$.
Also note that tuning of $\alpha=O(\sqrt{\Delta_1/\beta})$ yields a rate of $\smash{\wt{O}}\brk!{\ifrac{\sigma_0 \sqrt{\beta \Delta_1}}{\sqrt{T}}+\ifrac{(1+\sigma_1^2)\beta\Delta_1}{T}}$, matching  the results in expectation of \citet{ghadimi2013stochastic} and \citet[Theorem 4.8]{bottou2018optimization} up to few additional logarithmic factors being a result of the high probability analysis.

\subsection{Proof of \texorpdfstring{\cref{thm:known-non-convex-convergence}}{Theorem}}

In order to prove \cref{thm:known-non-convex-convergence} the next lemma establish a bound of $f(w_{t})-f^\star$ for all $t \leq T$ that holds with high probability, similarly to \cref{lem:fbound_affine} (proof given later in the section).
\begin{lemma}\label{lem:known-fbound-affine}
    With probability at least $1-\delta$, it holds that
    for all $t \leq T+1$,
    $f(w_{t})-f^\star \leq F$
    (with $F$ defined in
    \cref{thm:known-non-convex-convergence}
    ).
\end{lemma}
We continue to prove the main result using similar lines to the proof of \cref{thm:non_convex_convergence}.
\begin{proof}[Proof of \cref{thm:known-non-convex-convergence}]
    We start with a regret analysis.
    Using $\beta$-smoothness, we have for all $t \geq 1$ that
    \begin{align*}
        f(w_{t+1})
        \leq
        f(w_t) - \eta \nabla f(w_t) \boldsymbol{\cdot} g_t + \tfrac12 \beta \eta^2 \norm{g_t}^2
        ,
    \end{align*}
    which implies
    \begin{align*}
        \nabla f(w_t) \boldsymbol{\cdot} g_t
        &\leq
        \frac{1}{\eta} \brk!{f(w_t) - f(w_{t+1})}
        + \tfrac12 \beta \eta \norm{g_t}^2
        .
    \end{align*}
    Summing the above over $t=1,\ldots,T$, we obtain (with $\Delta_1=f(w_1)-f^\star$)
    \begin{align*}
        \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} g_t
        &\leq
        \frac{\Delta_1}{\eta}
        + \frac{\beta \eta}{2} \sum_{t=1}^T \norm{g_t}^2
        .
    \end{align*}
    Observe that
    \begin{align*}
        \norm{g_t}^2
        &\leq
        2\norm{\nabla f(w_t)}^2 + 2\norm{g_t - \nabla f(w_t)}^2
        \leq
        2(1+\sigma_1^2)\norm{\nabla f(w_t)}^2 + 2\sigma_0^2
        .
    \end{align*}
    Hence,
    \begin{align*}
        \sum_{t=1}^T \norm{g_t}^2
        &\leq
        2 \sum_{t=1}^T (1+\sigma_1^2) \norm{\nabla f(w_t)}^2
        + 2 \sigma_0^2 T
        .
    \end{align*}
    By the definition of $\eta$ (\cref{eq:known-eta}),
    \begin{align}\label{eq:known-regret}
        \nonumber
        \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} g_t
        &\leq
        \frac{\Delta_1}{\eta}
        + \beta \eta \sum_{t=1}^T (1+\sigma_1^2) \norm{\nabla f(w_t)}^2
        + \beta \eta \sigma_0^2 T
        \\
        &\leq
        4 \beta \Delta_1 (1+\sigma_1^2 \log \tfrac{T}{\delta})
        + \frac{1}{4} \sum_{t=1}^T \norm{\nabla f(w_t)}^2
        + \sigma_0 \sqrt{T} \brk2{\beta \alpha + \frac{\Delta_1}{\alpha}}
        .
    \end{align}
    We continue with a martingale analysis to replace $\sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} g_t$ with $\sum_{t=1}^T \norm{\nabla f(w_t)}^2$.
    Consider the sequence of R.V.s $(Z_t)_{t=1}^T$ defined by $Z_t = \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-g_t)$.
    Then $(Z_t)_{t=1}^T$ is a martingale difference sequence, as
    \begin{align*}
        \condE{Z_t}{t-1}
        &=
        \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-\condE{g_t}{t-1})
        =
        0
        .
    \end{align*}
    By Cauchy-Schwarz inequality and the noise assumption,
    \begin{align*}
        \abs{Z_t}
        &\leq
        \norm{\nabla f(w_t)}\norm{\nabla f(w_t)-g_t}
        \leq
        \norm{\nabla f(w_t)} \sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_t)}^2}
        .
    \end{align*}
    Thus, we obtain from \cref{lem:sub_gaussian} with
    $$
    \abs{Z_t} \leq \norm{\nabla f(w_t)} \sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_t)}^2}$$
    and $\lambda=1\big/3(\sigma_0^2 + 2 \beta \sigma_1^2 F)$ that with probability at least $1-\delta$,
    \begin{align*}
        \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-g_t)
        &\leq
        \frac{\sum_{t=1}^T (\sigma_0^2+\sigma_1^2 \norm{\nabla f(w_t)}^2) \norm{\nabla f(w_t)}^2}{4 (\sigma_0^2 + 2 \beta \sigma_1^2 F)}
        + 3 (\sigma_0^2+2 \beta \sigma_1^2 F) \log(1/\delta)
        .
    \end{align*}
    From \cref{lem:smooth_norm_bound,lem:known-fbound-affine}, with probability at least $1-\delta$, $\norm{\nabla f(w_t)}^2 \leq 2 \beta F$.
    Hence, under a union bound with \cref{lem:known-fbound-affine}, with probability at least $1-2\delta$,
    \begin{align}\label{eq:known-basic-martingale}
        \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-g_t)
        &
        \leq
        \frac{1}{4} \sum_{t=1}^T \norm{\nabla f(w_t)}^2
        + 3 (\sigma_0^2+2 \beta \sigma_1^2 F) \log \tfrac{1}{\delta}
        .
    \end{align}
    Summing \cref{eq:known-basic-martingale,eq:known-regret} and rearranging,
    \begin{align*}
        \sum_{t=1}^T \norm{\nabla f(w_t)}^2
        &\leq
        8 \beta \Delta_1 (1+\sigma_1^2 \log \tfrac{T}{\delta})
        + 6 (\sigma_0^2+2 \beta \sigma_1^2 F) \log \tfrac{1}{\delta}
        + 2 \sigma_0 \sqrt{T} \brk2{\beta \alpha + \frac{\Delta_1}{\alpha}}
        .
    \end{align*}
    From \cref{lem:known-fbound-affine,eq:known-eta},
    \begin{align*}
        F
        &\leq
        2 \Delta_1
        + 2 \beta \alpha^2
        + \frac{3 \sigma_0^2}{4 \beta \sigma_1^2}
        .
    \end{align*}
    Thus,
    \begin{align*}
        \sum_{t=1}^T \norm{\nabla f(w_t)}^2
        &\leq
        8 \beta \Delta_1 (1+\sigma_1^2 \log \tfrac{T}{\delta})
        + 24 \beta \sigma_1^2 \brk{\Delta_1 + \beta \alpha^2} \log \tfrac{1}{\delta}
        + 15 \sigma_0^2 \log \tfrac{1}{\delta}
        + 2 \sigma_0 \sqrt{T} \brk2{\beta \alpha + \frac{\Delta_1}{\alpha}}
        \\
        &\leq
        8 \beta \Delta_1 (1+ 4 \sigma_1^2) \log \tfrac{T}{\delta}
        + 24 \sigma_1^2 \beta^2 \alpha^2 \log \tfrac{1}{\delta}
        + 15 \sigma_0^2 \log \tfrac{1}{\delta}
        + 2 \sigma_0 \sqrt{T} \brk2{\beta \alpha + \frac{\Delta_1}{\alpha}}
        .\qedhere
    \end{align*}
\end{proof}

\subsection{Proof of \texorpdfstring{\cref{lem:known-fbound-affine}}{Lemma}}
\begin{proof}
    Consider the sequence of RVs $\brk{Z_t}_{t=1}^T$ defined by $Z_t = \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-g_t)$.
    Then $\brk{Z_t}_{t=1}^T$ is a martingale difference sequence, as
    \begin{align*}
        \condE{Z_t}{t-1}
        &=
        \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-\condE{g_t}{t-1})
        =
        0
        .
    \end{align*}
    From Cauchy–Schwarz inequality and the noise model,
    $\abs{Z_t}
    \leq \norm{\nabla f(w_t)} \sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_t)}^2}$. Thus, we obtain from \cref{lem:sub_gaussian} (together with a union bound over $t$) that for any fixed $\lambda>0$, with probability at least $1-\delta$, for all $1 \leq t \leq T$,
    \begin{align*}
        \sum_{s=1}^t \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s) - g_s)
        &
        \leq
        \frac{3 \lambda}{4} \sum_{s=1}^t \norm{\nabla f(w_s)}^2 (\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_s)}^2)
        + \frac{\log \frac{T}{\delta}}{\lambda}
        .
    \end{align*}
    We will proceed by induction.
    The claim is immediate for $t=1$ and we move to prove it for $t+1$.
    Using smoothness,
    \begin{align*}
        f(w_{s+1})
        &\leq f(w_{s}) - \eta \nabla f(w_s) \boldsymbol{\cdot} g_s + \frac{\beta \eta^2}{2} \norm{g_s}^2.
    \end{align*}
    Summing the above over $s=1,\ldots,t$,
    \begin{align}\label{eq:known-f-bound-sum}
        f(w_{t+1}) - f(w_{1})
        &\leq
        - \eta \sum_{s=1}^{t} \nabla f(w_s) \boldsymbol{\cdot} g_s
        + \frac{\beta \eta^2}{2} \sum_{s=1}^{t} \norm{g_s}^2.
    \end{align}
    The second term of the right hand side can be bounded by the definitions of $\eta$ and the noise model,
    \begin{align*}
        \frac{\beta \eta^2}{2} \sum_{s=1}^{t} \norm{g_s}^2
        &\leq
        \beta \eta^2 \sum_{s=1}^{t} \brk!{\norm{g_s-\nabla f(w_s)}^2 + \norm{\nabla f(w_s)}^2}
        \tag{$\norm{a+b}^2 \leq 2 \norm{a}^2+2\norm{b}^2$}
        \\
        &\leq
        \beta \eta^2 \sum_{s=1}^{t} \brk!{(1+\sigma_1^2) \norm{\nabla f(w_s)}^2 + \sigma_0^2}
        \tag{$\norm{\nabla f(w_s)-g_s}^2 \leq \sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_s)}^2$}
        \\
        &\leq
        \beta \alpha^2
        + \frac{\eta}{4} \sum_{s=1}^{t} \norm{\nabla f(w_s)}^2
        \tag{\cref{eq:known-eta}}
        .
    \end{align*}
    We move to bound the first term of the right hand side.
    We split the sum to obtain a martingale term,
    \begin{align*}
        - \eta \sum_{s=1}^t \nabla f(w_s) \boldsymbol{\cdot} g_s
        &=
        - \eta \sum_{s=1}^t \norm{\nabla f(w_s)}^2
        +
        \eta \sum_{s=1}^t \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s) - g_s)
        .
    \end{align*}
    From \cref{lem:smooth_norm_bound} and our induction, $\norm{\nabla f(w_t)}^2 \leq 2 \beta (f(w_t)-f^\star) \leq 2 \beta F$.
    Returning to our martingale bound, we have
    \begin{align*}
        \sum_{s=1}^t \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s) - g_s)
        &\leq
        \frac{3 \lambda}{4} \sum_{s=1}^t \norm{\nabla f(w_s)}^2 \brk{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_s)}^2}
        + \frac{\log \frac{T}{\delta}}{\lambda}
        \\
        &\leq
        \frac{3 \lambda}{4} \brk{\sigma_0^2 + 2 \beta \sigma_1^2 F} \sum_{s=1}^t \norm{\nabla f(w_s)}^2 + \frac{\log \frac{T}{\delta}}{\lambda}
        .
    \end{align*}
    Setting $\lambda = 1 \Big/ \brk{\sigma_0^2 + 2 \beta \sigma_1^2 F}$,
    \begin{align*}
        \sum_{s=1}^t \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s) - g_s)
        &\leq
        \frac{3}{4} \sum_{s=1}^t \norm{\nabla f(w_s)}^2
        + \brk{\sigma_0^2 + 2 \beta \sigma_1^2 F}\log \tfrac{T}{\delta}
        .
    \end{align*}
    Returning to \cref{eq:known-f-bound-sum} and combining the inequalities,
    \begin{align*}
        f(w_{t+1})-f(w_1)
        &\leq
        \beta \alpha^2
        + \eta\brk{\sigma_0^2 + 2 \beta \sigma_1^2 F}\log \tfrac{T}{\delta}
        \\
        &\leq
        \beta \alpha^2
        + \eta \sigma_0^2 \log \tfrac{T}{\delta}
        + \frac{1}{2} F
        =
        F - (f(w_1)-f^\star)
        ,
    \end{align*}
    where the last equality follows from the definition of $\eta$. Adding $f(w_1)-f^\star$ to both sides concludes the proof.
\end{proof}
\subsection*{Acknowledgements}
This work has received support from the Israeli Science Foundation (ISF) grant no.~2549/19, and from the Len Blavatnik and the Blavatnik Family foundation.

\bibliographystyle{abbrvnat}

\section{Proofs of \texorpdfstring{\cref{sec:non-convex}}{Non-convex Section}}
\label{sec:non_convex_appendix}
\subsection{Proof of \texorpdfstring{\cref{lem:sum_eta_g_t}}{Lemma}}
\begin{proof}%
    Define $\tilde{G}_t=\sqrt{G_t^2-G_0^2}$.
    Observe that for any $t \geq 1$,
    \begin{align*}
        \norm{g_t}^2
        =
        \sum_{s=1}^t \norm{g_s}^2 - \sum_{s=1}^{t-1} \norm{g_s}^2
        =
        \tilde{G}_{t}^2 - \tilde{G}_{t-1}^2
        =
        \brk!{ \tilde{G}_t + \tilde{G}_{t-1} }
        \brk!{ \tilde{G}_t - \tilde{G}_{t-1} }
        .
    \end{align*}
    Since $\tilde{G}_t \geq \tilde{G}_{t-1}$, this implies
    \begin{align*}
        \frac{\norm{g_t}^2}{\tilde{G}_t}
        =
        \brk*{\frac{\tilde{G}_t + \tilde{G}_{t-1}}{\tilde{G}_t}}
        \brk*{ \tilde{G}_{t} - \tilde{G}_{t-1} }
        \leq
        2 \brk*{ \tilde{G}_{t} - \tilde{G}_{t-1} }
        .
    \end{align*}
    Summing this over $t=1,\ldots,T$ and telescoping terms, we obtain
    \begin{align*}
        \sum_{t=1}^T \frac{\norm{g_t}^2}{G_t}
        \leq
        \sum_{t=1}^T \frac{\norm{g_t}^2}{\tilde{G}_t}
        \leq
        2 \tilde{G}_T
        =
        2 \sqrt{\sum_{t=1}^T \norm{g_s}^2}
        ,
    \end{align*}
    where the first inequality follows by $G_t \geq \tilde{G}_t$.
    On the other hand, using the inequality $1-x \leq \log(1/x)$,
    \begin{align*}
        \sum_{t=1}^T \frac{\norm{g_t}^2}{G_t^2}
        &=
        \sum_{t=1}^T \brk*{ 1 - \frac{G_{t-1}^2}{G_{t}^2} }
        \leq
        \sum_{t=1}^T \log\frac{G_{t}^2}{G_{t-1}^2}
        =
        2 \log\frac{G_T}{G_0}
        .\qedhere
    \end{align*}
\end{proof}
\subsection{Proof of \texorpdfstring{\cref{lem:smooth_norm_bound}}{Lemma}}
\begin{proof}%
    Let $w^+ = w-\frac{1}{\beta}\nabla f(w)$. Then from smoothness,
    \begin{align*}
        f(w^+)-f(w)
        &\leq
        \nabla f(w) \boldsymbol{\cdot} \brk{w^+-w}
        +\frac{\beta}{2}\norm{w^+-w}^2
        =
        -\frac{1}{2\beta}\norm{\nabla f(w)}^2
        .
    \end{align*}
    Hence,
    \begin{align*}
        \norm{\nabla f(w)}^2
        &\leq
        2 \beta \brk{f(w)-f(w^+)}
        \leq
        2 \beta \brk{f(w)-f^\star}
        .\qedhere
    \end{align*}
\end{proof}
\subsection{Proof of \texorpdfstring{\cref{lem:sum_eta_g_t_2_affine}}{Lemma}}
\begin{proof}%
    Using \cref{lem:sum_eta_g_t},
    \begin{align*}
        \sum_{t=1}^T \frac{\norm{g_t}^2}{G_t^2}
        &\leq 2 \log \frac{G_T}{G_0}.
    \end{align*}
    Since $\norm{g_t}^2 \leq 2 \norm{\nabla f(w_t)}^2 + 2 \norm{g_t-\nabla f(w_t)}^2 \leq 2(1+\sigma_1^2) \norm{\nabla f(w_t)}^2 + 2\sigma_0^2$,
    \begin{align*}
        \sum_{t=1}^T \frac{\norm{g_t}^2}{G_t^2}
        &\leq \log \brk*{G_0^2 + 2 \sigma_0^2 T + 2 (1+\sigma_1^2)\sum_{t=1}^T \norm{\nabla f(w_t)}^2}+2 \log(1/G_0).
    \end{align*}
    From smoothness,
    \begin{align*}
        \norm{\nabla f(w_t)}
        &\leq \norm{\nabla f(w_{t-1})} + \norm{\nabla f(w_{t})-\nabla f(w_{t-1})}
        \leq \norm{\nabla f(w_{t-1})} + \beta \norm{w_t-w_{t-1}}
        \\
        &=
        \norm{\nabla f(w_{t-1})} + \beta \eta_t \norm{g_t}
        \leq \norm{\nabla f(w_{t-1})} + \eta \beta
        .
    \end{align*}
    Thus,
    \begin{align*}
        \norm{\nabla f(w_t)}
        &\leq
        \norm{\nabla f(w_1)}
        + \eta \beta (t-1)
        \leq
        \norm{\nabla f(w_1)}
        + \eta \beta t
        .
    \end{align*}
    summing from $t=1,\ldots,T$,
    \begin{align*}
        \sum_{t=1}^T \norm{\nabla f(w_t)}^2
        &\leq 2 \sum_{t=1}^T (\eta^2 \beta^2 t^2 + \norm{\nabla f(w_1)}^2)
        \tag{$\norm{x+y}^2 \leq 2 \norm{x}^2 + 2 \norm{y}^2$}
        \\
        &\leq 2 \eta^2 \beta^2 T^3 + 2 \norm{\nabla f(w_1)}^2 T \tag{$\sum_{t=1}^T t^2 \leq T^3$} \\
        &\leq 2 \eta^2 \beta^2 T^3 + 4 \beta (f(w_1)-f^\star)T \tag{from smoothness, \cref{lem:smooth_norm_bound}}
        .
    \end{align*}
    Hence, with $G_0^2 = \gamma^2$,
    \begin{align*}
        \sum_{t=1}^T \frac{\norm{g_t}^2}{G_t^2}
        &\leq \log \brk*{1+\frac{2 \sigma_0^2 T + 4 (1+\sigma_1^2) \eta^2 \beta^2 T^3
        + 8 (1+\sigma_1^2) \beta (f(w_1)-f^\star)T}{\gamma^2}}
        \leq C_1
        .
        \qedhere
    \end{align*}
\end{proof}

\subsection{Proof of \texorpdfstring{\cref{lem:decorrelated_difference}}{Lemma}}
In order to prove the bound we also need the following technical lemma (proof follows).
\begin{lemma}\label{lem:decorrelated_step_size_diff}
    For all $1 \leq s \leq T$,
    \begin{align*}
        \abs{\tilde{\eta}_s-\eta_s}
        &\leq
        \frac{2 \tilde{\eta}_s \sqrt{\sigma_0^2 +\sigma_1^2 \norm{\nabla f(w_s)}^2}}{G_s}
        .
    \end{align*}
\end{lemma}

\begin{proof}[Proof of \cref{lem:decorrelated_difference}]
By \cref{lem:decorrelated_step_size_diff},
\begin{align*}
    \abs{\tilde{\eta}_s-\eta_s}
    &\leq
    \frac{2 \tilde{\eta}_s \sqrt{\sigma_0^2 +\sigma_1^2 \norm{\nabla f(w_s)}^2}}{G_s}
    .
\end{align*}
Hence, by applying Cauchy-Schwarz inequality and using $ab \leq \frac{a^2}{2} + \frac{b^2}{2}$ with $a=\sqrt{\tilde{\eta}_s}\norm{\nabla f(w_s)}$ and $b=2\sqrt{\tilde{\eta}_s(\sigma_0^2+\sigma_1^2 \norm{\nabla f(w_s)}^2)}\frac{\norm{g_s}}{G_s}$,
\begin{align*}
    \abs{\tilde{\eta}_s-\eta_s} \nabla f(w_s) \boldsymbol{\cdot} g_s
    &\leq
    2 \tilde{\eta}_s \norm{\nabla f(w_s)} \sqrt{\sigma_0^2 +\sigma_1^2 \norm{\nabla f(w_s)}^2} \frac{\norm{g_s}}{G_s}
    \\
    &\leq
    \frac{1}{2} \tilde{\eta}_s \norm{\nabla f(w_s)}^2 + 2 \tilde{\eta}_s \brk*{\sigma_0^2 +\sigma_1^2 \norm{\nabla f(w_s)}^2} \frac{\norm{g_s}^2}{G_s^2}
    \\
    &\leq
    \frac{1}{2} \tilde{\eta}_s \norm{\nabla f(w_s)}^2+ 2 \eta \sqrt{\sigma_0^2 +\sigma_1^2 \norm{\nabla f(w_s)}^2} \frac{\norm{g_s}^2}{G_s^2}
    \tag{$\tilde{\eta_s} \leq \eta/\sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_s)}^2}$}
    .
\end{align*}
Using the smoothness of $f$, $\norm{\nabla f(w_s)}^2 \leq 2 \beta (f(w_s)-f^\star)\leq 2 \beta \bar{\Delta}_t$ (\cref{lem:smooth_norm_bound}). Thus, summing for $s=1,\ldots,t$,
\begin{align*}
    \sum_{s=1}^t \abs{\tilde{\eta}_s-\eta_s} \nabla f(w_s) \boldsymbol{\cdot} g_s
    &\leq
    \frac{1}{2} \sum_{s=1}^t \tilde{\eta}_s \norm{\nabla f(w_s)}^2
    + 2 \eta \sqrt{\sigma_0^2 +2 \beta \sigma_1^2 \bar{\Delta}_t} \sum_{s=1}^t \frac{\norm{g_s}^2}{G_s^2}
    .
\end{align*}
Again, using $ab \leq \frac{a^2}{2} + \frac{b^2}{2}$ with $a=\sqrt{\bar{\Delta}_t/2}$ and $b=4\eta \sigma_1\sqrt{\beta}\sum_{s=1}^t \frac{\norm{g_s}^2}{G_s^2}$,
\begin{align*}
    \sum_{s=1}^t \abs{\tilde{\eta}_s-\eta_s} \nabla f(w_s) \boldsymbol{\cdot} g_s
    &\leq
    \frac{1}{2} \sum_{s=1}^t \tilde{\eta}_s \norm{\nabla f(w_s)}^2
    + 2 \eta \sigma_0 \sum_{s=1}^t \frac{\norm{g_s}^2}{G_s^2}
    + 2 \eta \sqrt{2 \beta \sigma_1^2 \bar{\Delta}_t} \sum_{s=1}^t \frac{\norm{g_s}^2}{G_s^2}
    \\
    &\leq
    \frac{1}{2} \sum_{s=1}^t \tilde{\eta}_s \norm{\nabla f(w_s)}^2
    + 2 \eta \sigma_0 \sum_{s=1}^t \frac{\norm{g_s}^2}{G_s^2}
    + \frac{\bar{\Delta}_t}{4}
    + 8 \eta^2 \beta \sigma_1^2 \brk*{\sum_{s=1}^t \frac{\norm{g_s}^2}{G_s^2}}^2
    . \qedhere
\end{align*}
\end{proof}
\begin{proof}[Proof of \cref{lem:decorrelated_step_size_diff}]
    Using $\frac{1}{\sqrt{a}}-\frac{1}{\sqrt{b}}=\frac{b-a}{\sqrt{a}\sqrt{b}(\sqrt{a}+\sqrt{b})}$,
    \begin{align*}
        \tilde{\eta}_s-\eta_s
        &=
        \frac{\eta}{\sqrt{G_{s-1}^2 + \sigma_0^2 + (1+\sigma_1^2) \norm{\nabla f(w_s)}^2}}-\frac{\eta}{\sqrt{G_{s-1}^2 + \norm{g_s}^2}}
        \\
        &=
        \frac{\eta(\norm{g_s}^2-\sigma_0^2 -(1+\sigma_1^2) \norm{\nabla f(w_s)}^2)}{G_{s} \sqrt{G_{s-1}^2 + \sigma_0^2 + (1+\sigma_1^2) \norm{\nabla f(w_s)}^2}(G_{s} + \sqrt{G_{s-1}^2 + \sigma_0^2 + (1+\sigma_1^2) \norm{\nabla f(w_s)}^2})}
        \\
        &=
        \frac{\tilde{\eta}_s}{G_s} \cdot \frac{\norm{g_s}^2-\sigma_0^2 -(1+\sigma_1^2) \norm{\nabla f(w_s)}^2}{G_{s} + \sqrt{G_{s-1}^2 + \sigma_0^2 + (1+\sigma_1^2) \norm{\nabla f(w_s)}^2}}
        .
    \end{align*}
    Thus,
    \begin{align*}
        \abs{\tilde{\eta}_s-\eta_s}
        &=
        \frac{\tilde{\eta}_s}{G_s}\cdot\frac{\abs*{\norm{g_s}^2-\sigma_0^2 -(1+\sigma_1^2) \norm{\nabla f(w_s)}^2}}{G_s + \sqrt{G_{s-1}^2 + \sigma_0^2 + (1+\sigma_1^2) \norm{\nabla f(w_s)}^2}}
        \\
        &\leq
        \frac{\tilde{\eta}_s}{G_s}\cdot \frac{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_s)}^2 + \norm{g_s-\nabla f(w_s)}\norm{g_s+\nabla f(w_s)}}{G_s + \sqrt{G_{s-1}^2 + \sigma_0^2 + (1+\sigma_1^2) \norm{\nabla f(w_s)}^2}}
        .
    \end{align*}
    Since
    $\norm{g_s+\nabla f(w_s)} \leq \norm{g_s} + \norm{\nabla f(w_s)} \leq G_s + \sqrt{G_{s-1}^2 + \sigma_0^2 + (1+\sigma_1^2) \norm{\nabla f(w_s)}^2}$,
    \begin{align*}
        \abs{\tilde{\eta}_s-\eta_s}
        &\leq
        \frac{\tilde{\eta}_s}{G_s} \cdot \brk*{\sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_s)}^2} + \norm{g_s-\nabla f(w_s)}}
        \leq
        \frac{2 \tilde{\eta}_s \sqrt{\sigma_0^2 +\sigma_1^2 \norm{\nabla f(w_s)}^2}}{G_s}
        .
        \qedhere
    \end{align*}
\end{proof}

\subsection{Proof of \texorpdfstring{\cref{lem:basic_martingale_affine}}{Lemma}}
\begin{proof}[Proof of \cref{lem:basic_martingale_affine}]
Consider the sequence of R.V.s $(Z_t)_{t=1}^T$ defined by $Z_t = \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-g_t)$.
Then $(Z_t)_{t=1}^T$ is a martingale difference sequence, as
\begin{align*}
    \condE{Z_t}{t-1}
    &=
    \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-\condE{g_t}{t-1})
    =
    0
    .
\end{align*}
By Cauchy-Schwarz inequality and the noise assumption,
\begin{align*}
    \abs{Z_t}
    &\leq
    \norm{\nabla f(w_t)}\norm{\nabla f(w_t)-g_t}
    \leq
    \norm{\nabla f(w_t)} \sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_t)}^2}
    .
\end{align*}
Thus, we obtain from \cref{lem:sub_gaussian} with
$$
\abs{Z_t} \leq \norm{\nabla f(w_t)} \sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_t)}^2}$$
and $\lambda=1\big/3(\sigma_0^2 + 2 \beta \sigma_1^2 F)$ that with probability at least $1-\delta$,
\begin{align*}
    \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-g_t)
    &\leq
    \frac{\sum_{t=1}^T (\sigma_0^2+\sigma_1^2 \norm{\nabla f(w_t)}^2) \norm{\nabla f(w_t)}^2}{4 (\sigma_0^2 + 2 \beta \sigma_1^2 F)}
    + 3 (\sigma_0^2+2 \beta \sigma_1^2 F) \log(1/\delta)
    .
\end{align*}
From \cref{lem:smooth_norm_bound,lem:fbound_affine}, with probability at least $1-\delta$, $\norm{\nabla f(w_t)}^2 \leq 2 \beta F$.
Hence, under a union bound with \cref{lem:fbound_affine}, with probability at least $1-2\delta$,
\begin{align*}
    \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-g_t)
    &
    \leq
    \frac{1}{4} \sum_{t=1}^T \norm{\nabla f(w_t)}^2
    + 3 (\sigma_0^2+2 \beta \sigma_1^2 F) \log(1/\delta)
    .\qedhere
\end{align*}
\end{proof}
\section{Analysis in the general non-convex case}
\label{sec:non-convex}

First, we focus on the general smooth case without convexity assumptions.
Our main result of this section is the following convergence bound for \adasgd.
\begin{theorem}\label{thm:non_convex_convergence}
Assume that $f$ is $\beta$-smooth.
Then for \adasgd, for any $\delta \in (0,1/3)$,
it holds with probability at least $1-\delta$ that
\begin{enumerate}[label=(\roman*)]
    \item%
    $\begin{aligned}[t]
        \frac{1}{T} \sum_{t=1}^T \norm{\nabla f(w_t)}^2
        &
        \leq
        O\brk*{ \tfrac{2F}{\eta} + \eta\beta }
        \frac{\sigma_0}{\sqrt{T}}
        + O\brk*{ \tfrac{\gamma}{\eta} \Delta_1 +(1+\sigma_1^2 \log\tfrac{1}{\delta})\brk!{\tfrac{2F}{\eta} + \eta\beta}^2 + \sigma_0^2 \log\tfrac{1}{\delta} } \frac{1}{T}
        ;
    \end{aligned}$
    \item%
    and $f(w_t)-f^\star \leq F$ for all $1 \leq t \leq T+1$,
\end{enumerate}
where $\Delta_1=f(w_1)-f^\star$ and
\begin{align}
\label{eq:fbound}
F
&=
\begin{aligned}[t]
    2 \Delta_1
    + \brk!{3 \log \tfrac{T}{\delta} + 4 C_1} \eta \sigma_0
    + \brk!{9 \log^2 \tfrac{T}{\delta} + 16 C_1^2} \eta^2 \beta \sigma_1^2
    + \eta^2 \beta C_1
    ,
\end{aligned}
\\
\label{eq:C1}
C_1
&=
\begin{aligned}[t]
    \log
    &
    \bigg(
        1
        \!+\!
        \frac{2 \sigma_0^2 T + 8 (1+\sigma_1^2) (\eta^2 \beta^2 T^3 + \beta \Delta_1 T)}{\gamma^2}
    \bigg)
    .\!
\end{aligned}
\end{align}
\end{theorem}
In addition to the features of this convergence result already highlighted in the introduction, we also note the high probability guarantee it provides that the maximal optimality gap $f(w_t)-f^\star$ remains $\smash{\wt{O}}(1)$ along the trajectory of \adasgd.
Comparing the result to the high probability analysis when the parameters are known (\cref{sec:known-non-convex}), the main difference is the bound of $f(w_t)-f^\star$ which is $O(\Delta_1)$ when the parameters are known (for large enough $T$), while the bound of \cref{thm:non_convex_convergence} has additional terms that depend on $\sigma_0$ and $\sigma_1$.
We remark that the prior results of~\citet{Faw2022ThePO} in this setting include a somewhat worse high-order noise term, of order $\ifrac{\sigma_1(\sigma_0^2+\sigma_1^{12})}{\sqrt{T}}$ in the noisy regime (for $\sigma_1 = \omega(1)$).

\subsection{Proof of~\texorpdfstring{\cref{thm:non_convex_convergence}}{Theorem}}

Henceforth we use the following notations: let $\Delta_t \triangleq f(w_t)-f^\star$ and $\bar{\Delta}_t \triangleq \max_{s \leq t} f(w_s)-f^\star$, for all $t$.
We also denote $G_t \triangleq \brk{\gamma^2 + \sum_{s=1}^t \norm{g_s}^2}^{1/2}$, thus, $\eta_t = \eta / G_t$.
As $\eta_t$ and $g_t$ are correlated random variables (indeed, $\eta_t$ depends on the norm of $g_t$), we define the following ``decorrelated stepsizes,''
\begin{align*}
    \forall ~ 1 \leq t \leq T :
    \tilde{\eta}_t \triangleq \frac{\eta}{\sqrt{G_{t-1}^2+ (1+\sigma_1^2)\norm{\nabla f(w_t)}^2 + \sigma_0^2}}
    .
\end{align*}
We begin by stating several lemmas that we use to prove \cref{thm:non_convex_convergence} (proofs are given later in the section).
The first lemma states a ``non-convex regret bound.''
\begin{lemma} \label{lem:non-convex-regret}
    If $f$ is $\beta$-smooth then for \adasgd we have
    \begin{align*}
        \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} g_t
        \leq
        \frac{\gamma \Delta_1}{\eta}
        + \brk*{\frac{2 \bar{\Delta}_T}{\eta}+\eta \beta} \sqrt{\sum_{t=1}^T \norm{g_t}^2}
        .
    \end{align*}
\end{lemma}
Note that the $\bar{\Delta}_T$ term in the right hand-side of the bound is a random variable.
In the noiseless setting with known smoothness and properly tuned stepsize, we know that $\bar{\Delta}_T = \Delta_1$ since $\Delta_{t+1} \leq \Delta_{t}$ for all $t$ and proving convergence would be much simplified.
The next lemma bounds $\bar{\Delta}_T = \max_{s \leq t} f(w_s)-f^\star$ with high probability in the noisy case.
\begin{lemma}\label{lem:fbound_affine}
With probability at least $1-\delta$, it holds that
$
    \bar{\Delta}_{T+1}
    \leq
    F
$
(with $F$ defined in \cref{eq:fbound}).
\end{lemma}
In order to translate the regret analysis of \cref{lem:non-convex-regret} to a high probability bound we need a bound on the sum of differences between $\norm{\nabla f(w_t)}^2$ and $g_t \boldsymbol{\cdot} \nabla f(w_t)$ which we obtain in the following lemma (see \cref{sec:non_convex_appendix} for a proof).
\begin{lemma}\label{lem:basic_martingale_affine}
    with probability at least $1 - 2 \delta$,
    it holds that
    \begin{align*}
        \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-g_t)
        &
        \leq
        \frac{1}{4} \sum_{t=1}^T \norm{\nabla f(w_t)}^2
        + 3 (\sigma_0^2+2 \beta \sigma_1^2 F) \log \tfrac{1}{\delta}
        ,
    \end{align*}
    with $F$ given in \cref{eq:fbound}.
\end{lemma}
Using the lemmas above we can now prove our main result.
\begin{proof}[Proof of \cref{thm:non_convex_convergence}]
All arguments below follow given the high probability events of \cref{lem:fbound_affine,lem:basic_martingale_affine}.
Hence, using a union bound, they follow with probability at least $1-3 \delta$.
We denote $C_2 = \ifrac{2 F}{\eta}+\eta \beta$ to simplify notation.
From \cref{lem:non-convex-regret,lem:fbound_affine},
\begin{align*}
    \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} g_t
    \leq
    \frac{\gamma \Delta_1}{\eta}
    + C_2 \sqrt{\sum_{t=1}^T \norm{g_t}^2}
    .
\end{align*}
Observe that
\begin{align*}
    \norm{g_t}^2
    &\leq
    2\norm{\nabla f(w_t)}^2 + 2\norm{g_t - \nabla f(w_t)}^2
    \leq
    2(1+\sigma_1^2)\norm{\nabla f(w_t)}^2 + 2\sigma_0^2
    .
\end{align*}
Hence, using $\sqrt{a+b} \leq \sqrt{a} + \sqrt{b}$ for $a,b \geq 0$,
\begin{align*}
    \sqrt{\sum_{t=1}^T \norm{g_t}^2}
    &\leq
    \sqrt{2\sum_{t=1}^T (1+\sigma_1^2) \norm{\nabla f(w_t)}^2} + \sigma_0 \sqrt{2 T}
    .
\end{align*}
By the elementary fact $ab \leq a^2/2 + b^2/2$ with
$a=\sqrt{\tfrac12 \sum_{t=1}^T \norm{\nabla f(w_t)}^2}$ and $b=2 C_2 \sqrt{1+\sigma_1^2}$,
\begin{align*}
    C_2 \sqrt{2(1+\sigma_1^2)\sum_{t=1}^T \norm{\nabla f(w_t)}^2}
    &\leq
    \frac{1}{4} \sum_{t=1}^T \norm{\nabla f(w_t)}^2
    + 2(1+\sigma_1^2) C_2^2
    .
\end{align*}
Hence,
\begin{align*}
    \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} g_t
    &\leq
    \frac{\gamma \Delta_1}{\eta}
    + 2(1+\sigma_1^2) C_2^2
    + C_2 \sigma_0 \sqrt{2 T}
    + \frac{1}{4} \sum_{t=1}^T \norm{\nabla f(w_t)}^2
    .
\end{align*}
Summing the above with \cref{lem:basic_martingale_affine} and rearranging,
\begin{align*}
    \sum_{t=1}^T \norm{\nabla f(w_t)}^2
    &\leq
    \frac{2 \gamma \Delta_1}{\eta}
    + 4(1+\sigma_1^2) C_2^2
    + 6(\sigma_0^2 + 2\beta \sigma_1^2 F) \log \tfrac{1}{\delta}
    + C_2 \sigma_0 \sqrt{8 T}
    .
\end{align*}
Substituting $C_2$ and dividing through by $T$ gives the theorem (due to the $O$ notation in the theorem we state a $1-\delta$ probability guarantee instead of $1-3\delta$).
\end{proof}
\subsection{Proof of \texorpdfstring{\cref{lem:non-convex-regret}}{Lemma}}
Before proving \cref{lem:non-convex-regret} we state a standard lemma (for completeness, we give a proof in \cref{sec:non_convex_appendix}).
\begin{lemma} \label{lem:sum_eta_g_t}
    Let $g_1,\ldots,g_T \in {\mathbb{R}}^d$ be an arbitrary sequence of vectors, and let $G_0 > 0$. For all $t \geq 1$, define
    \begin{align*}
        G_t
        =
        \sqrt{G_0^2 + \sum_{s=1}^t \norm{g_s}^2}
        .
    \end{align*}
    Then
    \begin{align*}
        \sum_{t=1}^T \frac{\norm{g_t}^2}{G_t}
        \leq
        2 \sqrt{\sum_{t=1}^T \norm{g_t}^2}
        ,
        \qquad
        \text{and}
        \qquad
        \sum_{t=1}^T \frac{\norm{g_t}^2}{G_t^2}
        \leq
        2\log\frac{G_T}{G_0}
        .
    \end{align*}
\end{lemma}
\begin{proof}[Proof of \cref{lem:non-convex-regret}]
    Using $\beta$-smoothness, we have for all $t \geq 1$ that
    \begin{align*}
        f(w_{t+1})
        \leq
        f(w_t) - \eta_t \nabla f(w_t) \boldsymbol{\cdot} g_t + \tfrac12 \beta \eta_t^2 \norm{g_t}^2
        ,
    \end{align*}
    which implies (with $\Delta_t=f(w_t)-f^\star$)
    \begin{align*}
        \nabla f(w_t) \boldsymbol{\cdot} g_t
        &\leq
        \frac{1}{\eta_t} \brk!{\Delta_t - \Delta_{t+1}}
        + \tfrac12 \beta \eta_t \norm{g_t}^2
        .
    \end{align*}
    Summing the above over $t=1,\ldots,T$, we obtain
    \begin{align*}
        \sum_{t=1}^T \! \nabla f(w_t) \boldsymbol{\cdot} g_t
        &\leq
        \frac{\Delta_1}{\eta_0}
        + \!\sum_{t=1}^T \brk3{ \frac{1}{\eta_t}
        - \frac{1}{\eta_{t-1}} } \Delta_t
        + \frac{\beta}{2} \!\sum_{t=1}^T\! \eta_t \norm{g_t}^2
        \!.
    \end{align*}
    Now, since $\eta_t \leq \eta_{t-1}$, we have for all $t$ that
    \begin{align*}
        \frac{1}{\eta_t} - \frac{1}{\eta_{t-1}}
        &
        \leq
        \eta_t \brk3{ \frac{1}{\eta_t^2} - \frac{1}{\eta_{t-1}^2} }
        =
        \frac{\eta_t}{\eta^2} \brk3{\gamma^2 + \sum_{s=1}^t \norm{g_s}^2 - \gamma^2 - \sum_{s=1}^{t-1} \norm{g_s}^2}
        =
        \frac{\eta_t}{\eta^2} \norm{g_t}^2
        ,
    \end{align*}
    thus
    \begin{align*}
        \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} g_t
        &\leq
        \frac{\Delta_1}{\eta_0}
        + \sum_{t=1}^T \eta_t \norm{g_t}^2 \brk*{\frac{\Delta_t}{\eta^2}+\frac{\beta}{2}}
        .
    \end{align*}
    Bounding $\Delta_t \leq \bar{\Delta}_T$ and applying \cref{lem:sum_eta_g_t} on the second sum
    yields
    \begin{align*}
        \sum_{t=1}^T \nabla f(w_t) \boldsymbol{\cdot} g_t
        &\leq
        \frac{\Delta_1}{\eta_0}
        + \brk*{\frac{2 \bar{\Delta}_T}{\eta}+\eta \beta} \sqrt{\sum_{t=1}^T \norm{g_t}^2}
        .
    \end{align*}
    Plugging in the expressions for $\eta_0$ yields the lemma.
\end{proof}

\subsection{Proof of \texorpdfstring{\cref{lem:fbound_affine}}{Lemma}}
For our high probability analysis we use the following martingale concentration inequality~\citep[][Lemma~1]{li2020high}.
\begin{lemma}
	\label{lem:sub_gaussian}
	Assume that $Z_1, Z_2, ..., Z_T$ is a martingale difference sequence with respect to $\xi_1, \xi_2, ..., \xi_T$ and $\mathbb{E}_t \left[\exp(Z_t^2/\sigma_t^2)\right] \leq \exp(1)$ for all $t$, where $\sigma_1,\ldots,\sigma_T$ is a sequence of random variables such that $\sigma_t$ is measurable with respect to $\xi_1, \xi_2, \dots, \xi_{t-1}$.
	Then, for any fixed $\lambda > 0$ and $\delta \in (0,1)$, with probability at least $1-\delta$, it holds that
	\[
	\sum_{t=1}^T Z_t \leq \frac{3}{4} \lambda \sum_{t=1}^T \sigma_t^2 + \frac{1}{\lambda} \log \frac{1}{\delta}~.
	\]
\end{lemma}
Note that it suffices to pick $\sigma_t \geq \abs{Z_t}$ with probability $1$ for the sub-Gaussian condition to hold.
We also need the following technical lemmas (see \cref{sec:non_convex_appendix} for proofs).
\begin{lemma}\label{lem:sum_eta_g_t_2_affine}
    For \adasgd we have:
    \begin{align*}
        \sum_{t=1}^T \frac{\norm{g_t}^2}{G_t^2}
        \leq
        C_1
        ,
    \end{align*}
    with $C_1$ given in \cref{eq:C1}.
\end{lemma}
\begin{lemma}\label{lem:decorrelated_difference}
    For all $t \geq 1$,
    \begin{align*}
        \sum_{s=1}^t \abs{\tilde{\eta}_s-\eta_s} (\nabla f(w_s) \boldsymbol{\cdot} g_s)
        &
        \leq
        \frac{\bar{\Delta}_t}{4}
        + \frac{1}{2} \sum_{s=1}^t  \tilde{\eta}_s \norm{\nabla f(w_s)}^2
        + 2 \eta \sigma_0 \sum_{s=1}^t \frac{\norm{g_s}^2}{G_s^2}
        + 8 \eta^2 \beta \sigma_1^2 \brk*{\sum_{s=1}^t \frac{\norm{g_s}^2}{G_s^2}}^2.
    \end{align*}
    \end{lemma}

Finally, we require a standard property of smooth functions.

\begin{lemma}\label{lem:smooth_norm_bound}
Let $f:{\mathbb{R}}^d \mapsto {\mathbb{R}}$ be $\beta$-smooth with minimum $f^\star$. Then
$
    \norm{\nabla f(w)}^2
    \leq
    2\beta\brk{f(w)-f^\star}
$
for all $w \in {\mathbb{R}}^d$.
\end{lemma}

\begin{proof}[Proof of \cref{lem:fbound_affine}]
First we handle the randomness of $\brk{g_t}_{t=1}^T$.
Consider the sequence of RVs $\brk{Z_t}_{t=1}^T$ defined by $Z_t = \tilde{\eta}_t \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-g_t)$.
Then $\brk{Z_t}_{t=1}^T$ is a martingale difference sequence, as
\begin{align*}
    \condE{Z_t}{t-1}
    &=
    \tilde{\eta}_t \nabla f(w_t) \boldsymbol{\cdot} (\nabla f(w_t)-\condE{g_t}{t-1})
    =
    0
    .
\end{align*}
From Cauchy–Schwarz inequality and the noise model,
$\abs{Z_t}
\leq \tilde{\eta}_t \norm{\nabla f(w_t)} \sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_t)}^2}$. Thus, we obtain from \cref{lem:sub_gaussian} (together with a union bound over $t$) that for any fixed $\lambda>0$, with probability at least $1-\delta$, for all $1 \leq t \leq T$,
\begin{align*}
    \sum_{s=1}^t \tilde{\eta}_s \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s) - g_s)
    &
    \leq
    \frac{3 \lambda}{4} \sum_{s=1}^t \tilde{\eta}_s^2 \norm{\nabla f(w_s)}^2 (\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_s)}^2)
    + \frac{\log \frac{T}{\delta}}{\lambda}
    .
\end{align*}
We will proceed by induction.
The claim is immediate for $t=1$ and we move to prove it for $t+1$.
Using smoothness,
\begin{align*}
    f(w_{s+1})
    &\leq f(w_{s}) - \eta_s \nabla f(w_s) \boldsymbol{\cdot} g_s + \frac{\beta}{2} \eta_s^2 \norm{g_s}^2.
\end{align*}
Summing the above over $s=1,\ldots,t$,
\begin{align}\label{eq:f_t_f_1}
    f(w_{t+1}) - f(w_{1})
    &\leq
    -\sum_{s=1}^{t} \eta_s \nabla f(w_s) \boldsymbol{\cdot} g_s + \frac{\beta}{2} \sum_{s=1}^{t} \eta_s^2 \norm{g_s}^2.
\end{align}
We move to bound the first term of the right hand side.
\begin{align*}
    - \sum_{s=1}^t \eta_s \nabla f(w_s) \boldsymbol{\cdot} g_s
    &
    =
    - \sum_{s=1}^t \tilde{\eta}_s \norm{\nabla f(w_s)}^2
    +
    \underbrace{
        \sum_{s=1}^t \tilde{\eta}_s \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s) - g_s)
    }_{(*)}
    +
    \underbrace{
        \sum_{s=1}^t (\tilde{\eta_s}-\eta_s) \nabla f(w_s) \boldsymbol{\cdot} g_s
    }_{(**)}
    .
\end{align*}
We will bound $(*)$ and $(**)$. From \cref{lem:smooth_norm_bound} and our induction, $\norm{\nabla f(w_t)}^2 \leq 2 \beta (f(w_t)-f^\star) \leq 2 \beta F$.
Returning to our martingale bound, with the fact that $\tilde{\eta}_s \leq \eta\Big/\sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_s)}^2}$, we have
\begin{align*}
    \sum_{s=1}^t \tilde{\eta}_s \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s) - g_s)
    &\leq
    \frac{3 \lambda \eta}{4} \sum_{s=1}^t \tilde{\eta}_s \norm{\nabla f(w_s)}^2 \sqrt{\sigma_0^2 + \sigma_1^2 \norm{\nabla f(w_s)}^2}
    + \frac{\log \frac{T}{\delta}}{\lambda}
    \\
    &\leq
    \frac{3 \lambda \eta}{4} \sqrt{\sigma_0^2 + 2 \beta \sigma_1^2 F} \sum_{s=1}^t \tilde{\eta}_s \norm{\nabla f(w_s)}^2 + \frac{\log \frac{T}{\delta}}{\lambda}
    .
\end{align*}
Setting $\lambda = 2 \Big/ 3 \eta \sqrt{\sigma_0^2 + 2 \beta \sigma_1^2 F}$,
\begin{align*}
    \sum_{s=1}^t \tilde{\eta}_s \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s) - g_s)
    &\leq
    \frac{1}{2} \! \sum_{s=1}^t \tilde{\eta}_s \norm{\nabla f(w_s)}^2
    + \frac{3 \eta}{2}\sqrt{\sigma_0^2 + 2 \beta \sigma_1^2 F}\log \tfrac{T}{\delta}
    \\
    &\leq
    \frac{1}{2} \! \sum_{s=1}^t \tilde{\eta}_s \norm{\nabla f(w_s)}^2
    + \frac{3 \eta}{2} \sigma_0 \log \tfrac{T}{\delta}
    + \frac{3 \eta}{2}\sqrt{2 \beta \sigma_1^2 F}\log \tfrac{T}{\delta}
    .
\end{align*}
Using the elementary fact that $ab \leq a^2/2+ b^2/2$ with $a=\sqrt{F/2}$ and $b=3 \eta \sqrt{\beta}\sigma_1\log \tfrac{T}{\delta}$,
\begin{align*}
    \sum_{s=1}^t \tilde{\eta}_s \nabla f(w_s) \boldsymbol{\cdot} (\nabla f(w_s) - g_s)
    &
    \leq
    \frac12 \sum_{s=1}^t \tilde{\eta}_s \norm{\nabla f(w_s)}^2
    + \frac32 \eta \sigma_0 \log \tfrac{T}{\delta}
    + \frac14F
    + \frac92 \eta^2 \beta \sigma_1^2 \log^2 \tfrac{T}{\delta}
    .
\end{align*}
In order to bound $(**)$, using \cref{lem:decorrelated_difference} with $\bar{\Delta}_t \leq F$ by induction,
\begin{align*}
    (**)
    &
    \leq
    \frac{F}{4}
    + \frac{1}{2}\sum_{s=1}^t  \tilde{\eta}_s \norm{\nabla f(w_s)}^2
    + 2 \eta \sigma_0 \sum_{s=1}^t \frac{\norm{g_s}^2}{G_s^2}
    + 8 \eta^2 \beta \sigma_1^2 \brk*{\sum_{s=1}^t \frac{\norm{g_s}^2}{G_s^2}}^2
    \\
    &\leq
    \frac{F}{4}
    + \frac{1}{2}\sum_{s=1}^t  \tilde{\eta}_s \norm{\nabla f(w_s)}^2
    + 2 \eta \sigma_0 C_1
    + 8 \eta^2 \beta \sigma_1^2 C_1^2
    ,
\end{align*}
where that last inequality follows from \cref{lem:sum_eta_g_t_2_affine}.
Hence, applying the bounds $(*) \text{ and } (**)$,
\begin{align*}
    - \sum_{s=1}^t \eta_s \nabla f(w_s) \boldsymbol{\cdot} g_s
    &
    \leq
    \tfrac12 F
    + \brk!{\tfrac32 \log \tfrac{T}{\delta} + 2 C_1} \eta \sigma_0
    + \brk!{\tfrac92 \log^2 \tfrac{T}{\delta} + 8 C_1^2} \eta^2 \beta \sigma_1^2
    .
\end{align*}
Thus, returning to \cref{eq:f_t_f_1} and applying \cref{lem:sum_eta_g_t_2_affine},
\begin{align*}
    f(w_{t+1})-f(w_1)
    &
    \leq
    \tfrac12 F
    + \brk!{\tfrac32 \log \tfrac{T}{\delta} + 2 C_1} \eta \sigma_0
    + \brk!{\tfrac92 \log^2 \tfrac{T}{\delta} + 8 C_1^2} \eta^2 \beta \sigma_1^2
    + \frac{\beta}{2} \sum_{s=1}^{t} \eta_s^2 \norm{g_s}^2
    \\
    &
    \leq
    \tfrac12 F
    + \brk!{\tfrac32 \log \tfrac{T}{\delta} + 2 C_1} \eta \sigma_0
    + \brk!{\tfrac92 \log^2 \tfrac{T}{\delta} + 8 C_1^2} \eta^2 \beta \sigma_1^2
    + \tfrac12 \eta^2 \beta C_1
    \\
    &= F - (f(w_1) - f^\star)
    .
\end{align*}
We conclude the induction by adding $f(w_1)-f^\star$ to both sides, obtaining $\bar{\Delta}_{T+1} = \max_{1 \leq t \leq T+1} \Delta_t \leq F$.
\end{proof}

\section{Preliminaries}

We consider unconstrained smooth optimization in $d$-dimensional Euclidean space ${\mathbb{R}}^d$ with the $\ell_2$ norm, denoted $\norm{\cdot}$.
A function $f$ is said to be $\beta$-smooth if for all $x,y \in {\mathbb{R}}^d$ it holds that $\norm{\nabla f(x)-\nabla f(y)} \leq \beta \norm{x-y}$.
This implies, in particular, that for all $x,y \in {\mathbb{R}}^d$, the following bound is true: $f(y) \leq f(x) + \nabla f(x) \boldsymbol{\cdot} (y-x)+\frac12 \beta\norm{y-x}^2$.
We additionally assume that $f$ is lower bounded over ${\mathbb{R}}^d$, and admits a minimum $f^\star = \min_{w \in {\mathbb{R}}^d} f(w)$.
We assume a standard stochastic first-order model, in which we are provided with an oracle $g(w)$ where, given a point $w \in {\mathbb{R}}^d$ returns a random vector $g(w) \in {\mathbb{R}}^d$ such that $\mathbb{E}[g(w) \mid w] = \nabla f(w)$.

\paragraph*{Affine noise model.}
We assume a general noise model, referred to in the literature as the ``affine variance noise'' (e.g., \citet{Faw2022ThePO} and references therein).
In this model, the gradient oracle $g$ satisfies
\begin{align*}
    \forall ~ w \in {\mathbb{R}}^d :
    \quad
    \norm{g(w)-\nabla f(w)}^2 \leq \sigma_0^2 + \sigma_1^2 \norm{\nabla f(w)}^2
\end{align*}
(deterministically), where $\sigma_0,\sigma_1 \geq 0$ are noise parameters.
The affine noise model permits the noise magnitude to grow as the norm of the true gradient increases (without the affine term, i.e., assuming $\sigma_1=0$, the gradients are effectively ``noiseless'' for gradients of large magnitude).
We refer to \citet{Faw2022ThePO} for an additional discussion of this model.

We remark that, for simplicity, we will assume throughout a deterministic bound on the noise magnitude, as stated above; in \cref{sec:sub_gaussian_reduction}, through a generic reduction, we extend our arguments to apply in a more general sub-Gaussian noise model with affine variance.

\paragraph*{Adaptive stepsize Stochastic Gradient Descent (\adasgd).}
Starting from an arbitrary point $w_1 \in {\mathbb{R}}^d$, for parameters $\eta,\gamma > 0$, the update rule of \adasgd is
\begin{align*} %
    w_{t+1}
    &=
    w_{t}-\eta_t g_t,
\end{align*}
where $g_t = g(w_t)$ is a stochastic gradient at $w_t$ and the step-size $\eta_t$ is given by
\begin{align}\label{eq:adaptive-eta}
    \forall ~ 1 \leq t \leq T :
    \qquad
    \eta_t = \frac{\eta}{\sqrt{\gamma^2 + \sum_{s=1}^t \norm{g_s}^2}}
    .
\end{align}
\paragraph*{Additional notation.}
Throughout, we use $\condE{\cdot}{t-1} \triangleq \mathbb{E}\brk[s]{\cdot \mid g_1,\ldots,g_{t-1}}$ to denote the conditional expectation with respect to the stochastic gradients queried up to (not including) step $t$.  We use $\log$ to denote the base $2$ logarithm.

\subsection{Additional related work}

\paragraph*{SGD for non-convex smooth optimization.}
The first to analyze SGD in the smooth non-convex setting and obtain tight convergence bounds were \citet{ghadimi2013stochastic}. They established that a \emph{properly tuned} $T$-steps SGD obtain a rate of $O\brk{\ifrac{1}{T}+\ifrac{\sigma}{\sqrt{T}}}$ assuming uniformly bounded variance of $\sigma^2$. A tight lower bound is given by \citet{arjevani2022lower}.
\citet[Theorem 4.8]{bottou2018optimization} extended the convergence result for the affine variance noise model.

\paragraph*{Adaptive methods for non-convex smooth optimization.}
An extensive line of work tackles convergence guarantees in expectation.%
\footnote{We include here results with polynomial dependence in the probability margin as they can be derived using Markov's inequality.}
\citet{ward2019adagrad} gave a rate of $\smash{\wt{O}}\brk{\ifrac{1}{\sqrt{T}}}$ for $T$-steps Adaptive SGD, assuming uniformly bounded gradients.
In parallel, \citet{li2019convergence} obtained a rate of $\smash{\wt{O}}\brk{\ifrac{1}{\sqrt{T}}}$ without assuming uniformly bounded gradients, assuming instead knowledge of the smoothness parameter and using a step-size that does not depend on the current stochastic gradient.
\citet{li2019convergence} also obtain a convergence guarantee for the smooth and convex case, again assuming knowledge of the smoothness.
A recent work by \citet{Faw2022ThePO} assumed affine variance noise and achieved the desired $\smash{\wt{O}}\brk{\ifrac{1}{\sqrt{T}}}$ in the noisy regime without assuming uniformly bounded gradients.

\paragraph*{High probability convergence bounds.}
Given knowledge of the smoothness and sub-Gaussian noise bounds, \citet{ghadimi2013stochastic} established convergence with high probability for non-adaptive SGD with properly tuned step-size.
Few high probability results of adaptive methods exist.
\citet{zhou2018convergence} and \citet{li2020high} established a high probability convergence result for AdaGrad and delayed AdaGrad respectively, assuming sub-Gaussian noise and knowledge of the smoothness to tune the step-size. \citet{kavis2021high} obtained a high probability convergence for Adaptive SGD, assuming uniformly bounded gradients and either sub-Gaussian noise or uniformly bounded stochastic gradients.
Assuming heavy-tailed noise, \citet{Cutkosky2021HighprobabilityBF} obtained convergence with high probability for normalized SGD with clipping and momentum with a tuned step-size according to a moment bound of the noise.

\paragraph{Unconstrained stochastic convex optimization.}

Also related to our work are \citet{harvey2019tight} and \citet{carmon2022making} who focused on the (non-smooth) unconstrained and convex setting.
\citet{harvey2019tight} established convergence with high probability of SGD using a generalized freedman's inequality.
\citet{carmon2022making} devised a stochastic parameter-free method which obtain high probability convergence guarantee and their method also obtain optimal rate for \emph{noiseless} convex and smooth objective.
Both papers used a careful martingale analysis in order to bound the distance of the iterates from a minimizer, a desired property for unconstrained stochastic optimization which we also establish for \adasgd.

\paragraph{Parameter-free methods for online optimization.}

The closely related literature on online optimization includes a plethora of adaptive methods.
Recent work in this context is mostly concerned with adaptivity to the norm of the comparator~\citep[e.g.,][]{orabona2016coin,cutkosky2018black,mhammedi2020lipschitz} and to the scale of the objective~\citep[e.g.,][]{orabona2018scale,mhammedi2020lipschitz}.

redpajama_set_name: RedPajamaArXiv