SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \langle \mathcal {T} (\gamma) _ {w ^ {a d v}} ^ {w} \operatorname {g r a d} \mathcal {L} _ {t} (w ^ {a d v}), \operatorname {g r a d} \mathcal {L} (w) \rangle_ {w} = \underbrace {\langle \mathcal {T} (\gamma) _ {w ^ {a d v}} ^ {w} \operatorname {g r a d} \mathcal {L} _ {t} (w ^ {a d v}) - \mathcal {T} (\zeta) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} _ {t} (w ^ {*}) , \operatorname {g r a d} \mathcal {L} (w) \rangle_ {w}} _ {T _ {1}} \\ + \underbrace {\langle \mathcal {T} (\zeta) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} _ {t} (w ^ {*}) , \operatorname {g r a d} \mathcal {L} (w) \rangle_ {w}} _ {T _ {2}} \\ \end{array}
$$
We will bound two terms,  $T_{1}$  and  $T_{2}$ , separately. Regarding the term  $T_{1}$ , we derive
$$
\begin{array}{l} - T _ {1} = - \left\langle \mathcal {T} (\gamma) _ {w ^ {a d v}} ^ {w} \operatorname {g r a d} \mathcal {L} _ {t} \left(w ^ {a d v}\right) - \mathcal {T} (\zeta) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} _ {t} \left(w ^ {*}\right), \operatorname {g r a d} \mathcal {L} (w) \right\rangle_ {w} \\ \leq \frac {1}{2} \left\| \mathcal {T} (\gamma) _ {w ^ {a d v}} ^ {w} \operatorname {g r a d} \mathcal {L} _ {t} \left(w ^ {a d v}\right) - \mathcal {T} (\zeta) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} _ {t} \left(w ^ {*}\right) \right\| _ {w} ^ {2} + \frac {1}{2} \left\| \operatorname {g r a d} \mathcal {L} (w) \right\| _ {w} ^ {2} \\ \leq \left\| \mathcal {T} (\gamma) _ {w ^ {a d v}} ^ {w} \operatorname {g r a d} \mathcal {L} _ {t} \left(w ^ {a d v}\right) - \operatorname {g r a d} \mathcal {L} _ {t} (w) \right\| _ {w} ^ {2} + \left\| \mathcal {T} (\zeta) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} _ {t} \left(w ^ {*}\right) - \operatorname {g r a d} \mathcal {L} _ {t} (w) \right\| _ {w} ^ {2} + \frac {1}{2} \left\| \operatorname {g r a d} \mathcal {L} (w) \right\| _ {w} ^ {2} \\ \leq L _ {\mathrm {R}} ^ {2} \| \rho \operatorname {g r a d} \mathcal {L} _ {t} (w) \| _ {w} ^ {2} + L _ {\mathrm {R}} ^ {2} \| \rho \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2} + \frac {1}{2} \| \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2} \\ \leq \rho^ {2} L _ {\mathrm {R}} ^ {2} \left(2 \| \operatorname {g r a d} \mathcal {L} _ {t} (w) - \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2} + 2 \| \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2}\right) + \left(\frac {1}{2} + \rho^ {2} L _ {\mathrm {R}} ^ {2}\right) \| \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2} \\ \leq \frac {2 \rho^ {2} L _ {\mathrm {R}} ^ {2} \sigma^ {2}}{b} + \left(\frac {1}{2} + 3 \rho^ {2} L _ {\mathrm {R}} ^ {2}\right) \| \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2} \\ \end{array}
$$
From the above inequality, we could finally bound the term  $T_{1}$  as
$$
T _ {1} \geq - \frac {2 \rho^ {2} L _ {\mathrm {R}} ^ {2} \sigma^ {2}}{b} - \left(\frac {1}{2} + 3 \rho^ {2} L _ {\mathrm {R}} ^ {2}\right) \| \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2}
$$
Regarding the term  $T_{2}$ , we just use the lemma as
$$
T _ {2} = \left\langle \mathcal {T} (\zeta) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} _ {t} \left(w ^ {*}\right), \operatorname {g r a d} \mathcal {L} (w) \right\rangle_ {w} \geq \left(1 - \rho L _ {\mathrm {R}}\right) \| \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2}
$$
Hence, we arrive at
According to Algorithm 1, we follow the notation as
$$
\operatorname {g r a d} \mathcal {L} _ {t} (w) = \frac {1}{b} \sum_ {i \in I _ {t}} \operatorname {g r a d} \ell_ {i} (w)
$$
$$
w _ {t} ^ {a d v} = \mathrm {R} _ {w _ {t}} \left(\rho \operatorname {g r a d} \mathcal {L} _ {t} (w _ {t})\right)
$$
We assume the stochastic  $m$ -SAM where the same batch is used for both inner and outer updates.
Lemma 3 (Descent inequality). Under the assumptions in Theorem 1, we have
$$
\mathbb {E} \big [ \mathcal {L} (w _ {t + 1}) \big ] \leq \mathbb {E} \big [ \mathcal {L} (w _ {t}) \big ] - \frac {3 \alpha}{8} \mathbb {E} \big [ \| \operatorname {g r a d} \mathcal {L} (w _ {t}) \| _ {w _ {t}} ^ {2} \big ] + \frac {\alpha^ {2} L _ {S} \sigma^ {2}}{b} + \frac {2 \alpha^ {2} L _ {S} ^ {3} \rho^ {2} \sigma^ {2}}{b} + \frac {2 \alpha \rho^ {3} L _ {\mathrm {R}} ^ {2} \sigma^ {2}}{b}
$$
Proof. Using the condition (C-4), we have
$$
\begin{array}{l} \mathcal {L} (w _ {t + 1}) = \mathcal {L} \left(\mathrm {R} _ {w _ {t}} \Big (- \alpha \mathcal {T} (\gamma) _ {w _ {t} ^ {a d v}} ^ {w _ {t}} \operatorname {g r a d} \mathcal {L} _ {t} (w _ {t} ^ {a d v}) \Big)\right) \\ \leq \mathcal {L} (w _ {t}) - \alpha \left\langle \operatorname {g r a d} \mathcal {L} (w _ {t}), \mathcal {T} (\gamma) _ {w _ {t} ^ {a d v}} ^ {w _ {t}} \operatorname {g r a d} \mathcal {L} _ {t} (w _ {t} ^ {a d v}) \right\rangle_ {w _ {t}} + \frac {\alpha^ {2} L _ {S}}{2} \left\| \mathcal {T} (\gamma) _ {w _ {t} ^ {a d v}} ^ {w _ {t}} \operatorname {g r a d} \mathcal {L} _ {t} (w _ {t} ^ {a d v}) \right\| _ {w _ {t}} ^ {2} \\ \end{array}
$$
For the last term in RHS, we can bound as
$$
\begin{array}{l} \left\| \mathcal {T} (\gamma) _ {w _ {t} ^ {a d v}} ^ {w _ {t}} \operatorname {g r a d} \mathcal {L} _ {t} (w _ {t} ^ {a d v}) \right\| _ {w _ {t}} ^ {2} = - \| \operatorname {g r a d} \mathcal {L} (w _ {t}) \| _ {w _ {t}} ^ {2} + \left\| \mathcal {T} (\gamma) _ {w _ {t} ^ {a d v}} ^ {w _ {t}} \operatorname {g r a d} \mathcal {L} _ {t} (w _ {t} ^ {a d v}) - \operatorname {g r a d} \mathcal {L} (w _ {t}) \right\| _ {w _ {t}} ^ {2} \\ + 2 \left\langle \mathcal {T} (\gamma) _ {w _ {t} ^ {a d v}} ^ {w _ {t}} \operatorname {g r a d} \mathcal {L} _ {t} (w _ {t} ^ {a d v}), \operatorname {g r a d} \mathcal {L} (w _ {t}) \right\rangle_ {w _ {t}} \\ \end{array}
$$
Again, we have
Taking the expectation on both sides, we have
$$
\begin{array}{l} \mathbb {E} \left[ \mathcal {L} \left(w _ {t + 1}\right) \right] \leq \mathbb {E} \left[ \mathcal {L} \left(w _ {t}\right) \right] - \frac {\alpha^ {2} L _ {S} \left(1 - 4 L _ {S} ^ {2} \rho^ {2}\right)}{2} \mathbb {E} \left[ \| \operatorname {g r a d} \mathcal {L} \left(w _ {t}\right) \| _ {w _ {t}} ^ {2} \right] + \frac {\alpha^ {2} L _ {S} \left(1 + 2 L _ {S} ^ {2} \rho^ {2}\right) \sigma^ {2}}{b} \\ - \alpha (1 - \alpha L _ {S}) \mathbb {E} \left[ \left\langle \mathcal {T} (\gamma) _ {w _ {t} ^ {a d v}} ^ {w _ {t}} \operatorname {g r a d} \mathcal {L} _ {t} \left(w _ {t} ^ {a d v}\right), \operatorname {g r a d} \mathcal {L} \left(w _ {t}\right) \right\rangle_ {w _ {t}} \right] \\ \mathbb {E} \left[ \mathcal {L} \left(w _ {t}\right) \right] - \frac {\alpha^ {2} L _ {S} \left(1 - 4 L _ {S} ^ {2} \rho^ {2}\right)}{2} \mathbb {E} \left[ \| \operatorname {g r a d} \mathcal {L} \left(w _ {t}\right) \| _ {w _ {t}} ^ {2} \right] + \frac {\alpha^ {2} L _ {S} \left(1 + 2 L _ {S} ^ {2} \rho^ {2}\right) \sigma^ {2}}{b} \\ - \alpha (1 - \alpha L _ {S}) \left[ \left(\frac {1}{2} - \rho L _ {\mathrm {R}} - 3 \rho^ {2} L _ {\mathrm {R}} ^ {2}\right) \mathbb {E} \left[ \| \operatorname {g r a d} \mathcal {L} (w _ {t}) \| _ {w _ {t}} ^ {2} \right] - \frac {2 \rho^ {2} L _ {\mathrm {R}} ^ {2} \sigma^ {2}}{b} \right] \\ \end{array}
$$
For sufficiently large number of total iteration  $T$ , the condition  $\rho \leq \frac{1}{4\widetilde{L}}$  is easily satisfied where  $\widetilde{L} = \max \{L_{\mathrm{R}}, L_{S}\}$  (defined in Theorem 1). Hence, we obtain
$$
\begin{array}{l} \frac {3 \alpha}{8} \mathbb {E} \left[ \| \operatorname {g r a d} \mathcal {L} (w _ {t}) \| _ {w _ {t}} ^ {2} \right] \leq \mathbb {E} \left[ \mathcal {L} (w _ {t}) \right] - \mathbb {E} \left[ \mathcal {L} (w _ {t + 1}) \right] + \frac {\alpha^ {2} L _ {S} (1 + 2 L _ {S} ^ {2} \rho^ {2}) \sigma^ {2}}{b} + \frac {2 \alpha (1 - \alpha L _ {S}) \rho^ {3} L _ {\mathrm {R}} ^ {2} \sigma^ {2}}{b} \\ \leq \mathbb {E} \left[ \mathcal {L} \left(w _ {t}\right) \right] - \mathbb {E} \left[ \mathcal {L} \left(w _ {t + 1}\right) \right] + \frac {\alpha^ {2} L _ {S} \sigma^ {2}}{b} + \frac {2 \alpha^ {2} L _ {S} ^ {3} \rho^ {2} \sigma^ {2}}{b} + \frac {2 \alpha \rho^ {3} L _ {\mathrm {R}} ^ {2} \sigma^ {2}}{b} \\ \end{array}
$$
By telescoping the above inequality from  $t = 0 \sim T - 1$ , we arrive at
$$
\mathbb {E} \left[ \| \operatorname {g r a d} \mathcal {L} (\tilde {w}) \| _ {\tilde {w}} ^ {2} \right] = \frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \left[ \| \operatorname {g r a d} \mathcal {L} (w _ {t}) \| _ {w _ {t}} ^ {2} \right] \leq \frac {8 \Delta}{3 \alpha T} + \frac {8 \alpha^ {2} L _ {S} \sigma^ {2}}{3 b} + \frac {1 6 \alpha^ {2} L _ {S} ^ {3} \rho^ {2} \sigma^ {2}}{3 b} + \frac {1 6 \alpha \rho^ {3} L _ {\mathrm {R}} ^ {2} \sigma^ {2}}{3 b}
$$
Under the step size condition  $\alpha_{t} = \frac{1}{\sqrt{T}\tilde{L}}$  and  $\rho_{t} = \frac{1}{T^{1 / 6}\tilde{L}}$ , we finally get
$$
\mathbb {E} \left[ \| \operatorname {g r a d} \mathcal {L} (\widetilde {w}) \| _ {\widetilde {w}} ^ {2} \right] \leq \frac {Q _ {1} \widetilde {L} \Delta}{\sqrt {T}} + \frac {Q _ {2} \sigma^ {2}}{b \sqrt {T}} + \frac {Q _ {3} \sigma^ {2}}{b T ^ {5 / 6}}
$$
for appropriate constants  $\{Q_i\}_{i=1}^3$ .
B Hyperparameter Details
We use the almost same hyperparameters in the study [52] and implement our experiments in Section 5 upon its official implementation. For completeness, we summarize the hyperparameter configurations in Table 3 and Table 4.
Table 3: Hyperparameter configurations for knowledge graph completion.
Dimension	WN18RR	FB15k-237
32	β	32	β
Batch Size	1000	1000	500	500
Negative Samples	50	50	50	50
Margin	8.0	8.0	8.0	8.0
Epochs	1000	1000	500	500
Max Norm	1.5	2.5	1.5	1.5
MaxScaler	3.5	2.5	2.5	2.5
Learning Rate	0.005	0.003	0.003	0.003
Gradient Norm	0.5	0.5	0.5	0.5
Table 4: Hyperparameter configurations for machine translation.
Hyperparameter	IWSLT'14	WMT'14
GPU Numbers	4	4
Embedding Dimension d	64	64
Feed-forward Dimension	256	256
Batch Type	Token	Token
Batch Size	10240	10240
Gradient Accumulation Steps	1	1
Training Steps	40000	200000
Dropout	0.0	0.1
Attention Dropout	0.1	0.0
Max Gradient Norm	0.5	0.5
Warmup Steps	8000	6000
Decay Method	noam	noam
Label Smoothing	0.1	0.1
Layer Number	6	6
Head Number	4	8
Learning Rate	5.0	5.0
Adam β2	0.998	0.998