SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbb {P} (\tau = \tau_ {2} \leq T) \leq \mathbb {P} (\tau_ {2} \leq T) \\ = \mathbb {P} \left(\bigcup_ {1 \leq t \leq T} \| \nabla f (x _ {t}) - \nabla f (x _ {t}, \xi_ {t}) \| > \sigma\right) \\ \leq \sum_ {1 \leq t \leq T} \mathbb {P} \left(\| \nabla f (x _ {t}) - \nabla f (x _ {t}, \xi_ {t}) \| > \sigma\right) \\ \leq \sum_ {1 <   t <   T} \mathbb {E} \left[ \mathbb {P} _ {t - 1} \left(\| \nabla f (x _ {t}) - \nabla f (x _ {t}, \xi_ {t}) \| > \sigma\right) \right] \\ \leq \sum_ {1 <   t <   T} \mathbb {E} \left[ 2 e ^ {- \frac {\sigma^ {2}}{2 R ^ {2}}} \right] \\ = 2 T e ^ {- \frac {\sigma^ {2}}{2 R ^ {2}}} \\ \leq \delta / 2, \\ \end{array}
$$
where the fourth inequality uses Assumption 4; and the last inequality uses  $\sigma = R\sqrt{2\log(4T / \delta)}$
Next, if  $\tau = \tau_{1} \leq T$ , by definition, we have  $f(x_{\tau}) - f^{*} > F$ , or equivalently,
$$
\frac {8 G}{\eta} (f (x _ {\tau}) - f ^ {*}) > \frac {8 G F}{\eta} = \frac {8 G ^ {3}}{3 \eta L} =: I _ {2}.
$$
On the other hand, since for any  $t < \tau$ , under the new definition of  $\tau$ , we still have
$$
f (x _ {t}) - f ^ {*} \leq F, \quad \| f (x _ {t}) \| \leq G, \quad \| \nabla f (x _ {t}) - \nabla f (x _ {t}, \xi_ {t}) \| \leq \sigma .
$$
Then we know that Lemma E.8 still holds because all of its requirements are still satisfied, i.e., there exists some event  $\mathcal{E}$  with  $\mathbb{P}(\mathcal{E})\leq \delta /2$ , such that under its complement  $\mathcal{E}^c$
$$
\begin{array}{l} \sum_ {t = 1} ^ {\tau - 1} \left\| \nabla f (x _ {t}) \right\| ^ {2} + \frac {8 G}{\eta} (f (x _ {\tau}) - f ^ {*}) \leq \frac {8 G}{\eta \lambda} \left(\Delta_ {1} \lambda + 8 \sigma^ {2} \left(\frac {\eta}{\beta} + \eta \beta T\right) + 2 0 \sigma^ {2} \eta \sqrt {(1 / \beta^ {2} + T) \iota}\right) \\ =: I _ {1}. \\ \end{array}
$$
By Lemma E.5, we know  $I_1 \leq I_2$ , which suggests that  $\mathcal{E}^c \cap \{\tau = \tau_1 \leq T\} = \emptyset$ , i.e.,  $\{\tau = \tau_1 \leq T\} \subset \mathcal{E}$ . Then we can show
$$
\mathbb {P} (\mathcal {E} \cup \{\tau \leq T \}) \leq \mathbb {P} (\mathcal {E}) + \mathbb {P} (\tau = \tau_ {2} \leq T) \leq \delta .
$$
$$
\mathbb {P} \left(\mathcal {E} ^ {c} \cap \{\tau = T + 1 \}\right) \geq 1 - \mathbb {P} \left(\mathcal {E} \cup \{\tau \leq T \}\right) \geq 1 - \delta ,
$$
and under the event  $\mathcal{E}^c\cap \{\tau = T + 1\}$  , we have  $\tau = T + 1$  and
$$
\frac {1}{T} \sum_ {t = 1} ^ {t} \| \nabla f (x _ {t}) \| ^ {2} \leq I _ {1} / T \leq \epsilon^ {2},
$$
F Convergence analysis of VRAdam
In this section, we provide detailed convergence analysis of VRAdam and prove Theorem 6.2. To do that, we first provide some technical definitions<sup>4</sup>. Denote
$$
\epsilon_ {t} := m _ {t} - \nabla f (x _ {t})
$$
4Note that the same symbol for Adam and VRAdam may have different meanings.
as the deviation of the momentum from the actual gradient. From the update rule in Algorithm 2, we can write
$$
\epsilon_ {t} = (1 - \beta) \epsilon_ {t - 1} + W _ {t}, \tag {25}
$$
$$
W _ {t} := \nabla f (x _ {t}, \xi_ {t}) - \nabla f (x _ {t}) - (1 - \beta) (\nabla f (x _ {t - 1}, \xi_ {t}) - \nabla f (x _ {t - 1})).
$$
Let  $G$  be the constant defined in Theorem 6.2 and denote  $F := \frac{G^2}{3(3L_0 + 4L_\rho G^\rho)}$ . We define the following stopping times as discussed in Section 6.1.
$$
\tau_ {1} := \min  \left\{t \mid f \left(x _ {t}\right) - f ^ {*} > F \right\} \wedge (T + 1),
$$
$$
\tau_ {2} := \min  \left\{t \mid \| \epsilon_ {t} \| > G \right\} \wedge (T + 1), \tag {26}
$$
$$
\tau := \min  \left\{\tau_ {1}, \tau_ {2} \right\}.
$$
It is straight forward to verify that  $\tau_{1},\tau_{2},\tau$  are all stopping times. Then if  $t <   \tau$  , we have
$$
f (x _ {t}) - f ^ {*} \leq F, \quad \| \nabla f (x _ {t}) \| \leq G, \quad \| \epsilon_ {t} \| \leq G.
$$
Then we can also bound the update  $\| x_{t + 1} - x_t\| \leq \eta D$  where  $D = 2G / \lambda$  if  $t < \tau$  (see Lemma F.3 for the details). Finally, we consider the same definition of  $r$  and  $L$  as those for Adam. Specifically,
$$
r := \min  \left\{\frac {1}{5 L _ {\rho} G ^ {\rho - 1}}, \frac {1}{5 \left(L _ {0} ^ {\rho - 1} L _ {\rho}\right) ^ {1 / \rho}} \right\}, \quad L := 3 L _ {0} + 4 L _ {\rho} G ^ {\rho}. \tag {27}
$$
F.1 Useful lemmas
We first list several useful lemmas in this section without proofs. Their proofs are deferred later in Appendix F.3.
To start with, we provide a lemma on the local smoothness of each component function  $f(\cdot, \xi)$  when the gradient of the objective function  $f$  is bounded.
Lemma F.1. For any constant  $G \geq \sigma$  and two points  $x \in \operatorname{dom}(f), y \in \mathbb{R}^d$  such that  $\| \nabla f(x) \| \leq G$  and  $\| y - x \| \leq r/2$ , we have  $y \in \operatorname{dom}(f)$  and
$$
\left\| \nabla f (y) - \nabla f (x) \right\| \leq L \| y - x \|,
$$
$$
\| \nabla f (y, \xi) - \nabla f (x, \xi) \| \leq 4 L \| y - x \|, \forall \xi ,
$$
$$
f (y) \leq f (x) + \left\langle \nabla f (x), y - x \right\rangle + \frac {1}{2} L \| y - x \| ^ {2},
$$
where  $r$  and  $L$  are defined in (27).
With the new definition of stopping time  $\tau$  in (26), all the quantities in Algorithm 2 are well bounded before  $\tau$ . In particular, the following lemma holds.
Lemma F.2. If  $t < \tau$ , we have
$$
\| \nabla f (x _ {t}) \| \leq G, \quad \| \nabla f (x _ {t}, \xi_ {t}) \| \leq G + \sigma , \quad \| m _ {t} \| \leq 2 G,
$$
Next, we provide the following lemma which bounds the update at each step before  $\tau$ .
Lemma F.3. if  $t < \tau$ ,  $\|x_{t+1} - x_t\| \leq \eta D$  where  $D = 2G / \lambda$ .
The following lemma bounds  $\| W_t\|$  when  $t\leq \tau$
Lemma F.4. If  $t \leq \tau$ ,  $G \geq 2\sigma$ , and  $\eta \leq \frac{r}{2D}$
$$
\left\| W _ {t} \right\| \leq \beta \sigma + \frac {5 \eta L}{\lambda} \left(\| \nabla f (x _ {t - 1}) \| + \| \epsilon_ {t - 1} \|\right).
$$
Finally, we present some inequalities regarding the parameter choices, which will simplify the calculations later.
Lemma F.5. Under the parameter choices in Theorem 6.2, we have
$$
\frac {2 \Delta_ {1}}{F} \leq \frac {\delta}{4}, \quad \frac {\lambda \Delta_ {1} \beta}{\eta G ^ {2}} \leq \frac {\delta}{4}, \quad \eta \beta T \leq \frac {\lambda \Delta_ {1}}{8 \sigma^ {2}}, \quad \eta \leq \frac {\lambda^ {3 / 2}}{4 0 L} \sqrt {\frac {\beta}{G}}.
$$
F.2 Proof of Theorem 6.2
Before proving the theorem, we will need to present several important lemmas. First, note that the descent lemma still holds for VRAdam.
Lemma F.6. If  $t < \tau$ , choosing  $G \geq \sigma + \lambda$  and  $\eta \leq \min \left\{\frac{r}{2D}, \frac{\lambda}{6L}\right\}$ , we have
Proof of Lemma F.6. The proof is essentially the same as that of Lemma E.6.
Lemma F.7. Choose  $G \geq \max \{2\sigma, 2\lambda\}$ ,  $S_1 \geq \frac{1}{2\beta^2T}$ , and  $\eta \leq \min \left\{\frac{r}{2D}, \frac{\lambda^{3/2}}{40L} \sqrt{\frac{\beta}{G}}\right\}$ . We have
$$
\mathbb {E} \left[ \sum_ {t = 1} ^ {\tau - 1} \frac {\beta}{2} \| \epsilon_ {t} \| ^ {2} - \frac {\lambda \beta}{1 6 G} \| \nabla f (x _ {t}) \| ^ {2} \right] \leq 4 \sigma^ {2} \beta^ {2} T - \mathbb {E} [ \| \epsilon_ {\tau} \| ^ {2} ].
$$
Proof of Lemma F.7. By Lemma F.4, we have
$$
\begin{array}{l} \left\| W _ {t} \right\| ^ {2} \leq 2 \sigma^ {2} \beta^ {2} + \frac {1 0 0 \eta^ {2} L ^ {2}}{\lambda^ {2}} \left(\| \nabla f (x _ {t - 1}) \| ^ {2} + \| \epsilon_ {t - 1} \| ^ {2}\right) \\ \leq 2 \sigma^ {2} \beta^ {2} + \frac {\lambda \beta}{1 6 G} \left(\| \nabla f (x _ {t - 1}) \| ^ {2} + \| \epsilon_ {t - 1} \| ^ {2}\right), \\ \end{array}
$$
where in the second inequality we choose  $\eta \leq \frac{\lambda^{3/2}}{40L} \sqrt{\frac{\beta}{G}}$ . Therefore, noting that  $\frac{\lambda\beta}{16G} \leq \beta / 2$ , by (25), we have
$$
\begin{array}{l} \left\| \epsilon_ {t} \right\| ^ {2} = (1 - \beta) ^ {2} \left\| \epsilon_ {t - 1} \right\| ^ {2} + \left\| W _ {t} \right\| ^ {2} + (1 - \beta) \left\langle \epsilon_ {t - 1}, W _ {t} \right\rangle \\ \leq (1 - \beta / 2) \| \epsilon_ {t - 1} \| ^ {2} + \frac {\lambda \beta}{1 6 G} \| \nabla f (x _ {t - 1}) \| ^ {2} + 2 \sigma^ {2} \beta^ {2} + (1 - \beta) \Big \langle \epsilon_ {t - 1}, W _ {t} \Big \rangle . \\ \end{array}
$$
Taking a summation over  $2 \leq t \leq \tau$  and re-arranging the terms, we get
$$
\sum_ {t = 1} ^ {\tau - 1} \frac {\beta}{2} \| \epsilon_ {t} \| ^ {2} - \frac {\lambda \beta}{1 6 G} \| \nabla f (x _ {t}) \| ^ {2} \leq \| \epsilon_ {1} \| ^ {2} - \| \epsilon_ {\tau} \| ^ {2} + 2 \sigma^ {2} \beta^ {2} (\tau - 1) + (1 - \beta) \sum_ {t = 2} ^ {\tau} \left\langle \epsilon_ {t - 1}, W _ {t} \right\rangle .
$$
Taking expectations on both sides, noting that
$$
\mathbb {E} \left[ \sum_ {t = 2} ^ {\tau} \left\langle \epsilon_ {t - 1}, W _ {t} \right\rangle \right] = 0
$$
by the Optional Stopping Theorem (Lemma C.2), we have
$$
\mathbb {E} \left[ \sum_ {t = 1} ^ {\tau - 1} \frac {\beta}{2} \| \epsilon_ {t} \| ^ {2} - \frac {\lambda \beta}{1 6 G} \| \nabla f (x _ {t}) \| ^ {2} \right] \leq 2 \sigma^ {2} \beta^ {2} T + \mathbb {E} [ \| \epsilon_ {1} \| ^ {2} ] - \mathbb {E} [ \| \epsilon_ {\tau} \| ^ {2} ] \leq 4 \sigma^ {2} \beta^ {2} T - \mathbb {E} [ \| \epsilon_ {\tau} \| ^ {2} ],
$$
where in the second inequality we choose  $S_{1}\geq \frac{1}{2\beta^{2}T}$  which implies  $\mathbb{E}[\| \epsilon_1\|^2 ]\leq \sigma^2 /S_1\leq 2\sigma^2\beta^2 T$
Lemma F.8. Under the parameter choices in Theorem 6.2, we have
$$
\mathbb {E} \left[ \sum_ {t = 1} ^ {\tau - 1} \| \nabla f (x _ {t}) \| ^ {2} \right] \leq \frac {1 6 G \Delta_ {1}}{\eta}, \quad \mathbb {E} [ f (x _ {\tau}) - f ^ {*} ] \leq 2 \Delta_ {1}, \quad \mathbb {E} [ \| \epsilon_ {\tau} \| ^ {2} ] \leq \frac {\lambda \Delta_ {1} \beta}{\eta}.
$$
Proof of Lemma F.8. First note that according to Lemma F.5, it is straight forward to verify that the parameter choices in Theorem 6.2 satisfy the requirements in Lemma F.6 and Lemma F.7. Then by Lemma F.6, if  $t < \tau$ ,
Taking a summation over  $1 \leq t < \tau$ , re-arranging terms, multiplying both sides by  $\frac{8G}{\eta}$ , and taking an expectation, we get
$$
\mathbb {E} \left[ \sum_ {t = 1} ^ {\tau - 1} 2 \| \nabla f (x _ {t}) \| ^ {2} - \frac {8 G}{\lambda} \| \epsilon_ {t} \| ^ {2} \right] \leq \frac {8 G}{\eta} \mathbb {E} [ f (x _ {1}) - f (x _ {\tau}) ] \leq \frac {8 G}{\eta} (\Delta_ {1} - \mathbb {E} [ f (x _ {\tau}) - f ^ {*} ]). \tag {28}
$$
By Lemma F.7, we have
$$
\mathbb {E} \left[ \sum_ {t = 1} ^ {\tau - 1} \frac {8 G}{\lambda} \| \epsilon_ {t} \| ^ {2} - \| \nabla f (x _ {t}) \| ^ {2} \right] \leq \frac {6 4 G \sigma^ {2} \beta T}{\lambda} - \frac {1 6 G}{\lambda \beta} \mathbb {E} [ \| \epsilon_ {\tau} \| ^ {2} ] \leq \frac {8 G \Delta_ {1}}{\eta} - \frac {1 6 G}{\lambda \beta} \mathbb {E} [ \| \epsilon_ {\tau} \| ^ {2} ], \tag {29}
$$
where the last inequality is due to Lemma F.5. Then (28) + (29) gives
$$
\mathbb {E} \left[ \sum_ {t = 1} ^ {\tau - 1} \| \nabla f (x _ {t}) \| ^ {2} \right] + \frac {8 G}{\eta} \mathbb {E} [ f (x _ {\tau}) - f ^ {*} ] + \frac {1 6 G}{\lambda \beta} \mathbb {E} [ \| \epsilon_ {\tau} \| ^ {2} ] \leq \frac {1 6 G \Delta_ {1}}{\eta},
$$
With all the above lemmas, we are ready to prove the theorem.
Proof of Theorem 6.2. First note that according to Lemma F.5, it is straight forward to verify that the parameter choices in Theorem 6.2 satisfy the requirements in all the lemmas for VRAdam.
Then, first note that if  $\tau = \tau_{1} \leq T$ , we know  $f(x_{\tau}) - f^{*} > F$  by the definition of  $\tau$ . Therefore,
$$
\mathbb {P} (\tau = \tau_ {1} \leq T) \leq \mathbb {P} (f (x _ {\tau}) - f ^ {*} > F) \leq \frac {\mathbb {E} [ f (x _ {\tau}) - f ^ {*} ]}{F} \leq \frac {2 \Delta_ {1}}{F} \leq \frac {\delta}{4},
$$
where the second inequality uses Markov's inequality; the third inequality is by Lemma F.8; and the last inequality is due to Lemma F.5.
Similarly, if  $\tau_{2} = \tau \leq T$ , we know  $\| \epsilon_{\tau}\| > G$ . We have
$$
\mathbb {P} (\tau_ {2} = \tau \leq T) \leq \mathbb {P} (\| \epsilon_ {\tau} \| > G) = \mathbb {P} (\| \epsilon_ {\tau} \| ^ {2} > G ^ {2}) \leq \frac {\mathbb {E} [ \| \epsilon_ {\tau} \| ^ {2} ]}{G ^ {2}} \leq \frac {\lambda \Delta_ {1} \beta}{\eta G ^ {2}} \leq \frac {\delta}{4},
$$
where the second inequality uses Markov's equilibriat; the third inequality is by Lemma F.8; and the last inequality is due to Lemma F.5. where the last inequality is due to Lemma F.5. Therefore,
$$
\mathbb {P} (\tau \leq T) \leq \mathbb {P} (\tau_ {1} = \tau \leq T) + \mathbb {P} (\tau_ {2} = \tau \leq T) \leq \frac {\delta}{2}.
$$
Also, note that by Lemma F.8
$$
\begin{array}{l} \frac {1 6 G \Delta_ {1}}{\eta} \geq \mathbb {E} \left[ \sum_ {t = 1} ^ {\tau - 1} \| \nabla f (x _ {t}) \| ^ {2} \right] \\ \geq \mathbb {P} (\tau = T + 1) \mathbb {E} \left[ \sum_ {t = 1} ^ {T} \| \nabla f (x _ {t}) \| ^ {2} \middle | \tau = T + 1 \right] \\ \geq \frac {1}{2} \mathbb {E} \left[ \sum_ {t = 1} ^ {T} \| \nabla f (x _ {t}) \| ^ {2} \middle | \tau = T + 1 \right], \\ \end{array}
$$
where the last inequality is due to  $\mathbb{P}(\tau = T + 1) = 1 - \mathbb{P}(\tau \leq T) \geq 1 - \delta / 2 \geq 1 / 2$ . Then we can get
$$
\mathbb {E} \left[ \frac {1}{T} \sum_ {t = 1} ^ {T} \| \nabla f (x _ {t}) \| ^ {2} \Bigg | \tau = T + 1 \right] \leq \frac {3 2 G \Delta_ {1}}{\eta T} \leq \frac {\delta \epsilon^ {2}}{2}.
$$
Let  $\mathcal{F} \coloneqq \left\{\frac{1}{T}\sum_{t=1}^{T}\|\nabla f(x_t)\|^2 > \epsilon^2\right\}$  be the event of not converging to stationary points. By Markov's inequality, we have
$$
\mathbb {P} (\mathcal {F} | \tau = T + 1) \leq \frac {\delta}{2}.
$$
$$
\mathbb {P} (\mathcal {F} \cup \{\tau \leq T \}) \leq \mathbb {P} (\tau \leq T) + \mathbb {P} (\mathcal {F} | \tau = T + 1) \leq \delta ,
$$
i.e., with probability at least  $1 - \delta$ , we have both  $\tau = T + 1$  and  $\frac{1}{T}\sum_{t=1}^{T}\|\nabla f(x_t)\|^2 \leq \epsilon^2$ .
F.3 Proofs of lemmas in Appendix F.1
Proof of Lemma F.1. This lemma is a direct corollary of Corollary 5.2. Note that by Assumption 6, we have  $\| \nabla f(x,\xi)\| \leq G + \sigma \leq 2G$ . Hence, when computing the locality size and smoothness constant for the component function  $f(\cdot ,\xi)$ , we need to replace the constant  $G$  in Corollary 5.2 with  $2G$ , that is why we get a smaller locality size of  $r / 2$  and a larger smoothness constant of  $4L$ .
Proof of Lemma F.2. The bound on  $\| m_t\|$  is by the definition of  $\tau$  in (26). All other quantities for VRAdam are defined in the same way as those in Adam (Algorithm 1), so they have the same upper bounds as in Lemma E.2.
Proof of Lemma F.3.
$$
\left\| x _ {t + 1} - x _ {t} \right\| \leq \eta \left\| m _ {t} \right\| / \lambda \leq 2 \eta G / \lambda = \eta D,
$$
where the first inequality uses the update rule in Algorithm 2 and  $h_t \preceq \eta / \lambda$  by Lemma F.2; the second inequality is again due to Lemma F.2.
Proof of Lemma F.4. By the definition of  $W_{t}$ , it is easy to verify that
$$
W _ {t} = \beta (\nabla f (x _ {t}, \xi_ {t}) - \nabla f (x _ {t})) + (1 - \beta) \delta_ {t},
$$
$$
\delta_ {t} = \nabla f (x _ {t}, \xi_ {t}) - \nabla f (x _ {t - 1}, \xi_ {t}) - \nabla f (x _ {t}) + \nabla f (x _ {t - 1}).
$$
Then we can bound
$$
\begin{array}{l} \left\| \delta_ {t} \right\| \leq \left\| \nabla f (x _ {t}, \xi_ {t}) - \nabla f (x _ {t - 1}, \xi_ {t}) \right\| + \left\| \nabla f (x _ {t}) - \nabla f (x _ {t - 1}) \right\| \\ \leq 5 L \left\| x _ {t} - x _ {t - 1} \right\| \\ \leq \frac {5 \eta L}{\lambda} \left(\| \nabla f (x _ {t - 1}) \| + \| \epsilon_ {t - 1} \|\right), \\ \end{array}
$$
where the second inequality uses Lemma F.1; and the last inequality is due to  $\| x_{t} - x_{t - 1}\| \leq \eta \| m_{t - 1}\| /\lambda \leq \eta \left(\| \nabla f(x_{t - 1})\| +\| \epsilon_{t - 1}\|\right) / \lambda$ . Then, we have
$$
\| W _ {t} \| \leq \beta \sigma + \frac {5 \eta L}{\lambda} \left(\| \nabla f (x _ {t - 1}) \| + \| \epsilon_ {t - 1} \|\right).
$$
Proof of Lemma F.5. These inequalities can be obtained by direct calculations.