SlowGuess/NeurIPS_2023_OCR huggingface dataset:
sign is negative then it turns to remove from its active half-space all training points that were initially inside. Moreover, those training points cross the activation boundary in the same order as they cross the half-space boundary of the corresponding yardstick trajectory  $\alpha_{j}^{t}$ , and at approximately the same times (cf. Proposition 2).
Lemma 3. For all  $j \in J_{+} \cup J_{-}$  there exist unique  $0 = t_j^0 < t_j^1 < \ldots < t_j^{n_j}$  such that for all  $\ell \in [n_j]$ :
(i)  $I_{s_j}(\pmb{w}_j^t) = I_{s_j}(\pmb{z}_j) \cup \{i_j^1, \dots, i_j^{\ell - 1}\}$  for all  $t \in (t_j^{\ell - 1}, t_j^\ell)$ ;
(ii)  $I_0(\pmb{w}_j^t) = \emptyset$  for all  $t \in (t_j^{\ell - 1}, t_j^\ell)$ , and  $I_0\left(\pmb{w}_j^{t_j^\ell}\right) = \{i_j^\ell\}$ ;
(iii)  $|\tau_j^\ell -t_j^\ell |\leq \lambda^{1 - \left(1 + \frac{3\ell - 1}{3n_j}\right)\varepsilon}.$
The preceding lemma is proved by establishing, for this first phase of the training, non-asymptotic upper bounds on the Euclidean norms of the hidden neurons and hence on the absolute values of the network outputs, and inductively over the stage index  $\ell$ , on the distances between the unit-sphere normalisations of  $\alpha_{j}^{t}$  and  $\pmb{w}_{j}^{t}$ . Based on that analysis, we then obtain that each negative-sign hidden neuron does not grow from its initial length and deactivates by time  $T_0 \coloneqq \max_{j \in J_{+} \cup J_{-}} \tau_{j}^{n_j} + 1$ .
Lemma 4. For all  $j \in J_{-}$  we have:
$$
\left\| \boldsymbol {w} _ {j} ^ {T _ {0}} \right\| \leq \lambda \left\| \boldsymbol {z} _ {j} \right\| \quad \boldsymbol {w} _ {j} ^ {t} = \boldsymbol {w} _ {j} ^ {T _ {0}} \quad f o r a l l t \geq T _ {0}.
$$
We also obtain that, up to a later time  $T_{1} \coloneqq \varepsilon \ln (1 / \lambda) / \| \gamma_{[n]} \|$ , each positive-sign hidden neuron: grows but keeps its length below  $2\| z_{j}\| \lambda^{1 - \varepsilon}$ , continues to align to the vector  $\gamma_{[n]}$  up to a cosine of at least  $1 - \lambda^{\varepsilon}$ , and maintains bounded by  $\lambda^{1 - 3\varepsilon}$  the difference between the logarithm of its length divided by the initialisation scale and the logarithm of the corresponding yardstick vector length.
Lemma 5. For all  $j \in J_{+}$  we have:
$$
\| \boldsymbol {w} _ {j} ^ {T _ {1}} \| <   2 \| \boldsymbol {z} _ {j} \| \lambda^ {1 - \varepsilon} \quad \overline {{\boldsymbol {w}}} _ {j} ^ {T _ {1} ^ {\top}} \overline {{\gamma}} _ {[ n ]} \geq 1 - \lambda^ {\varepsilon} \quad | \ln \| \boldsymbol {\alpha} _ {j} ^ {T _ {1}} \| - \ln \| \boldsymbol {w} _ {j} ^ {T _ {1}} / \lambda \| | \leq \lambda^ {1 - 3 \varepsilon}.
$$
5 Second phase: growth and convergence
We next analyse the gradient flow subsequent to the deactivation of the negative-sign hidden neurons by time  $T_0$  and the alignment of the positive-sign ones up to time  $T_1$ , and establish that the loss converges to zero at a rate which is exponential and does not depend on the initialisation scale  $\lambda$ .
Theorem 6. Under Assumptions 1 and 2, there exists a time  $T_{2} < \ln (1 / \lambda)(4 + \varepsilon)d\Delta^{2} / \delta^{6}$  such that for all  $t \geq 0$  we have  $L(\pmb{\theta}^{T_2 + t}) < 0.5\Delta^2\mathrm{e}^{-t\cdot 0.4\delta^4 /\Delta^2}$ .
In particular, for  $\varepsilon = 1/4$  and  $\lambda = \left((mn^{n})^{9\Delta^{2}/\delta^{3}}\right)^{-3/\varepsilon}$  (cf. Assumption 2), the first bound in Theorem 6 becomes  $T_{2} < (\ln m + n\ln n)d \cdot 17 \cdot 27\Delta^{4}/\delta^{9}$ .
The proof of Theorem 6 is in large part geometric, with a key role played by a set  $S \coloneqq S_1 \cup \dots \cup S_d$  in predictor space, whose constituent subsets are defined as
$$
\mathcal {S} _ {\ell} := \left\{\boldsymbol {v} = \sum_ {k = 1} ^ {d} \nu_ {k} \boldsymbol {u} _ {k} \left| \bigwedge_ {1 \leq k <   \ell} \Omega_ {k} \wedge \Phi_ {\ell} \wedge \bigwedge_ {\ell \leq k <   k ^ {\prime} \leq d} \left(\Psi_ {k, k ^ {\prime}} ^ {\downarrow} \wedge \Psi_ {k, k ^ {\prime}} ^ {\uparrow}\right) \wedge \Xi \right. \right\},
$$
where the individual constraints are as follows (here  $\eta_0\coloneqq \infty$  so that e.g.  $\frac{\eta_1}{2\eta_0} = 0$ ):
$$
\begin{array}{l} \Omega_ {k} \colon 1 <   \frac {\nu_ {k}}{\nu_ {k} ^ {*}} \qquad \Phi_ {\ell} \colon \frac {\eta_ {\ell}}{2 \eta_ {\ell - 1}} <   \frac {\nu_ {\ell}}{\nu_ {\ell} ^ {*}} \leq 1 \qquad \Psi_ {k, k ^ {\prime}} ^ {\downarrow} \colon \frac {\eta_ {k ^ {\prime}}}{2 \eta_ {k}} \frac {\nu_ {k}}{\nu_ {k} ^ {*}} <   \frac {\nu_ {k ^ {\prime}}}{\nu_ {k ^ {\prime}} ^ {*}} \\ \Xi \colon \overline {{\boldsymbol {v}}} ^ {\top} \overline {{\boldsymbol {X} \boldsymbol {X} ^ {\top} (\boldsymbol {v} ^ {*} - \boldsymbol {v})}} > \lambda^ {\varepsilon / 3} \qquad \qquad \qquad \Psi_ {k, k ^ {\prime}} ^ {\uparrow} \colon \frac {\nu_ {k ^ {\prime}}}{\nu_ {k ^ {\prime}} ^ {*}} <   1 - \left(1 - \frac {\nu_ {k}}{\nu_ {k} ^ {*}}\right) ^ {\frac {1}{2} + \frac {\eta_ {k ^ {\prime}}}{2 \eta_ {k}}}. \\ \end{array}
$$
Thus  $S$  is connected, open, and constrained by  $\Xi$  to be within the ellipsoid  $\pmb{v}^{\top}\pmb{X}\pmb{X}^{\top}(\pmb{v}^{*} - \pmb{v}) = 0$  which is centred at  $\frac{v^*}{2}$ , with the remaining constraints slicing off further regions by straight or curved boundary surfaces.
In the most complex component of this work, we show that, for all  $t \geq T_1$ , the trajectory of the sum  $\boldsymbol{v}^t \coloneqq \sum_{j \in J_+} a_j^t \boldsymbol{w}_j^t$  of the active hidden neurons weighted by the last layer stays inside  $S$ , and the cosines of the angles between the neurons remain above  $1 - 4\lambda^\varepsilon$ . This involves proving that each face of the boundary of  $S$  is repelling for the training dynamics when approached from the inside; we remark that, although that is in general false for the entire boundary of the constraint  $\Xi$ , it is in particular true for its remainder after the slicing off by the other constraints. We also show that all points in  $S$  are positively correlated with all training points, which together with the preceding facts implies that, during this second phase of the training, the network behaves approximately like a linear one-hidden layer one-neuron network. Then, as the cornerstone of the rest of the proof, we show that, for all  $t \geq T_2$ , the gradient of the loss satisfies a Polyak-Lojasiewicz inequality  $\| \nabla L(\boldsymbol{\theta}^t) \|^2 > \frac{2\eta_d \| \gamma_{[n]} \|}{5\eta_1} L(\boldsymbol{\theta}^t)$ . Here  $T_2 \coloneqq \inf \{ t \geq T_1 \mid \nu_1^t / \nu_1^* \geq 1/2 \}$  is a time by which the network has departed from the initial saddle, more precisely when the first coordinate  $\nu_1^t$  of the bundle vector  $\boldsymbol{v}^t$  with respect to the basis consisting of the eigenvectors of the matrix  $\frac{1}{n} \boldsymbol{X} \boldsymbol{X}^\top$  crosses the half-way threshold to the first coordinate  $\nu_1^*$  of the teacher neuron.
The interior of the ellipsoid in the constraint  $\Xi$  actually consists of all vectors that have an acute angle with the derivative of the training dynamics in predictor space, and the "padding" of  $\lambda^{\varepsilon /3}$  is present because the derivative of the bundle vector  $\pmb{v}^{t}$  is "noisy" due to the latter being made up of the approximately aligned neurons. The remaining constraints delimit the subsets  $S_{1},\ldots ,S_{d}$  of the set  $S$ , through which the bundle vector  $\pmb{v}_{t}$  passes in that order, with each unique "handover" from  $S_{\ell}$  to  $S_{\ell +1}$  happening exactly when the corresponding coordinate  $\nu_{\ell}^{t}$  exceeds its target  $\nu_{\ell}^{*}$ . The non-linearity of the constraints  $\Psi_{k,k'}^{\uparrow}$  is needed to ensure the repelling for the training dynamics.
6 Implicit bias of gradient flow
Let us denote the set of all balanced networks by
$$
\Theta := \left\{\left(\boldsymbol {a}, \boldsymbol {W}\right) \in \mathbb {R} ^ {m} \times \mathbb {R} ^ {m \times d} \mid \forall j \in [ m ] \colon | a _ {j} | = \| \boldsymbol {w} _ {j} \| \right\}
$$
and the subset in which all non-zero hidden neurons are positive scalings of  $\boldsymbol{v}^*$ , have positive last-layer weights, and have lengths whose squares sum up to  $\| \boldsymbol{v}^* \| = 1$ , by
$$
\Theta_ {\boldsymbol {v} ^ {*}} := \left\{\left(\boldsymbol {a}, \boldsymbol {W}\right) \in \Theta \mid \sum_ {j = 1} ^ {m} \| \boldsymbol {w} _ {j} \| ^ {2} = 1 \wedge \forall j \in [ m ]: \boldsymbol {w} _ {j} \neq \mathbf {0} \Rightarrow \left(\overline {{\boldsymbol {w}}} _ {j} = \boldsymbol {v} ^ {*} \wedge a _ {j} > 0\right) \right\}.
$$
Our main result establishes that, as the initialisation scale  $\lambda$  tends to zero, the networks with zero loss to which the gradient flow converges tend to a network in  $\Theta_{\mathbf{v}^*}$ . The explicit subscripts indicate the dependence on  $\lambda$  of the parameter vectors. The proof builds on the preceding results and involves a careful control of accumulations of approximation errors over lengthy time intervals.
Theorem 7. Under Assumptions 1 and 2,  $L\left(\lim_{t \to \infty} \pmb{\theta}_{\lambda}^{t}\right) = 0$  and  $\lim_{\lambda \to 0^{+}} \lim_{t \to \infty} \pmb{\theta}_{\lambda}^{t} \in \Theta_{\pmb{v}^{*}}$ .
7 Interpolators with minimum norm
To compare the set  $\Theta_{v^*}$  of balanced rank-1 interpolator networks with the set of all minimum-norm interpolator networks, in this section we focus on training datasets of cardinality  $d$ , we assume the network width is greater than 1 (otherwise the rank is necessarily 1), and we exclude the threshold case of Lebesgue measure zero where  $\mathcal{M} = 0$ . The latter measurement of the training dataset is defined below in terms of angles between the teacher neuron and vectors in any two cones generated by different generators of the dual of the cone of all training points.
Let  $[\chi_1,\dots ,\chi_d]^\top \coloneqq X^{-1}$  and
$$
\mathcal {M} := \max  \left\{\cos \angle (\boldsymbol {p}, \boldsymbol {q}) - \sin \angle (\boldsymbol {p}, \boldsymbol {v} ^ {*}) \left| \begin{array}{c} \emptyset \subsetneq K \subsetneq [ d ] \\ \wedge \mathbf {0} \neq \boldsymbol {p} \in \operatorname {c o n e} \{\boldsymbol {\chi} _ {k} \mid k \in K \} \\ \wedge \mathbf {0} \neq \boldsymbol {q} \in \operatorname {c o n e} \{\boldsymbol {\chi} _ {k} \mid k \notin K \} \end{array} \right. \right\}.
$$
Assumption 3.  $n = d$ ,  $m > 1$ , and  $\mathcal{M} \neq 0$ .
We obtain that, surprisingly,  $\Theta_{\mathbf{v}^*}$  equals the set of all interpolators with minimum Euclidean norm if  $\mathcal{M} < 0$ , but otherwise they are disjoint.
Figure 1: Dependence of the maximum angle between active hidden neurons on the initialisation scale  $\lambda$ , for two generation schemes of the training dataset and a range of input dimensions, at the end of the training. Both axes are logarithmic, and each point plotted shows the median over five trials.
Theorem 8. Under Assumptions 1 and 3:
(i) if  $\mathcal{M} < 0$  then  $\Theta_{\pmb{v}^*}$  is the set of all global minimisers of  $\| \pmb{\theta} \|^{2}$  subject to  $L(\pmb{\theta}) = 0$ ;
(ii) if  $\mathcal{M} > 0$  then no point in  $\Theta_{\pmb{v}^*}$  is a global minimiser of  $\| \pmb{\theta}\|^2$  subject to  $L(\pmb{\theta}) = 0$ .
For each of the two cases, we provide a family of example datasets in the appendix. We remark that a sufficient condition for  $\mathcal{M} < 0$  to hold is that the inner product of any two distinct rows  $\chi_{k}$  of the inverse of the dataset matrix  $X$  is non-positive, i.e. that the inverse of the Gram matrix of the dataset (in our setting this Gram matrix is positive) is a  $Z$ -matrix (cf. e.g. Fiedler and Pták [1962]). Also, if the training points were orthogonal then all the  $\cos \angle (\pmb{p}, \pmb{q})$  terms in the definition of  $\mathcal{M}$  would be zero and consequently we would have  $\mathcal{M} < 0$ ; this is consistent with the result that, per sign class of the training labels in the orthogonal setting, minimising the Euclidean norm of interpolators coincides with minimising their rank [Boursier et al., 2022, Appendix C].
We consider two schemes for generating the training dataset, where  $\mathbb{S}^{d - 1}$  is the unit sphere in  $\mathbb{R}^d$ .
Centred: We sample  $\pmb{\mu}$  from  $\mathcal{U}(\mathbb{S}^{d - 1})$ , then sample  $x_{1},\ldots ,x_{d}$  from  $\mathcal{N}(\boldsymbol {\mu},\frac{\rho}{d}\boldsymbol {I}_d)$  where  $\rho = 1$ , and finally set  $\pmb{v}^{*} = \pmb{\mu}$ . This distribution has the property that, in high dimensions, the angles between the teacher neuron  $\pmb{v}^{*}$  and the training points  $\pmb{x}_i$  concentrate around  $\pi /4$ . We exclude rare cases where some of these angles exceed  $\pi /2$ .
Uncentred: This is the same, except that we use  $\rho = \sqrt{2} - 1$ , sample one extra point  $x_0$ , and finally set  $\pmb{v}^{*} = \overline{\pmb{x}}_{0}$ . Here the angles between  $\pmb{v}^{*}$  and  $\pmb{x}_i$  also concentrate around  $\pi /4$  in high dimensions, but the expected distance between  $\pmb{v}^{*}$  and  $\pmb{\mu}$  is  $\sqrt{\rho}$ .
For each of the two dataset schemes, we train a one-hidden layer ReLU network of width  $m = 200$  by gradient descent with learning rate 0.01, from a balanced initialisation such that  $\pmb{z}_j\stackrel {\mathrm{i.i.d.}}{\sim}\mathcal{N}(\mathbf{0},\frac{1}{dm}\pmb {I}_d)$  and  $s_j\stackrel {\mathrm{i.i.d.}}{\sim}\mathcal{U}\{\pm 1\}$ , and for a range of initialisation scales  $\lambda$  and input dimensions  $d$ .<sup>2</sup>
We present in Figure 1 some results from considering initialisation scales  $\lambda = 4^2, 4^1, \ldots, 4^{-12}, 4^{-13}$  and input dimensions  $d = 4, 16, 64, 256, 1024$ , where we train until the number of iterations
2We are making code to run the experiments available at https://github.com/englert-m/shallow_ReLU_dynamics.
Figure 2: Evolution of the training loss, and of an outside distribution test loss, during training for an example centred training dataset in dimension 16 and width 200. The horizontal axes, logarithmic for the training loss and linear for the test loss, show iterations. The vertical axes are logarithmic.
reaches  $2 \cdot 10^{7}$  or the loss drops below  $10^{-9}$ . The plots are in line with Theorem 7, showing how the maximum angle between active hidden neurons at the end of the training decreases with  $\lambda$ .
Figure 2 on the left illustrates the exponentially fast convergence of the training loss (cf. Theorem 6), and on the right how the implicit bias can result in good generalisation. The test loss is computed over an input distribution which is different from that of the training points, namely we sample 64 test inputs from  $\mathcal{N}(\mathbf{0},\mathbf{I}_d)$ . These plots are for initialisation scales  $\lambda = 4^{-2},4^{-3},\ldots ,4^{-7},4^{-8}$
We provided a detailed analysis of the dynamics of training a shallow ReLU network by gradient flow from a small initialisation for learning a single neuron which is correlated with the training points, establishing convergence to zero loss and implicit bias to rank minimisation in parameter space. We believe that in particular the geometric insights we obtained in order to deal with the complexities of the multi-stage alignment of hidden neurons followed by the simultaneous evolution of their norms and directions, will be useful to the community in the ongoing quest to understand implicit bias of gradient-based algorithms for regression tasks using non-linear networks.
A major direction for future work is to bridge the gap between, on one hand, our assumption that the angles between the teacher neuron and the training points are less than  $\pi /4$ , and the other, the assumption of Boursier et al. [2022] that the training points are orthogonal, while keeping a fine granularity of description. We expect this to be difficult because it seems to require handling bundles of approximately aligned neurons which may have changing sets of training points in their active half-spaces and which may separate during the training. However, it should be straightforward to extend our results to orthogonally separable datasets and two teacher ReLU neurons, where each of the latter has an arbitrary sign, labels one of the two classes of training points, and has angles less than  $\pi /4$  with them; the gradient flow would then pass close to a second saddle point, where the labels of one of the classes have been nearly fitted but the hidden neurons that will fit the labels of the other class are still small. We report on related numerical experiments in the appendix.
We also obtained a condition on the dataset that determines whether rank minimisation and Euclidean norm minimisation for interpolator networks coincide or are distinct. Although this dichotomy remains true if the  $\pi /4$  correlation bound is relaxed to  $\pi /2$ , the implicit bias of gradient flow in that extended setting is an open question. Other directions for future work include considering multi-neuron teacher networks, student networks with more than one hidden layer, further non-linear activation functions, and gradient descent instead of gradient flow; also refining the bounds on the initialisation scale and the convergence time.
We acknowledge the Centre for Discrete Mathematics and Its Applications at the University of Warwick for partial support, and the Scientific Computing Research Technology Platform at the University of Warwick for providing the compute cluster on which the experiments presented in this paper were run.
Gal Vardi, Gilad Yehudai, and Ohad Shamir. Gradient Methods Provably Converge to Non-Robust Networks. In NeurIPS, 2022. 2
Mingze Wang and Chao Ma. Early Stage Convergence and Global Convergence of Training Mildly Parameterized Neural Networks. In NeurIPS, 2022.
Yifei Wang and Mert Pilanci. The Convex Geometry of Backpropagation: Neural Network Gradient Flows Converge to Extreme Points of the Dual Convex Program. In ICLR, 2022. 3
Blake E. Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and Rich Regimes in Overparametrized Models. In  $COLT$ , pages 3635-3673, 2020. 2
Weihang Xu and Simon Du. Over-Parameterization Exponentially Slows Down Gradient Descent for Learning a Single Neuron. In  $COLT$ , pages 1155-1198, 2023.
Gilad Yehudai and Ohad Shamir. Learning a Single Neuron with Gradient Methods. In  $COLT$ , pages 3756-3786, 2020.
Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. A unifying view on implicit bias in training linear neural networks. In ICLR, 2021. 2
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Commun. ACM, 64(3):107-115, 2021. 1