Generalization Guarantees of Gradient Descent for Multi-Layer Neural Networks

** Q1. One of the drawback of this analysis is that since the network width $m$ is usually chosen first in practice, this result (such as theorem 2,3) basically says the training can't be too long and $\eta T$ need to be upper bounded by some functions over the network width. Also, this work prove convergence via lower bounding the minimum eigenvalue of the Hessian and uses the Hessian at the initialization as a reference point. This limits how far the network weights can travel from the initialization value. For $L^2$ loss, I think this is OK but for losses like cross-entropy where the network achieve zero-loss at infinity, it is not clear how far this analysis can be generalized.

Thanks for your thoughtful comment. In this work, we answer a theoretical question of how wide the network is required for GD to achieve the desired excess risk bound $O({1 / \sqrt{n}})$. Our results also shed light on how to choose the suitable $\eta T$ and $m$ related to data size $n$ to achieve the $O({1 / \sqrt{n}})$ risk rate under the different settings.

In addition, if the network width is chosen first, our results (especially Corollaries 4 and 8 in the revised version) reflect the learning capability of this network. For example, let us consider the case $c=1$ and $\mu=1/2$. Corollary 4 implies, at the theoretical aspect, that if $n\lesssim m^{4/3}$, then GD for two-layer NNs with at least $\sqrt{n}$ iterations can achieve the desired error rate $O(1/\sqrt{n})$. While the estimation of the minimum eigenvalue of the Hessian $\lambda_{\min}(\ell(\mathbf{W};z))$ depends on $||\mathbf{W}-\mathbf{W}_0||_2$, we can show that along the trajectory of GD, $||\mathbf{W}_t-\mathbf{W}_0||_2\le c\sqrt{\eta T}$. Hence, our analysis does not limit how far the network weights can travel from the initialization value. We mainly focus on the least square loss in this work. We agree with you that extending our results to the cross-entropy loss is very interesting and important, which we leave as an important future direction.

** Q2. How are the network weights initialized? It seems the only requirement on initialization is assumption 2.

Yes, we can initialize the network weights randomly as the previous works did or use a fixed initialization as long as Assumption 2 holds. Indeed, for sigmoid activation or hyperbolic tangent activation considered in our paper, it can be easily verified that Assumption 2 holds for any initialization if the inputs and labels are bounded. That is, for any network initialization $\mathbf{W}_0$ and any data $z=(\mathbf{x},y)$ with $||\mathbf{x}||_2\le c_x$ and $|y|\le c_y$, we can always find a constant $c_0>0$ such that the assumption $\ell(\mathbf{W}_0;z)\le c_0$ holds. We have added the clarification below Assumption 2 in the revised version.

** Q4. Do you see a way to argue that the dependence on width (in equation (4)) is tight/necessary?

Thanks for the thoughtful comment. The dependence on width $m$ in equation (4) relies on estimating the maximum/minimum eigenvalue of the Hessian matrix of the empirical risk. Therefore, it may not be precisely tight. It remains a challenging problem to explore what is the tight bound for width $m$. However, as mentioned in Remark 1, [1] provided a similar generalization bound with an assumption $m\gtrsim (\eta T)^5/n^2 + (\eta T)^2$, and our result with $c=1/2$ is consistent with their result. Further, [2] derived a similar excess risk bound for $c=1/2$ (refer to Theorem 3 in the revised version) with $m\gtrsim (\eta T)^5$. We relax this condition to $m\gtrsim (\eta T)^3$ by providing a better estimation of the smallest eigenvalue of a Hessian matrix.

[1] Y. Lei, R. Jin, and Y. Ying. Stability and generalization analysis of gradient methods for shallow neural networks. In Advances in Neural Information Processing Systems, volume 35, 2022.

[2] D. Richards and I. Kuzborskij. Stability \\& generalisation of gradient descent for shallow neural networks without the neural tangent kernel. In Advances in Neural Information Processing Systems, volume 34, 2021.

** Q1. The proof thus far only works for three layers. [Although it is believable that with more work, one may be able to extend the analysis for multiple layers, e.g., by arguing a similar weak convexity condition on the loss function.] / The abstract (and the paper's title) mentions multi-layer neural networks, but the analysis currently holds for three layers. This is a mismatch.

Thanks for your constructive comment. Following your excellent suggestion, we have changed our title to "Generalization guarantees of gradient descent for shallow neural networks", and have revised the abstract accordingly.

** Q2. It is by now widely accepted that the NTK analysis is more of an analogy rather than a realistic characterization of what happens in reality. Thus, the practical relevance of this paper, which follows similar lines of arguments as the NTK (e.g., by rescaling the loss and expanding the network width), is very limited.

Thank you for your comment. In this work, we provide an alternative approach in a kernel-free regime to study the generalization guarantees of GD with neural networks, using the concept of algorithmic stability. One noteworthy merit of our stability analysis lies in its ability to bypass tailored assumptions frequently necessitated by prior works [1,2]. Notably, these assumptions encompass aspects like random initialization and data distribution/kernel matrices which have a positive eigenvalue. In this regard, we consider our stability analysis perspective to be both complementary to existing research and at the same time, sheds new light on understanding GD’s generalization in neural networks.

Moreover, our generalization results show that for both two-layer and three-layer neural networks, GD can achieve the dimension-independent generalization bound with only very milder overparameterization (e.g., $c=1/2$ and $\mu=1/2$, $m\asymp n^{3/2}$, $c=1$ and $\mu=1/3$, $m\asymp n^{3/2}$ for two-layer neural networks. More details can be found in Corollaries 4 and 8).

For instance, Theorem 5.1 of [1] established generalization bounds of GD for two-layer ReLU neural networks with random initialization under an assumption on data distribution, very high overparameterization and the least eigenvalue of the kernel matrix. [2] introduced the neural tangent random feature (NTRF) model, and proved the generalization bounds of one-pass SGD for deep ReLU neural networks with random initialization in the over-parameterization regime, which requires the network weights to stay very close to their initialization throughout training. Their results imply with high overparameterization $m\ge n^7$ that if the data can be classified by a function in the NTRF function class with a small training error, a small classification error can be achieved with the over-parameterized ReLU network.

To conclude, our stability-based results do not require conditions on the random initialization or the data distribution, and provide generalization bounds under milder overparameterization. In this sense, the perspective from our stability analysis sheds new light on comprehending GD's generalization in neural networks.

[1] S. Arora, S. Du, W. Hu, et al. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. ICML 2019.

[2] Y. Cao and Q. Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. NeurIPS 2019.

** Q3. How much more difficulty would it take to extend the analysis to multiple layers?

The key challenge of extending our results from three layers to multiple layers is to control the maximum and minimum eigenvalues of a Hessian matrix of the empirical risk, which rely on the norm of the coupled weights of different layers. Once having the estimates of eigenvalues, we can establish the almost co-coercivity of the gradient operator, which is crucial for the stability analysis. However, it's not easy to estimate the coupled weights of different layers by using the existing technique. We acknowledge the importance of considering the applicability of our findings to networks with multiple layers. Since the existing stability analysis mainly focused on two-layer neural networks, we regard our work as an important starting point to study multiple layers for future research directions. We have added some discussion in the Conclusion to illustrate the main technical challenges in studying the multi-layer case via algorithmic stability.

** Q1. It seems that the authors primarily concentrate on three-layer neural networks. It would be beneficial if the authors explicitly discuss the potential applicability of their findings to networks with multiple layers.

Great point! We acknowledge the importance of considering the applicability of our findings to networks with multiple layers. Since the existing stability analysis mainly focused on two-layer neural networks, we regard our work as an important starting point to study multiple layers for future research directions. Indeed, extending our results from three layers to multiple layers is not straightforward. The main difficulty lies in estimating the maximum and minimum eigenvalues of a Hessian matrix of the empirical risk, which rely on the interconnected weights across layers. Controlling these interconnected weights using existing techniques presents considerable difficulty. In the Conclusion section, we've included discussions highlighting the principal technical obstacles in investigating multiple layers through algorithmic stability.

** Q2. The uniformity of scaling factors across layers may appear somewhat unconventional. It would be helpful if the authors provided further insight into this choice and its implications.

Thank you for the constructive suggestion. To more explicitly illustrate the relationship between the scaling parameter $c$ and the over/under-parameterization of the network (as shown in Corollaries 4 and 8, and Figure 1), we make the assumption that both the width and scaling factors of each hidden layer are identical. Nevertheless, our findings can readily extend to scenarios where the widths and scaling factors vary across layers. In the revised version, we've included additional clarifications to accommodate this aspect.

** Q3. It would be beneficial if the authors could offer more intuitive explanations regarding how the upper bound on the W-norm improves from a t-scale to a $\sqrt{t}$-scale, shedding light on the underlying mechanisms that drive this improvement.

Thank you for the constructive suggestion. As mentioned in Remark 4, for three-layer neural networks, the technical challenge in estimating $||\mathbf{W}_t||_2$ is closely tied to the smoothness of the loss function. Simultaneously, the smoothness parameter of the loss function relies on the norm of $\mathbf{W}_t$. We overcame this difficulty by initially providing a rough bound $||\mathbf{W}_t-\mathbf{W}_0||_2\le \eta t m^{2c-1}$ by induction. Once having the estimate of $||\mathbf{W}_t||_2$, we can upper bound the maximum eigenvalue of the Hessian by a constant $\hat{\rho}$ if $m$ satisfies (6). It implies that the loss is $\hat{\rho}$-smooth along the trajectory of GD. Leveraging this smoothness property, we deduce that

$\eta(1-\eta \hat{\rho}/2)\sum\_{j=0}^t||\nabla L_S(\mathbf{W}_j)||_2^2\le\sum\_{j=0}^t L_S(\mathbf{W}\_j)-L_S(\mathbf{W}\_{j+1})\le L_S(\mathbf{W}\_0).$

Additionally, considering the update rule of GD and $\eta \hat{\rho}\le 1/2$, we conclude that $||\mathbf{W}_t-\mathbf{W}_0||_2^2= \eta^2 ||\sum\_{j=0}^t \nabla L_S(\mathbf{W}_j)||_2^2 \le 2\eta t L_S(\mathbf{W}_0)$.

We have added more details in Remark 4 in the revised version.

**Q1. In Assumption 2, one of the conditions seem to depend on a specific initialization of the weights $\mathbf{W}_0$, where the loss should be less than $c_0$ when evaluated under the initialization. However, nowhere in the paper mentions (at least to that point) how the networks are initialized or if we are assuming a fixed initialization – I am not sure if the bound holds for $c_0$ uniformly over any initialization of the weights or not. Please, clarify.

Thank you for this thoughtful comment. We would like to emphasize that our results hold for both a random initialization setting and a fixed initialization setting as long as $\ell(\mathbf{W}_0;z)$ is uniformly bounded by a constant $c_0$. Indeed, one can verify that such uniform bound $c_0$ holds true for some widely used activations including sigmoid activation and hyperbolic tangent activation for any initialization. Since we assume that $\mathbf{x}$ and $y$ are bounded, $f_{\mathbf{W}_0}(\mathbf{x})$ is uniformly bounded for both two-layer and three-layer cases by noting that the sigmoid and hyperbolic tangent activation functions $\sigma(\mathbf{w}_k\^\top \mathbf{x})\le 1$ for any $k$. Consequently, for any $\mathbf{W}_0$, the least square loss $\ell(\mathbf{W}_0;z)$ is uniformly bounded by a constant $c_0$.

In addition, the assumption of the loss is standard in the stability and generalization analysis of GD with neural networks [1, 2]. We have added the clarification below Assumption 2 and above Theorem 3 in the revised version.

[1] D. Richards and I. Kuzborskij. Stability \\& generalisation of gradient descent for shallow neural networks without the neural tangent kernel. In Advances in Neural Information Processing Systems, volume 34, 2021.

[2] Y. Lei, R. Jin, and Y. Ying. Stability and generalization analysis of gradient methods for shallow neural networks. In Advances in Neural Information Processing Systems, volume 35, 2022.

** Q2. In Theorem 3, it says that equation (4) should hold, but then show another bound on $m$ in equation (5). How should we interpret this? Should we just take the maximum of both lower bounds on $m$?

Yes, we take the maximum of both lower bounds on $m$ to ensure that our main result for two-layer NNs (Theorem 3 in the revised version) holds. This is what we did in Corollary 4, i.e., choosing $m\gtrsim (\eta T)^{\frac{3}{2c}}$ such that the maximum of the lower bounds in (4) and (5) hold. Notably, the generalization error bounds in Theorem 1 hold under condition (4). Introducing condition (5), we further establish the optimization error bounds in Theorem 2. Combining these two results together, we derive the excess risk bounds in Theorem 3 when (4) and (5) hold simultaneously.