SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbb {E} [ F _ {\mu} (\mathbf {J _ {T}}) - F _ {S} (\mathbf {J _ {T}}) ] \leq - 2 \gamma^ {\prime} \mathbb {V} _ {m} \mathbb {E} \int_ {t} \frac {d F _ {S} (\mathbf {J _ {t}})}{\sqrt {n}} \sqrt {1 + \frac {\mathrm {T r} (\Sigma (\mathbf {J _ {t}}))}{\| \nabla F _ {S} (\mathbf {J _ {t}}) \| _ {2} ^ {2}}} \\ \underbrace {- 2 \gamma^ {\prime} \mathbb {V} _ {m} \mathbb {E} \int_ {t} \frac {d F _ {S} ^ {l i n} (\mathbf {J} _ {\mathbf {t}}) - d F _ {S} (\mathbf {J} _ {\mathbf {t}})}{\sqrt {n}} \sqrt {1 + \frac {\operatorname {T r} (\Sigma (\mathbf {J} _ {\mathbf {t}}))}{\| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) \| _ {2} ^ {2}}} _ {(A)} \tag {45} \\ + \underbrace {\beta \mathbb {E} \sum_ {t = 0} ^ {T - 1} \left\| \mathbf {J} _ {\mathbf {t} + \mathbf {1}} - \mathbf {J} _ {\mathbf {t}} \right\| ^ {2}} _ {(B)}. \\ \end{array}
$$
The term " $(A)$ " is caused by using  $F_{S}(\mathbf{J}_{\mathbf{t} + 1}) - F_{S}(\mathbf{J}_{\mathbf{t}})$  to replace  $F_{S}^{lin}(\mathbf{J}_{\mathbf{t} + 1}) - F_{S}^{lin}(\mathbf{J}_{\mathbf{t}})$ . The term " $(B)$ " is induced by  $\mathrm{gen}^{nl}(\mathbf{J}_{\mathbf{T}})$ . Then, we want to give a upper bound of  $(A)$  using  $M_4^4$ :
$$
\begin{array}{l} (A) \stackrel {(\star)} {\leq} 2 \gamma^ {\prime} \mathbb {V} _ {m} \mathbb {E} \sum_ {t = 0} ^ {T - 1} \frac {\beta \| \mathbf {J} _ {\mathbf {t} + \mathbf {1}} - \mathbf {J} _ {\mathbf {t}} \| ^ {2}}{\sqrt {n}} \sqrt {1 + \frac {\operatorname {T r} (\Sigma (\mathbf {J} _ {\mathbf {t}}))}{\| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) \| _ {2} ^ {2}}} \\ \stackrel {(\star \star)} {\leq} 2 c ^ {2} \gamma^ {\prime} \mathbb {V} _ {m} \mathbb {E} \sum_ {t = 0} ^ {T - 1} \frac {\| \nabla F _ {S} (\mathbf {J} _ {t}) + \epsilon (\mathbf {J} _ {t}) \| ^ {2}}{\beta \sqrt {n} (t + 1) ^ {2}} \sqrt {1 + \frac {\mathrm {T r} (\Sigma (\mathbf {J} _ {t}))}{\| \nabla F _ {S} (\mathbf {J} _ {t}) \| _ {2} ^ {2}}} \\ = 2 c ^ {2} \gamma^ {\prime} \mathbb {V} _ {m} \sum_ {t = 0} ^ {T - 1} \mathbb {E} _ {t - 1} \frac {\left\| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) + \epsilon (\mathbf {J} _ {\mathbf {t}}) \right\| ^ {2}}{\beta \sqrt {n} (t + 1) ^ {2}} \sqrt {1 + \frac {\operatorname {T r} \left(\Sigma (\mathbf {J} _ {\mathbf {t}})\right)}{\left\| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) \right\| _ {2} ^ {2}}} (46) \\ \stackrel {(\star \star \star)} {\leq} 2 c ^ {2} \gamma^ {\prime} \mathbb {V} _ {m} \sum_ {t = 0} ^ {T - 1} \sqrt {\frac {\mathbb {E} _ {t - 1} \| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) + \epsilon (\mathbf {J} _ {\mathbf {t}}) \| ^ {4}}{\beta^ {2} n (t + 1) ^ {4}} \mathbb {E} _ {t - 1} \left(1 + \frac {\operatorname {T r} (\Sigma (\mathbf {J} _ {\mathbf {t}}))}{\| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) \| _ {2} ^ {2}}\right)} (46) \\ \leq 2 c ^ {2} \gamma^ {\prime} \mathbb {V} _ {m} \sum_ {t = 0} ^ {T - 1} \sqrt {\frac {M _ {4} ^ {4}}{\beta^ {2} n (t + 1) ^ {4}} \mathbb {E} _ {t - 1} \left(1 + \frac {\mathrm {T r} (\Sigma (\mathbf {J} _ {t}))}{\| \nabla F _ {S} (\mathbf {J} _ {t}) \| _ {2} ^ {2}}\right)} \\ \leq 2 c ^ {2} \gamma^ {\prime} \mathbb {V} _ {m} M _ {4} ^ {2} \sqrt {\mathbb {E} \sum_ {t = 0} ^ {T - 1} \frac {1}{\beta^ {2} n (t + 1) ^ {4}} \left(1 + \frac {\operatorname {T r} (\Sigma (\mathbf {J} _ {\mathbf {t}}))}{\| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) \| _ {2} ^ {2}}\right)}. \\ \end{array}
$$
where  $(\star)$  is due to the Equation 44,  $(\star \star)$  is due to the update rule of  $\mathbf{J}_{\mathbf{t}}$  and  $(\star \star \star)$  is due to Hölder's inequality. In the following, we use  $M_2^2$  to give a upper bound for  $(B)$ :
$$
\begin{array}{l} (B) \leq \frac {c ^ {2}}{\beta} \sum_ {t = 0} ^ {T - 1} \frac {1}{(t + 1) ^ {2}} \mathbb {E} \| \nabla F (\mathbf {J} _ {\mathbf {t}}) + \epsilon (\mathbf {J} _ {\mathbf {t}}) \| ^ {2} \\ \leq \frac {c ^ {2}}{\beta} \sum_ {t = 0} ^ {T - 1} \frac {1}{(t + 1) ^ {2}} M _ {2} ^ {2} \\ \leq \frac {c ^ {2}}{\beta} \left(M _ {2} ^ {2} + \sum_ {t = 1} ^ {T - 1} \frac {1}{(t + 1) ^ {2}} M _ {2} ^ {2}\right) \tag {47} \\ \leq \frac {c ^ {2}}{\beta} \left(M _ {2} ^ {2} + \sum_ {t = 1} ^ {T - 1} \left(\frac {1}{t} - \frac {1}{t + 1}\right) M _ {2} ^ {2}\right) \\ \leq \frac {c ^ {2}}{\beta} \left(2 M _ {2} ^ {2} - \frac {1}{T} M _ {2} ^ {2}\right) \\ \leq 2 c ^ {2} \frac {M _ {2} ^ {2}}{\beta} \\ \end{array}
$$
Taking the upper bound value of "(A)" and "(B)" into Equation 45, we obtain the result.
B Relaxed Assumption and Corresponding Bound
Assumption B.1. There is a value  $\gamma$ ,  $T_0$  and  $\zeta$ , so that for all  $\mathbf{w} \in \{\mathbf{J}_{\mathbf{t}} | t \in \mathbb{N} \land t < T_0\}$ , we have  $\|\nabla F_{\mu}(\mathbf{w})\| \leq \gamma \| \nabla F_S(\mathbf{w})\|$  and for all  $\mathbf{w} \in \{\mathbf{J}_{\mathbf{t}} | t \in \mathbb{N} \land t \geq T_0\}$ , we have  $\|\nabla F_{\mu}(\mathbf{w})\| \leq \gamma \| \nabla F_S(\mathbf{w})\| + \zeta$ .
Theorem B.2. Under Assumption B.1, given  $S \sim \mu^n$ , let  $\mathbf{J} = \mathcal{A}(S)$ , where  $\mathcal{A}$  denotes the SGD or GD algorithm training with  $T$  steps, we have:
$$
\begin{array}{l} \mathbb {E} \left[ F _ {\mu} \left(\mathbf {J} _ {\mathbf {T}}\right) - F _ {S} \left(\mathbf {J} _ {\mathbf {T}}\right) \right] \leq - 2 \gamma^ {\prime} \mathbb {V} _ {m} \mathbb {E} \int_ {t} \frac {d F _ {S} \left(\mathbf {J} _ {\mathbf {t}}\right)}{\sqrt {n}} \sqrt {1 + \frac {\operatorname {T r} \left(\Sigma \left(\mathbf {J} _ {\mathbf {t}}\right)\right)}{\| \nabla F _ {S} \left(\mathbf {J} _ {\mathbf {t}}\right) \| _ {2} ^ {2}}} + \frac {1}{2} \sum_ {t = T _ {0}} ^ {T} \eta_ {t} \| \nabla F _ {S} \left(\mathbf {J} _ {\mathbf {t}}\right) \| \zeta + \mathcal {O} (\eta_ {m}), \tag {48} \\ \begin{array}{l l} \text {w h e r e} \quad \mathbb {V} (\mathbf {w}) & = \frac {\| \nabla F _ {S} (\mathbf {w}) \|}{\mathbb {E} _ {U \subset S} \| \frac {| U |}{n} \nabla F _ {U} (\mathbf {w}) - \frac {n - | U |}{n} \nabla F _ {S / U} (\mathbf {w}) \|}, \quad \mathbb {V} _ {m} \\ \max  \{1, \max  _ {U \subset S; t} \frac {| U | \| \nabla F _ {U} (\mathbf {J} _ {\mathbf {t}}) \|}{n \| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) \|} \} \gamma . \end{array} \\ \end{array}
$$
Proof. Most of the proofs in this part are the same as those in Appendix A, except for Equation 30. The Equation 30 is replaced by:
$$
\begin{array}{l} k \mathbb {E} _ {\sigma} \sup  _ {h \in \mathcal {L} _ {\mathbf {J} | S}} \left(\frac {1}{n} \sum_ {i} ^ {n} \sigma_ {i} h (z _ {i})\right) - \mathbb {E} _ {S ^ {\prime}, \sigma} \sup  _ {h \in \mathcal {F} _ {\mathbf {J} | S}} \left(\frac {1}{n} \sum_ {i} ^ {n} \sigma_ {i} (h (z _ {i}))\right) \\ = k \mathbb {E} _ {\sigma} \left(\frac {1}{n} \sum_ {t = 0} ^ {T} \delta_ {t} \| g _ {S _ {+}} (\mathbf {J} _ {\mathbf {t}}) - g _ {S _ {-}} (\mathbf {J} _ {\mathbf {t}}) \|\right) - \mathbb {E} _ {S ^ {\prime}, \sigma} \left(\frac {1}{n} \sum_ {t = 0} ^ {T} \delta_ {t} \frac {g _ {S} (\mathbf {J} _ {\mathbf {t}})}{\| g _ {S} (\mathbf {J} _ {\mathbf {t}}) \|} \left(g _ {S _ {+} ^ {\prime}} (\mathbf {J} _ {\mathbf {t}}) - g _ {S _ {-}} (\mathbf {J} _ {\mathbf {t}})\right)\right) \\ = k \mathbb {E} _ {\sigma} \left(\frac {1}{n} \sum_ {t = 0} ^ {T} \delta_ {t} \| g _ {S _ {+}} (\mathbf {J} _ {\mathbf {t}}) - g _ {S _ {-}} (\mathbf {J} _ {\mathbf {t}}) \| - \frac {1}{n} \sum_ {t = 0} ^ {T} \delta_ {t} \frac {g _ {S} (\mathbf {J} _ {\mathbf {t}})}{\| g _ {S} (\mathbf {J} _ {\mathbf {t}}) \|} (| S _ {+} | \nabla F _ {\mu} (\mathbf {J} _ {\mathbf {t}}) - g _ {S _ {-}} (\mathbf {J} _ {\mathbf {t}}))) \right. \tag {49} \\ \geq k \mathbb {E} _ {\sigma} \left(\frac {1}{n} \sum_ {t = 0} ^ {T} \delta_ {t} \| g _ {S _ {+}} (\mathbf {J} _ {\mathbf {t}}) - g _ {S _ {-}} (\mathbf {J} _ {\mathbf {t}}) \| - \frac {1}{n} \sum_ {t = 0} ^ {T} \delta_ {t} \| | S _ {+} | \nabla F _ {\mu} (\mathbf {J} _ {\mathbf {t}}) - g _ {S _ {-}} (\mathbf {J} _ {\mathbf {t}}) |\right) \\ \geq k \frac {1}{n} \sum_ {t = 0} ^ {T} \delta_ {t} \mathbb {E} _ {\sigma} \left(\| g _ {S _ {+}} (\mathbf {J} _ {\mathbf {t}}) - g _ {S _ {-}} (\mathbf {J} _ {\mathbf {t}}) \|\right) - \sum_ {t = 0} ^ {T} \delta_ {t} \gamma^ {\prime} \| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) \| - \frac {1}{2} \sum_ {t = T _ {0}} ^ {T} \delta_ {t} \zeta \\ \geq - \frac {1}{2} \sum_ {t = T _ {0}} ^ {T} \eta_ {t} \| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) \| \zeta . \\ \end{array}
$$
Remark B.3. Compared of Theorem 3.6, we have a extra term  $\sum_{t = T_0}^T\eta_t\| \nabla F_S(\mathbf{J}_t)\| \zeta$  here. Since the unrelaxed assumption  $\| \nabla F_{\mu}(\mathbf{w})\| \leq \gamma \| \nabla F_S(\mathbf{w})\|$  is not satisfied only when  $\| \nabla F_S(\mathbf{w})\|$  is relative small, the term  $\sum_{t = T_0}^T\eta_t\| \nabla F_S(\mathbf{J}_t)\| \zeta$  is small value.
C.1 Calculation of  $\mathcal{C}(\mathbf{J})$
To reduce the calculation, we construct a randomly sampled subset  $S_{\mathrm{sp}} = \{z_1^{\mathrm{sp}}, \dots, z_n^{\mathrm{sp}}\} \subset S$ .
$$
\begin{array}{l} \int_ {t} \frac {d F _ {S} (\mathbf {J} _ {\mathbf {t}})}{\sqrt {n}} \sqrt {1 + \frac {\operatorname {T r} (\Sigma (\mathbf {J} _ {\mathbf {t}}))}{\| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) \| _ {2} ^ {2}}} = \int_ {t} \frac {d F _ {S} (\mathbf {J} _ {\mathbf {t}})}{\sqrt {n}} \sqrt {\frac {\sum_ {i = 1} ^ {n} \| \nabla f (\mathbf {J} _ {\mathbf {t}} , z _ {i}) \| _ {2} ^ {2}}{n} \frac {1}{\| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) \| _ {2} ^ {2}}} \\ \approx \int_ {t} \frac {d F _ {S} (\mathbf {J} _ {\mathbf {t}})}{\sqrt {n}} \sqrt {\frac {\sum_ {i = 1} ^ {n _ {\mathrm {s p}}} \| \nabla f (\mathbf {J} _ {\mathbf {t}} , z _ {i} ^ {\mathrm {s p}}) \| _ {2} ^ {2}}{n _ {\mathrm {s p}}}} \frac {1}{\| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t}}) \| _ {2} ^ {2}} \\ \end{array}
$$
Denote the weights after  $t$ -epoch training as  $\mathbf{X}_{\mathbf{t}}$ . We can roughly calculated  $\mathcal{C}(\mathbf{J}_{\mathbf{t}})$
$$
\sum_ {t = 1} ^ {T} \frac {F _ {S} (\mathbf {X} _ {\mathbf {t}}) - F _ {S} (\mathbf {X} _ {\mathbf {t} - 1})}{\sqrt {n}} \sqrt {\frac {\sum_ {i = 1} ^ {n _ {\mathrm {s p}}} \| \nabla f (\mathbf {X} _ {\mathbf {t}} , z _ {i} ^ {\mathrm {s p}}) \| _ {2} ^ {2}}{n _ {\mathrm {s p}}} \frac {1}{\| \nabla F _ {S} (\mathbf {X} _ {\mathbf {t}}) \| _ {2} ^ {2}}}
$$
Here, we give a detail setting of the experiment for each figure.
Figure 1 The learning rate is fixed to 0.05 during all the training process. The batch size is 256. All experiments is trained with 100 epoch. The test accuracy for CIFAR-10, CIFAR-100, and SVHN are  $87.64\%$ ,  $55.08\%$ , and  $92.80\%$ , respectively.
Figure 2 The initial learning rate is set to 0.05 with the batch size of 1024. We use the Cosine Annealing LR Schedule to adjust the learning rate during training.
Figure 3 Each point is an average of three repeated experiments. We stop training when the training loss is small than 0.2.
C.3 Experimental exploration of  $\mathrm{gen}^{nl}(\mathbf{J}_{\mathbf{T}})$
In this section, our aim is to investigate the conditions under which  $\mathrm{gen}^{nl}(\mathbf{J}_{\mathbf{T}})\approx 0$ . Since directly calculating the difference  $|R_{\mu}(\mathbf{J}_{\mathbf{T}}) - R_{S}(\mathbf{J}_{\mathbf{T}})|$  is challenging, we concentrate on the upper bound value  $|R_{\mu}(\mathbf{J}_{\mathbf{T}})| + |R_{S}(\mathbf{J}_{\mathbf{T}})|$ .
We conduct the experiment using cifar10-5k dataset and fc-tanh network, following the setting of paper [9]. Cifar10-5k[9] is a subset of cifar10 dataset. Building upon the work of [1], we compute the Relative Progress Ratio (RP) and Test Relative Progress Ratio (TRP) throughout the training process. We initially consider the case of gradient descent. The definitions of RP and TRP for gradient descent are as follows:
$$
\mathrm {R P} \left(\mathbf {J} _ {\mathbf {t}}\right) \triangleq \frac {F _ {S} \left(\mathbf {J} _ {\mathbf {t} + \mathbf {1}}\right) - F _ {S} \left(\mathbf {J} _ {\mathbf {t}}\right)}{\eta \| \nabla F _ {S} \left(\mathbf {J} _ {\mathbf {t}}\right) \| ^ {2}} \tag {50}
$$
$$
\operatorname {T R P} \left(\mathbf {J} _ {\mathbf {t}}\right) \triangleq \frac {F _ {S ^ {\prime}} \left(\mathbf {J} _ {\mathbf {t} + 1}\right) - F _ {S ^ {\prime}} \left(\mathbf {J} _ {\mathbf {t}}\right)}{\eta \nabla F _ {S} \left(\mathbf {J} _ {\mathbf {t}}\right) ^ {\mathrm {T}} \nabla F _ {S ^ {\prime}} \left(\mathbf {J} _ {\mathbf {t}}\right)}. \tag {51}
$$
$$
\begin{array}{l} F _ {S} (\mathbf {J _ {0}}) + \sum_ {t = 1} ^ {T} (F _ {S} (\mathbf {J _ {t}}) - F _ {S} (\mathbf {J _ {t - 1}})) - F _ {S} (\mathbf {J _ {0}}) - \sum_ {t = 0} ^ {T - 1} (\mathbf {J _ {t}} - \mathbf {J _ {t - 1}}) ^ {\mathrm {T}} \nabla F _ {S} (\mathbf {J _ {t - 1}}) \\ = \sum_ {t = 1} ^ {T} \left[ \left(F _ {S} \left(\mathbf {J} _ {\mathbf {t}}\right) - F _ {S} \left(\mathbf {J} _ {\mathbf {t} - \mathbf {1}}\right)\right) - \left(\mathbf {J} _ {\mathbf {t}} - \mathbf {J} _ {\mathbf {t} - \mathbf {1}}\right) ^ {\mathrm {T}} \nabla F _ {S} \left(\mathbf {J} _ {\mathbf {t} - \mathbf {1}}\right) \right] \tag {52} \\ = \sum_ {t = 1} ^ {T} \left[ \left(F _ {S} (\mathbf {J} _ {\mathbf {t}}) - F _ {S} (\mathbf {J} _ {\mathbf {t} - \mathbf {1}})\right) + \eta \| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t} - \mathbf {1}}) \| ^ {2} \right] \\ = \sum_ {t = 1} ^ {T} \left[ \eta_ {t} \left(1 + \operatorname {R P} \left(\mathbf {J} _ {\mathbf {t} - \mathbf {1}}\right)\right) \| \nabla F _ {S} \left(\mathbf {J} _ {\mathbf {t} - \mathbf {1}}\right) \| ^ {2} \right] \\ \end{array}
$$
Following the same way, we have:
$$
\begin{array}{l} F _ {S ^ {\prime}} (\mathbf {J} _ {\mathbf {0}}) + \sum_ {t = 1} ^ {T} \left(F _ {S ^ {\prime}} \left(\mathbf {J} _ {\mathbf {t}}\right) - F _ {S ^ {\prime}} \left(\mathbf {J} _ {\mathbf {t} - \mathbf {1}}\right)\right) - F _ {S ^ {\prime}} \left(\mathbf {J} _ {\mathbf {0}}\right) - \sum_ {t = 0} ^ {T - 1} \left(\mathbf {J} _ {\mathbf {t}} - \mathbf {J} _ {\mathbf {t} - \mathbf {1}}\right) ^ {\mathrm {T}} \nabla F _ {S ^ {\prime}} \left(\mathbf {J} _ {\mathbf {t} - \mathbf {1}}\right) \tag {53} \\ = \sum_ {t = 1} ^ {T} \left[ \eta_ {t} \left(1 + \operatorname {T R P} \left(\mathbf {J} _ {\mathbf {t} - \mathbf {1}}\right)\right) \nabla F _ {S} \left(\mathbf {J} _ {\mathbf {t} - \mathbf {1}}\right) ^ {\mathrm {T}} \nabla F _ {S ^ {\prime}} \left(\mathbf {J} _ {\mathbf {t} - \mathbf {1}}\right) \right] \\ \end{array}
$$
Combining Equation (52) and Equation (53), we have:
$$
\begin{array}{l} | \operatorname {g e n} ^ {n l} (\mathbf {J} _ {\mathbf {T}}) | \leq \\ \sum_ {t = 1} ^ {T} \eta_ {t} \left[ (1 + \operatorname {T R P} (\mathbf {J} _ {\mathbf {t} - 1})) | \nabla F _ {S} (\mathbf {J} _ {\mathbf {t} - 1}) ^ {\mathrm {T}} \nabla F _ {S ^ {\prime}} (\mathbf {J} _ {\mathbf {t} - 1}) | + (1 + \operatorname {R P} (\mathbf {J} _ {\mathbf {t} - 1})) \| \nabla F _ {S} (\mathbf {J} _ {\mathbf {t} - 1}) \| ^ {2} \right] \tag {54} \\ \end{array}
$$
Therefore, if we have for all  $t,\mathrm{RP}(\mathbf{J}_{\mathbf{t}})\approx -1$  and  $\mathrm{TRP}(\mathbf{J}_{\mathbf{t}})\approx -1$  , then  $|\operatorname {gen}^{nl}(\mathbf{J}_{\mathbf{T}})|\approx 0$
From Figure 4 we find that in stable regime, where the sharpness is below the  $\frac{2}{\eta}$ , we have  $\mathrm{TRP} \approx \mathrm{RP} \approx -1$ . Under small learning rate, the gradient descent doesn't enter the regime of edge of stability and we have  $\mathrm{TRP} \approx \mathrm{RP} \approx -1$  during whole training process and  $\mathrm{gen}^{nl}(\mathbf{J}_{\mathbf{T}}) \approx 0$ .
Next, we consider the case of Stochastic Gradient Descent (SGD). Due to the stochastic estimation of the gradient, we need to rely on some approximations. Let  $\mathbf{X}_{\mathrm{t}}^{\mathrm{i}}$  represent the weights after the  $t$ -epoch and  $i$ -th iteration of training. We assume a constant learning rate  $\eta$  for SGD. The gradient is approximated as follows:
$$
\eta \nabla F _ {S} \left(\mathbf {X} _ {\mathbf {t}} ^ {\mathbf {i}}\right) \approx \frac {B}{n} \left(\mathbf {X} _ {\mathbf {t}} - \mathbf {X} _ {\mathbf {t} + 1}\right) = \frac {B}{n} \sum_ {i = 1} ^ {\frac {n}{B}} \nabla F _ {S} \left(\mathbf {X} _ {\mathbf {t}} ^ {\mathbf {i}}\right), \tag {55}
$$
$\eta = 0.01$
Figure 4: Exploration of  $\mathrm{gen}^{nl}(\mathbf{J}_{\mathbf{T}})$  on Gradient descent. Experiments is conducted on CIFar10-5k dataset with cross entropy loss. The blue dash line in fourth row denotes  $\frac{2}{\eta}$ . Gradient descent enter the EoS regime when the sharpness is above  $\frac{2}{\eta}$ . Both RP and TRP have values around -1 when sharpness is below the  $\frac{2}{\eta}$ .
and we approximate  $\nabla F_{S'}(\mathbf{X}_{\mathbf{t}}^{\mathbf{i}})$  as:
$$
\eta \nabla F _ {S} \left(\mathbf {X} _ {\mathbf {t}} ^ {\mathbf {i}}\right) \approx \eta \nabla F _ {S} \left(\mathbf {X} _ {\mathbf {t}}\right). \tag {56}
$$
$$
\operatorname {R P} \left(\mathbf {X} _ {\mathbf {t}}\right) \approx \frac {\eta \left(F _ {S} \left(\mathbf {X} _ {\mathbf {t} + \mathbf {1}}\right) - F _ {S} \left(\mathbf {X} _ {\mathbf {t}}\right)\right)}{\left\| \mathbf {X} _ {\mathbf {t} + \mathbf {1}} - \mathbf {X} _ {\mathbf {t}} \right\|} \tag {57}
$$
$$
\operatorname {T R P} \left(\mathbf {X} _ {\mathbf {t}}\right) \approx \frac {F _ {S ^ {\prime}} \left(\mathbf {X} _ {\mathbf {t} + \mathbf {1}}\right) - F _ {S ^ {\prime}} \left(\mathbf {X} _ {\mathbf {t}}\right)}{\left(\mathbf {X} _ {\mathbf {t}} - \mathbf {X} _ {\mathbf {t} + \mathbf {1}}\right) ^ {\mathrm {T}} \nabla F _ {S ^ {\prime}} \left(\mathbf {X} _ {\mathbf {t}}\right)}. \tag {58}
$$
Figure 5: Exploration of  $\mathrm{gen}^{nl}(\mathbf{J}_{\mathbf{T}})$  on SGD case. Here, the effective learning rate is defined as  $\eta_{ef}\triangleq \frac{n}{B}\eta$ . We still have  $\mathrm{gen}^{nl}(\mathbf{J}_{\mathbf{T}})\approx 0$  under small learning rate.
We calculated the effect learning rate for SGD as  $\eta_{ef} \triangleq \frac{n}{B}\eta$ . Figure 5 shows that the conclusions of SGD are similar as GD, except that the conditions of entering EoS are different.
Table 4: Comparison of trajectory based generalization bounds. Only our proposed method can apply to the SGD with rich trajectory related information.
Method	Conditions	T.R.T
Nikolakakis et al. [25]	Gradient Descent, ηt ≤ c/t ≤ 1/β, β-smooth	ΣTt=1Tηt/n ∑i=1n ||∇f(Jt, zi)||2
Neu et al. [24]	β-smooth, E[||∇f(w, z) - ∇Fμ(w)||] ≤ v, f(·) is subgaussian distribution	√Tη2
Park et al. [27]	Weak Lipschitz continuity, Piecewise β′-smooth, f(·) is bounded, η < 2/β	T
Ours	Small Learning Rate, ||∇Fμ(w)|| ≤ γ||∇F(S)(w)||	∫t dFS(Jt)√1 + Tr(Σ(Jt))/||∇F(S(Jt)||2
D Other Related Work
This part compares the works that is not listed in Table 2. Table 4 gives other trajectory based generalization bounds. [25] is a stability based work designed mainly for generalization of gradient descent. It removes the Lipschitz assumption, and replaced by the term  $\sum_{t=1}^{T} \eta_t \frac{1}{n} \sum_{i=1}^{n} \|\nabla f(\mathbf{J}_t, z_i)\|^2$  in the generalization bounds. This helps enrich the trajectory information in the bounds. The limitation of this work is that it can only apply to the gradient descent and it is hard to extend to the stochastic gradient descent. Neu et al. [24] adapt the information-theoretical generalization bound to the stochastic gradient descent. The Theorem 1 in Neu et al. [24] contains rich information about the learning trajectory, but most is about  $\nabla F_\mu(\mathbf{w})$ , which is unavailable for us. Therefore, we mainly consider the result of Corollary 2 in Neu et al. [24], which removes the term  $\nabla F_\mu(\mathbf{w})$  by the assumption listed in Table 4. For this Collory, the remained information within trajectory is merely the  $\sqrt{T}\eta^2$ . Although Neu et al. [24] doesn't require the assumption of small learning rate, the bound contains the dimension of model, which is large for deep neural network. Compared with these work, our proposed method has advantage in that it can both reveal rich information about learning trajectory and applied to stochastic gradient descent.
Chandramoorthy et al. [8] analyzes the generalization behavior based on statistical algorithmic stability. The proposed generalization bound can be applied into algorithms that don't converge. Let  $S^{(i)}$  be the dataset obtained by replace  $z_{i}$  in  $S$  with another sample  $z_{i}^{\prime}$  draw from distribution  $\mu$ . The generalization bound relies on the stability measure  $m \triangleq \sup \left\{\frac{1}{T}\sum_{t = 0}^{T - 1}f(\mathbf{J}_{\mathbf{t}}|S,z) - \frac{1}{T}\sum_{t = 0}^{T - 1}f(\mathbf{J}_{\mathbf{t}}|S^{(i)},z)|z\in \mathcal{Z},i\in [n]\right\}$ . We don't directly compare with this method because the calculation of  $m$  relies on  $S^{(i)}$  which contains sample outside of  $S$ . Therefore, we treat this result as intermediate results. More assumption is needed to remove this dependence of the information about the unseen samples, i.e., the samples outside set  $S$ .
E Effect of Learning Rate and Stochastic Noise
In this part, we want to analyze how learning rate and the stochastic noise jointly affect our proposed generalization bound. Specifically, we denote  $p_t(\mathbf{w})$  as the distribution of the  $\mathbf{J}_{\mathbf{t}}$  during the training with multiple training steps. Following the work [12], we consider the SDE function as an approximation, which is shown as below:
$$
\mathrm {d} \mathbf {w} = - \nabla F _ {S} (\mathbf {w}) \mathrm {d} t + \sqrt {\eta} C ^ {\frac {1}{2}} \mathrm {d} \mathbf {W} (t). \tag {59}
$$
The SDE can be regarded as the continuous counterpart of Equation(3) when sets the distribution of noise term  $\epsilon^{\prime}$  in Equation(3) as Gaussian distribution. The influence of the noise  $\epsilon$  on  $p_t(\mathbf{w})$  is shown in the following theorem.
Theorem E.1. When the updating of the weight  $w$  follows Equation (59), the covariance matrix  $C$  is a hessian matrix of a function with a scalar output, then we have:
$$
\frac {\partial p _ {t} (\mathbf {w})}{\partial t} = - \sum_ {i = 1} ^ {d} \frac {\partial}{\partial \mathbf {w} _ {\mathbf {i}}} [ \nabla F _ {S} (\mathbf {w}) p _ {t} (\mathbf {w}) - \frac {\eta}{2} [ \nabla \operatorname {T r} (C (\mathbf {w})) + \underbrace {C (\mathbf {w}) \nabla_ {w} \log (p _ {t} (\mathbf {w}))} _ {\text {d a m p l i n g f a c t o r}} ] p _ {t} (\mathbf {w}) ]. \tag {60}
$$
Remark E.2. Previous studies [45, 32, 12] tell that the covariance matrix  $C$  is proximately equal to the hessian matrix of the loss function with respect to the parameters of DNN. Thus, the above condition that the covariance matrix  $C$  is a hessian matrix of a function with scalar output is easy to be satisfied. Formula (60) contains three parts. The item  $F_{S}(\mathbf{w})p_{t}(\mathbf{w})$  enlarge the probability of parameters being located in the parameter space with low  $F_{S}(\mathbf{w})$ .  $\nabla \mathrm{Tr}(C(\mathbf{w}))$  and  $C(\mathbf{w})\nabla_w\log (p_t(\mathbf{w}))$  usually
contradict with each other.  $\nabla \mathrm{Tr}(C(\mathbf{w}))$  enlarge the probability of parameters being located in the parameter space with low  $\operatorname {Tr}(C(\mathbf{w}))$  value, while  $C(\mathbf{w})\nabla_w\log (p_t(\mathbf{w}))$  serves as a damping factor to prevent the probability from concentrating on a small space. Therefore, setting larger learning rate gives stronger force for the weight to the area with lower  $\operatorname {Tr}(C(\mathbf{w}))$  values. According to Equation 5, we also have a lower  $\Sigma (\mathbf{w})$ . As a result, large learning rate causes a small lower bound in Theorem 3.6
Proof. Based on the condition described above, we can infer that  $C(\mathbf{w}) = \nabla \nabla G(\mathbf{w})$ , where  $G$  is a function with a scalar output.
We first prove that  $\nabla \cdot C(\mathbf{w}) = \nabla \operatorname {Tr}(C(\mathbf{w}))$  as below: