SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Theorem 3.  $\hat{\mathcal{R}}_S(\mathcal{G}_T)\leq \min \{U_1,U_2\}$  .Here
$$
U _ {1} = \frac {B}{n} \sqrt {\sup  _ {\mathsf {K} (\cdot , \cdot ; \mathcal {S} ^ {\prime}) \in \mathcal {K} _ {T}} \operatorname {T r} \left(\mathsf {K} \left(\mathbf {Z} , \mathbf {Z} ; \mathcal {S} ^ {\prime}\right)\right) + \sum_ {i \neq j} \Delta \left(z _ {i} , z _ {j}\right)},
$$
$$
U _ {2} = \inf  _ {\epsilon > 0} \left(\frac {\epsilon}{n} + \sqrt {\frac {2 \ln \mathcal {N} (\mathcal {G} _ {T} ^ {S} , \epsilon , \| | | _ {1})}{n}}\right),
$$
where  $\mathcal{G}_T^S = \{g(\mathbf{Z}) = (g(z_1),\dots ,g(z_n)):g\in \mathcal{G}_T\}$ $\mathcal{N}(\mathcal{G}_T^S,\epsilon ,\| \cdot \| _1)$  is the covering number of  $\mathcal{G}_T^S$  with the  $\ell_1$  -norm and
$$
\Delta (\boldsymbol {z} _ {i}, \boldsymbol {z} _ {j}) = \frac {1}{2} \left[ \sup _ {\mathsf {K} (\cdot , \cdot ; \mathcal {S} ^ {\prime}) \in \mathcal {K} _ {T}} \mathsf {K} (\boldsymbol {z} _ {i}, \boldsymbol {z} _ {j}; \mathcal {S} ^ {\prime}) - \inf _ {\mathsf {K} (\cdot , \cdot ; \mathcal {S} ^ {\prime}) \in \mathcal {K} _ {T}} \mathsf {K} (\boldsymbol {z} _ {i}, \boldsymbol {z} _ {j}; \mathcal {S} ^ {\prime}) \right].
$$
The term  $U_{1}$  is composed by two components  $\sup_{\mathsf{K}(\cdot ,\cdot ;\mathcal{S}^{\prime})\in \mathcal{K}_{T}}\mathrm{Tr}(\mathsf{K}(\mathbf{Z},\mathbf{Z};\mathcal{S}^{\prime}))$  and  $\Delta (z_i,z_j)$ . The first component, according to the definition of LPK, quantifies the maximum magnitude of the loss gradient in  $\kappa_T$  evaluated with the set  $s$  throughout the training trajectory. The second component assesses the range of variation of LPK within the set  $\kappa_T$ . The term  $U_{2}$  is obtained from analyzing the covering number of  $\mathcal{G}_T$ . It shows that if the variation of the loss dynamics of gradient flow with different training data is small, then the complexity of  $\mathcal{G}_T$  will also be small. The norm constraint  $\frac{1}{n^2}\sum_{i,j}\mathsf{K}_T(z_i',z_j';\mathcal{S}')\leq B^2$  balances a tradeoff between the tightness of the bound and the expressiveness of the set  $\mathcal{G}_T$  (the number of datasets covered). Combining these two bounds with Theorem 1, we obtain the following generalization bound.
Corollary 1 (Generalization bound for NN). Fix  $B > 0$ . Let  $\hat{\mathcal{R}}_{\mathcal{S}}^{gf}(\mathcal{G}_T) = \min(U_1, U_2)$  where  $U_1$  and  $U_2$  are defined in Theorem 3. For any  $\delta \in (0,1)$ , with probability at least  $1 - \delta$  over the draw of an i.i.d. sample set  $\mathcal{S} = \{\pmb{z}_i\}_{i=1}^n$ , if  $\frac{1}{n^2} \sum_{i,j} \mathsf{K}_T(\pmb{z}_i, \pmb{z}_j; \mathcal{S}) \leq B^2$ , the following holds for  $\ell(\pmb{w}_T, \pmb{z})$  that trained from  $\mathcal{S}$ ,
$$
L _ {\mu} (A _ {T} (\mathcal {S})) - L _ {S} (A _ {T} (\mathcal {S})) \leq 2 \hat {\mathcal {R}} _ {\mathcal {S}} ^ {g f} (\mathcal {G} _ {T}) + 3 \sqrt {\frac {\log (2 / \delta)}{2 n}},
$$
where  $\mathbf{w}_T = A_T(\mathcal{S})$  is the output from the gradient flow (1) at time  $T$  by using  $\mathcal{S}$  as input.
Our result owns many compelling properties.
- First, our bound holds in a general setting as it does not hedge on a special NN architecture. In contrast, existing works [2, 11] only consider fully connected NNs and require NN to be ultra-wide.
- Our bound depends on the data distribution through the quantities in  $U_{1}$  and  $U_{2}$ . This property not only significantly tightens our bound but can also help explain some empirical observations of NNs. For example, different from classical generalization theory, e.g. VC dimension, our complexity bounds depend on the labels directly, which helps explain the random label phenomenon [59] as shown in Figure 3 in Sec. 7.
- Our experiments in Sec. 7 (Figure 2) demonstrate the tightness of the generalization bound. Intuitively, our bound is tight because (1) instead of considering the entire hypothesis class, we focus on the subset of interest characterized by running gradient flow from a starting point  $w_0$ ; (2) we get the bound from an equivalence between NNs and general kernel machines, whose generalization bounds are tighter. Finally, we compare our generalization bound with two existing NTK-based bounds in Table 1.
5 Stochastic Gradient Flow
In the previous section, we derived a generalization bound for NNs trained from full-batch gradient flow. Here we extend our analysis to stochastic gradient flow and derive a corresponding generalization bound. To start with, we recall the dynamics of stochastic gradient flow (SGD with infinitesimal step size). Let  $S_{t} \subseteq \{1, \dots, n\}$  be the indices of batch data used in time interval  $[t, t + 1]$  and  $|S_{t}| = m$  be the batch size. We establish a new connection between the loss dynamics of stochastic gradient flow and a general kernel machine. Then we investigate the complexity of the collection of such kernel machines that can be induced by various training sets.
Theorem 4. Suppose  $\mathbf{w}(T) = \mathbf{w}_T$  is a solution of stochastic gradient flow at time  $T \in \mathbb{N}$  with initialization  $\mathbf{w}(0) = \mathbf{w}_0$ . Then for any  $z \in \mathcal{Z}$ ,
$$
\ell (\boldsymbol {w} _ {T}, \boldsymbol {z}) = \sum_ {t = 0} ^ {T - 1} \sum_ {i \in \mathcal {S} _ {t}} - \frac {1}{m} \mathrm {K} _ {t, t + 1} (\boldsymbol {z}, \boldsymbol {z} _ {i}; \mathcal {S}) + \ell (\boldsymbol {w} _ {0}, \boldsymbol {z}),
$$
where  $\mathsf{K}_{t,t + 1}(\pmb {z},\pmb {z}_i;\mathcal{S}) = \int_t^{t + 1}\bar{\mathsf{K}} (\pmb {w}(t);\pmb {z},\pmb {z}_i)dt$  with  $\bar{\mathsf{K}}$  defined in Definition 3.
The above theorem shows that the loss of the NN in stochastic gradient flow dynamics can be characterized by a sum of general kernel machines. In particular, when we use the full batch at each time interval (i.e.,  $m = n$ ), the above result recovers Theorem 2. To study its generalization behavior, we introduce the class of kernel machines induced by different training sets  $\mathcal{S}' \in \mathrm{supp}(\mu^{\otimes n})$  with constrained RKHS norms. Specifically, given  $B_{t} > 0$  for  $t = 0,\dots ,T - 1$ , we define
$$
\mathcal {K} _ {T} = \left\{\left(\mathsf {K} _ {0, 1} (\cdot , \cdot ; \mathcal {S} ^ {\prime}), \dots , \mathsf {K} _ {T - 1, T} (\cdot , \cdot ; \mathcal {S} ^ {\prime})\right): \mathcal {S} ^ {\prime} \in \operatorname {s u p p} \left(\mu^ {\otimes n}\right), \frac {1}{m ^ {2}} \sum_ {i, j \in \mathcal {S} _ {t}} \mathsf {K} _ {t, t + 1} \left(\boldsymbol {z} _ {i} ^ {\prime}, \boldsymbol {z} _ {j} ^ {\prime}; \mathcal {S} ^ {\prime}\right) \leq B _ {t} ^ {2} \right\}.
$$
Note this set includes the kernel induced by the training set  $S$  if it satisfies the constraints. Then  $\ell(\boldsymbol{w}_T, \boldsymbol{z})$  trained from all feasible  $S' \in \mathrm{supp}(\mu^{\otimes n})$  form a function class
$$
\mathcal {G} _ {T} \triangleq \left\{\sum_ {t = 0} ^ {T - 1} \sum_ {i \in \mathcal {S} _ {t}} - \frac {1}{m} \mathrm {K} _ {t, t + 1} \left(\boldsymbol {z}, \boldsymbol {z} _ {i} ^ {\prime}; \mathcal {S} ^ {\prime}\right) + \ell \left(\boldsymbol {w} _ {0}, \boldsymbol {z}\right): \mathrm {K} (\cdot , \cdot ; \mathcal {S} ^ {\prime}) \in \mathcal {K} _ {T} \right\}. \tag {3}
$$
Next, we upper bound the Rademacher complexity of the function class  $\mathcal{G}_T$ . This bound can naturally translate into a generalization bound by equipping with Theorem 1.
Theorem 5. The Rademacher complexity of  $\mathcal{G}_T$  defined in (3) has an upper bound:
$$
\hat {\mathcal {R}} _ {\mathcal {S}} (\mathcal {G} _ {T}) \leq \sum_ {t = 0} ^ {T - 1} \frac {B _ {t}}{n} \sqrt {\sup _ {\mathrm {K} (\cdot , ; \mathcal {S} ^ {\prime}) \in \mathcal {K} _ {T}} \operatorname {T r} \left(\mathrm {K} _ {t , t + 1} (\mathbf {Z} , \mathbf {Z}) ; \mathcal {S} ^ {\prime}\right) + \sum_ {i \neq j} \Delta_ {t} (z _ {i} , z _ {j})}.
$$
where  $\Delta_t(\pmb{z}_i, \pmb{z}_j) = \frac{1}{2} \left[ \sup_{\mathsf{K}(\cdot, \cdot; \mathcal{S}') \in \mathcal{K}_T} \mathsf{K}_{t,t+1}(\pmb{z}_i, \pmb{z}_j; \mathcal{S}') - \inf_{\mathsf{K}(\cdot, \cdot; \mathcal{S}') \in \mathcal{K}_T} \mathsf{K}_{t,t+1}(\pmb{z}_i, \pmb{z}_j; \mathcal{S}') \right]$ .
We assumed that mini-batches indices  $S_{t}$  are chosen before training. However, our analysis can be extended to accommodate random mini-batch selections of any sampling strategy by enumerating all potential  $S_{t}$  in  $\mathcal{K}_T$ .
6 Case Study & Use Case
In the previous sections, we derived generalization bounds for NNs trained with (stochastic) gradient flow. These bounds may initially appear complex due to their dependence on the training process. Here we show that these bounds can be significantly simplified by applying them to infinite-width NNs (and stable algorithms in Appendix E.2, norm-constraint NNs in Appendix E.3). Moreover, we demonstrate that our generalization bounds maintain a high correlation with the true generalization error. As a result, we use them to guide the design of NAS, and our experimental results demonstrate that this approach has a favorable performance compared with state-of-the-art algorithms.
6.1 Infinite-width NN
In this subsection, we consider a special case of infinite-width NNs trained by gradient flow and derive pre-computed generalization bounds. We focus on gradient flow to simplify the presentations but our results can be directly extended to stochastic gradient flow. For an infinite-width NN, under certain conditions, the neural tangent kernel keeps unchanged during training:  $\hat{\Theta}(\boldsymbol{w}_t; \boldsymbol{x}, \boldsymbol{x}') \to \Theta(\boldsymbol{x}, \boldsymbol{x}') \cdot \mathbf{I}_k$ . Consider a  $\rho$ -Lipschitz loss function, i.e.  $\| \nabla_f \ell(\boldsymbol{w}, \boldsymbol{z}) \| \leq \rho$ . The Rademacher complexity in Theorem 3 can be bounded by  $\hat{\mathcal{R}}_S(\mathcal{G}_T) \leq U_\infty$ , where
$$
U _ {\infty} = \frac {\rho B \sqrt {T}}{n} \sqrt {\sum_ {i , j} | \Theta (\boldsymbol {x} _ {i} , \boldsymbol {x} _ {j}) |}. \tag {4}
$$
In this infinite-width regime, our bound has no dependence on the initialization  $\pmb{w}_0$  since the NTK converges to a deterministic limit and has no dependence on the parameters. That means the bound holds for all possible  $\pmb{w}_0$  of infinite-width NNs trained by gradient flow from initialization. Compared with the bound  $\tilde{O}(L \cdot \sqrt{\frac{\mathbf{Y}^\top (\Theta)^{-1} \mathbf{Y}}{n}})$  in [11],  $U_{\infty}$  has several advantages: (1) it has no dependence on the number of layers  $L$ ; (2) it holds for NNs with multiple outputs.
6.2 Correlation Analysis and NAS
As a practical application, we apply our generalization bounds to guide the design of NAS. We first introduce a quantity  $U_{\mathrm{sgd}}$  simplified from the bound in Theorem 5, defined as
$$
U _ {\mathrm {s g d}} = \sum_ {t = 0} ^ {T - 1} \frac {1}{n} \sqrt {\frac {1}{m ^ {2}} \sum_ {i , j \in S _ {t}} K _ {t , t + 1} \left(\boldsymbol {z} _ {i} , \boldsymbol {z} _ {j} ; \mathcal {S}\right)} \sqrt {\operatorname {T r} \left(\mathrm {K} _ {t , t + 1} (\mathbf {Z} , \mathbf {Z}) ; \mathcal {S}\right)}.
$$
$U_{\mathrm{sgd}}$  can be computed along with training a NN via SGD on a training set  $S$ . Combining it with the training loss, we define the following quantity as an estimate of the population loss:
$$
\operatorname {G e n e} (\boldsymbol {w}, \mathcal {S}) = L _ {\mathcal {S}} (\boldsymbol {w}) + 2 U _ {\mathrm {s g d}}. \tag {5}
$$
We analyze the correlation between  $\mathrm{Gene}(\boldsymbol {w},\mathcal{S})$  and the true generalization error by randomly sampling 100 NN architectures from NAS-Bench-201 [19]. For each, we compute both  $\mathrm{Gene}(\boldsymbol {w},\mathcal{S})$  and the true generalization error. Since solving the gradient flow ODE is computationally infeasible for the large NNs in NAS-Bench-201, we apply a trapezoidal rule to approximate the integration in LPK  $K_{t,t + 1}$ . This approximation enables us to compute  $U_{\mathrm{sgd}}$  efficiently. Figure 1 demonstrates the correlation between  $\mathrm{Gene}(\boldsymbol {w},\mathcal{S})$  and the test error. The left figures plot the test error at epoch 1 or 2 against  $\mathrm{Gene}(\boldsymbol {w},\mathcal{S})$  of the respective epochs, showing a strong positive correlation between them. The right figures plot the test error at convergence against  $\mathrm{Gene}(\boldsymbol {w},\mathcal{S})$  at epoch 1 or 2, which also demonstrate a positive correlation. The outlier is caused by some architecture with large loss gradients. This experiment shows that  $\mathrm{Gene}(\boldsymbol {w},\mathcal{S})$  at the initial training stage can predict the performance of NNs at convergence. Based
Figure 1: Correlation between  $\mathrm{Gene}(\boldsymbol {w},\mathcal{S})$  and the test error on CIFAR-100 at epoch 1 and epoch 2. Kendall's tau shows they have a strong positive correlation.
on this observation, we use  $\mathrm{Gene}(\vec{w},\mathcal{S})$  as a metric in NAS for selecting architectures at the initial training stage (see Table 2). This approach significantly reduces computational costs compared with training-based NAS algorithms [47, 34, 18, 33].
Figure 2: Experiment (I). (a) shows the dynamics of logistic loss for 5 randomly selected training samples for NN trained by gradient flow (NN GF), NN trained by GD (NN GD), and the equivalent general kernel machine (KM) in Theorem 2. The dynamics of NN GF and KM overlap and thus verify the equivalence in Theorem 2. The dynamics of NN GF and NN GD are consistently close throughout the training process. (b) shows NN GF's training loss, test loss, test error, and upper bound for  $L_{\mu}(\pmb{w}_T)$  in Corollary 1. (c) shows that the complexity bound  $\hat{\mathcal{R}}_S^{gf}(\mathcal{G}_T)$  in Corollary 1 captures the generalization gap  $L_{\mu}(\pmb{w}_T) - L_S(\pmb{w}_T)$  well. It first increases and then converges after sufficient training time.
7 Numerical Experiments
We conduct comprehensive numerical experiments to demonstrate our generalization bounds. We observe that our complexity bounds are tight with respect to the generalization gap and can capture how noisy label influences the generalization behaviors of NNs. Moreover, we apply  $\mathrm{Gene}(\boldsymbol{w},\mathcal{S})$  in (5) to NAS and demonstrate favorable performance compared with state-of-the-art algorithms.
(I) Generalization bound in Corollary 1. In Figure 2 (more detailed in Figure A.4), we use a logistic loss to train a two-layer NN with 100 hidden nodes for binary classification on MNIST 1 and 7 [31] by full-batch gradient flow and compute its generalization bound. Due to the computational cost of solving the gradient flow ODE and computing the kernel, we only train and calculate the bound on  $n = 1000$  training samples. The bound would be tighter with more training samples. The NN is initialized using the NTK parameterization [28]. We use the Softplus activation function, defined as  $\mathrm{Softplus}(x) = \frac{1}{\beta} \ln(1 + e^{\beta x})$ . This function is continuously differentiable and serves as a smooth approximation to the ReLU activation function. In our experiments, we set  $\beta = 10$ . To train the NN via gradient flow, we solve the gradient flow ODE given by (1) to decide the NN parameters. For the equivalent general KM, we compute the LTK using the NN parameters and integrate it to get the LPK  $\mathcal{K}_T$ . These ODEs are computed with torchdiffeq [12]. To estimate the generalization bound, we train the NN on (20, 50, 100) independently sampled training sets  $S'$  to estimate the  $\mathcal{K}_T$  and  $\mathcal{G}_T$ , and the supremum in the bound  $U_1$  is estimated by taking the maximum over the finite set  $\mathcal{K}_T$ . For  $U_2$ , we compute an upper bound of it by setting  $\epsilon$  as the largest  $\ell_1$  distance between any two  $g(\mathbf{Z}) \in \mathcal{G}_T^S$  and  $\mathcal{N}(\mathcal{G}_T^S, \epsilon, \| \cdot \|_1) = 1$  because in this case any  $g(\mathbf{Z}) \in \mathcal{G}_T^S$  will satisfy as an  $\epsilon$ -cover. We run each experiment five times and plot the mean and standard deviation. The numerical experiments demonstrate that our complexity bound is tight and can capture the generalization gap well. As the number of  $S'$  increases, our bound converges to the true supremum value. To estimate the true supremum value, we apply the extreme value theory in Appendix A.3 and show the gap between the finite maximum and supremum is small, validating using a finite maximum as an estimate for our bound.
We train the NN using GD with a finite learning rate  $\eta = 10$  to compare with the NN trained by gradient flow. The training time  $t = \eta \times$  training steps. In Figure 2 (a) and Figure A.4 (a)(b), we observe that the loss of the NN trained by GD and that trained by gradient flow are consistently close throughout the entire training process. Consequently, while we established upper bounds for NNs trained by gradient flow, these results can serve as an (approximate) upper bound for NNs trained by GD with a finite learning rate.
Notably, for the NN in this experiment, the VC dimension [7] is 55957.3, the norm-based bound in [6] is 140.7 at  $T = 1000$ , and the NTK-based bound for an ultra-wide NN in [11] is 1.44, which are all vacuous (larger than 1), while our bound is tight (0.03 at  $T = 1000$ ). See a detailed comparison in Appendix A.2. We also conduct an experiment of three-layer NN (3072-100-100-1) trained on binary CIFAR-10 (cat and dog) [30] in Figure A.5, where there is a larger generalization gap.
Figure 3: Experiment (II). Generalization bound with label noise at  $T = 20000$
Table 2: Experiment (III). Comparison with state-of-the-art training-free/minimum-training NAS methods on NAS-Bench-201. Test accuracy with mean and deviation are reported. "Best" is the best accuracy over the four runs. "Optimal" indicates the best test accuracy achievable in NAS-Bench-201 search space. RS: randomly sample 100 architectures and select the one with the best metric value.
Algorithm	CIFAR-10 Accuracy	Best	CIFAR-100 Accuracy	Best
Baselines
TENAS [13]	93.08±0.15	93.25	70.37±2.40	73.16
RS + LGA3 [39]	93.64	69.77
Ours
RS + Gene(w,S)1	93.68±0.12	93.84	72.02±1.43	73.15
RS + Gene(w,S)2	93.79±0.18	94.02	72.76±0.33	73.15
Optimal	94.37	73.51
(II) Generalization bound with label noise. The settings are similar to Experiment (I) except we corrupt the labels with different portions of noise and calculate the bound after training NN until  $T = 20000$ . We estimate the bound with 20 training sets  $S'$ . The results in Figure 3 show that our bound has a strong correlation with the generalization gap and increases with the portion of label noise. Unlike classical generalization theory, e.g. VC dimension, our generalization bound can help explain the random label phenomenon [59].
(III) Neural architecture search (NAS). We apply  $\mathrm{Gene}(\boldsymbol{w},\mathcal{S})$  in Eq. (5) to guide the design of NAS. The results are shown in Table 2. We use a simple random search (RS) with  $\mathrm{Gene}(\boldsymbol{w},\mathcal{S})$ , where 100 architectures are sampled from the search space for evaluation, and the architecture with smallest  $\mathrm{Gene}(\boldsymbol{w},\mathcal{S})$  is selected.  $U_{\mathrm{sgd}}$  is estimated with a batch of data of size 600. Build upon Sec. 6.2, we apply  $\mathrm{Gene}(\boldsymbol{w},\mathcal{S})_1$  and  $\mathrm{Gene}(\boldsymbol{w},\mathcal{S})_2$  ( $\mathrm{Gene}(\boldsymbol{w},\mathcal{S})$  after training 1 and 2 epochs) to select NN architectures at the initial training stage in order to reduce computational costs. We compare our method with state-of-the-art training-free/minimum-training NAS algorithms [13, 39]. We run the experiments four times with different random seeds and report the mean and standard deviation. We reproduce the results in Chen et al. [13] using their released code and directly adopt the results reported in Mok et al. [39] as they did not release the code. The results show our approach of  $\mathrm{RS} + \mathrm{Gene}(\boldsymbol{w},\mathcal{S})$  can achieve favorable performance compared with state-of-the-art training-free/minimum-training NAS algorithms.
8 Conclusion and Future Work
In this paper, we establish a new connection between the loss dynamics of (stochastic) gradient flow and a general kernel machine. Building upon this result, we introduce generalization bounds for NNs trained from (stochastic) gradient flow. Our bounds hold for any continuously differentiable NN architectures (both finite-width and ultra-wide) and are generally tighter than existing bounds. Moreover, for infinite-width NNs, we obtain a pre-computed generalization bound for the whole training process. Finally, we apply our results to NAS and demonstrate favorable performance compared with state-of-the-art NAS algorithms.
There are several directions for future research. First, evaluating our generation bounds relies on the loss gradient, which may contain private and sensitive information. One potential fix would be accessing such information in a differentially private manner and it would be interesting to investigate how this "noisy" observation of gradient information influences our generalization bounds. Second, it is worth exploring how other optimization algorithms and different model architectures influence the generalization bounds. Finally, our bounds provide worst-case guarantees to the generalization of NNs and it would be interesting to extend our results to obtain expected bounds for further sharpening the results.
9 Acknowledgement
We thank the anonymous reviewers for valuable suggestions to improve the paper. We also thank the San Diego Supercomputer Center and the MIT-IBM Watson AI Lab for computing resources. Y. Chen and T.-W. Weng are supported by National Science Foundation under Grant No. 2107189 and 2313105.