Convergence of the Gradient Flow for Shallow ReLU Networks on Weakly Interacting Data

Response: Thank you for your review and questions, which we address below. Please feel free to follow up if any of our responses do not fully resolve your concerns.

1. Theorem 1. The convergence rate decreasing with $n$ is due to the orthogonality constraint. Note that $n$ could not grow larger than $d$, and hence the limit $n \to \infty$ (which would break down the result due to the infinitely slow convergence) cannot happen. Moreover, imagine the case $p=1$, we have for all $i$ that $y = ||w||<w|x_i>$, where the $x_i$ are orthonormal. This implies $||w||^2 = \sqrt{n}y$, and thus the PL is at most of order $\frac{1}{\sqrt{n}}$.

2. Gradient Descent. The gradient flow corresponds to the limit of a gradient descent with the step size going to 0. In the small but positive step size limit, any equation of the form $\frac{d f(t)}{dt} = c(t)$ becomes $\frac{f(t + s) - f(t)}{s} = c(t) \sim c(t + s)$ with $s$ the step size. Now, we can carry the same computations and inequalities as before to find a valid discrete equivalent to Theorem 1, up to a constant depending on $s$. Taking $s$ small enough guarantees the convergence, but $s$ might need to depend on $p, n, d$ or the constants $C_{x,y}^{+,-}$.

3. Gradient Flow Equation. This is indeed quite confusing (even for the author up until writing the paper), and comes from the normalization of the network. We wanted to find a way to normalize both the network $h$ and the loss $L$ in order to have consistent definition if $p$, $n$ or $d$ would go to infinity. We thus introduce the $\frac{1}{p}$ before the summation in $h$, which introduces a temporal dilatation: the more neurons, the slower each neuron learns. To have learning in the regime $p\rightarrow+\infty$, we thus rescale the gradient descent equation by a factor $p$. Note that this is strictly equivalent to having the normal equation without $p$, a network without the normalization, and weights initialized with norm $\frac{1}{\sqrt{p}}$, which is the default equations used in pytorch.

4. Local Minima. Local minima exist in this setup, but the paper claims that in average, we will avoid them all. This works in two steps: first, with enough neurons, we put ourselves into a high probability zone of the parameter space which do not trap us into local minima (example: all the $|a_j|$ are positive, hence no negative label can be fitted), and second, the near-orthogonality hypothesis says that we don’t have local minima created from the weights moving around and interacting (see the example in appendix C.1). In particular, the local-PL viewpoint means that although the local-minima exist, they are avoided on the trajectory starting from the high probability zone.

Response: Thank you for the review, we addressed your questions below.

1. Proposition 2. Orthogonal data enables us to integrate exactly the system, and have the tight bound in equation (10), of order $\frac{1}{n^{\alpha}}$. We need this tightness to show that all speeds in the interval $[\frac{1}{n}, \frac{1}{\sqrt{n}}]$ can be achieved by appropriate weight initialization.

2. Group Initialization. We think that there exists a small enough variance for which the Polyak-Łojasiewicz (PL) curvature still remains of order $\frac{1}{n^{\alpha}}$, but the lack of differentiability of the problem makes this a non-trivial claim which we should further verify. Note that the main point of the group initialization is to show that the speed interval $[\frac{1}{n}, \frac{1}{\sqrt{n}}]$ are achieved by some initialization of the weights.

3. Implicit Bias. The problem of the implicit biases are intertwined with our Conjecture: if the weights are approximately equal in magnitude (i.e., implicit bias of an L2-norm), then the speed should be of order $\frac{1}{\sqrt{n}}$; and if the weights are very unequal (i.e., implicit bias of an L1 norm), then the speed will be of order $\frac{1}{n}$. The experiments in section 5.2 suggests a speed of order $\frac{1}{\sqrt{n}}$, and thus an implicit bias minimizing the L2 norm. We did not discuss why convergence speed would be related to implicit bias, and think this would be an interesting direction for future work.

4. Residual Definition. Thank you for pointing this out — we will add the definition to the paper. For completeness, we provide it here $r_i(t) = y_i - h_{\theta(t)}(x_i)$.

Response: Thank you for your positive review and thoughtful questions. We’re glad you appreciated the paper. Please find our responses to your comments below:

1. Necessity of assumptions. The paper shows that the convergence holds with the asymmetric assumption, i.e.for all $j, \, |a_j| > ||w_j||$, and that it may not hold when $|a_j| < ||w_j||$ for all $j$, as shown in Appendix C.1. This makes this assumption necessary for the result to hold. Now, the assumption-free case makes room for edge cases that look like the proof from Proposition 3. However, these edge cases often require initializations which are unlikely, and thus can be excluded in a reasonning in high probability. While we believe this assumption is unnecessary, we were unable to rigorously prove this within the scope of the present work.

2. Scaling improvements. We do think that the gap between Theorem 1 and the experiments can be improved. In the proof of Theorem 1, we need a constraint on the dimension in two critical steps: 1. the lower bound in equation (23) which requires only $d > n$, this is a rank constraint, and cannot be improved with our method, 2. the lower bound between equations (24) and (25) requires the constraint $n \sim \sqrt(d)$, and comes from an upper bound on the norm of the residuals $|r_i(t)|$. The second constraint could be losened up by having a better control on the residuals. In this case, the theoretical scaling would become $d \sim n$.

3. Extension to deeper networks. Multi-layer networks are more difficult to control due to the possibility of intermediate layers collapsing before fitting is complete. In our setup, collapse is avoided thanks to Lemma 1 (which prevents $a_j$ from approaching zero) and by controlling the weights $w_j$. However, in deeper architectures or with multi-dimensional outputs, our result no longer holds, as it becomes difficult to simultaneously control multiple sets of weights. Understanding why deep networks avoid collapse in our regime is an interesting and important question, which warrants further investigation and will likely require new techniques for managing groups of neurons.

4. Practical implications. Our goal in this paper was to understand convergence of neural networks in the most natural training setup. Unfortunately, this requires unrealistic assumptions on the data, which should be lifted by future work. In practice, since one-layer networks are not really used, it is hard to find examples to compare ourselves with. To illustrate, if we compare our model in terms of parameters-data ratio, and cite Phi-2 with a 519:1 or the ResNet-50 trained on Cifar with a 460:1 ratio. Given the sizes of the models, this is still far from the scaling $d * p \sim n^2\log(n)$. But should Theorem 1 be extended to the empirical scaling $d\sim n$, the ratios would be very close for some architectures, with a $\sim d*\log(d):d$ ratio.

5. Comparison with existing methods. While this would certainly be valuable, we did not compare with existing methods in this paper. This is because the theoretical understanding of neural networks remains far from practical setups. We focused our efforts on understanding convergence rather than generalization. Nonetheless, studying test-time behavior and the properties of the minima found is a natural and important direction for our future work.

Response: Thank you for your review and interesting questions. We address them below.

1. Width of the Network. The large input dimensionality $d$ and the large network width $p$ are two distinct components in our proof. The dimensionality is crucial for ensuring that the weights do not interact too strongly, which allows us to establish the Polyak-Łojasiewicz (PL) lower bounds. The width of the network, on the other hand, is mainly used to guarantee a good initialization—specifically, that for each input $x_i$, there exists some neuron $w_j$ that is positively correlated with it, with high probability. The scaling $p = \log(n)$ primarily arises due to the ReLU non-linearity; different activation functions require different scalings. For instance, non-zero Leaky ReLU requires only a constant number of neurons. That said, increasing the network width could potentially strengthen the results, but doing so would require a different analysis that is compatible with, yet distinct from, the PL framework.

2. Group Initialization. The group initialization setup is a toy example intended solely to demonstrate that the lower and upper bounds are tight for orthogonal data. Its purpose is not to model any real-world data distribution. As stated in Claim 5.2 of our experimental section, we expect that Gaussian data should converge at a rate of $\frac{1}{\sqrt{n}}$.

3. Theorem 2. We apologize for the lack of clarity in the original statement. Theorem 2 aims to show that under group initialization, the network learns one group at a time in the large-$n$ limit. This is demonstrated by rescaling the loss, which becomes a piecewise constant and decreasing function in the limit. Consequently, within a short interval of order $\log(n)$, the network fully learns a group. We also quantify the width of this transition using the times $t_n^j(\epsilon)$ and $t_n^j(1 - \epsilon)$, following the methodology used in cutoff phenomena for Markov chains. This analysis reveals that the transition width decreases as $\frac{1}{\log(n)}$, which is very slow and hence not visible in practice. We will make that discussion clearer in the revised version of the paper.

4. Residual. Thank you for pointing this out—we indeed omitted the definition of the residual. We will clarify that $r_i(t) = y_i - h_{\theta(t)}(x_i)$ denotes the $i$-th residual of the loss.

5. Activation Function. It is possible to use other activation functions and still obtain convergence results. A simple example is any monomial function of the form $x \mapsto \alpha x^\beta$, with $\beta > 1$ and $\alpha \neq 0$. Such functions satisfy Lemma 1 up to a constant, and the proofs adapt easily. We believe other monotonic activation functions could also work well, though computing their convergence rates might be more challenging. While various activation functions are indeed applicable, we chose not to discuss them in the main text as they do not significantly contribute to the intuition. However, we will include an appendix discussing how different activation functions affect the convergence probability and rate.

6. Mean-Field Scaling. We adopted mean-field scaling to normalize the network and to clearly distinguish our regime from the Neural Tangent Kernel (NTK) regime, even as $d$, $n$, or $p$ become large. Using other scalings often hides key constants (e.g., $\frac{1}{p}$) inside initialization terms, as is done in PyTorch, which makes the limit $p \to \infty$ less transparent. We chose this scaling for the sake of clarity and interpretability.

7. Multi-Layer Networks. Multi-layer networks are more difficult to control due to the possibility of intermediate layers collapsing before fitting is complete. In our setup, collapse is avoided thanks to Lemma 1 (which prevents $a_j$ from approaching zero) and by controlling the weights $w_j$. However, in deeper architectures or with multi-dimensional outputs, our result no longer holds, as it becomes difficult to simultaneously control multiple sets of weights. Understanding why deep networks avoid collapse in our regime is an interesting and important question, which warrants further investigation and will likely require new techniques for managing groups of neurons.

Thank you for your response. I now understand that your approach can indeed handle scenarios where plateaus are present.

As an additional question, several works, such as [1] and [2], derive local PL conditions and prove convergence in the teacher-student setting, which you also briefly mention in your conclusion. Could you please clarify how your results relate to these works?

[1] Zhou, Mo, Rong Ge, and Chi Jin. "A local convergence theory for mildly over-parameterized two-layer neural network." COLT 2021.
[2] Chizat, Lenaic. "Sparse optimization on measures with over-parameterized gradient descent." Mathematical Programming 194.1 (2022): 487-532.