Generalization Bound of Gradient Flow through Training Trajectory and Data-dependent Kernel

We thank the reviewer for the positive feedback, thoughtful comments, and for appreciating the novelty and value of the work! In response to the reviewer’s inquiries, we present the following detailed explanation.

---

Q1. The exact order of the upper bound.
A1. A naive bound of $\Gamma$ is $L\sqrt{T/n}$ using the Lipschitz property based on our assumption. Combining with the $\epsilon$ term, our bound is at most $O(\sqrt{T/n})$ in the worst case. However, when the model converges fast, $\Gamma$ may not increase with $\sqrt{T}$, and the $\epsilon$ can also achieve a faster rate of $O(n^{-3/4})$ when $T$ is small. Our Table 1 shows when $\Gamma$ will dominate our generalization bound. For some special cases, e.g., NTK regime, we refer to Section 6 for more details.

Q2. The bound is a post-hoc analysis tool, and cannot be used to predict generalization or guide the training process beforehand.

A2. In the case studies in Section 6, we can predict the generalization before the training. See Corollary 6.1, 6.2, 6.3. In these case studies, we have made additional assumptions regarding the dataset, optimizer, and training dynamics. Consequently, we can further simplify our generalization error bound and obtain a more precise analysis of the leading $\Gamma$ terms. If we do not have detailed assumptions about the learning dynamics, we may need to carry out the training process sufficiently to extract enough information to determine our generalization error bound. Potentially, a practical application of our framework is to train a model with a smaller size on a smaller dataset and then use its observed training trajectory to compute the LPK‐based generalization bound. This “mini‐experiment” provides an early, data‐driven estimate of how well the full‐scale model is likely to generalize without the computational cost of training a very large network.

By monitoring the evolution of our bound $\Gamma$ during training, we can predict the overfitting for overparameterized NNs and identify optimal stopping time for training without access to test data. Our bound can also be used as a regularization in the training. Moreover, our bound can also serve as a proxy to compare model architectures in Neural Architecture Search (NAS), e.g., [1, 2, 3, 4].

[1] Mellor, J., Turner, J., Storkey, A., and Crowley, E. J. Neural architecture search without training. ICML 2021.

[2] Chen, W., Gong, X., and Wang, Z. Neural architecture search on imagenet in four gpu hours: A theoretically inspired perspective. ICLR 2021.

[3] Mok, J., Na, B., Kim, J.-H., Han, D., and Yoon, S. Demystifying the neural tangent kernel from a practical perspective: Can it be trusted for neural architecture search without training? CVPR 2022.

[4] Chen, Y., Huang, W., Wang, H., Loh, C., Srivastava, A., Nguyen, L. M., and Weng, T.-W. Analyzing generalization of neural networks through loss path kernels. NeurIPS 2023.

Q3. Computing this term is very expensive as it requires the full gradient history, making it hard to use in practice.

A3. The computation of our bound is tractable. In the $\Gamma$ term (Eq. (3)), as Reviewer xevc, mentioned $L_S(w_0) - L_S(w_T)$ can be easily computed from the training loss. For $\sum_{i=1}^n \int_0^T ||\nabla_w \ell(w_t, z) ||^2 dt$, we can compute the per-sample gradients (which can be computed in parallel for all the samples) at each training step, calculate their 2-norm, and sum over the training steps. The computation cost of the bound is similar to that of gradient descent training of the model. For large datasets, we can estimate $\sum_{i=1}^n \int_0^T ||\nabla_w \ell(w_t, z) ||^2 dt$ with a subset of training samples, such as the batch data used at each training step in SGD. We will add a detailed discussion of the computational complexity of evaluating the LPK term $\Gamma$.

Q4. Do we have a way to bound the term $\int_0^T ||\nabla_w \ell(w_t, z)|| dt$?

A4. A naive bound of $\int_0^T ||\nabla_w \ell(w_t, z)|| dt$ is $L^2 T$ using the Lipschitz property. It is hard to give a tighter bound for general cases. But for more specific models, such as the case studies in Sec 6, we can derive tighter bounds for this quantity.

For the NTK case, in line 730, we bound it by $\sum_i^n \int_0^T ||\nabla_w \ell(w_t, z_i)|| dt \leq \frac{n \lambda_{max} \|\|f(w_0, X) - y\|\|^2 }{2 \lambda_{min}} ( 1 - e^{-\frac{2 \lambda_{min} T}{n}})$.

For the kernel regression, in line 744, we bound it by $\sum_i^n \int_0^T ||\nabla_w \ell(w_t, z_i)|| dt \leq \frac{K_{max} n}{2}||w_0 - w^*||^2$.

Q5. Do Gradient Descent and Gradient Flow have a difference in the order of the resulting bound?

A5. In the early phase of this paper, we computed the gradient flow and found that the difference between the gradient flow and gradient descent with a small learning rate is very small, as also shown in [5, 6]. Therefore, the bounds of gradient flow and GD should not differ much under certain conditions of the learning rate. We will present a simulation to control the difference between the GD and gradient flow dynamics for two-layer neural networks on a portion of the MNIST dataset.

Our experimental findings demonstrate a strong correlation between our bound and gradient descent, particularly when employing reasonable learning rates. In our simulation, two-layer neural networks (NNs) and ResNets were trained with learning rates of 0.01 and 0.001, respectively, which are widely adopted learning rate settings in practical applications. These experimental results also suggest the potential extension of our theory to discrete SGD with a constant learning rate, which we intend to explore further in future research endeavors. Furthermore, it is noteworthy that [5] indicates that deep neural networks with homogeneous activations can be approximated by gradient descent. Consequently, we can potentially transfer our analysis of gradient flow for neural networks to a gradient descent training framework based on this observation. We will elaborate on this limitation in the revised version of our paper.

[5] Continuous vs. Discrete Optimization of Deep Neural Networks. NeurIPS 2021.

[6] Chen, Y., Huang, W., Wang, H., Loh, C., Srivastava, A., Nguyen, L. M., and Weng, T.-W. Analyzing generalization of neural networks through loss path kernels. NeurIPS 2023.

---

We are happy to address any further questions and feedback.

We express our gratitude for the positive feedback, thoughtful review, and the reviewer’s appreciation for our work. In response to the reviewer’s suggestions and inquiries, we will incorporate a comparative analysis with the results of Chen et al. and provide a concise summary of the improvements made to our paper.

---

Q1. The Liptichitzness and smoothness assumptions are quite strong and make the bound computationally intractable, since the Lipschitz or smoothness parameters are hard to estimate in practice.

A1. Thanks for the question. However, we would like to note that the $\Gamma$ term in our bound does not depend on the Lipschitz or smoothness constant. As we have compared in Table 1, we conclude that in many cases, $\Gamma$ is the dominant term compared with $\epsilon$. In these cases, $\epsilon$ will be small compared with $\Gamma$ as $n\to\infty$. Hence, we can use the $\Gamma$ to represent our generalization bound. For large-scale simulations, our $\Gamma$ can be applied to bound the generalization gap, and it empirically tracks the actual generalization gap well.

To improve the results of Chen et al. and obtain a better LPK-type generalization bound, we need to include the Lipschitz and smoothness assumptions. Chen et al. do not require these assumptions because they did not perform any optimization analysis. Their bound calculates the supremum over an infinite set, which cannot be accurately estimated. Additionally, their bounds must be evaluated on (infinitely many) datasets different from the training set, which is impossible in practice. By explicitly enforcing Lipschitz and smoothness conditions, one transforms an intractable supremum‐based LPK guarantee into a finite-sample, data-dependent, and computationally verifiable bound. This modification not only tightens the theoretical rates but also aligns the analysis with the practical realities of deep‐network training. We will add a detailed comparison with Chen et al. in the revision of our paper.

Q2. Compare the tightness between the bounds of this paper and the previous one by Chen et al. Miss some comparison between these bounds and those ones derived solely based on stability or uniform convergence approaches.

A2. Our bound is clearly tighter than the previous one by Chen et al. Our bound has a rate of $O(\sqrt{T/n})$ in the worst case, while the bound in Chen et al can be as large as $O(\sqrt{T})$ since the $\sum_i \delta(z_i, z_j)$ term in their bound can be $O(L^2 T n^2)$.

Stability-based bounds are exponential with $T$ for non-convex loss with constant learning rate, as we showed in Lemma 4.3, while our bound is at most $O(\sqrt{T/n})$ in the worst case. Uniform convergence bounds such as [1, 2, 3, 4] listed below are usually based on the norm of the weight matrices. These bounds are independent of the algorithm and cannot capture the inductive bias of the learning algorithm. In contrast, our bound is data-dependent and algorithm-dependent. Empirically, our bound is also orders of magnitude tighter than [3], as shown in Fig. 2.

[1] Bartlett, P. L. and Mendelson, S. Rademacher and gaussian complexities: Risk bounds and structural results. JMLR 2002.

[2] Neyshabur, B., Tomioka, R., and Srebro, N. Norm-based capacity control in neural networks. COLT 2015.

[3] Bartlett, P. L., Foster, D. J., and Telgarsky, M. J. Spectrally-normalized margin bounds for neural networks. NeurIPS 2017.

[4] Neyshabur, B., Li, Z., Bhojanapalli, S., LeCun, Y., and Srebro, N. The role of over-parametrization in generalization of neural networks. ICLR 2019.

Q3. Impact of the big O terms on the tightness and computational tractability of these bounds? How are these bounds empirically compared to the previous one by Chen et al?

A3. As we compared in Table 1, in many cases, $\Gamma$ is the dominant term compared with $\epsilon$. In these cases, the big-O terms are small as $n$ increases. Hence, we can use $\Gamma$ to track the trend of the generalization error bound through different training dynamics.

Our bound is tighter than the previous one by Chen et al. As we mentioned above in A2, our bound has a rate of $O(\sqrt{T/n})$ in the worst case, while the bound in Chen et al can be as large as $O(\sqrt{T})$. Numerically, we will add an empirical estimate of the big-O terms for the MNIST dataset and small neural networks, and compare our bound with Chen et al in the revised version.

Q4. Can these bounds provide some new insights for the design of learning algorithms?

A4. Thank you for the insightful question. The primary motivation of this paper is to investigate the generalization capabilities of neural networks. It would be more insightful if our theory could provide practical insights for training neural networks. We concur with the reviewer’s observation that our theory indicates gradient clipping. Moreover, by monitoring the evolution of our bound during training, we can potentially predict and analyze the overfitting for overparameterized NNs and identify optimal stopping time for training without access to test data. We could potentially utilize a smaller model and a smaller dataset to train and approximate our generalization bound. Subsequently, we can employ this bound to predict the generalization capability of larger models. Our bound can also be used as a regularization in the training. Moreover, our bound can also serve as a proxy to compare model architectures in Neural Architecture Search (NAS), e.g., [5, 6, 7, 8].

[5] Mellor, J., Turner, J., Storkey, A., and Crowley, E. J. Neural architecture search without training. ICML 2021.

[6] Chen, W., Gong, X., and Wang, Z. Neural architecture search on imagenet in four gpu hours: A theoretically inspired perspective. ICLR 2021.

[7] Mok, J., Na, B., Kim, J.-H., Han, D., and Yoon, S. Demystifying the neural tangent kernel from a practical perspective: Can it be trusted for neural architecture search without training? CVPR 2022.

[8] Chen, Y., Huang, W., Wang, H., Loh, C., Srivastava, A., Nguyen, L. M., and Weng, T.-W. Analyzing generalization of neural networks through loss path kernels. NeurIPS 2023.

---

We are happy to address any further questions and feedback.

We appreciate the positive feedback, thoughtful review, and the reviewer’s appreciation for the work! We will add a limitation section, provide further discussion in the numerical experiments section, and address the points the reviewer mentioned in our revised paper. We also thank the reviewer for mentioning the paper, “A Generalization Bound for Nearly-Linear Networks”. We will have a discussion of this paper in our literature review.

---

Q1. The Lipschitzness and smoothness assumptions.

A1. Thanks for the question and comments. The Lipschitzness and smoothness assumptions are standard assumptions in optimization and generalization error bound analysis of neural networks [1, 2, 4]. There is also considerable empirical evidence that neural networks are Lipschitz, and many efforts have been made to estimate the Lipschitz constant (see e.g. [3, 4]). In practice, training with weight decay or early stopping confines the weights to a compact subset of parameter space, on which the loss is indeed (locally) Lipschitz.

For a two-layer network with smooth activation functions, as long as the norms of the weight matrices are bounded, there should be a finite Lipschitz constant. For the toy model the reviewer is concerned with, the Lipschitz constant could be large. But the Lipschitz constant will still be finite as long as the weight matrices are bounded and the activation function is Lipschitz. This is why in many previous generalization bound papers, the final bound is determined by the Frobenius norm of the weight matrices, e.g., [4, 5, 6, 7, 8]. However, this Lipschitz constant can be significantly large depending on the model’s size, although the trend is consistent with the true generalization gap. This is because, as the reviewer mentioned, the bound in the ResNet18 simulation is much larger.

We will add additional discussion about the limitations of Lipschitzness and smoothness assumptions in the conclusion section. In the simulations, we will also add an experiment to estimate the Lipschitz constant and weight matrix norm during training. This will enable us to ascertain that these terms do not exhibit explosive growth and are bounded by a constant. In the event that this is the case, it will provide empirical evidence of our simulation’s consistency with our theoretical framework. Furthermore, our simulations show that our $\Gamma$ in the theorems has already served as a good bound of the generalization gap across diverse architectures and datasets. Even when a network may not strictly satisfy the global Lipschitz or smoothness conditions, we find that our bound still tracks its behavior remarkably well. These results suggest that, with modest technical extensions, our theory could be broadened to cover a wider class of models. This is one of the reasons why various simulations on real-world data and various neural networks have been presented in this paper. We will conclude by outlining the challenges and potential strategies for relaxing these regularity assumptions in future work in the revision of our paper.

[1] Hardt, M., Recht, B., and Singer, Y. Train faster, generalize better: Stability of stochastic gradient descent. ICML 2016.
[2] Optimization Methods for Large-Scale Machine Learning.
[3] Efficient and Accurate Estimation of Lipschitz Constants for Deep Neural Networks. NeurIPS 2019.
[4] Generalization bounds of stochastic gradient descent for wide and deep neural networks. NeurIPS 2019.
[5] Bartlett, P. L. and Mendelson, S. Rademacher and gaussian complexities: Risk bounds and structural results. JMLR 2002.
[6] Neyshabur, B., Tomioka, R., and Srebro, N. Norm-based capacity control in neural networks. COLT 2015.
[7] Bartlett, P. L., Foster, D. J., and Telgarsky, M. J. Spectrally-normalized margin bounds for neural networks. NeurIPS 2017.
[8] Neyshabur, B., Li, Z., Bhojanapalli, S., LeCun, Y., and Srebro, N. The role of over-parametrization in generalization of neural networks. ICLR 2019.

Q2. An intuitive explanation why the epsilon term in your Theorem 5.2 is a minimum between two terms, not a maximum.

A2. The epsilon term is the minimum between the two terms because the two bounds both hold at the same time, and we can take the minimum one. The mechanism is different from what the reviewer mentioned. Here we have two different methods for the epsilon term. So we take the minimum to get a sharper bound.

Q3. Have you tried to compute your bound up to a number? How close is it to gamma on your experiments? Does it also track the actual generalization gap well?

A3. Thanks for your suggestion! We have not tried to do that since estimating the Lipschitz/smoothness constant, and the constant in the big-O terms is hard. In the revision, we will numerically estimate our bound, including the $\epsilon$ term for a two-layer neural network during training to see how it tracks or affects the generalization gap.

However, as we have compared in Table 1, we want to emphasize that in many cases, $\Gamma$ is the dominant term compared with $\epsilon$. In these cases, $\epsilon$ will be small compared with $\Gamma$ as $n\to\infty$. Hence, we can use the $\Gamma$ to represent our generalization bound. For large-scale simulations, our $\Gamma$ can be applied to bound the generalization gap, and it empirically tracks the actual generalization gap well. Also, for future work, we hope to develop a new method to estimate our LPK term so that it could be useful for predicting the optimal early stopping time.

Q4. Missing related work.

A3. Thanks for the related work! We will add and discuss the paper in the revised version.

---

We are happy to address any further questions and feedback.

Thanks for your follow-up question! We really appreciate you taking the time to review our paper.

We apologize for the ambiguous references in the rebuttal. You are correct that [3] indeed estimates the Lipschitz constant of the input. [4] does not estimate the Lipschitz constant but rather derives a generalization bound using the Lipschitz condition of the loss function and activation function. On the other hand, we want to mention [9], which estimates the upper bound of the Lipschitz constant of NN parameters. They show the Lipschitz constant is upper-bounded as long as the parameters and input are bounded.

In general, it could be hard to guarantee the global Lipschitzness and boundedness of parameters in general settings. But for some specific setting, such as the NTK case, the parameters can be bounded in a ball [10, 11, 12]. Additionally, we can show the boundedness of the parameter space under the gradient flow with an $\ell_2$ weight decay. Suppose that we consider loss function $J(w)=L_S(w)+\frac{\lambda}{2}\|\|w\|\|^2$, and gradient flow $\dot w(t)=-\nabla J(w(t)),\quad w(0)=w_0.$ We can check that $J(w(t))\le J(w_0)$ for all $t\ge0$ because of $\frac{d}{dt}J(w(t))=\langle\nabla J(w),\dot w\rangle=-\|\|\nabla J(w(t))\|\|^2\le0$. Since $L_S(w)\ge0$, $\frac{\lambda}{2}\|\|w(t)\|\|^2\le J(w(t))\le J(w_0)=L_S(w_0)+\frac{\lambda}{2}\|\|w_0\|\|^2.$ Hence, we can show that $\||w(t)\||\le\sqrt{\||w_0\||^2+\frac{2}{\lambda}L_S(w_0)}\equiv R$ which depends only on $w_0$, $\lambda$, and $L_S(w_0)$ (the initial loss). $R$ is bounded whenever the loss is bounded.

In deep learning, there are many other ways to potentially guarantee the global Lipschitzness and boundedness of training parameters, such as gradient clipping, projected GD [13], spectral normalization [14], and explicitly parametrizing weight matrices as orthogonal matrices [15]. We leave the extension of our assumption to local Lipschitzness as future work.

We will add a detailed discussion on the Lipschitzness and smoothness assumptions in our limitations section. Thank you for your thoughtful suggestions and constructive feedback!

---

[9] Local Lipschitz Bounds of Deep Neural Networks.

[10] Gradient Descent Provably Optimizes Over-parameterized Neural Networks. ICLR 2019.

[11] Gradient Descent Finds Global Minima of Deep Neural Networks. ICML 2019.

[12] A Convergence Theory for Deep Learning via Over-Parameterization.ICML 2019.

[13] Stability of Stochastic Gradient Descent on Nonsmooth Convex Losses. NeurIPS 2020.

[14] Optimization and Generalization Guarantees for Weight Normalization

[15] Full-Capacity Unitary Recurrent Neural Networks. NIPS 2016.

We thank the reviewer for the positive feedback, thoughtful review, and for appreciating the merits of the work!

---

Q1. The theoretical analysis is limited to full-batch gradient flow with infinitesimal step size, which does not align with practical stochastic optimization scenarios in deep learning.

A1. Although our analysis is based on gradient flow, our bound correlates well with the true generalization gap even for constant learning rates, as shown in our experiments. Also, we have the theory of stochastic gradient flow in section 5.2, which extends full-batch gradient flow to mini-batch stochastic gradient flow. In our simulation, the two-layer neural networks (NNs) and ResNets are trained with learning rates of 0.01 and 0.001, which are common learning rate setups in practice. These experimental results also show that our theory can be potentially extended to discrete SGD with a constant learning rate. We will leave this to future research work. We want to emphasize that using gradient flow to do theoretical analysis is easier and does not lose the general properties of the training dynamics for NNs. The gradient flow dynamics provides a clean, continuous-time model that captures the core features of neural-network training without sacrificing generality. Many qualitative phenomena observed under SGD (e.g., phases of feature learning, generalization plateaux, late overfitting) have been shown to persist in the gradient-flow regime. Moreover, even the training dynamic of gradient flow for two-layer NNs is still underexplored and an ongoing research direction currently, e.g., [43] ("Dynamical decoupling of generalization and overfitting in large two-layer networks") in our reference. Additionally, we want to mention that [1] shows that deep neural networks with homogeneous activations, gradient flow trajectories can be approximated by gradient descent, so we can potentially transfer our analysis of gradient flow for neural networks into a gradient descent training based on this result. We will have a further discussion of this limitation in the revision of our paper.

[1] Continuous vs. Discrete Optimization of Deep Neural Networks. NeurIPS 2021.

Q2. Computing LPK in practice may be computationally demanding.

A2. The computation of our bound is tractable. In the $\Gamma$ term (Eq. (3)), $L_S(w_0) - L_S(w_T)$ can be easily computed from the training loss. For $\sum_{i=1}^n \int_0^T ||\nabla_w \ell(w_t, z) ||^2 dt$, we can compute the per-sample gradients (which can be computed in parallel for all the samples) at each training step, calculate their 2-norm, and sum over the training steps. The computation cost of the bound is similar to that of gradient descent training of the model. For large datasets, we can estimate $\sum_{i=1}^n \int_0^T ||\nabla_w \ell(w_t, z) ||^2 dt$ with a subset of training samples, such as the batch data used at each training step in SGD. We will add a detailed discussion of the computational complexity of evaluating the LPK term $\Gamma$.

Q3. When training time T is large, the asymptotic generalization bound aligns with standard bounds via Radamacher complexity analysis and poorer than that via covering number analysis.

A3. Standard bounds via Radamacher complexity and covering number analysis are usually based on the norm of the weight matrices [2, 3, 4, 5]. These bounds are independent of the algorithm and cannot capture the inductive bias of the learning algorithm. In contrast, our bound is data-dependent and algorithm-dependent. Empirically, our bound is also orders of magnitude tighter than [4], as shown in Fig. 2.

We would appreciate it if the reviewer could clarify which Rademacher complexity bound and covering number bound they are referring to.

[2] Bartlett, P. L. and Mendelson, S. Rademacher and gaussian complexities: Risk bounds and structural results. JMLR 2002.
[3] Neyshabur, B., Tomioka, R., and Srebro, N. Norm-based capacity control in neural networks. COLT 2015.
[4] Bartlett, P. L., Foster, D. J., and Telgarsky, M. J. Spectrally-normalized margin bounds for neural networks. NeurIPS 2017.
[5] Neyshabur, B., Li, Z., Bhojanapalli, S., LeCun, Y., and Srebro, N. The role of over-parametrization in generalization of neural networks. ICLR 2019.

Q4. How robust is the LPK-based generalization bound to architectural variations?

A4. Thanks for the question! We will test our bounds for transformers or larger language models later. Vision transformers (ViT) satisfy our theory. For language models, we can treat each sentence as a data point $x_i$ and its corresponding response as a label $y_i$. Our analysis does not require the output dimension of the model to be one dimension. The same bound will hold for the language model and ViT. As a remark, our LPK-based generalization bound exhibits intricate and multifaceted dependencies on the dataset, training dynamics, and the architecture. In Figures 1 and 4 (Appendix A), we compare the predicted bound against the actual generalization gap for ResNet-18 versus ResNet-34 on CIFAR-10. Despite the difference in depth and width, the bound tracks the observed gap closely for both architectures, demonstrating its quantitative stability under architectural scaling.

Q5. Extend the results to discrete-time gradient descent or SGD with fixed learning rate

A5. Our results could be extended to discrete-time GD and SGD using an idea similar to [6], by considering the taylor expansion of GD at each time step: $\ell(w_{t+1}, z) - \ell(w_{t}, z) = \nabla_w \ell(w_{t}, z)^\top (w_{t+1} - w_{t}) + O(||w_{t+1} - w_{t}||)$, where $\nabla_w \ell(w_{t}, z)^\top (w_{t+1} - w_{t})$ can be analyzed like a linear function class and $O(||w_{t+1} - w_{t}||)$ can be controlled with learning rate. We defer this to future work, as it requires additional effort and lies beyond the scope of the present study.

[6] Learning Trajectories are Generalization Indicators. NeurIPS 2023.

Q6. Computational complexity of evaluating the LPK term

A6. As we replied in Q2, the computational complexity of computing our bound is similar to the auto-gradient in GD training of the model. In our simulation, it won’t cost too much since we can compute the gradient at each step of the GD training and compute the LPK term along the training process. We will have a detailed discussion of the computational complexity in our revised paper.

Q7. How sensitive is the bound to initialization?

A7. Thanks for the insightful question! The choice of initialization could indeed influence the training trajectory and the generalization bound. Our analysis assumes a fixed random initialization. The final bound does not depend directly on the initialization, but rather indirectly through the training trajectory. It would be very interesting to consider the influence of the initialization on the bound, which may require more involved analysis of optimization. Intuitively, the initialization close to the stationary points may induce a smaller bound and better generalization ability.

In our case study of feature learning in Sec 6.3, we show that a mean-field type initialization enables feature learning and can obtain better sample complexity compared with NTK initialization from our generalization bound. This may indicate how initialization affects the generalization error. Our generalization bound captures this effect.

---

We are happy to address any further questions and feedback.