Learning Trajectories are Generalization Indicators

The proof of this theorem is placed in Appendix A. Here, we give the sketch for this proof.

Beginning with Equation (8), we decomposite the $F_{\mu}(\mathbf{\mathbf{J_{T}}})-F_{S}(\mathbf{J_{T}})$ into a linear part ($\operatorname{gen}^{lin}(\mathbf{J_{T}})$) and nonlinear part($\operatorname{gen}^{nl}(\mathbf{J_{T}})$). We have $\operatorname{gen}^{lin}(\mathbf{J_{T}})=\sum_{t=1}^{T}(ii)^{lin}_{t}$, where $(ii)^{lin}_{t}\triangleq(\mathbf{J_{t}}-\mathbf{J_{t-1}})^{\mathrm{T}}(\nabla F_{\mu}(\mathbf{J_{t-1}})-\nabla F_{S}(\mathbf{J_{t-1}}))$. The nonlinear part is $\operatorname{gen}^{nl}(\mathbf{J_{T}})=F_{\mu}(\mathbf{\mathbf{J_{T}}})-F_{S}(\mathbf{J_{T}})-\operatorname{gen}^{lin}(\mathbf{J_{T}})$. We takle these two parts differently. Here, we focus on analyzing $\operatorname{gen}^{lin}(\mathbf{J_{T}})$ because it dominates under small learning rate. Detail discussion of $\operatorname{gen}^{nl}(\mathbf{J_{T}})$ is given in Appendix (Propositon A.1 and Subsection C.3)

We construct the addictive linear space $\mathcal{L}_{\mathbf{J}|S}\triangleq\{\sum_{t=0}^{T-1}\mathbf{w_{t}}^{\mathrm{T}}\nabla f(\mathbf{J_{t}})\ |\ \|\mathbf{w_{t}}\|\leq\Delta_{t}\}$, where $\Delta_{t}\triangleq\|\eta_{t}\nabla F_{S}(\mathbf{J_{t}})\|$. Then $\mathbb{E}[\operatorname{gen}^{lin}(\mathbf{J_{T}})]\leq 2\gamma^{\prime}\mathbb{V}_{m}\mathbb{E}R_{S}(\mathcal{L}_{\mathbf{J}|S})$, where $R_{S}(\mathcal{L}_{\mathbf{J}|S})\triangleq\mathbb{E}_{\sigma}\sup\limits_{h\in\mathcal{L}_{\mathbf{J}|S}}(\frac{1}{n}\sum_{i=1}^{n}\sigma_{i}h(z_{i}))$.

Finally, we compute the upper bound of $R_{S}(\mathcal{L}_{\mathbf{J}|S})$, which follows same techniques used in Radermacher Complexity theory. By combining this with Proposition A.1, we establish the theorem.

Directly applying the Rademacher complexity to calculate the generalization error bound fails because the large complexity of neural network’s function space leads to trival bound[40]. In this work, we want to calculate the complexity of the function space that can be explored during the training process. However, there are two challenges here. First, the trajectory of neural network is a "line", instead of a function space that can be calculated the complexity. To solve this problem, we indroduce the addictive linear space $\mathcal{L}_{\mathbf{J}|S}$. This space contains the local information of learning trajectory, and can serve as the pseudo function space. Second, the function space $\mathcal{L}_{\mathbf{J}|S}$ has a dependent on the sample set $S$, while the theory of Rademacher complexity requires that the function space is independent with training samples. To decouple this dependence, we adapt the Rademacher complexity and we obtain that $\mathbb{E}[\operatorname{gen}^{lin}(\mathbf{J_{T}})]\leq 2\gamma^{\prime}\mathbb{V}_{m}\mathbb{E}R_{S}(\mathcal{L}_{\mathbf{J}|S})$. Here, $\gamma^{\prime}$ is indroduced to decouple the dependent fact mentioned above.

In Table 1, we present a summary of stability-based methods, while other methods are outlined in Appendix D. We focus on generalization bounds from previous works that eliminate terms dependent on extra information about data distribution $\mu$, apart from the training data $S$, using assumptions such as smoothness or Lipschitz continuity. Analyzing Table 1 reveals that most prior works primarily depend on the learning rate $\eta$ and the total number of training steps $T$. This suggests that we can achieve the same bound by using an identical learning rate schedule and total training steps, which does not align with our practical experience. Our proposed generalization bound considers the evolution of function values, gradient covariance, and gradient norms throughout the training process. As a result, our bounds encompass more comprehensive information about the learning trajectory.

We will first analyze the dependent of $n$ for $\mathbb{V}$. The $\mathbb{V}$ is calculated as $\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})\|}$. Obviously, the gradient of individual sample is unrelated to the sample size $n$. And $|U|\sim n$. Therefore, $\mathbb{V}=\mathcal{O}(1)$. Similarly, we have $\mathbb{E}\int_{t}\frac{dF_{S}(\mathbf{J_{t}})}{\sqrt{n}}\sqrt{1+\frac{\operatorname{Tr}(\Sigma(\mathbf{J_{t}}))}{\|\nabla F_{S}(\mathbf{J_{t}})\|_{2}^{2}}}=\mathcal{O}(\frac{1}{\sqrt{n}})$. As for the $\mathcal{O}(\eta_{m})$ term in Theorem 3.6, we have $\lim\limits_{n\to\infty}\mathcal{O}(\eta_{m})=0$ according to Proposition A.1. We simply assume that $\mathcal{O}(\eta_{m})=\mathcal{O}(\frac{1}{n^{c}})$. Therefore, our bound has $\mathcal{O}(\frac{1}{n^{\text{min}\{0.5,c\}}})$

Next, in order to draw a clearer comparison with the stability-based method, we present the following corollary. This corollary employs the $\beta$-smooth assumption to bound $\operatorname{gen}^{nl}(\mathbf{J_{T}})$ and leverages a similar learning rate setting to that found in stability based works.

1) Some values in the bounds are hard to be calculated. Calculating $\beta$ (under the $\beta$-smooth assumption) and $L$ (under the $L$-Lipschitz assumption) in stability-based work, as well as the values of $\mathbb{V}$ and $\gamma$ in our proposed bound, are challenging. 2) Stability-based methods require a batch size of 1. The training is hard for batch size of 1 with learning rate setting $\eta_{t}=\frac{1}{t}$ in complex datasets.

In the following, we discuss the construction of the toy dataset used to compare the tightness of the generalization bounds. The training data is $X_{tr}=\{x_{i}\}_{i=1}^{n}$. All the data $x_{i}$ is sampled from Guassian distribution $\mathcal{N}(0,\mathbf{I}_{d})$. Sampling $\tilde{\mathbf{w}}\sim\mathcal{N}(0,\mathbf{I}_{d})$,the ground truth is generated by $y_{i}=1\ \ \text{if}\ \ \tilde{\mathbf{w}}^{\mathrm{T}}x_{i}>0\ \ \text{else}\ \ 0$. The weights for learning is denoted as $\mathbf{w}$. The predict $\tilde{y}$ is calculated as $\tilde{y}_{i}=\mathbf{w}^{\mathrm{T}}x_{i}$. The loss for a simple data point is $l_{i}=\left\|y_{i}-\mathbf{w}^{\mathrm{T}}x_{i}\right\|_{2}$. The training loss is $\mathcal{L}=\sum_{i=1}^{n}l_{i}$. The test data is $X_{te}=\{x^{\prime}_{i}\}$, where $x^{\prime}_{i}=\tilde{x}^{\prime}_{i}$ and $\tilde{x}^{\prime}_{i}\sim\mathcal{N}(0,\mathbf{I}_{d})$. We use 100 samples for training and 1,000 samples for evaluation. The model is trained using SGD for 200 epochs.

We evaluate the tightness of our bound by comparing our results with those in references Hardt et al. [11] and Zhang et al. [42] from the original paper. We set the learning rate as $\eta_{t}=\frac{1}{\beta t}$. Our reasons for comparing with these two papers are: 1) Hardt et al. [11] is a representative study, 2) Both papers have theorems using a learning rate setting of $\eta_{t}=\mathcal{O}(\frac{1}{t})$, which aligns with Corollary 3.8 in our paper, and 3) They do not assume convexity. The generalization bounds we compare include Corollary 3.8 from our paper, Theorem 3.12 from Hardt et al. [11], and Theorem 5 from Zhang et al. [42].

Our results are given in Table 3. Our bound is tighter under this setting.

To explore the Assumption 3.4, we define $\gamma_{t}\triangleq\frac{\|\nabla F_{\mu}(\mathbf{J_{t}})\|}{\|\nabla F_{S}(\mathbf{J_{t}})\|}$ and $\widetilde{\gamma}_{t}\triangleq\frac{\|\nabla F_{S^{\prime}}(\mathbf{J_{t}})\|}{\|\nabla F_{S}(\mathbf{J_{t}})\|}$, where $S^{\prime}$ is another data set i.i.d sampled from distribution $\mu$. Because $S^{\prime}$ is independent with $S$, we have $\widetilde{\gamma}_{t}\approx\gamma_{t}$. We found that $\widetilde{\gamma}_{t}$ is stable around 1 during the early stage of training(Figure 1). When the training loss is reaching a relative small value, $\widetilde{\gamma}_{t}$ increases as we continue training. This phenomenon remain consistant aross the Cifar10, Cifar100 and SVHN datasets. The $\gamma$ in Assumption 3.4 can be assigned as $\gamma=\max_{t}\widetilde{\gamma}_{t}$. We can always find such $\gamma$ if the optimizer is not extreme overfitting. Under the extremely overfitting case, we can use the relaxed theorem in Appendix B to bound the generalization error.

The generalization error and the $\mathcal{C}(\mathbf{J_{t}})$ both changes as the training continues. Therefore, we want to verify whether they correlate with each other during the training process. Here, we use the term $\nabla F_{S^{\prime}}(\mathbf{J_{t}})-\nabla F_{S}(\mathbf{J_{t}})$ to approximate the generalization error. We find that $\nabla F_{S^{\prime}}(\mathbf{J_{t}})-\nabla F_{S}(\mathbf{J_{t}})$ has similar trend with $\mathcal{C}(\mathbf{J_{t}})$ (Figure 2 Center). What’s more, we also find that the curve of $\nabla F_{S}(\mathbf{J_{t}})+\mathcal{C}(\mathbf{J_{t}})$ exhibits a comparable pattern with the curve $F_{S^{\prime}}(\mathbf{J_{t}})$ (Figure 2 Left). To explore whether $\mathcal{C}(\mathbf{J_{t}})$ reveals influence of the change of $F_{S}(\mathbf{J_{t}})$ to the generalization error, we plot $\frac{d\mathcal{C}(\mathbf{J_{t}})}{dF_{S}(\mathbf{J_{t}})}$ (Figure 2 Right) during the training process. $\frac{d\mathcal{C}(\mathbf{J_{t}})}{dF_{S}(\mathbf{J_{t}})}$ increases slowly during the early stage of training, but surge rapidly afterward. This discovery is aligned with our intuition about the overfitting.

In Figure 3, we carry out experiments under various settings. Each data point in the figure represents the average of three repeated experiments. The results demonstrate that both the generalization error and $\mathcal{C}(\mathbf{J_{t}})$ increase as the level of label noise is raised (Figure 3 Left). The another experiments measure $\mathcal{C}(\mathbf{J_{t}})$ and generalization error for different learning rate and discover that $\mathcal{C}(\mathbf{J_{t}})$ can capture the trend generalization error. The reasons behind a larger learning rate resulting in a smaller generalization error have been explored in Li et al. [17], Barrett and Dherin [3]. Additionally, Appendix E discusses why a larger learning rate can lead to a smaller $\mathcal{C}(\mathbf{J_{t}})$.