Dynamical Decoupling of Generalization and Overfitting in Large Two-Layer Networks

$**Questions**$

1. In our setting we do observe benign overfitting. At the end of the "feature unlearning" phase, the test error is the one achieved by the neural tangent model. This model has benign overfitting [see, e.g. Ghorbani et al 2016; Montanari, Zhong, 2020]. Namely, for $n\to\infty$ the excess error converges to $0$. However this happens only for $n$ superpolynomial in $d$. (The phenomenology established in [Montanari, Zhong, 2020] for neural tangent models matches the general framework for benign overfitting in [Bartlett et al. 2019] where is known that the excess risk vanishes only very slowly with sample size) On the other in the mean field/feature learning regime of training the excess error vanishes already when $n\gg d$.

2. There are several possible routes to making the analysis of Section 2. While a full discussion would be too complex, a possible route would be:
   (i) Use the DMFT equations for actual 2-layer networks that were proven in Celentano, Cheng, Montanari, arXiv preprint arXiv:2112.07572, 2021 (instead of the simplified DMFT we use in our paper).
   (ii) Carry out the singular perturbation analysis of these equations.
   A first step towards (ii) would be to rigorize the singular perturbation analysis in the simplified DMFT that we use in our paper. Even in this simplified setting, step (ii) is the most challenging because it requires to control the solutions of complex non-linear PDEs in the long time limit. The rigorous results in Section 3 partially address this point.

3. There are no particular complications that make the analysis more complicated for SGD. Indeed, for $m=1$ and single index data, the DMFT analysis of SGD dynamics was performed in:
   Kamali, Urbani, arXiv preprint arXiv:2309.04788, 2023. We do expect a similar behavior.

4. The rigorous results in Section 3 imply that the results obtained from the DMFT analysis have a larger domain of applicability. In particular, no-overfitting take place for $t=o(m^{1/4})$ and long as $n/md = \alpha$ is kept constant (regardless of the order of limits). Similarly, the results in the first regime of DMFT are correct for $t=O((\log m)^c)$, and a small $c$ (for $n/md = \alpha$ constant regardless of the order of limits). In particular, these claims apply to cases in which $m/d$ is kept constant.
   While we do not characterize the full domain of applicability of our analysis, these results imply that it is broader than could be assumed from the original derivation. We plan to make this point clearer by making explicit the dependence of the results in section 3 on $n$, $d$.

5. The reviewer is correct. In line 159, we fix the second layer weights at initialization and do not train them. Since we use a symmetric initialization, see Sec. C.4, all the second layer weights are the same.

6. The Bayes optimal value of the test error is $\tau^2/2$ (in the limit $n/d\to\infty$ considered here). This is achieved during the feature learning phase of training. In the lazy phase, the test error increases because the model converges to a neural tangent kernel model that only learns the linear component of the target.

$**Limitations**$

We appreciate the referee comments. As discussed above, we will try our best to address them and discuss limitations/extensions.

$**Weaknesses**$

We will address the problem of readability of the manuscript. As other reviewers requested we plan to:

1. Clarify the notations and terminology.

2. State formally a the most important conjectures/results arising from the DMFT analysis.

3. We will connect better the main text to the appendix.

4. Try to make the best selection for the material in main text.

$**Questions**$

1. In the revised version of the manuscript we will clarify the role of the assumption $\mathbb{E}(\varphi(G)G)\neq 0$ and refer to earlier literature concerning this point.

2. The fact that $a_i$ is independent of $i$ is a consequence of two facts. First, under the initialization, neurons are exchangeable (in the sense of probability theory), and hence they remain exchangeable at all subsequent times. As a consequence of exchangeability ${\mathbb E} a_i(t)$ does not depend on $i$. Second, second layer weights $a_i(t)$ concentrate around their expectation for any $t=O(1)$ as $n,d\to\infty$. Within DMFT, this is discussed in sec. C.4. We could expand on this point from a mathematical perspective.'

The claim that (under the symmetric initialization that we use) the $a_i$ concentrate around a value independent of $i$ for $t=O(1)$ as $n,d\to\infty$ can in fact be proven rigorously.

3. $z$ is the scaled time. We will make the notation clearer in the revised version.

$**Weaknesses**$

• $*Writing.*$ We are going to improve the revised version of the manuscript in two ways:

1. Clarify the notations and terminology.

2. We will connect better the main text to the appendix.

3. Try to make the best selection for the material in main text.

• $*Regularization*$. We point out that we regularize explicitly the first layer weights by constraining them to have unit norm. Instead we do not regularize the second layer weights. From a technical viewpoint, adding a regularization for second-layer weights leads to a simple modification of the same DMFT equations. The final result will be that such a regularization stops the growth of second-layer weights at a value depending on the regularization parameter. For suitable choices of this parameter the unlearning/overfitting dynamical regime will become milder or totally eliminated. Even larger regularization strength will affect the first phase of training and hence induce underfitting. We regard this as a natural (and quite simple) follow-up project.

• $*Measure of complexity.*$ The fact that the $\ell_1$ norm of second layer weights is a bound on complexity is a classical result by Bartlett, see Ref.[6] (cf lines 145-150 of the manuscript). In principle it could be that the complexity is significantly smaller because of some special structure in the first layer weights. However we see no evidence of this in our analysis of the dynamics:

1. First layer weights remain isotropically distributed (apart from the projection on the direction of the signal). Also, by construction their norm is fixed.
2. Increase of the generalization error is accurately tracked by the increase in norm of second layer weights. This is shown very clearly in Figure 5 right frame.

$**Questions**$

1. The definition of the network in Eq.(1.1) has a factor $1/m$ in front of the sum. Therefore if one defines the second layer weights as $A_i=a_i/m$, then our "lazy initialization" corresponds to He initialization $A_i = \mathcal{O}(1/\sqrt m)$, while our "mean field initialization" corresponds to $A_i\sim \mathcal{O}(1/m)$. (We point out that the distinction between "lazy" and "mean field" is standard in the theory literature.)

2. See discussion above for the use of the norm of second layer weights as a measure of complexity.

3. We will add definitions for this terminology: we agree that while common in the theory literature it is best to explain it for a broader audience.

4. The connection between activation function and the function $h(z)$ is straightforward the coefficients of polynomial expansion of the latter are the square of the Hermite coefficients of the former. A similar relation holds for $\widehat\varphi$. We will clarify this point. We also plan to add the details of the numerical simulations.

5. The choice of the functions $h(z)$ and $\widehat \varphi(z)$ in the figures was dictated by desire to select examples such that (i) The qualitative behavior of learning curves is the same as for generic such functions; (ii) The relevant phenomena are clearly visible. We underline that the singular perturbation theory analysis applies to general $h(z)$ and $\widehat \varphi(z)$

6. We apologize for these misprints. We will change "SM" to "the appendix" and replace the "see the appendix" present in the appendix with the corresponding sections.

7. Throughout the paper, it is understood that the activation function does not change with the width $m$. We will clarify this in the text.

8. We will add a discussion the effect of adding an explicit regularization term in the dynamics.

$**Minor comments.**$

Thank you for the detailed comments. We will do our best to address them.

The reviewer points out a general weakness of our work. Admittedly, this is a theoretical work that focuses on:

1. A simple model of data (multi-index model with Gaussian features)

2. A simple neural network (2-layer fully connected).

Focusing on a simple example is the price we pay for being able to carry out an analytical study of the gradient descent dynamics.

$*Earlier theoretical work was limited to significantly simpler cases*$ (eg linear models, networks with 1 neuron) or special regimes (eg neural tangent limit).

Despite the simplicity of the setting, there is a lot that is poorly understood and a lot to learn from this setting. As we explain starting with lines 63 and 64, several of the fundamental question in this model generalize to more complex ML models and the conceptual picture we develop generalizes as well.

Concerning the empirical validation, we do validate our results against numerical simulations on two layer neural networks, see for example Fig. 2 and 4.

We limit ourselves to simulations with synthetic data because our theory can only capture the behavior on real data at a qualitative level. In machine learning, this is nearly always the case for a completely analytical theory. The purpose of our simulations is to verify that various approximations made in our analysis are indeed accurate.

$**Questions.**$

1. Both the general DMFT theory and the rigorous analysis results of Section 3 apply to general multi-index models. Again, for moderate time-scales and large $m$, the dynamics is captured by the mean field theory of [Mei, Montanari, Nguyen, 2018], [Chizat, Bach, 2018]. If the latter converges to vanishing training error in (rescaled) times of order one, then our analysis extends verbatimly. However, in multi-index models convergence can take time slowly diverging with $d$. We do not expect that this will change the qualitative picture we derive, but the analysis becomes technically more complicated, to account for those cases.

2. We assume isotropic Gaussian feature vectors ${\boldsymbol x}_i$ Universality results can be used to show that the same DMFT captures the dynamics when the ${\boldsymbol x}_i$'s have independent entries (see, e.g. [Celentano, Cheng, Montanari, 2021]). Certain distributions with non-identity covariance can be also treated (e.g. linear transforms of iids), but the analysis becomes more intricate and we do not expect qualitatively new insights.

3. A general method that is supported by general theory is to monitor the error on a validation set and stop GF when this is minimal. Interestingly, our theory indicates that ---for large networks--- the validation error will stay close to the minimum for a large interval of time. This suggests that large networks are naturally stable with respect to the choice of stopping time.

$**Weaknesses.**$

• The Gaussian approximation was mentioned in line 123 of the manuscript. However this a short statement and we plan to expand on it.

• The expectation in Eq. A.4 is over the data distribution. We will clarify this and related points.

• We will add a reference to the figure showing that approximation error is vanishing on constant time scales (this can be seen, for instance, in Fig 4 right frame, by noting that the asymptote corresponds to the Bayes error.)

• In the revised version of the manuscript we will specify the activation used in simulations.

• Indeed, because of the symmetric initialization, second layer weights concentrate around a common deterministic function $a(t)$ for all $t$ that remain bounded as $n,d\to\infty$. This is further discussed in Appendix C.4, but we will make it more explicit in the main text.

• $\gamma_{GF}$ is the threshold defined via Eq.(2.3) and lines 166-167: we will clarify this threshold phenomenon.

• We agree that it is remarkable that our theory is capable to predict the unlearning phenomenon. Moving figure 30 and 31 to the main text is probably hard because of space constraints. Note that Figure 5 also illustrate clearly the same phenomenon.

$**Questions.**$

• $*Order of limits.*$ As the reviewer writes, in the DMFT analysis we take first $n$ and $d$ to infinity and then we study the large $m$ limit, allowing possibly the time to diverge with $m$. However, the rigorous results in Section 3 imply that the results obtained from the DMFT analysis have a larger domain of applicability. In particular, no-overfitting take place for $t=o(m^{1/4})$ and long as $n/md = \alpha$ is kept constant (regardless of the order of limits). Similarly, the results in the first regime of DMFT are correct for $t=O((\log m)^c)$, and a small $c$ (for $n/md = \alpha$ constant regardless of the order of limits). While we do not characterize the full domain of applicability of our analysis, these results imply that it is broader than could be assumed from the original derivation. In the revised version we make the dependence on $m$, $d$ and $n$ more explicit in Sec. 3.

• $*Figure 10.*$ This figure refers to a setting in which second-layer weights are fixed to $\gamma\sqrt{m}$. When the control parameters (the overparametrization ratio $\alpha$ or the scale $\gamma$) are such that gradient flow coverges to zero train error, this convergence is exponentially fast in time. The exponential convergence rate can be used to define a relaxation timescale that we estimate as detailed in lines 858-868. When either $\gamma$ or $\alpha$ cross the interpolation threshold, the relaxation time diverges. In the inset of Figure 10 we plot the behavior of the relaxation time as a function of the control parameters. This allows to determine the interpolation threshold $\alpha_{GF}$. In the main panel we plot the relaxation time as a function of $\alpha_{GF}-\alpha$, to verify that the relaxation time diverges as $(\alpha_{GF}-\alpha)^{-\nu}$. The data suggest that $\nu$ is close to $2$. In the revised version of the appendix we will add a plot of the train error in linear-logarithmic scale. This will confirm the exponential relaxation to zero train error and clarify Figure 10.

• $*Section G.2.1*.$ The ansatz used on this timescale is presented in Eqs.~(G.39)-(G.41). Among the scaling forms used in various dynamical regimes, we consider this particularly simple. We simply posit that all quantities have a finite limit as $m\to \infty$, without time rescaling. An indication that this should be the case is already contained in the mean field theory [Mei, Montanari, Nguyen, 2018; Chizat, Bach 2019] which obtains a finite non-trivial limit as $m\to\infty$ and $t=\Theta(1)$. Notice however that the out-of-diagonal response $R_o$ turns out to be of order $1/m$ in this ansatz. This is intuitive because $R_o(t,s)$ captures the effect of changing the drift for the $j$-th neuron at time $s$ on the $i$-th neuron, $i\neq j$ at time $t$. It is reasonable to think that this effect is of order $1/m$ because each neuron is multiplied by a factor $1/m$ in the overall function.