Paper: Generalization of Scaled Deep ResNets in the Mean-Field Regime (ICLR-2024)

Reviewer comment (Theory):
Perhaps this is a necessity for using such learning-theory tools, but I found the use of r-regions somewhat misleading. Fixing the r-region, it seems to be that, in one way or another, one is limiting the training time. Thus qualitatively, one is limiting the number of data points as training for a very short time on infinitely many data points is not much different than training for such a time on a finite large number of data points. If so what is the meaning of taking large n here if we don't give the network enough time to learn these points?

Author response:
**A1**: In Theorem 4.8, we prove that when $\beta$ is large enough, for any $t>0$, the KL divergence is bounded by a constant order for **infinite** time. Therefore, our results have no constraint on the training time.

Reviewer comment (Theory):
Similarly confusing for me was the claim that $R_n(F) = \beta \sqrt{r/n}$ when \beta scales larger than $\sqrt{n}$. If so it seems that $R_n(F)=O(n^0)$ and not $1/\sqrt{n}$.

Author response:
**A2**:  We have made the proof clearer in our updated version. Concretely, in proving the generalization bound, we take $r = O(1/\beta^2)$ in our result and thus $R_n(F) = \beta \sqrt{r/n}$ scales $1/\sqrt{n}$ as $\beta \sqrt{r}$ cancels out.

Reviewer comment (Evaluation):
The authors' results do not reflect the complexity of the task of the dataset measure. While this can also be viewed as a positive point, I find it hard to believe that one can provide tight or even order-of-magnitude estimates on actual neural networks (in the regime where they generalize well) by putting so little data in. While this is not a criticism of the specific work under consideration, it does weaken its broader appeal in my mind.

Author response:
**A3**: We follow data assumptions on datasets as used in previous mean-field paper, such as "A mean field view of the landscape of two-layer neural networks". We remove the minimum eigenvalue assumption by explicitly estimating the lower bound, which is independent of $n$. 
We agree with the reviewer that a stronger data assumption would allow us for certain results under some specific settings, e.g., $k$-sparse parity problem [1].

Reviewer comment (Experiments):
With relation to point number 3, some numerical experiments (to the very least on toy data sets) showing how accurate these bounds are would be useful.

Author response:
**A4**: According to your suggestion, we have already added experiments in the updated paper, see Sec. 4.5. 

Here we briefly introduce here. We validate our findings on the toy datasets [Two Spirals](https://i.imgur.com/FTjqjuV.jpg), where the data dimension $d=2$. We use a neural ODE model ([torchdyn](https://github.com/DiffEqML/torchdyn)) to approximate the infinite depth ResNets, we take the discretization $L=10$. The neural ODE model and the output layer are both parametrized by a two-layer network with the tanh activation function, and the hidden dimension is $M=K=20$. The parameters of ResNet encoder and the output layer are jointly trained by Adam optimizer with an initial learning rate $0.01$. We perform full-batch training for 1,000 steps on the training dataset of size $n_{\rm train}$, and test the resulting model on the test dataset of size $n_{\rm test}=1024$ by the 0-1 classification loss. We run experiments over 3 seeds and report the mean. We fit the results (after logarithm) by ordinary least squares, and obtain the slope is $-1.02$ with p-value $10^{-5}$, as shown in [this figure](https://i.imgur.com/ZJ0eAAJ.jpg).

Reviewer comment (Experiments):
Can the authors provide some supporting experiments showing how tight are their bounds?

Author response:
**A4**: According to your suggestion, we have already added experiments in the updated paper, see Sec. 4.5. 

Here we briefly introduce here. We validate our findings on the toy datasets [Two Spirals](https://i.imgur.com/FTjqjuV.jpg), where the data dimension $d=2$. We use a neural ODE model ([torchdyn](https://github.com/DiffEqML/torchdyn)) to approximate the infinite depth ResNets, we take the discretization $L=10$. The neural ODE model and the output layer are both parametrized by a two-layer network with the tanh activation function, and the hidden dimension is $M=K=20$. The parameters of ResNet encoder and the output layer are jointly trained by Adam optimizer with an initial learning rate $0.01$. We perform full-batch training for 1,000 steps on the training dataset of size $n_{\rm train}$, and test the resulting model on the test dataset of size $n_{\rm test}=1024$ by the 0-1 classification loss. We run experiments over 3 seeds and report the mean. We fit the results (after logarithm) by ordinary least squares, and obtain the slope is $-1.02$ with p-value $10^{-5}$, as shown in [this figure](https://i.imgur.com/ZJ0eAAJ.jpg).

Reviewer comment (Theory):
The proof of global convergence relies on lower-bounding the smallest eigenvalue of the Gram matrix of the gradients of the two-layer neural network in the end of the architecture. There is no guarantee on the Gram matrix of the gradients of the infinite-depth ResNet part of the architecture. For this reason, I am wondering if the proof would still hold if the ResNet part of the architecture was fixed to its initial value or even altogether removed. If this is the case, I think this should be stated more clearly, and in my opinion questions a bit the narrative of the paper which gives an important focus on the ResNet part of the architecture, whereas I do not fully understand at this point the role it plays in the proof.

Author response:
Thanks for your support on this work on increasing the score. We agree with the reviewer that our model degenerates to the two-layer NN setting if we set the residual part as an identity mapping. According to the reviewer's suggestion, we clearly mentioned this in the remark after Lemma 4.3  in our updated paper.

Reviewer comment (Theory):
The generalization bound is proven in the case where the radius \( r \) of the KL ball decreases in \( O(1/n) \). The paper proves that, with the appropriate scaling of \( \beta \), the trained network remains in this ball, and thus the generalization bound holds for the trained network. This means that the distribution of the weights is required to be increasingly close to the distribution at initialization as the sample size increases. For this reason, I am wondering if the regime where the generalization bound holds can really be thought of as feature learning. While this is not a problem in itself, it questions in my opinion the statement in the conclusion of the paper that “this is the first paper to build a generalization bound for deep neural networks in the so-called feature learning regime”. In my opinion, either more caution in the phrasing or more explanations would be highly welcome.

Author response:
We agree with the reviewer that the stationary distribution after training will be close to its initialization as suggested by our theory.
Our results provide an upper bound of the KL divergence, which allows for a decent movement in terms of distribution, but we cannot provide a lower bound to the KL divergence to ensure a constant order distance. 
According to the reviewer's suggestion, we have rephrased the expression to "beyond NTK regime" to avoid any confusion.

Reviewer comment (Theory):
A less important remark regards Theorem 4.11. The rate of convergence of the empirical risk in Lemma 4.7 actually depends on \( n \) through \( \beta \) (which appears both in the exponential and implicitly in the loss at initialization, since the output of the network at initialization is proportional to \( \beta \)). I don’t think this dependence is taken into account by the authors in Theorem 4.11, or at least it is not stated since the authors state that the empirical risk decays as \( O(e^{-t}) \). While I believe that the stated theorem is correct, the result could probably be sharpened (since the rate of convergence is actually \( O(e^{-nt}) \), although the magnitude of the loss at initialization also depends on \( n \) and would have to be taken into account). In any case, a proof would be beneficial.

Author response:
Thanks for pointing it out. We have already updated Theorem 4.10, such that it is able to describe the limiting distribution more accurately.

Reviewer comment (Novelty):
There are several papers which show global convergence of deep residual networks and which are not cited. In particular, the beginning of Section 4 suggests that this paper is the first one to prove a rate of convergence with mild assumptions. There have actually been other papers proposing this, see e.g. https://arxiv.org/abs/1910.02934, https://arxiv.org/abs/2204.07261, https://arxiv.org/abs/2309.01213. Please rephrase accordingly the beginning of Section 4. Also, I believe Barboni et al. do not “assume that the model parameters are close to the initialization”, they prove it to be the case in their setting, which is a very different statement.

Author response:
Thanks for pointing out these related works. We have already rephrased our expression for a better understanding in the revised version. We summarize the difference here. 

[1] obtains an algorithm-dependent bound for the overparameterized deep residual networks. [2] proves linear convergence of gradient descent to a global optimum for the training of deep residual networks with constant layer width. 
Both of these two works focus on the finite width setting, instead of the infinite mean-field limit in our work. The employed techniques are accordingly different. 
[3] discusses the implicit regularization of ResNets converging to ODEs by virtue of the PL condition that is different from our KL divergence technique.
Besides, the generalization properties are not characterized in [3].
We are thankful to the reviewer for pointing out these references. Notice that [3] was posted on arXiv in September 2023.

Regarding the work from "Barboni et al." [4], we agree with the reviewer that they do not assume that the model parameters are close to the initialization. We also do not assume this in our work. 
In Theorem 2 of [4], the condition of radius $R$  is relatively difficult to verify, since the order of $\lambda(\cdot)$ function in the $N$-universality assumption (Assumption 2) is hard to estimate. 

---

[1] Frei, Spencer, Yuan Cao, and Quanquan Gu. "Algorithm-dependent generalization bounds for overparameterized deep residual networks." Advances in neural information processing systems 32 (2019).

[2] Cont, Rama, Alain Rossier, and RenYuan Xu. "Convergence and Implicit Regularization Properties of Gradient Descent for Deep Residual Networks." arXiv preprint arXiv:2204.07261 (2022).

[3] Marion, Pierre, et al. "Implicit regularization of deep residual networks towards neural ODEs." arXiv preprint arXiv:2309.01213 (2023).

[4] Barboni et al. "On global convergence of ResNets: From finite to infinite width using linear parameterization." NeurIPS 2022.

Reviewer comment (Writing):
Page 2, second paragraph: \( \tau \) et \( \nu \) should be swapped.

Author response:
We are thankful to the reviewer for the attentive study of our work. We have fixed the typos that the reviewer pointed out. Specifically, 
- (1) The sentence, “zero initialization of the network’s prediction”, means that the initial prediction of the limiting network is zero, which will be useful in the later proofs. We notice that it is confusing here, and we delete this sentence. 
- (2 )$\Lambda$ is $\Lambda(d)$, we omit the dependence on $d$ for simplicity. 
- (3) We have changed Chizat and Bach's duplicated entries. For Ding, et. al, there are actually two different papers: "*On the Global Convergence of Gradient Descent for multi-layer ResNets in the mean-field Regime*" and "*Overparameterization of Deep ResNet: Zero Loss and Mean-field Analysis*". Therefore, there is no duplication here. 
- (4) We have added the definition of the distribution’s infinite norm in Lemma B.7 and B.8.

Reviewer comment (Writing):
Page 3: “nerual” -> neural.

Author response:
We are thankful to the reviewer for the attentive study of our work. We have fixed the typos that the reviewer pointed out. Specifically, 
- (1) The sentence, “zero initialization of the network’s prediction”, means that the initial prediction of the limiting network is zero, which will be useful in the later proofs. We notice that it is confusing here, and we delete this sentence. 
- (2 )$\Lambda$ is $\Lambda(d)$, we omit the dependence on $d$ for simplicity. 
- (3) We have changed Chizat and Bach's duplicated entries. For Ding, et. al, there are actually two different papers: "*On the Global Convergence of Gradient Descent for multi-layer ResNets in the mean-field Regime*" and "*Overparameterization of Deep ResNet: Zero Loss and Mean-field Analysis*". Therefore, there is no duplication here. 
- (4) We have added the definition of the distribution’s infinite norm in Lemma B.7 and B.8.

Reviewer comment (Theory):
Page 3: in your ResNet model, you assume a scaling in \( 1/L \). While this is necessary to obtain an ODE limit and usually assumed in the corresponding literature, it departs from the actual practice. I think it would be interesting to mention this. See https://arxiv.org/abs/2105.12245, https://arxiv.org/abs/2206.06929 for a related discussion.

Author response:
Thanks for pointing it out. We have added these references to the paper. [1,2] discuss the influence of scaling on the ResNets training. [1] points out the $1/L$ scaling is inconsistent with the practical observation, which should be $1/L^{0.7}$ instead, which will lead to different infinite depth limit such as rough path. We leave this as future work. [2] derive a limiting neural SDE for the $1/\sqrt{L}$ scaling, and discuss the stability around the critical scaling $1/\sqrt{L}$ and $1/L$.  Different scaling leads to different continuous limits varying from Ito process to rough differential equation. We believe that our approach not only serves effectively in the Neural ODE regime but also plays a crucial role in establishing generalization bounds in other continuous limits, for our proof lies in the stability exhibited by the continuous model under consideration. We thank the reviewer for pointing this out, we'll leave this for future work. 

In this paper, we still adopt the $1/L$ scaling, following the traditional neural ODE community. One reason is that [1] is not the complete story, the $1/L$ scaling is also observed in practice, such as Section 5.3 in [3], which observes the approximately $1/L$ scaling in the residual connected networks (ReZero); and Section 5 in [4] also validates the $1/L$ scaling both in theory and experiment. 

[1] Cohen, Alain-Sam, et al. "Scaling properties of deep residual networks." International Conference on Machine Learning. PMLR, 2021.

[2] Marion, Pierre, et al. "Scaling resnets in the large-depth regime." arXiv preprint arXiv:2206.06929 (2022).

[3] Bachlechner, Thomas, et al. "Rezero is all you need: Fast convergence at large depth." Uncertainty in Artificial Intelligence. PMLR, 2021.

[4] Marion, Pierre, et al. "Implicit regularization of deep residual networks towards neural ODEs." arXiv preprint arXiv:2309.01213 (2023).

Reviewer comment (Writing):
Page 4: the sentence “We introduce a trainable MLP …” is not very clear. What do you mean by “zero initialization of the network’s prediction”?

Author response:
We are thankful to the reviewer for the attentive study of our work. We have fixed the typos that the reviewer pointed out. Specifically, 
- (1) The sentence, “zero initialization of the network’s prediction”, means that the initial prediction of the limiting network is zero, which will be useful in the later proofs. We notice that it is confusing here, and we delete this sentence. 
- (2 )$\Lambda$ is $\Lambda(d)$, we omit the dependence on $d$ for simplicity. 
- (3) We have changed Chizat and Bach's duplicated entries. For Ding, et. al, there are actually two different papers: "*On the Global Convergence of Gradient Descent for multi-layer ResNets in the mean-field Regime*" and "*Overparameterization of Deep ResNet: Zero Loss and Mean-field Analysis*". Therefore, there is no duplication here. 
- (4) We have added the definition of the distribution’s infinite norm in Lemma B.7 and B.8.

Reviewer comment (Writing):
Pages 8 and 9: who is \( \Lambda \)? Should it be \( \Lambda(d) \)?

Author response:
We are thankful to the reviewer for the attentive study of our work. We have fixed the typos that the reviewer pointed out. Specifically, 
- (1) The sentence, “zero initialization of the network’s prediction”, means that the initial prediction of the limiting network is zero, which will be useful in the later proofs. We notice that it is confusing here, and we delete this sentence. 
- (2 )$\Lambda$ is $\Lambda(d)$, we omit the dependence on $d$ for simplicity. 
- (3) We have changed Chizat and Bach's duplicated entries. For Ding, et. al, there are actually two different papers: "*On the Global Convergence of Gradient Descent for multi-layer ResNets in the mean-field Regime*" and "*Overparameterization of Deep ResNet: Zero Loss and Mean-field Analysis*". Therefore, there is no duplication here. 
- (4) We have added the definition of the distribution’s infinite norm in Lemma B.7 and B.8.

Reviewer comment (Presentation):
There are duplicated entries in the references, e.g., Chizat and Bach, and (probably) Ding, Chen, Li, and Wright. Please double-check this.

Author response:
We are thankful to the reviewer for the attentive study of our work. We have fixed the typos that the reviewer pointed out. Specifically, 
- (1) The sentence, “zero initialization of the network’s prediction”, means that the initial prediction of the limiting network is zero, which will be useful in the later proofs. We notice that it is confusing here, and we delete this sentence. 
- (2 )$\Lambda$ is $\Lambda(d)$, we omit the dependence on $d$ for simplicity. 
- (3) We have changed Chizat and Bach's duplicated entries. For Ding, et. al, there are actually two different papers: "*On the Global Convergence of Gradient Descent for multi-layer ResNets in the mean-field Regime*" and "*Overparameterization of Deep ResNet: Zero Loss and Mean-field Analysis*". Therefore, there is no duplication here. 
- (4) We have added the definition of the distribution’s infinite norm in Lemma B.7 and B.8.

Reviewer comment (Writing):
Lemma B.7: I think the infinite norm is undefined. It is possible to infer what it means from context, but please define it.

Author response:
We are thankful to the reviewer for the attentive study of our work. We have fixed the typos that the reviewer pointed out. Specifically, 
- (1) The sentence, “zero initialization of the network’s prediction”, means that the initial prediction of the limiting network is zero, which will be useful in the later proofs. We notice that it is confusing here, and we delete this sentence. 
- (2 )$\Lambda$ is $\Lambda(d)$, we omit the dependence on $d$ for simplicity. 
- (3) We have changed Chizat and Bach's duplicated entries. For Ding, et. al, there are actually two different papers: "*On the Global Convergence of Gradient Descent for multi-layer ResNets in the mean-field Regime*" and "*Overparameterization of Deep ResNet: Zero Loss and Mean-field Analysis*". Therefore, there is no duplication here. 
- (4) We have added the definition of the distribution’s infinite norm in Lemma B.7 and B.8.

Reviewer comment (Novelty):
Missing reference. The following work is also closely relevant to the paper. 

Z. Chen, E. Vanden-Eijnden, and J. Bruna. On feature learning in neural networks with global convergence guarantees.

Author response:
**A4**. Thanks for pointing out these related works. We have already rephrased our expression for a better understanding in the revised version. 

We summarize the difference here. [1] obtains an algorithm-dependent bound for the overparameterized deep residual networks. [2] proves linear convergence of gradient descent to a global optimum for the training of deep residual networks with constant layer width. 
Both of these two works focus on the finite width setting, instead of the infinite mean-field limit in our work. The employed techniques are accordingly different. 
[3] discusses the implicit regularization of ResNets converging to ODEs by virtue of the PL condition that is different from our KL divergence technique.
Besides, the generalization properties are not characterized in [3].
We are thankful to the reviewer for pointing out these references. Notice that [3] was posted on arXiv in September 2023.

Regarding the work from "Barboni et al." [4], we agree with the reviewer that they do not assume that the model parameters are close to the initialization. We also do not assume this in our work. 
In Theorem 2 of [4], the condition of radius $R$  is relatively difficult to verify, since the order of $\lambda(\cdot)$ function in the $N$-universality assumption (Assumption 2) is hard to estimate. 

[1] Frei, Spencer, Yuan Cao, and Quanquan Gu. "Algorithm-dependent generalization bounds for overparameterized deep residual networks." Advances in neural information processing systems 32 (2019).

[2] Cont, Rama, Alain Rossier, and RenYuan Xu. "Convergence and Implicit Regularization Properties of Gradient Descent for Deep Residual Networks." arXiv preprint arXiv:2204.07261 (2022).

[3] Marion, Pierre, et al. "Implicit regularization of deep residual networks towards neural ODEs." arXiv preprint arXiv:2309.01213 (2023).

[4] Barboni et al. "On global convergence of ResNets: From finite to infinite width using linear parameterization." NeurIPS 2022.v