Early Neuron Alignment in Two-layer ReLU Networks with Small Initialization
Authors: Hancheng Min
Enrique Mallada
Rene Vidal
Venue: ICLR 2024 poster (2024)
Keywords: Neural Networks
Small Initialization
Gradient Flow
Neuron Alignment
Primary area: learning theory
TL;DR: We provide a quantitative analysis of neuron dynamics in two-layer ReLU networks under small initialization

Abstract: This paper studies the problem of training a two-layer ReLU network for binary classification using gradient flow with small initialization. We consider a training dataset with well-separated input vectors: Any pair of input data with the same label are positively correlated, and any pair with different labels are negatively correlated. Our analysis shows that, during the early phase of training, neurons in the first layer try to align with either the positive data or the negative data, depending on its corresponding weight on the second layer. A careful analysis of the neurons' directional dynamics allows us to provide an $\mathcal{O}(\frac{\log n}{\sqrt{\mu}})$ upper bound on the time it takes for all neurons to achieve good alignment with the input data, where $n$ is the number of data points and $\mu$ measures how well the data are separated. After the early alignment phase, the loss converges to zero at a $\mathcal{O}(\frac{1}{t})$ rate, and the weight matrix on the first layer is approximately low-rank. Numerical experiments on the MNIST dataset illustrate our theoretical findings.

Published as a conference paper at ICLR 2024
EARLY NEURON ALIGNMENT IN TWO-LAYER RELU NETWORKS WITH SMALL INITIALIZATION
Hancheng Min
Enrique Mallada Johns Hopkins University ma adaę jhu whhu.edu.
René Vidal
University of Pennsylvania n@seas. annchmin@seas.o
University of Pennsylvania vi da rr wpesn.upenn... edu
upenn. eedu hanchmi
ABSTRACT
This paper studies the problem of training a two-layer ReLU network for binary classification using gradient flow with small initialization. We consider training dataset with well-separated input vectors: Any pair of input data with the same label are positively correlated, and any pair with different labels are negatively correlated. Our analysis shows that, during the early phase of training, neurons in the first layer try to align with either the positive data or the negative data, depending on its corresponding weight on the second layer. A careful analysis of the neurons' directional dynamics allows us to provide an O (Log 2) upper bound on the time it
takes for all neurons to achieve good alignment with the input data, where n is the number of data points and μ measures how well the data are separated. After the early alignment phase, the loss converges to zero a (+t) rate, and the weight matrix on the first layer is approximately low-rank. Numerical experiments on the MNIST dataset illustrate our theoretical findings.
INTRODUCTION
Neural networks have shown excellent empirical performance in many application domains such vision (Krizhvskyy et al., 2012). speech (Hinton al., 2012) and video games (Silver al.) 2016). Despite being highly overparametrized, networks trained by gradient descent with random initialization and without explicit regularization enjoy good generalization performance. One possible explanation for this phenomenon is the implicit bias or regularization induced by first-order algorithms under certain initialization assumptions. For example, first-order methods applied to (deep) matrix factorization models may produce solutions that have low nuclear norm (Gunasekar et al.] (2117) and/or low rank (Arora et al., (201)), and similar phenomena have been observed for deep tensor factorization (Razin al., 2022). Moreover, prior work such as (Saxe et al. [2014] SSögee & Soltanolkotabi (021)) has found that deep linear networks sequentially learn the dominant singular values of the ouput- correlation matrix.
biases can often be achieved by small initialization. This has motivated series of works that theoretically analyze the training dynamics of first-order methods for neural networks with small initialization. For linear networks, the implicit bias of small
is widely known that these
initialization has been studied in the context of linear regression (Saxe et al.)
2014
Gidel et al.
2019
Min et al. 2021 Varre et al. (003) and matrix factorization (Gunasekar et al. 2017 Arora et al
2019
Li al. 2018 |021] SSögerr& & Soltanolkotabi 2021 Yaras al. 2023 Soltanolkotabi et al.
2023).
Recently, the effect of small initialization has been studied for two-layer ReLU networks (Maennel
al. 2018 Lyu et al.) [202]] Phuong & Lampert 2021 Boursier et al. 2022]. For example, Maennel ct al. (2018) observes that during the early stage of tramiig, neurons in the first layer converge to
one out of finitely many directions determined by the dataset. Based on this observation, (Punnn,
Lampert (202T) shows that in the case of well-separated data, where any pair of input data with the same label are positively correlated and any pair with different labels are negatively correlated, there are only two directions the neurons tend to converge to: the positive data center and the negative one. Moreover, Phuong & Lampert (202T) shows that such directional convergence holds, then the loss converges, and the resulting first-layer weight matrix is low-rank. However, directional convergence is assumed in their analysis; there is no explicit characterization of how long it takes to achieve directional convergence and how the time to convergence depends on the initialization scale.

---

Published as a conference paper at ICLR 2024
Paper contributions: In this paper, provide complete analysis of the dynamies of gradient flow for training a two-layer ReLU network on well-separated data with small initialization. Specifically, show that if the initialization is sufficiently small, during the early phase of training the neurons in the first layer try to align with either the positive data or the negative data, depending on its corresponding weight on the second layer. Moreover, through a careful analysis of the neuron's directional dynamics we show that the time it takes for all neurons to achieve good alignment with the input data is upper bounded by O log_) ), where n. is the number of data points and measures how well the data separated. We also show that after the early alignment phase the loss converges to zero at a (()) rate and that the weight matrix on the first layer is approximately low-rank
vR
Notation: We denote the Euclidean norm of vector by ||ع||, the inner product between the vectors & and by (ar y) = x y, and the cosine of the angle between them as cos(1 (1, y) ( T Tell
x y ). For
an n X m matrix A, we let A denote its transpose. We also let ||A|| and || denote the spectral norm and Frobenius norm of A, respectively. For scalar-valued or matrix-valued function of time, d (t) denote its time derivative. Furthermore, we define 1 to be the is true and 1 = 0 otherwise. We also let I denote the
F(t), we let F = F()) dt indicator for statement A identity matrix, and N N(u, o²) denote the normal distribution with mean μ and variance o².
2 PRELIMINARIES
In this section, we first discuss problem setting. We then present some key ingredients for analyz- ing the training dynamics of ReLU networks under small initialization, and discuss some of the weaknesses/issues from prior work.
2.1 PROBLEM SETTING
We are interested in a binary classification problem with dataset [ar,,... ,In] E RDx (input data) and [....... yo]" E +1 (labels). For the classifier, f IRD R, we consider two-layer
ReLU network:
W,)
= = =
E u, x),
parametrized by network weights W = [w],.... w.] E RDxh , = [u1, vh T E Rhx1 where (-) max(-,0] is the ReLU activation function. We aim to find the network weights that minimize the training loss L(W, v) = En_ E(y,, I (xx; W,u), u)). where E
R is either the exponential loss E 1) exp(-wú) or the logistic loss Cly, = log(1 exp(- The network is trained via the gradient flow (GF) dynamics
W E Ow DWW,,, à O,L(W, v),
where Ow L, O.L are Clark sub-differentials of L. Therefore, 2 a differential inclusion (Bolte et al., 2010). For simplicity of presentation, instead of directly working on this differential inclusion, our theoretical results will be stated for the Caratheodory solution (Reid 197I) of ( when the ReLU subgradient is fixed as o' (x)) 1a>d") T In Appendix E show that under the data assumption of our interest (to be introduced later), the Caratheodory solution (s) {W()), (W)))v(() v()} exists globally for all t E [0, oo), which we call the solution (s) of (2) throughout this paper.
To initialize the weights, we consider the following initialization scheme. First, we start from a weight matrix Wo Dxx and then and then initialize the weights
W(0)= eWo v0( E {llw (0)] Ilo,, (0) Vj [h).
That is, the weight matrix Wo determines the initial shape of the first-layer weights W()) and we use c to control the initialization scale and we are interested in the regime where is sufficiently small. For the second layer weights v(0), each Uy (0) has magnitude || (0)]] and we only need to decide its sign. Our results in later sections are stated for deterministic choice of e, Wo. and v(0), then we comment on the case where Wo is chosen randomly via some distribution.
The resulting weights in 3 are always "balanced". i.e., of (0) ey (0) can only take two values: either [[w]( - -1 (0)). More importantly, under GF (2. this
(0)]1² = 0, Vj [/.] because
In Appendix
F we discuss how our results can he extended to the solution differential inclusion.
2

---

Published as a conference paper at ICLR 2024
balancedness is preserved (Du et al. 2018: : (t) [t,(1)) 0, 2 0, E [b.. In addition. it is shown in Boursier et al. (2022) that sign oy (t)) (1 Vt 2 0, E [h]. and the dynamical behaviors of neurons will be divided into two types, depending on sign(r;(0))
Remark 1. For our theoretical results, the balancedness condition assumed for technical purposes:
it simplifies the dynamics of GF and thus the analysis. It is common assumption for many existing
works on both linear (Arora et al. 2018b) and nonlinear @Phuong & Laamper 2027 Boursier et al.
2022) neural networks. For the experiments in Section + we use a standard Gaussian initialization (not balanced) with a small variance to validate our theoretical findings.
consider the case where all columns of Wo are nonzero. Le.. (10y(00]] > 0, Vj E [h]. We make this assumption because whenever wy(0) = O, we also have yi vy(0)= from the balancedness, which together would imply ún = 0, ز = 0 under gradient flow. As a result, Wj and Uy would remain zero and thus they could be ignored in the convergence analysis.
Remark 2. Without loss of generality,
Remark 3. Our main results will depend on both max, [00]( and min, [0)((0]], as shown in our proofs in Appendices [
[and(D) Therefore, whenever speak of small initialization, we will say that e is small without worrying about the scale of Wo, which is already considered in our results.
2.2 NEURAL ALIGNMENT WITH SMALL INITIALIZATION: AN OVERVIEW
Prior work argues that the gradient flow dynamics (2) under small initialization 3 i.e., when is sufficiently small, can be roughly described as "align then fit" (Maennel et al. 2018 Boursier al. 2022) Durinn the early phase of training, every neuron Wy. E [h] keeps small norm [uy]
Twu significantly in order to locally maximize "signed coverage" (Maennel et al. 2018) of itself w.r.t. the training data. After the alignment phase, part of (potentially all) the neurons grow their norms in order to fit the training data, and the loss decreases significantly. The analysis for the fitting phase generally depends on the resulting neuron directions at the end of the alignment phase (Phuong & Lampert [2021] Boursier et al., (022). However, prior analysis of the alignment phase either is based on a vanishing initialization argument that can not be directly translated into the case finite but small initialization (Maennel et al. 2018) or assumees some stringent assumption on the data (Boursier et al. 2022). In this section, we provide a brief overview of the existing analysis for neural alignment and then point out several weaknesses in prior work.
I while changing their directions "i
Since during the align- ment phase all neurons have small norm, prior work mainly d wy focuses on the directional dynamics, i.e., ar [T::T of the neu- rons. The analysis relies on the following approximation of the dynamics of every neuron Wj, j E [A]:
Prior analysis of the alignment phase:
Pa, (11z)
taa (wy
22
d w/ sign( (0)) (r)))) (w)))))) dt [wyl
322
(4)
37%
sign(r, (0))
where P.. = 1 ww orthogonal to and
is the projection onto the subspace T
Figure 1: Illustration of d w, ing the early alignment phase. 21 (5) has +1 label, and 22,42 22 have -1 labels, 21 22 lie inside the halfs- pace (x, wz)
dr ||wy] dur-
xa(w)= E 2: u yeu;
denotes the signed combination of the data points activated by w. First of all, 44) implies that the dynamics wy [Te are approxi- w mately decoupled, and thus one can study each |mu! separately. Moreover, as illustrated in Figure [ if sign(u, (0)) > O, the flow (4) pushes w, towards x. (,(w)) (w;). since wi is attracted by its currently activated positive data and repelled by its currently activated negative data. Intuitively, during the alignment phase, neuron Wy with sign (v, (0)) > 0 would try to find direction where can activate as much positive data and as less negative data as possible. If (ign(r;(0)) < O. the opposite holds.
0 (gray shaded),
thus xa (W,) = I1 22. Since
sign(r,(0) > 0, GF pushes wj to- wards xz
Indeed, Maennel et al. (2018) claims that the neuron Wj would be aligned with some "extreme vectors". defined as vector w E SD-1 that locally maximizes Σiein) y:o ((x,, w)) (similarly, wy with sign(u, (0)) < 0 would be aligned with the local minimizer). and there are only finitely many such vectors. The analysis is done under the limit + where the approximation in (4) is exact.
3

---

Published as a conference paper at ICLR 2024
Weakness in prior analyses: Although Maennel et al. (2018) provides great insights into the dynamical behavior of the neurons in the alignment phase, the validity of the aforementioned approximation for finite but small remains in question. First, one needs to make sure that the error
is sufficiently small when is finite in order to justify (4) as a good approximation. Second, the error bound needs to hold for the entire alignment phase. Maennel et al. 2018) assumes + O: hence there is no formal error bound. In addition, prior analyses on small initialization (Stöger & Soltanolkotabi 202T: Boursier et al. (022) suggest the alignment phase only holds for (log 2) time. Thus, the clamm in Maennel et al. (2018) would only hold if good alignment is achieved before the alignment phase ends. However, Maennel et al. (2018) provides no upper bound on the time it takes to achieve good alignment. Therefore, without a finite analysis, [Maennel et al.] (2118) fails to fully explain the training dynamics under small initialization. Understanding the alignment phase with finite requires additional quantitative analysis. To the best of our knowledge, this has only been studied under a stringent assumption that all data points are orthogonal to each other (Bourster et al. 2022), or that there are effectively two data points Wang & Ma (2033)
ai; de [aal sign ,
Goal of this paper: In this paper, we want to address some of the aforementioned issues by developing a formal analysis for the early alignment phase with a finite but small initialization scale e. We first discuss our main theorem that shows that directional convergence can be achieved within bounded time under data assumptions that are less restrictive and have more practical relevance. Then, we discuss the error bound for justifying (4) in the proof sketch of the main theorem.
3. CONVERGENCE OF TWO-LAYER RELU NETWORKS WITH SMALL INITIALIZATION
Our main results require the following data assumption:
Assumption 1. Any pair of data with the same (different) label is positively (negatively) correlated, i.e., min:,, T| ||I
(,,9),,,;;). := f O.
Given training dataset, we define S. :- E ((,,)>> 19,00, Vi} to be the cone in R" such that whenever neuron us E S4. w is activated exclusively by every x; with positive label (see Figure 2). Similarly, for x; with negative labels, we define S_ ==
Figure 2: Neuron alignment un- der Assumption T For neurons in V4, ① if it lies inside S_, then it gets repelled by x_ and escapes S.; Once outside S_, it may 2 get repelled by some negative data and eventually en- ters Saead or may 3 gain some activation on positive data and eventually enter S4, then get constantly attracted by x..
I (2,,)>00 = 1g, Vij. Finally, we define Stead {a E RD (z) & 0, Vi} to be the cone such that whenever data activates w. Given Assumption [ it can be shown (see Appendix C that S+ (S.) ) non-empty, convex cone that contains all positive data xy E I+ (negative data 24, i E I_). Saead is a convex cone as well, but not necessarily non-empty. We illustrate these cones in Figure 2 given some training data (red solid arrow denotes positive data and blue denotes negative ones).
{z E RD
E Saead,
Moreover, given some initialization from 3). we define I+
{i E [m] : y/4 > 0} to be the set of indices of positive data, and
I_ = (i E [n] Yi < 0) for negative data. We also define
0} to be the set of indices of neurons with positive second-layer entry and V- := {j E [6] sign by (t) < <0} for neurons with negative second-layer entry. Note that, discussed in Section 2.1 (t) does not change under balanced initialization, thus V4,V= are time invariant. Further, as we discussed in Section 2.2 about the early alignment phase, expect that every neuron in V4 will drift toward the region where positive data concentrate and thus eventually reach S+ or Saead, as visualized in Figure 2 (2+ x_ shown in the figure are defined in Assumption 2 Similarly, all neurons in V_ would chase after negative data and thus reach S_ or Stead- Our theorem precisely characterizes this behavior.
V+ = (j [h] sign(e; (t)
3.1
MAIN RESULTS
Our main results are stated for solutions to the GF dynamics 2 However, in rare cases, solutions to could be non-unique and there are potentially "irregular solutions" (please refer Appendix E.4
4

---

Published as a conference paper at ICLR 2024
for details) that allow some neurons regain activation even after becoming completely deactivated in Stead- We deem such irregular solutions of little practical relevance since when implementing gradient descent algorithm in practice, neurons in Sacad would receive zero update and thus stay in Sdead- Therefore, our main theorem concerns some regular solutions to (2) (the existence of such solutions is shown in Appendix (.2). as defined below.
Definition 1. A solution (I),, v(t) to (2) is regular if satisfy that wy(1o) E Stead for some j E [b)] and some to 2 0 implies wy(t) E Sdead, Vt 2 lo-
Before present main theorem, also need the following assumption on the initialization, for technical reasons, essentially asking the neuron wy (0), E V- (or ()()). E V-, resp.) to not be completely aligned with 24 (or x-, resp.).
Assumption 2. The initialization from
satisfies that maxjev,
(r,(0)
< 1, and
maxjev_
[wy(
*n (0)(0)
1, where x+ EieI, 24 and x_ DiEI_ Σz-
We are now ready to present our main result (given Assumption [and Assumption 2
Theorem 1. Given some initialization from
3 if 0 VR exp(- vu log n)), then for any regular
solution to the gradient flow dynamics 2 we have
1. (Directional convergence in early alignment phase) 3t = 0 log
), such that
vp
Mj E V+, either (y((1) E S+ or woy(t)) E SSeend Moreover, if maxiez, (mn{(0),z;) > 0, then my(ti) S_. Vj E V_, either wy(ti) E S_ or wy((t) E Stead- Moreover, if maxxe_ (w,(0). x;) > then wr(ti) S..
2. (Final convergence and low-rank bias) Vt 2 t and Vjj E [h], neuron wz(t) stays within S_ (S.. or Steea) if wf,(t)) E S+ (S_, or Stead resp.). Moreover, if both S+, and S_ contains at least one neuron at time tg then
3a > 0 and 3t, with t t2 = 2 log Vhe such that [(t) S (()=)a(t-ta)+1 Vt m22tz 1z
Lit))
l oo, [[((III + oo and (1) < 2/I1( + O(e). Thus, the stable rank of (() satisfies lim supt- /WOLP//WWW (1)B S
We provide proof sketch that highlights the technical novelty of our results in Section 3 Our O-)) notations hide additional constants that depend on the data and initialization, for which we refer readers to the complete proof of Theorem in Appendix D We make the following remarks:
Early neuron alignment: The first part of the Theorem I describes the configuration of all neurons at the end of the alignment phase. Every neuron in V_ reaches either S. Saend by f,, and stays there for the remainder of training. Obviously, we care about those neurons reaching S+, as any neuron in Steadd does contribute the final convergence all. Luckily, Theorem suggests that any neuron in V4 that starts with some activation on the positive data, i.c.. it is initialized in the union of halfspaces Usez, {w (w, 24) > 0}, will eventually reach S4. A similar discussion holds for neurons in V_ We argue that randomly initializing Wo ensures that with high probability, there will be least a pair of neurons reaching S+ and S_ by time & (please see the next remark). Lastly, we note that it is possible that Saead 0, in which case every neuron reaches either S_. or S_.
Merits of random initialization: Our theorem is stated for a deterministic initialization 3 given an initial shape Wo. In practice, one would use random initialization to find a Wo. for example,
[Woks s.idd N (0, 1).. First, our Theorem applies to this Gaussian initialization: Assumption is x- | = and { [ ) ||=|| 1 = wy [[)](()0]] ||
satisfied with probability one because the events
(0) x+
have probability zero. Moreover, any neuron in V4 has at least probability 1/2 of being initialized
within the union of halfspaces USET {w (w, xr > 0}. which ensures that this neuron reaches S4. Thus when there are m neurons in V4, the probability that S. has at least one neuron at time is lower bounded by 1 2 (same argument holds for S_), Therefore, with only very mild overparametrization on the network width h, one can make sure that with high probability there is at least neuron in both S, and 8_, leading to final convergence.
5

---

Published as a conference paper at ICLR 2024
Importance of quantitative bound on f1: The analysis for neural alignment relies on the ap- proximation in 4.. which, through our analysis (see Lemma [T). is shown to only hold before T = achieved within 4 time, then the initialization scale must be chosen to be O
6 log vhe ) thus if one proves, through the approximation in ④ that good alignment is
yh. exp(-nl1)) so
that 1 - T, i.e. the proved alignment should finish before the approximation (4) fails. Therefore, without an explicit bound on E1, one does not know a prior how small should be. Our quantitative analysis shows that under e directional convergence is achieved within 1 = O
logn time. This vF
bound, in return, determines the bound for initialization scale c. Moreover, our bound quantitatively reveals the non- non-trivial dependency on the "data separation" y for such directional convergence to occur. Indeed, through numerical illustration in Appendix A.2 we show that the dependence on the data separability fa > O is crucial in determining the scale of the initialization: As approaches zero, the time needed for the desired alignment increases, necessitate the use of a smaller e.
Refined aligument within S_, Once neuron in V+ reaches S+, it never leaves S+. Moreover, it always gets attracted by x.. Therefore, every neuron gets well aligned with x4. cos(w,, xT) ~ 1, Ywu E S+. A similar argument shows neurons in S_ get attracted by x_. We opt not to formally state it in Theorem a the result would be similar that in (Boursier et al., 2022). and alignment with ,,, x- is not necessary to guarantee convergence. Instead, we show this refined alignment through our numerical experiment Section
Final convergence and low-rank bias: We present the final convergence results mostly for the completeness of the analysis. GF after i can be viewed as fitting positive data xz, i I+, with a subnetwork consisting of neurons in S4. and fitting negative data with neurons in S_. By the fact that all neurons in S+ activate all z,, i E I+, the resulting subnetwork is linear, and so is the subnetwork for fitting 24, i E I_. The convergence analysis reduces to establishing (1/tt) convergence for two linear networks (Arora et al., Sat Min et al.) [021 Yun et al.) 2020). The non-trine non-trivial and novel part is to show that right after the alignment phase ends, one expect substantial decrease of the loss (starting from time 22 = @ A log vhe )). An alternative way of proving convergence is
by observing that at f,, all data has been correctly classified (w.r.t. sign of f). which is sufficient for showing O (Tiggz) convergence (Lyu & Li. 2019 Ji & Telgarsky 2020 of the loss, but this asymptotic rate does not suggest a time after which the loss start to decrease significantly. As for the stable rank, our result follows the analysis in Le & Jegelka (2022). but in a simpler form since ours is for linear networks. Although convergence is established partially by existing results, we note that these analyses are all possible because we have quantitatively bound t in the alignment phase.
3.2 COMPARISON WITH PRIOR WORK
Our results provide a complete (from alignment to convergence), non-asymptotic (finite e), quantita- tive (bounds on f1, f2) analysis for the GF in (2) under small initialization. Similar neural aligmment has been studied in prior work for orthogonally separable data (same as ours) and for orthogonal data, and we shall discuss them separately.
Alignment under orthogonally separable data: Phuong & Lampert 2021 assumes that there exists a time f such that at f1, the neurons are in either S_ or Saeud and their main contribution is the analysis of the implicit bias for the later stage of the training. they justify their assumption by the analysis in Maennel et al. (2018), which does not necessarily apply to the case of finite e, as we discussed in Section 2.2 Later [Wang & Pilancil (2022) shows L exists, provided that the initialization scale is sufficiently small, but still with no explicit analysis showing how t1 depends on the data separability μ and the size of the training data n. Moreover, there is no quantification how small should be. In our work, all the results are non-asymptotic and quantitative: we show that good alignment is achieved within t = O( log ) time and provide an explicit upper bound on Moreover, our results highlight the dependence on the separability μ >> 0, (Further illustrated in Appendix A..2) which is not studied in Phuong & Lampert (202I); Wang & Pilanci (2022).
Alignment under orthogonal data: In Boursier et al. (2022) the neuron alignment is carefully analyzed for the case all data points are orthogonal to each other, i.e., 2y) = 0, Vi + j E [n]
We point out that neuron behavior is different under orthogonal data (illustrated in Appendix A.3) only the positive (negative) neurons initially activate all the positive (negative) data will end up in S. (S_). ). In our case, all positive (negative) neurons will arrive at S+ (S_), unless they become dead

---

Published as a conference paper at ICLR 2024
neuron. Moreover, due to such distinction, the analysis is different: Boursier et al. (2022) restrict their results to positive (negative) neurons w, that initially activate all the positive (negative) data, and there is no need for analyzing neuron activation. However, since our analysis is on all positive neurons, regardless of their initial activation pattern, it utilizes novel techniques to track the evolution of the activation pattern (see Section 3.3).
Other related work: Convergence of two-layer (leaky- ReeU networks are also studied under non-
small initialization settings, mainly for gradient descent Wang & Ma (2022) and for training only the first-layer weights Frei et al. (2022); ; Kou et al. (2023). There is no direct comparison to them as they study the convergence in other regimes. Nonetheless, the analyses of neural alignment remain essential in these works but are done through different tools (one no longer has an approximation in (4). We note that such analyses also require certain restrictive data assumptions. For example, Wang & Ma
(2022) assumes orthogonal separability, together with some geometric constraint on the data: Frei et al. 2022); Kou et al. (2033) assumes high-dimensional near-orthogonal data.
3.3 PROOF SKETCH FOR THE ALIGNMENT PHASE
In this section, we sketch the proof for our Theorem I First of all, it can be shown that S4, Sdead are trapping regions for all (t))), E V-, that is. whenever mty(t) gets inside S+ (or Saeud), it never leaves S+ (or Stea)). Similarly, S are trapping regions for all ( E V_ The alignment phase analysis concerns how long it takes for all neurons to reach one of the trapping regions, followed by the final convergence analysis on fitting data with +1 label by neurons S., and fitting data with -1 label by those in S_. We have discussed the final convergence analysis in the remark "Final convergence and low-rank bias", thus we focus on the proof sketch for the early alignment phase here, which is considered as main technical contribution.
Approximating dt T Our analysis for the neural alignment is rooted in the following Lemma:
Lemma 1. Given some initialization from 3 ife = O (T) then there exists T = O L log Vhe
such that any solution to the gradient flow dynamics ② scatiffies that Vt = T,
d wx(t) dt [1]
sign (o (envh))
max
(6)
This Lemma shows that the error between de [[y] and sign(u,(0))P (t)) (t,(()) can be arbi- trarily small with some appropriate choice of (to be determined later). This allows one analyze the true directional dynamics [[]T
d (t)()
mnx(t) using some property of (a) (w)(()))) (t)), which lcads to a
1 = ( log 12) upper bound on the time it takes for the neuron direction to converge to the sets S+,, be made sufficiently small so that the error bound holds until the directional convergence is achieved, i.e. T 2 f1. We will first illustrate the analysis for directional convergence, then close the proof sketch with the choice of sufficiently small e.
v/z S_, or Sdead- Moreover, it also suggests
Activation pattern evolution: Given a sufficiently small e, one can show that under Assumption for every neuron wy that is not in Saead we have:
d
W
> 0, Vi [n],ifje
dt
||w,]]
|x>||
d
w
< [n)
(8)
dt ||y]l
|>>:||
This is because if neuron satisfies (1ra, wyy) = O for some i, and is not in Stead. GF moves wy towards Fa (w, Eite) Wiylie Interestingly. Assumption [
implies (17,)) xo (w;)) > 0, Vi E [n].
which makes d wy = sign(u; (O)) (w;) point inward (or outward) the halfspace (1a7), wy) > Zt TTe 0, if sign(v; (0))
(or sign(1;(0)) < O. respectively). See Figure 3 for illustration.
As consequence, neuron can only change its activation pattern in particular manner: a neuron in V4, whenever activated by some z; with y/s = +1, never loses the activation on s,; thereafter.
because implies that GF pushes [wyl towards I; at the boundary (wy. O. Moreover,
7

---

Published as a conference paper at ICLR 2024
also shows that a neuron in V+ will never regain activation on 2; with y -1 once it loses the activation because GF pushes ||ws]] V- never loses activation on negative data and never gains activation on positive data.
wi, against z at the boundary (uy, xz) = O. Similarly, a neuron in
JE V+
2i (u(()) (uy)
!
Saad
S,.
Fa (uy)
101
impossible for (wo))
St
Figure 4: Illustration of the activation pattern evolution. The epochs on the time axis denote the time Wi changes its ac- tivation pattern by either losing one negative data (denoted by "+") or gaining one positive data (denoted by "-"). The markers are colored if currently activates w. During the alignment phase 0 I S f1. neuron w. E V++ starts with activation on all negative data and no positive data, every O (1/ (1/na) time, it must change its activation, unless either ! reaches Saead. or 2 it activates some positive data at some epoch then eventually reaches S4.
Figure 3: For j E V+. Assump- tion enforces (Tilk, Fa (w,)) O, thus GF pushes w/ inward the half- space (1il/s wyy > 0 at (1,,, wy) = (i.e. towards gaining activation on I, if y/4 +1, or losing activation on Vi. if y/ = 1.). Sz and S denotes the subspace orthogonal to T2 and wy, respectively.
Equations 7 and (8) are key in the analysis of alignment because they limit how many times neuron can change its activation pattern: a neuron in V+ can only gain activation on positive data and lose activation on negative data, thus at maximum, neuron Wjj j E V+, can start with full activation on all negative data and no activation on any positive one (which implies w,) E S_) then lose activation every negative data and gain activation every positive data as GF training proceeds (which implies (t1) E S4), taking at most n changes on its activation pattern. See Figure 4 for illustration. Then, since it is possible to show that neuron w with j E V that has cos( W0, x-- < 1 (guaranteed by Assumption 2 and is not in S+ or Saead. must change its activation pattern after O V time (that does not depend on c), where na is the number of data that currently activates wy. one can upper bound the time for wy to reach S4 or Stead by some has activation on at least one positive data, i.e., max,gr, (m,(0), x;) >> 0 since it cannot lose this activation. A similar argument holds for wj. j V_ that they reaches either S_ or Ssead before t1.
Bound on activation transitions and duration:
O log
v/ constant independent of e. Moreover, Wj must reach S+, if it initially
Choice of e: All the aforementioned analyses rely on the assumption that the approximation in equation holds with some specific error bound. We show in Appendix (htat the desired bound
is d
(0) (1) (0g (1) (wy (t)) < ((VT), which, by Lemma can be achieved
N0
by sufficiently small initialization scale = O( yTh
Moreover, the directional convergence
Vin
(which takes O logn time) should be achieved before the alignment phase ends, which happens V at T { log Vhe E = O vे exp( n opt to present 62 in our main theorem because e2 beats when is large.
This is ensured by choosing another sufficiently small initialization scale
log n)). Overall, the initialization scale should satisfy e S min(fr,e)) We
vi
4 NUMERICAL EXPERIMENTS
We use toy example in Appendix A. 1 to clearly visualize the neuron alignment during training (due to space constraints). In the main body of this paper, we validate our theorem using binary classification task for two MNIST digits. Such training data do not satisfy Assumption [ every data vector is grayscale image with non-negative entries, making the inner product between any pair of data non-negative, regardless of their labels. However, we can preprocess the training data by centering: 2; + I; - F. where ă = Σiejn 2t/n. The preprocessed data, then, approximately satisfies our assumption (see the left-most plot in Figure 55) a pair of data points is very likely to have a positive correlation if they have the same label and have negative correlation if they have
8

---

Published as a conference paper at ICLR 2024
Correlation
Aligiment with
Convergence and low-rank bias
Dsta/iteuron Visualization
1.0
140
0.8
mate,
[Mg
0.6
5000
10:00
15000
1.0
0.4
-2
-10
5000 10:00
15000
5000
10000 15000
Figure 5: Training two-layer ReLU network under small initialization for binary classification on MNIST digits 0 and 1. (First Plot) Data correlation [xx,, 2,ylay heatmap, where the data reordered by their label (digit first, then digit 0); (Second Plot) Alignment between neurons and the aggregate positive/negative data x+ Liez, Iz. H- = Dier_ T4. (Third Plot) The loss L, the stable rank and the squared spectral norm of W during training; (Fourth Plot) Visualizing neuron centers , wh and data centers ž+, ___ (at iteration 15000).
different labels. Thus we expect our theorem to make reasonable predictions on the training dynamics with preprocessed data. For the remaining section, we use 2;, i [n], denote the preprocessed d
(centered) data and use E to denote the mean of the original data.
We build two-layer ReLU network with h 50 neurons and initialize all entries of the weights as d.d.d. [W], ~ A (0, (",a²) aa))) Uy i.ddd N (0, Vi E [n].] E [h] with a 10 10-6 Then we run gradient descent on both W and with step size 2 x 10-3. Notice that here the weights are not initialized to be balanced as (3). The numerical results are shown in Figure
Alignment phase: Without balancedness. one no longer has sign( (t)) = sign(u; (0)). With little abuse of notation, we denote V- (1) = (j E [h] sign(v,(()) > 0} and (1) = E [b] sign(1z(1) (4) > 0), and we expect that at the end of the alignment phase, neurons in V+ are aligned with 244 = Σiez, I4, and neurons in V- with x_ = Σier_ 24. The second plot in Figure shows such an alignment between neurons and x+, x_. In the top part, the red solid line shows cos(w,,24 during training, where w+ Σjev, wy/VV-1, and the shaded region defines the range between minjev, cos(w;, and maxjev_ cos( (1, x+). Similarly, the bottom part, the green solid line shows cos(w,, during training, where w_ = Ejev_ w;/DV-1, and the shaded region delineates the range between minjey cos(w;, and maxjev cos(w;, .)). Initially, every neuron is approximately orthogonal to Z+, x_ due to random initialization. Then all neurons in V+ (V-) start to move towards x+ (x and achieve good alignment after 2000 itcrations. When the is not exactly
loss starts to decrease, the alignment drops. We conjecture that because Assumption I satisfied, neurons in V4 have to fit some negative data, for which x. not the best direction.
Final convergence: After ~ 3000 iterations, the norm || 12 starts to grow and the loss decreases, as shown the third plot in 5igure Moreover, the stable rank || / [2 decreases below For this experiment, we almost have cos(x4,
1,, thus the neurons in V++ (aligned with x..) and those in V_ (aligned with x-)) )) are almost co-linear. Therefore, the stable rank / is almost 1, as seen from the plot. Finally, at iteration 15000, visualize the mean neuron V. = Ejev, V+ T_ = Σjev_ w as grayscale images, and compare them with F4 r II
x x | good alignment.
Comparison with other training schemes: For two-layer ReLU networks, there another line of
work (Brutzkus et al. 2018 Wang et al. 2019) that studies GD/SGD only on the first-layer weights fixed throughout training. In Appendix A.5] we compare our
W and keeping the second-layed training schemes to those in Brutzkus et al. (2018); Wang et al. (2019), and show that while both schemes achieve small training loss, the aforementioned two-phase training (alignment then final convergence) does no happen if only the first-layer in trained.
5 CONCLUSION
This paper studies the problem of training a binary classifier via gradient flow two-layer ReLU networks under small initialization. We consider a training dataset with well-separated input vectors. A careful analysis of the neurons' directional dynamics allows to provide an upper bound on the time it takes for all neurons to achieve good alignment with the input data. Numerical experiment on classifying two digits from the MNIST dataset correlates with our theoretical findings.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENT
The authors thank the support of the NSF-Simons Research Collaborations on the Mathematical and Scientific Foundations of Deep Learning (NSF grant 2031985), the NSF HDR TRIPODS Institute for the Foundations of Graph and Deep Learning (NSF grant 1934979), the ONR MURI Program (ONR grant 503405-78051), and the NSF CAREER Program (NSF grant 1752362). The authors thank Ziqing Xu and Salma Tarmoun for the insightful discussions.
REFERENCES
Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. In International Conference on Learning Representations, 2018a.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterizzaion. In 35th International Conference on Machine Learning. 2018b.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
Jérôme Bolte, Aris Daniilidis, Olivier Ley, and Laurent Mazet. Characterizations of tojasiewicz inequalities: subgradient flows. talweg, convexity. Transactions of the American Mathematical Society, 362(6)33319-336 2010.
Etienne Boursier, Loucas Pullaud-Vive Vivien, and Nicolas Flammarion. Gradient flow dynamics of shallow relu networks for square loss and orthogonal inputs. In Advances in Neural Information Processing Systems, volume 35, pp. 20105-20118 18, 2022.
Alon Brutzkus, Amir Globerson, Eran Malach, and Shai Shalev-Shwartz. SGD learns over- parameterized networks that provably generalize on linearly separable data. International Conference on Learning Representations, 2018.
Simon Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
Zafer Erean. Extension and separation of vector valued functions. Turkish Journal of Mathematics, 21(4), 1997.
A. Filippov. The existence of solutions of generalized differential equations. Mathematical notes of the Academy of Sciences of the USSR, 10(3):608-611 611, September 1971. ISSN 1573-8876. doi: 10.1007/BF004444722.
Spencer Frei,
Gal Vardi, Peter Bartlett, Nathan Srebro, and Wei Hu. Implicit bias in leaky relu networks trained on high-dimensione data. In The Eleventh International Conference on Learning Representations, 2022.
Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien Implicit regularization of discrete gradient dynamics linear neural networks. In Advances in Neural Information Processing Systems, volume 32, pp. 2202-3211 Curran Associates, Inc., 2019.
Suriya Gunasekar, Blake Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nathan Srebro. Implicit regularization in matrix factorization. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp. 6152 60,,6 2017.
Geoffrey Hinton, Li Deng, Dong Yu, George E Dahl, Abdel-rahman Mohamed, Navdeep Jaitly, Andrew Senior, Vincent Vanhoucke, Patrick Nguyen, Tara N Sainath, et al. Deep neural networks for acoustic modeling in speech recognition: The shared views four research groups. 1EEE Signal processing magazine, 29(6)82-97,
2012. 2012.
10

---

Published as a conference paper at ICLR 2024
Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS' 20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.
Yiwen Kou, Zixiang Chen, and Quanquan Gu. Implicit bias of gradient descent for two-layer reLU and leaky reLU networks on nearly-orthogonal data. In Thirly-seventh Conference on Neural Information Processing Systems, 2023.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. In Advances in neural information processing systems, pp. 1097-1105, 2012.
Thien Le and Stefanie Jegelka. Training invariances and the low-rank phenomenon: beyond linear networks. In International Conference on Learning Representations, 2022.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang, Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet (eds.), Proceedings of the st Conference On Learning Theory, volume of Proceedings of Machine Learning Research, pp. 2-47. PMLR, 06-09 Jul 2018.
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference on Learning Representations, 2021.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2019.
Kaifeng Lyu, Zhiyuan Li, Runzhe Wang, and Sanjeev Arora. Gradient descent on two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 14:12298-12999 2021.
Hartmut Maennel, Olivier Bousquet, and Sylvain Gelly. Gradient descent quantizes relu network features. arXiv preprint arXiv:/803.08360 2018.
Hancheng Min, Salma Tarmoun, René Vidal, and Enrique Mallada. On the explicit role of initializa- tion on the convergence and implicit bias of overparametrized linear networks. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 7760-7768. PMLR, 18-24 Jul 2021.
Jan Persson. A generalization of carathéodory's 's existence theorem for ordinary differential equations. Journal of Mathematical Analysis and Applications, 49(2):496-503, 1975. ISSN 0022-247X. doi:
https://doi.org otg//0...........o 10. 1016/0022--00 1016/0022-2477(755777 247X(75)90192-4,
Mary Phuong and Christoph H Lampert. The inductive bias of relu networks on orthogonally separable data. In International Conference on Learning Representations, 2021.
Noam Razin, Asaf Maman, and Naday Cohen. Implicit regularization in hierarchical tensor factoriza- tion and deep convolutional neural networks. In International Conference on Machine Learning. pp. 18462. PMLR, 2022.
W. T. Reid. Ordinary Differential Equations. Wiley, New York, 1971.
Andrew M Saxe, James L. Mcclelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural network. In International Conference on Learning Representations, 2014.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Joannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484, 2016.
Mahdi Soltanolkotabi, Dominik Stöger, and Changzhi Xie. Implicit balancing and regularization: Generalization and convergence guarantees for overparameterized asymmetric matrix sensing. arXiv preprint ar0iv: 2303. 2023.
11

---

Published as a conference paper at ICLR 2024
Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction. Advances in Neural Information Processing Systems, 34, 2021.
A.J. van der Schaft and J.M. Schumacher. An Introduction Hybrid Dynamical Systems. Number 251 in Lecture Notes in Control and Information Sciences. Springer Verlag, Germany, 2000.
Aditya Vardhan Varre, Maria-Luiza Vladarean, Loucas Pillaud- Pillaud-Vivien, and Nicolas Flammarion. On the spectral bias of two-layer linear networks. Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Gang Wang, Georgios B. Giannakis, and Jie Chen. Learning relu networks linearly separable data: Algorithm, optimality, and generalization. IEEE Transactions on Signal Processing, 67(
2370, 2019. doi: 10.1 10.1009SS...0992220 2904921..
Mingze Wang and Chao Ma. Early stage convergence and global convergence of training mildly parameterized neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Mingze Wang and Chao Ma. Understanding multi-phase optimization dynamics and rich nonlinear behaviors of relu networks. arXiv preprint arXiv:2305.12467 2023.
Yifei Wang and Mert Pilanci. The convex geometry of backpropagation: Neural network gradient flows converge to extreme points of the dual convex program. In International Conference on Learning Representations, 2022.
Can Yaras, Peng Wang, Wei Hu, Zhihui Zhu, Laura Balzano, and Qing Qu. The law of parsimony in gradient descent for learning deep linear networks, 2023.
Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. A unifying view on implicit bias in training linear neural networks. International Conference on Learning Representations, 2020.
12

---

Published as a conference paper at ICLR 2024
A ADDITIONAL EXPERIMENTS
A.1 ILLUSTRATIVE EXAMPLE
We illustrate our theorem using a toy example: we train a two-layer ReLU network with = 50 neurons under toy dataset in R² (See Figure. that satisfies Assumption | and initialize all
entries of the weights [W]; s.i.d. ~ N (0,a²) 1i sid. . (0,a²) vi E [n]. E [h] with a 10 10-6.
Then we run gradient descent on both W and with step size 2 x 10-3, Our theorem well predicts the dynamics neurons at the early stage of the training: aside from neurons that ended up in Saaad, neurons in V+ reach S+ and achieve good alignment with 2++, and neurons in V_ are well aligned with x in S_. Note that after alignment, the loss experiences two sharp decreases before gets close zero, which is studied and explained in Boursier et al. 2022)
lol-heration
[b)-ieration 7000
[c).iteratien 21000
12
Loss Curve
-1.0
-1.0
-1.0
ang
0.8
3.0
0.0
0.0
U.0
0.2
1.0
1.0
1.0
-10
-0.5
0.0
0.3
1.0
-1.0
-0.5
0.0
0.5
1.0
-1.
-2.5
0.0
0.5
1.0
1000D 20:00 30300 40000
Eatation
Figure 6: Illustration gradient descent on two-layer ReLU network with small initialization. The marker represents either a data point neuron. Solid lines represent the directions of neurons. (a) at initialization, all neurons have small norm and are pointing in different directions; (b) around the end of the alignment phase, all neurons well aligned with x+ (x-);; (c) With good alignment, neurons in S_, S+ start to grow in norm and the loss decreases. When the loss is close to zero, the resulting network has its first-layer weight approximately low-rank.
in S_, Sdead Moreover, neurons in S+ (S_) are
A.2 EFFECT OF DATA SEPARABILITY
This section investigates the effect of data separability y on the time required to achieve the desired alignment as in Theorem [ through simple example. we consider a similar setting as in A and explore the cases when data separability μ < 1. We expect that as separability decreases, the time for neurons to achieve the desired alignment as in Theorem [ increases, necessitating a smaller initialization scale. For simplicity, we consider dataset with only two positive data (x1, V1 +1), (22, 1/2 +1).
In Figure we first set = (x1, = s0.10. and the neuron alignment is consistent with Theorem positive neurons (that are not dead) eventually enters S4. activating both data points, and then final convergence follows.
Ial-iteration
iorattr 7000
(c)-lteration 21000
Loss
1.0
-0.5
0.5
0055
0.8
Nocoppdt 351
015
0.5
0.5
04
9O YERT
00
Karation
[1, 0), +1, X2 (sin(0.1). co.(().))))) (0.1)) V2 +1 The experimental setting is exactly the same as illustra- tive example in Appendix A.1 (initialization scale 10 10-6). The marker represents either data point or neuron. Solid lines represent the directions of neurons. In the alignment phase, positive neurons aligned with F4, and then grow their norm for final convergence.
Figure 7: Neural alignment under two data points with small positive correlation:
13

---

Published as a conference paper at ICLR 2024
8 decrease the separability μ sin(0.001) (other settings remain changed), the neural alignment becomes slower: 1) at iteration 7000, there are still neurons (that not dead) outside S4, namely those aligned with either I1 22. while in previous setting (u= (0)))) all neurons (that are not dead) have reached S-: 2) In this particular instance of the experiment, also see neuron remains outside S+ the late stage of the training (at iteration 21000). This clearly shows that as data separability μ. decreases, the time needed for all neurons (that are not dead) reach S+ increases, and if the initialization scale is not small enough for the alignment phase to hold for long time, there will be neurons remains outside S+.
However, in Figure
Ial.teration
Ib).iteration 7000
[c)ilteration
1.2
piccid
Ly (030000
0.5
0.5
0.5
0.8
0.0
0.6
0.5
022
1.0
EooAetoo 258-21
-10
-0.5
0.0
0.5
1.0
-1.0 -0.5
1.0
-16
-0.
0.0
10000 20:00 30300 40000
Figure Neural alignment under two data points with tiny positive correlation: 21 [1, y1 +1, X2 = [sin(O... (0.001,, cos(0.00)))) (,, y2 +1. The experimental setting is exactly the same as the illustrative example in Appendix A.1 (initialization scale 10-6). The marker represents either data point or neuron. Solid lines represent the directions of neurons. In the alignment phase, positive neurons are aligned with 4. but the alignment is slower
A.3 NEURON DYNAMICS UNDER ORTHOGONAL DATA
We have seen in the last section how small i affects the neuron dynamics. The orthogonal data assumption studied in Boursier et al. (2022) is precisely the extreme case of + 0, where the neuron behavior changes substantially. We follow exactly the same setting in Appendix A.2 and consider the case of O.
In Figure 9 we see that S+ is no longer the region that contains all (non-dead) positive neurons at the end of the alignment phase. Depending on where each neuron is initialized, it could end up being S4, aligned with æ1, or aligned with T2. Moreover, for final convergence, only the neurons ended up in S_ grow their norms and fit the data, whose number is clearly less than that in the case of μ > 0.
lajiiteration
[b).iteration 7000
[c).lteration 21000
12
Loss Curve
-1.0
-1.0
-1.0
1.0
3.0
0.0
6.0
0.6
See
0.5
8A
1.0
10
1.0
-10
-0.5
1.0
-1.3
-0.5
-14
-25
80
1000B 20300
30200 40200 Earation
Figure 9: Neural alignment under two orthogonal data points: = [1,0] +1, 00, 1) U2
+1. The experimental setting is exactly the same as the illustrative example in Appendix A. The marker represents either a data point or neuron. Solid lines represent the directions of neurons. In the alignment phase, positive neurons are aligned with one of these directions: I1, I2, T2, T., and only those aligned with 24 grow their norm for final convergence.
This difference in neurons' dynamical behavior makes the analysis in Boursier et al. 2002 different than ours: First, Boursier et al. (2022) only studies the dynamics of the positive (negative) neurons that initially activate all positive (negative) data, which will end up in S+ (S_) and fit the data, and the analysis does not evolve the changes in their activation pattern. In our case, any positive (negative) neurons could potentially end up in S+ (S-), and in particular, will if it initially activates least
14

---

Published as a conference paper at ICLR 2024
positive (negative) data, thus it becomes necessary to track the evolution of the activation pattern of all these neurons (novelty in our analysis). Moreover, consider the case that neurons are being randomly initialized, Boursier et al. (2022) requires the set of positive (negative) neurons that initially activate all positive (negative) data being non-empty, which needs the number of neurons h to scale exponentially in number training data (extremely overparametrized). In our case, we only require h = (1) (See Merits of overparametrization after Theorem D a mild overparamerization.
In summary, while Boursier al. (2022) also provides quantitative analysis neural alignment under small initialization, it is done under the assumption that all data are orthogonal to each other, leading to a different neuron dynamical behavior than ours. Due to such differences, their analysis cannot be directly applied to the case of orthogonally separable data (ours), for which we develop novel analyses on the evolution of neuron activation patterns (See proof sketch in Section 3.3
A.4 ADDITIONAL EXPERIMENTS ON MNIST DATASET
We use exactly the same experimental setting as in the main paper and only use a different pair of digits. The results are as follows:
Data Corretation
Alignment with x
Convergence and low-rank bias
Data/Neuran Visualization
40
1.0
Loss C
10
0.5
3
8
...)
[M]
V+
0
5000
10000
15000
0.6
0.5
0.4
-00
Co0 x.
0.2
n.
W. +x
5000
10000
15000
5000
10000
15000
Keraticn
toration
Figure 10: Binary classification on MNIST Digits and 8.
Data Correlation
Allgament with X+,
Convergence and low-rank blas E pWii-wwn
Data/Neuron Visualizattin
3.0
40
7
0.8
20
006/1 )
0.0
0.6
X
X=
X,
4
x.
12
15000
-20
0.5
-40
0000
N-
0
5000
10000
15000
21
5200
10000
15000
Keration
Eeration
Figure 11: Binary classification on MNIST Digits 1 and 7.
A.5 DISCUSSION ON THE TWO-PHASE CONVERGENCE
With the same two-digit MNIST dataset in Section 4 further discuss the two-phase convergence under small initialization. We use a two-layer ReI.U network with = 50 neurons and initialize all
entries of the weights [W]y sidd N (0, .22)),. Uy 1idd N a2).).. Vi E [m]]. [h] with a 10 Then we run stochastic gradient descent (SGD) with batch size 2000 on both W and with step size 17 = 2 x 10-3 For comparison, we also consider the training schemes studied in Brutzkus ct al. 2018); Wang et al. (2019), where only the first-layer weight W is trained starting from a small
initialization [W]s i.i.d. N (0, a²). and vy are chosen to be either +1 or - with equal probability, then fixed throughout training.
We consider the changes in neuron norms and directions separately. In particular, these quantities are defined as
d
Σ
=> 2 C,
=
(changes neuron norms)
dt
VuyB
15

---

Published as a conference paper at ICLR 2024
1.0
Loss
0.8
Changes in neuron narms
Changes directions
10
5
0.0 250
0
500
750
1000
1250
1500
1750
2000
epochs
Figure 12: Two-phase training under small initialization (SGD with batch size of 2000, step size of 2 x 10-3). At the early phase of the training, only neuron directions are changing while neurons' second stage, neurons start grow their norms and loss starts to decrease.
norms do not grow. See Appendix A.5 for the precise definitions of "changes in neuron norms" and "changes in neuron directions"
0.6
Loss
Cha ngange
200
0.5
Cha nges in neuron directions
0.4
150
0.3
100
0.2
0
250
500
750
1000
1250
1500
1750
2000
epochs
Figure 13: No two-phase training when only the first layer is trained (SGD with batch size of 2000, step size of 10-3).
Σ d W3
Σ Pes (T1wy)11 1
(changes in neuron directions)
dt |0y|
and they measure, at the end of every epoch, how much the neuron norms and directions will change one uses a one-step full gradient descent with small step size.
we show the changes in neuron norms and directions over the training trajectory when we run stochastic gradient descent (SGD) on both first- and second-layer weights. The two-phase (alignment phase then final convergence) is clearly shown by comparing the relative scale changes in neuron norms and directions different phases the training.
Training both layers: In Figure 12
Training only the first layer: In Figure 13 we show the changes in neuron norms and directions over the training trajectory when we run stochastic gradient descent (SGD) on ONLY the first-layer weights (Brutzkus et al. 2018 Wang et al. 2019). The plot indicates that two-phase training does not happen in this case.
16

---

Published as a conference paper at ICLR 2024
B PROOF OF LEMMA T NEURON DYNAMICS UNDER SMALL INITIALIZATION
The following property of (exponential loss E(y, i)) = exp( yÿ) or logistic loss fyy 1) = 2 log 1+ exp(-v0)) will be used throughout the Appendix for proofs of several results:
Lemma 2. For l. we have
Vzl(v.0) - yl S 2(01. vy E (+1,-11
-
11
(9)
Proof. Exponential loss: when fyy, 0) exp(- -y1):
V Vv y) 1 (exp(-yj) y/
I |ull exp( (yp) 1
I eee(p= - 1| 2|
where the last inequality due the fact that max exp(- exp(x)) 1}, Ve [0.
Logistic loss: when E(y. ù) = 2 log(1 + exp( vū)):
2y exp(-yύ)
Vgl(v.0)
exp(-xti y0)
y
yexp(- y
exp(-v0)
< /ul/exp( (vv)
exp(-yy) - 1| S 2
Remark 4. More generally, our subsequent results regarding neuron dynamics under small initial- ization hold for any loss function that satisfies the condition stated in Lemma 2 which includes the l2 loss Ely, 9) (y 2) studied in Boursier et al. (022))
B. FORMAL STATEMENT
Our results for neuron direction dynamics during the early phase of the training will be stated for networks with any a-leaky ReL.U activation (xx) = max (a,, ax) with a E [0, 1]. In particular, the ReLU activation when a O, which is the activation function we considered in the main paper, and it is the linear activation when a 1.
Denote: Xmax max; |as]], Wmax max, Il The formal statement of Lemma is as follow:
Lemma 1. Let the activation function be a-leaky ReLU activation (xx) maxfæ, az).. Given some initialization from B). for any e I 4VhXmox W then any solution to the gradient flow dynamics
satisfies that VI < T In Inxmaax max
log
Vhe
wy(t (t)
max
di [1](
(O)P
< JenvTx2
max max
W
where
La (ww)
Σ Iivoo' ((z,, w/))
Σ
I&li + a
Tili
(0 50
With Lemma and set 0, we obtain the results stated in the main paper. Lemma is direct result of the following two lemmas.
Lemma 3. Let the activation function be o-leaky ReLU activation o (z)) maxfa, az). Given some initialization in (3), then for any e S 4VToXmma any solution to the gradient flow dynamics satisfies
2
max (1)(1))1
max (f( (1), s 2eVh. Kax W max
(10)
Vh
Vt
log Vhe
17

---

Published as a conference paper at ICLR 2024
Lemma 4. Let the activation function be a-leaky ReLU activation o(z) max Con- sider any solution to the gradient flow dynamic (2) starting from initialization (3) max; (2 W, u)) - 1, we have, Vi E [n)]
Whenever
d wj
sign(u,(0))
(1-
Ia (a(w))
2nx max (x4; W)) (11)
|wy|l
w
where
Ka (w))
Σ aatha' ((ri, mz))
E
Tilk +
Σ
Filk-
Remark 5. By stating our approximation results for neuron directional dynamics with any a-leaky ReLU activation function, we highlight that even for some networks with other activation functions than ReLU, there is similar notion of neuron alignment al the early stuge of the training, and the analytical wols used this paper can be upplied to them. However, we note that our main results (Theorem 7) will not directly apply as the neuron directional dynamics have changed as consider an activation function different than ReLU (see the general definition of z (a(ww)) and additional efforts required establish the directional convergence for general leaky-ReLU functions.
B.2 PROOF OF LEMMA 3 BOUNDS NEURON NORMS
Proof of Lemma 3
Under gradient flow, we have
dt " Σ V f(x;; W. u)),
(12)
Balanced initialization enforces vy = sign(u; (0g) hence
dt " Σ o' ((zr, w/) v V f(x W, (cr)::/:).)))rsssmm(((()))))))
(13)
Let T := inf max, [f(z; ((),;;(()), 2e W2 W then S J E [M]. we have
|,yl1? = ( w/. dt Wij
di
a 2
o' d'(,, work,, wy)) f(zn W. u)) (ar, wyy sign(e, (g)n( (0))|
V Vollyi, f (2 W, b) my)] [|yl||
< 2
(lval + (x:: (z,, my)س [[ [yyl
(by Lemma
<
Σ
(1 + 4ev hxx max W2 2 (2,, wyyس [2yy||
(Since S T)
i=1
< >
(1 1eeeeehhe h..hhamaa W ||= ||
i=1
< 2n(Xmax + 4c hxx2 max W w1
(14)
Let : inf((: 110;(11)
2eW2
mpa ), and let j arg miny Tj, then Ty- miny Tj S T due vh
the fact that
[f W, w.
Σ JE[A]
o' ((2,, wr))e, (wy, xz)
||an || < h.X h.mmax max ||,,
jE(h
JE
18

---

Published as a conference paper at ICLR 2024
which implies (If ((),, x(])] > Tex max Whax 7 (/,8.1.[l,,(1) Then for t S Ty,, we have
BeW
vh
d |w 12 S 2n (x(xxxx + devhX2 max W maxx) 1| 11
(15)
By Grönwall's inequality, we have Vt S Tz*
I|wy- (t)l z exp (2mxX + 4EV hvvx max W Wmxx)t
[00 (0)113,
= exp
22m(X + 4e 4..44e max
W2 max )6
[ Wo 11?
I exp
(2mxX + 4ev hxx max W W 2² W?
max
Suppose Tj.
log
1 Vho
then by the continuity of || (f)]l we have
4m max
2 max
max (T=-)]1 ك exp (2n(X + 4e he max W m
2 w
hh
max
1
I exp
2n(Xmux + 4c hex2222
max W Wmax) 4nXmaxx
log
2w2
Vhe
max
< exp
1 W2
max log
1
2
Vhe
max
1
eW2
S exp
log
ew2
max
Vhe
Vh
which leads to contradiction 2c < 6. Therefore, one must have T 2 Tj. 2 InXmax
log Vhe
This finishes the proof.
B.3 PROOF OF LEMMA 4 DIRECTIONAL DYNAMICS OF NEURONS
Proof of Lemma 4
As we showed in the proof for Lemma 3 under balanced initialization,
d
Wy
(e V ( (0)]|
(16)
di
Then for any E [n]. d Wj
n dt ||w]|| sign(u,(0)) Σ
{rz, wz)
1 (2 Volly, W.
rz
|1 112 ws
{2r wi)
-sign(v;(0))
Vellyr [(x;; W, w))
wis
||0 112
=(a),1))>0
si) 1 ||00
(( w f(1; W. W.)))a
Therefore, whenever max; [x W, W) S 1,
Wy sign(1;(0) (1 dt ||w
Na (w3))
|
sign(u;(0)
Σ o' ((a,, (V W. + v1) x4
n1
V Volly,, f(x; W. vvl ||
<
2)5 W, w)1 -|x=a| < 2nM, max (f(; W, v)
(17)
19

---

Published as a conference paper at ICLR 2024
20

---

Published as a conference paper at ICLR 2024
C PROOF FOR THEOREM I EARLY ALIGNMENT PHASE
We break the proof of Theorem into two parts: In Appendix we prove the first part regarding directional convergence. Then in Appendix D we prove the remaining statement on final convergence and low-rank bias.
C
C.1 AUXILIARY LEMMAS
The first several Lemmas concern mostly some conic geometry given the data assumption:
Consider the following conic hull
K CH(( [n]])
Σ a,r.y : a; 0, [n]
(18)
It is clear that 2,y; E K, Vi, and Fa (w) K, V.. The following lemma shows any pair of vectors in K is y-coherent. Lemma 5, cos(21, 22) 2 Hr v0 + 21, 22 K.
Proof. Since 21, 22 E K we let 21 = Ei= F4yaa1 and22 = Σド= Uzj where ais 0zj 2 0
but not all of them.
cos(21, 22)
(21,
01j22j (1.31i.
| 22|
1 1221
Ensep
| | |1102211
2 14:
2
22
where the last inequality is due to
2 221 S
|a 1o1s
|
||a2]
Σ
|ل
The following lemma is some basic results regarding S4 and S-: Lemma 6. S+ and S are convex cones (excluding the origin).
Proof. Since ((a,,)) a2). Vi [n], a > 0, S are cones. Moreover, (x,, 21) and (m,, 22) implies (I,, (1:,0121 21 + 2222) Vai,as 02 thus S+, S_ are convex cones.
Now we consider the complete metric space D (w.r.t. arccos( 1)) and we are interested in its subsets K n S nSP-1, and S_ 1-1 First, we have (we use Int(S) to denote interior of S)
Lemma 7, K Kns c Int(( nsP-l), and - K K c Int D-P-1)
Proof. Consider any Ie Σラ= ajozUy E K For any 24, ye, E [n], we have
n
Tンリj Kiyii | ||ا,||' ||>;||
(Fc, Σ الو1امم
Vi
pXmin >> 0,
D1 0 V: 0
S -uXmin 0,
Depending on the sign of ys, we have either
|
|2
n
(Ze, = Σ 2y11
Σ ولا|ا 2 pXmiin >0, (v: ++1)
||>
[
M/i
21

---

Published as a conference paper at ICLR 2024
0:) = Σ ااو|اروl
زززي TzYi
|
S |1
|
Σ
ayl وارراري1 & -xXmi <0, (V: = -1)
||
|
y/s
where the fact that 1 || To E ES+ n 15-1
= اررر This already tells us
Since fi(=) = (z, 2.) is continuous function of = E SD-1 There exists open ball B (are 6:) centered at zy. with some radius 8, > 0, such that V2 E (Iie, o:), one have [f((z) - Sa (re) S pXmin, which implies
(2, 2 x min /2> 0, y1
15 -uX min / 0, y/
Hence 1 B de 8:) E est nSD-1, Therefore, Ie E Int n SD-1). This suffices to show
||=||| K n SD-1 C Int(S, n SSD-1). The other statement - nsP-1 C Int(S_ n SP-1) is proved
similarly.
The following two lemmas are some direct results of Lemma 7 Lemma 8. 3G > such that
SC S++ 2t
85 C
(19)
where S = ERD cos(2, x) 2 vI [
Proof. By Lemma s+ll E c Tnt (S+). Since SD-1 is a complete metric space (w.rt arccos (s,,,, there exists open ball centered at =t of some radius arccos(vI- (1) that is || subset of S+, from which one can show S C S4. The other statement St. C S_ simply comes and Int(S+) - Int(S_).
from the fact that I+
Lemma 9. 3 > O, such that
sup #1 mLEKniD-1 nSD-))nnns - cos((1,, I2)1 S -I-
(20)
(St here is defined to be DD-1 S, the set complement w.r.. complete space DD1)
Proof. Notice that
sup ZI =IEKn (nt(s (x1, 2,EKr5p-L (ag.:P--______n((S____ (00 arecos (I1,
inf
Since DD 1 a complete metric space (w.rt arccos (., () and Kn5P-1 and x2 E (Int (In-() are two of its compact subsets. Suppose
inf 1ee
arcoos (I1, 22) = =,,
then Sr, E Knn 22 E (Int(S, n such that arccos (I), x2) = O, i.e., xi = Ø2, which contradiets the fact that K n D-D C Int(S PP-1) (Lemma Therefore, have the infimum strictly larger than zero, then
sup
(2))
sup - Lec (01, < 1.
(21)
Similarly, one can show that
- - (11, < 1.
sup
(22)
Finally, find such that
max sup CiEKnSP-L,aee$$1
(x1,22),
sup zi ( (1), = VI-E.
22

---

Published as a conference paper at ICLR 2024
then for any x1 E KnS and 22 (6__33-)lnn$_nn we have
- -1-55 (21, - / 1-
which is the desired result.
The remaining two lemmas technical but extensively used in the main proof. Lemma 10. Consider any solution to the gradient flow dynamic starting from initialization 3 Let By S be some reference direction, we define
Vry = (xr |w,|1 ) Vrr ( (se) va(W)) [
Wi
) Day
<
wy
Fa((1))
(23)
||2yy11 (w;()
where za (a(())) = 2i((a,,0))>> Ye
Whenever max; (f W, v) S 1, we have
d
dt Wro sign(v;(0)) (fhra - ())))az) | < 2mx max f(z;; W, w)).
(24)
Proof. A simple application of Lemma 4 together with Souuhy-Schwartz::
d dl
Vij - sign(u;(0)) (Uraa - var))as) (w;)ll
(
))
Fr
d wj
sign(1;(0)) (1-
Σ yiurs
2n max max (f W, v)]-
dt
|10ll
Lemma 11.
[ (w)ll z Vuna
( Xmin
(25)
where na(w) Hi [n] (Z w) > 0}}.
Proof. Let (ww) denote {i [n]] (ze,, w) > 0}, then
|I,[w)]|=
Σ
Iith
Σ
||
Σ
|
|x:i||
Tzlhi
Zyyy |
(=(2,,0)20
IEI. (w)
(w),n(w)
Σ (ET,(w)
|| 12y?
Σ
||ا 1| M1ballusl
(w).
W
na (a(w) min yana(w) (Na(w) 1) X2
IN Vanua (ann(ww).. X
no(w)(1 p(na(w) 1))Xmin
C.2 PROOF FOR EARLY ALIGNMENT PHASE
Proof of Theorem First Part. Given some initialization in 3 by Assumption 2 that
3C2 0, such
max cos(w; (0), x / C2, JEV+
max cos(w, x4) < JEV_
Co
(26)
We define : = max (S1: (2). where Ç is from Lemma In addition, by Lemma that
E 0, such
sup mieKnsP-L,age wessssnssmm nsP-1 ns-
cos(2 S VI-E.
(27)
23

---

Published as a conference paper at ICLR 2024
We pick initialization scale that satisfies:
min(uu, (,))
1
64nX max
)
1
< min
exp
log min V/Xmmin
4v X
W
Vh
4v Avhxmma
W max
max
max
(28)
By Lemma
Vt <T=
log
we have
InXmax
Vice
C.in/,,C.E)VV max ((a; W. W)l S
(29)
4n
which is the key to analyzing the alignment phase. For the sake of simplicity, we only discuss the analysis of neurons in V+ here, the proof for neurons in V- is almost identical.
Activation pattern evolution:
Pick any wy in V+ and pick xr, Tzys for some i E [n]. and consider
the case when (wy, 24) = 0. From Lemma 10 we have
d Vrj (tira Wry Was) I (wy)))
S max max (x(;; w)]
dt
(wg, = 0 implies Why =
saye |Taaa TMwill ) = O, thus we have
thes k0)0 - Vraal @raal=aa(0 (2a))=) I Xmx max f(zs v)l. dt
Then whenever wy & Stend, we have
dt Ury W 2 Ura (la 2n max [f(; W,, w)
Mll (w;)) 2n.X max max (x;; W, v)
(by Lemma 5
N M
11 3/2 Xin -nxX max [x W.
(by Lemma
3/2 Xmin/ > 0.
(by 29
This is precisely 7 in Section 3.3
Bound on activation transitions and duration: Next we show that if at time to < T, wy (to) $ S+U
Saead, and the activation pattern of wy is (to))) then 1 (2,, + > * 1 1 ww (to))>0+ where At (ainCC&E)V/Xaaa (ws (fo3] and na (wy (to)) ) defined Lemma aa as long as to + At
T as well. That is, during the alignment phase [0, 7]. wy must change its activation pattern within At time. There are two cases:
The first case is when (to) E St n S Si In this case, suppose that (ta
1 (mr,ns (to))>00 A0 < T < At, i.e. wy fixes its activation during [to, to + At). then we have za (w, (to +T) = Za (w)))),,,,, 4o < r At. Let us pick I, Ia (t))) then Lemma 10 leads to
dt co F.a (w;)) (1 cos2 (wij xa (mz))) ||,a(()))]|| < 2nXmma max [I;;; w)]-
Since za (wy) is fixed, we have Vt E [to, to + At).
d
cos(wj, Fo (w)(())) -1-cos "wos(w,, J. ( (w)((()))) (15(6(1
< 2nx max [f W, v)
di
d dt
cos(w,, Fa (w, (to))) 2
cos" (ws,. (1-cos((u;, Da (w, (to)))) [I
2n max max (x w)
E ar (w, (lo)) 2n.X max max (# W, v)
(by 27
E Vumaa (wy (to)) Xmin - 2nX max (I;; v)
(by Lemma
Ev Evima (w Xmin//2
(min(())))mmn/2
(by 29
min 5) VTnna (wy (to)) X
24

---

Published as a conference paper at ICLR 2024
which implies that, by the Fundamental Theorem of Calculus, cos(w,((a (fo + At), (w(g(()))
(to))) / St d cos(wy(to + Fa (((y((o)))
dt
2 cos(w (tu))tt)::10)) (fo)) At minf5. St/mma (to) X (to)) + 1,
cos(w,
which leads to cos(w, (to + At), Za (w, (to))) 1. This would imply wy(to + At) E S. because En (w,((1)) E S4. which contradicts our original assumption that wy fixes the activation pattern. Therefore, 30
S Al such that 1l (x,,,, (to+to)) o 1 (e, (t))>00 due to the restriction on how wy can change its activation pattern, it cannot return to its previous activation pattern, then must have 1 (to, Ad))) * (0r,,
The other case is when wy(t E S_. For this case, we need first show that wto + r) & $5_, VO s & Åt, or more generally, S does not contain any wy in V+ during [0, Tj. To see this, let us pick H, = x_, then Lemma T0 suggests that
Dij (Vra (ar
2m. max (a, W;
Consider the case when cos(wy, x-) VT- S, i.e. wy the boundary of S We know that in this case, wy E SSS S S_ thus (w())) = xx_, and
d cos(w,,
(1 cos (oop (w,, x) || 2nXmax max (0 W,v),,
di
cos(w,n==)==D)
which is
d cos(w;. x-)
S 2nX max max (f(;;; W₂ + 51|2-11
i $ -SVTXmmiu + 2mx max f( W, v)
dt
-oo(u)j.-) =-VI-
d cos(wi; x
dt
con(m),--)=VTT-
S -cl|s- + 2nX max max (1, w)
(by Lemma I1
S -SVTXmin 2 1200 0.
(by (29)
Therefore, during [0, T). neuron wy in V+ cannot enter S5. if at initialization, wyy(0) d S which is guaranteed by (26).
With the argument above, we know that wy(to + +7) S$_ vo < S At. Again we suppose that wi(t) (t) E S_ S$__, vt E [to to + At). i.e.,wy fixes its activation during (to, to + At). Let us pick
2r x_, then Lemma 10 suggests that
d cos(u,, x)) + (1 (ws, x))) ||--1| < 2n max (s W,r),,
dt
which leads V/. [40, lo + At).
d
cos( cos(w;,x=_) x-)
5 - cos ( x ||- + 2nX max (f ([;WW,,) W,+)
dt
s -110=11 || 2nx max max (x:; W,
(wy $ S$_)
s -CVtina (u (w, (lo)) Xmin + 2n X max max (1 W,
(by Lemma
s -V
(by
min(e, (/imma (to)) X
25

---

Published as a conference paper at ICLR 2024
Similarly, by FTC, we have
cos(wy(to + ++)
This would imply wy(to + At) E S+ because -I- Ia (wy(lo) E S4, which contra- dicts our original assumption that wy fixes its activation pattern. Therefore, one must have 1 (m (to Ad))) * (t)))()))))
In summary, have shown that, during [0, a neuron in V4 can not keep fixed activation pattern for a time longer than At min((,5) piXiinn where na is the number of data points that activate wy under the fixed activation pattern.
Bound on total travel time until directional convergence As we have discussed Section 3.3 and also formally proved here, during alignment phase [0, T", neuron in V4 must change its activation pattern within A v
time unless is in either S+ or Saead- And the new activation negative data points and must keep all
it is transitioning into must contain no new activation existing activation on positive data points, together it shows that a neuron must reach either S+ or Stead within fixed amount of time, which is the remaining thing we need to formally show here.
For simplicity of the argument, we first assume T = oc, i.e., the alignment phase lasts indefinitely, 16 log and we show that neuron in V4 must reach S+ or Saead before 4 minkfSEEv/uX... Lastly, such directional convergence can be achieved if t < 'T, which is guaranteed by our choice of e in 28
For neuron V+ that reaches Sdend. the analysis is easy: It must start with no activation on positive data and then lose activation negative data one by one until losing all of its activation. Therefore, it must reach Stead before
na(my (0)) min V mink
4
16 log min /aXmin = 41-
S
min(S. E) MuXmin
k
For neuron in V4 that reaches S4, there is no difference conceptually, but it can switch its activation pattern in many ways before reaching S+, so it is not straightforward to see its travel time until S+ upper bounded by f1.
To formally show the upper bound on the travel time, we need some definition of a path that keeps a record of the activation patterns of neuron (y( before reaches S+.
Let = -I+1, n- I__ be the number of positive, negative data respectively, then we call P(KC (.1))...,, ,(()) a path of length-L. if
1. 4O 11 < L, we have k(1) = (k = N x N with 0 S k < n4. 0 <n.;
2. For k((1), k((1), with I 12, have either k(11) ) > k(h)) k.01)) < (1))))
3. ((L)) = (n+, 0);
4. 1(0) + (). 40 L.
Given all our analysis on how a neuron Wy (t) can switch its activation pattern in previous parts, we know that for any wy that reaches S4, there is associated P P(asoo ,.),....... that keeps an ordered record of encountered values of
(l I (z,, wy(()) 0H. Hi e I_ (o,, (t) 0})
before wy reaches S4. That is, a neuron wy starts with some activation pattern that activates 4 (0) positive data and (0) negative data, then switch its activation pattern (by either losing negative data or gaining positive data) to one that activates (+(1) positive data and k (1) negative data. By
keep doing so, it reaches S4, that activates k2 (L) n4 positive data and k (L) 0 negative data. Please see Figure 14 for an illustration of a path.
Given path Paa.o,, ((2) of neuron Wj, define the travel time of this path as
4
((P(Moo,A AM).. k(.))
t min(s. Xmm (k + k0))
26

---

Published as a conference paper at ICLR 2024
P Praco 20, AC01
Picor ,O,.. (210))
(A
k+
PIax
k4
Figure 14: Illustration of path of length-10.
10..
Figure 15: Illustration of path and the maxi- mal path
Each dot the grid represents one k0))
which exactly the traveling time from k:()) to k((4) if one spends min(C.&)y/TTXNn (M (42+ on the
edge between (0()) and b((+1).
Our analysis shows that if wy reaches S+, then
inf(t: (ty()) S+} S T(P (1:0) (100)) ,00 x00)))
Now define the maximal path Pmax as a path that has the maximum length n244 + n which is uniquely determined by the following trajectory of g(l))
(0, nn_), (0, n - 1), (0, n_ - 2).... ... (0, 1), (1, (1,0), ... (n+ 1, 1,0), (n4, 0).
Please see Figure 15 for illustration.
The traveling time for P max is
T(P max) min { / Ppppppmii Xmin
k
k=1
4
16 log
min(c,E) VpXXmin
The proof is complete by the fact that any path satisfies
T(P (PPa)))),)).... pr),.. ... ma2s)) S T(P
This is because there is one-to-one correspondence between the edges (x,0,,(+1)) k(+11) in
Pieco (c0)....,.o)) and subset of edges in Pmaxx and the travel time from of edge (K(1), ((+1)) is shorter than the corresponding edge in Pmmax Formally stating such correspondence is tedious and visual illustration in Figure T6 and i7] more effective (Putting all correspondence makes clustered plot thus we split them into two figures):
Therefore, wy reaches S4, then it reaches S+ within f1;
inf(t: (t) E S+} < ((P w0),..
(.. maar)) S T(Pmax) S f1.
So far we have shown when the alignment phase lasts long enough, i.e., T large enough, the directional convergence is achieved by t1. We simply pick such that
1 Vhe
16 logn
T 4nXmaax
log
t min(g.E) V/pX
and 28 suffices.
27

---

we will use it as fact, then prove the Published as a conference paper at ICLR 2024
k
k:
Paro 20
P
Pmax
k4
Pmax
k+
Figure 16: Correspondence between edges in Figure 17: Correspondence between edges in
P(Kto ,00, M002) and Pmax- (Part 1)
Pueoi,a A(D) and Pmax. (Part 2)
D PROOF FOR THEOREM 1 FINAL CONVERGENCE
Since we have proved the first part of Theorem in Section C remaining part of Theorem
D.1 AUXILIARY LEMMAS
First, we show that S4, S_, Socd are trapping regions.
Lemma 12. Consider any solution to the gradient flow dynamic 2 have the following:
If at some time t 2 0, we have wy(tt)= E Sdead, then wy(ty + T) E Sdead, V7 0:
If at some time 1 2 0, we have wy (f1) E S+ for some j E V+, then wy + E S++ VT 2 0;
If some time t 2 0, we have wy (t1) E S- some V_, then wy (t1 + 7) E S_, Wr 2 0;
Proof. The first statement is simple, if wy E Sdead, then one have ú 0, thus wy remains in Stead For the second statement, we have, since j e V4.
dt w = Σ 120 f(z; W, u) =/
When is the exponential loss, by the Fundamental Theorem of Caleulus, writes, Vr 2 0,
wr(t1 + +T)= =w](tt) + S"d d mey(11)
dt w,dr
1e Vge(y, (1x W, v)) ol)allluusll
eey(t)) + 00
1 (a (er.y))008 exp(-xipffff fif (ara W, 2aa lll
(
my(tt) Σ
exp( exp(-wi. 1- (x,; W, W.)))) ww)
SEI
Here (t,tt) E S. by our assumption, ž, E C S, because ã, is a conical combination of T, E It. Since S+ convex cone, we have my(h + r) E S4, as well.
28

---

Published as a conference paper at ICLR 2024
When is the logistic loss, have, similarly,
wy(ti +) = w,(tz)+ + /-² Σ 1 (2),)))000 1 exp( exp(-Mii (x;; v))
2 2exp(-viJ(ae W,,)
I rillwylldr
2 1x W. v))
= wsy(tt +
1+ exp( (2, W. u)) |1m |
24 E S.+
VEI
The proof of the third statement is almost identical (we only show the case of exponential loss here): when j E V-, we have
d
n 1 ( Vellu,, f(y;; W. w))sl//ullll
W
dt
and
(L sep(-yf(az;: W, (ll xx iEI_
Σ
=
Again, here wy (41) ) E S_ by our assumption, 7_ E - K C S_ because ã_ is a conical combination of nyg i I__ Since S_ is convex cone, have wy (11 7) S+ well.
Then the following Lemma provides a lower bound on neuron norms upon tj.
Lemma 13. Consider any solution to the gradient flow dynamic ( starting from initialization (3) Let t be the time when directional convergence is achieved, as defined in Theorem 7 and we define V4 Li wy((1) E S and V_- {jj wj:ws(h) (1())) E ES... If both V. and V._ are non-emply, we have
Σ
[ey 2 exp( -44xmaxt Σ |1y(0))12
JeV.
jeY
Σ [woy(t))11 2 exp( exp(-4nXmaxe (1)
Σ Ilw,
jEV_
jev_
Proof. We have shown that
d
||,,11 = -2 D 1 Lenm Vol(u, f(x W, b)) (z,, wy) sign(r;(0))]len)
dt
Then before
we have Vi [h]
di 12 = -2
Σ 1 (1) Vel(vs. f(x W; v)) (r,, ws) sign(u, (1)
-2
Σ
(lval + 2 2max [( W.
i=1
| 2 AnXmmax|lul
where the second last inequality is because max; [f W, w) S we have
before t,. Summing over j E V4.
d
Σ |u, 2 - -nnX max
Σ ||wr|12
dt
JEV,
jEV,
Therefore, we have the following bound:
Σ
11,((1)11 > exp( exp(-AnXmacee
Σ
(0)]
JEL
SED
29

---

Published as a conference paper at ICLR 2024
Moreover, after 1,, the neuron norms non-decreasing, as suggested by
Lemma 14. Consider any solution to the gradient flow dynamic (2) starting from initialization (3). Let to be the time when directional convergence achieved, as defined Theorem | and we define
V+ : {j : (t) E S+} and V_ : {j: wy(t)) S-}. If both V+ and V- are non non-empty, we have
WT 2 0 and t2 2 ti,
Σ (0,(122 + + Σ (1,(1
110,(12 +
Σ (1ws(12) JEV_
(30)
JEV+
JEV,
JEV_
Proof. It suffices to show that after L1, the following derivatives: d
d
Σ
Σ (t)s(
dt JEV+
dt JEV_
are non-negative.
For E V4, wj stays in S+ by Lemma
12 and we have
d
||,11 = -2 Σ
v Vol(v, f(z;; W, o)) (oi, wyy ||wll-
di
IEI+
" Eier, V1 exp(pp ( W, o)) (me, wy) ||yll
(6 is exponential)
2exp (o..W.)) (1i, my) ||yl|| 2 LiEr, y1 1+expl -i (2):W.))
(E is logistic)
2 0.
Summing over j E V4, we have de Ejer, (t)l 2 0. Similarly one has d at Ejev. (t)1
IV
0.
Finally, the following lemma used for deriving the final convergence. Lemma 15. Consider the following loss function
Liia (W, 6) = Σ Evi T (Txx))
if (xt, y E [n] are linearly separable, i.e., By > 0 and 2 E SD-1 such that yı (2, 2 y, Vi e [n]. then under the gradient flow Lim (W, v), whenever you' W 2 O. Vi, we have
Lim S - 2 c2?2
(31)
Proof. For l. being exponential loss, we have:
Ł= = - - IVV - V C1
Σ villys, v. W 22)x
- ||0||
vellye v. T...TTT.. WTTTT)or x=)
12
s - ||012
2, yel v T.T xx)
s ||ol1
Σ Elyi, TW x2))
< - ||0| C2 < 1 c²
30

---

Published as a conference paper at ICLR 2024
For being logistic loss, have:
-
-IV
2 2exp( you W ni)
Σ
Iz2
+ exp( 1:pp-y111' W x.,)
12
2 exxp( y:u0 xzi)
||0||
Vi
exp( y:0 W TTx))) 26
2 2exp( - y;u W (Txx)
$ ||0|
V/ exp( - y,u² W xx))
24
2 eexp you WT r;)
||01
WTTa. Iz)) ?
+ exp( 3/22" W
2 exxp Y:u WT 1;)
- ?
+ exp( 4/21 W (WTx) x1.)
i=1
2
S -1
log(1 + exp(-yeU WT [xx)) Σ
4 11011/22
where the last inequality uses the fact that 2 2 2 log(1 + 2) when [0, 1].
D.2 PROOF OF FINAL CONVERGENCE
Proof of Theorem 7 Second Part. By Lemma 12
we know that after 4, neurons in S_ (S_) stays
in S+ (S-). Thus the loss can be decomposed as
L Σ 1EI1
l
Vis Σ
Uy (Wg. vi)
+
Σ SEI
Viv Σ Uy (w0,.
(32)
JEV+
JEV_
where V4 : D w (t1) E S+} and v_ : Ej: wy(ti) E S_J. Therefore, the training after 4 is decoupled into 1) using neurons in D4, to fit positive data in I, and 2) using neurons in ii_ to fit positive data I__.
We define f+(x; w, = EjED, oy (W), 2z) and let to inf maxiez, [f+ W,v) > }. Similarly, we also define f- fz; w,v) = Ejet, Uj x.)) and let to = inf(t max,ez (x; W, w) }}. Then I min ). by Lemma 3
O (1/t) convergence after f:z: We first show that when both t t are finite, then implies O(1/t) convergence on the loss. Then we show that they are indeed finite and to = max(t; t, O log :).
At 12 max t by definition, 3i E I such that
1
W) 5 Σ
Us (Wz, xza))
Σ
||,,
Il
(33)
jEN+
jEN4
which implies, by Lemma 14
Vt > t
Σ (w,(1)]
Σ
(1))12))1
(34)
4 1Xmax
jEV
jev
31

---

Published as a conference paper at ICLR 2024
Similarly, we have Vit 2 t2.
(,) (t)ll
(35)
max
JED
Under the gradient flow dynamics
2 we apply Lemma 15 the decomposed loss (32)
4
Σ JeV,
- (uXmin)?
Σ JEV
of
L3 (pXMia))
Here, pick the same 7 4uX for both L+ and L_ because {z,, i I+ is linearly separable with x = 912 || 2,96) 2 ul| 2 pXmin by Assumption And similarly, { ya), i I_ is linearly separable with (2, 2xUs) 2 al|=aal pXXinn- Replace of by ||w,]12 from balancedness,
together with 34 335) we have
45 S - Σ 1k0y|12
L2 (2x
Σ
||wsy1?
c (uXXinn)
JEV+
JEV+
(uXmin)) )2
(uXmin))* )²
(uXmin)² C²
(C3
(C+ 6_)²
4x
8Xmax
8.mmax
which is
C (HXmii))
2
32Xmuxx
Integrating both side from l to any
2 12, we have
((xmmin)²
(t -2),
C t2
32Xmax
which leads
((I) S L(t2)a(t 12) +1 where (p.mmia)
C(t2)
32X max
Showing 12 O log t): The remaining thing is to show to is ( log 4
Since after t, the gradient dynamics are fully decoupled into two gradient flow dynamics (on L.-
and on L.), suffices to show 1 = ( log 2 and L O log {) separately, then combine them to show 1/2 = max ( = O 1 log ) The proof almost identical for L4 and C.,, thus we only prove tx = O 1 log !) here.
Suppose
(og 22 VTT W2 min
t2
2 t + X Xmin
+ A4nxmttt
(36)
where n [ J. It takes two steps to show contradiction: First, we show that for some t, 2 a refined alignment cos(w,(t) (4 ta), x4) 2 +. Vj E V., is achieved, and such refined alignment maintained until at least t²: (t) x.) 2 4. vjj E V4 for all t + to It t Then, keeping this refined alignment leads to contradiction.
For j E D4. have
Wi di |w,]l
w Tray ||
Σ
Vollyt f (xzi W, W)
Then
d cos(x w,) (cos(x+, 2x) cos(244. wyy cos(T.. w))) ||
(cos(#+ Ha) (os(z,,0))) |Ball-
32

---

Published as a conference paper at ICLR 2024
//im,-mmm! We can show that cos(24. 2 3/ and | 2 vjm X / when St E (we defer the proof the end it breaks the flow), thus within [41, tgj, we have
cos(x w,) 2 { cos(x4+,0))) v0;)) Vunn+
(37)
di
We use (37) two ways: First, since
d cos(x4, wi)
Vun X
> 0,
dt
coa(x)
24
cos(14, my) 2 is trapping region for wy during [1,, Define to = inf 2 4 min cos(a-, (t)) 2 1), then clearly, if ta S t then cos( (t), 1. Vj V+ for all 1 ta <15 1
Now we use 37) again to show that ta S ti + Jum, X Suppose that ta 2 t1 + then 7j" such that cos(14, my (1)) < 1. 1I E 11144 Xam x. and we have
6
6 Jum X
1 cos(x+,, Wz) Xmin/
Vun Xumn
di cos(x1,,0y))
(38)
3
24
This shows
1 cos(x+, wy- (ti + 1)) 2 cos(x4, wz- + - 4 4
which contradicts that cos(24, Wy- (t))
4' Hence we know la S 4 +
VKM4Xmmn
summary, we have cos( w. (t), x+)
4 Vj E V4 for all t + Vin+Xmin S st t
Now we check the dynamics of Σje 11w, (t)l during t + vin, Xmin St S E For simplicity. we denote 1 + Xmni : --
For j E V4, we have, for 12 St
1|,]11== 2 Σ
- -Vglli f(zss W, e)) (zi, wyy) |eyll
di
iEI
Σ VEI+
(
ws) w)))
||wi|l
(by (40)
(x+, wz) ||wyy|l
|| || w))
cos(s+,,
1 - |0,
(Since 2 4
/um+ Xmin
Ju |1,113
(by Lemma
4
which leads (summing over E (4)
d
Vun+ Xmiu
Σ
|
Σ ||,,112
dt
4
JEV
JED.
33

---

Published as a conference paper at ICLR 2024
By Gronwall's inequality, we have
Σ
(t2))
JEV.
Viin+Xminn (tr (ti) Σ
exp
I1wy((
JEV+
Vun Xmin
2 exp
(1/) - (i)
Σ
[1,r((1))l
(By Lemma 14
JEV+
Vun+ Xmin
2 exp
(1/ - (1) exp ( (-4nXmaxli)
Σ
[m]]
(By Lemma 13
4
JEV-
exp
Vum Xminn
(1) - (i) exp (-4nXmaxhh (1) 2
2
min
VJXXmnn min
2
(by 36
However, at tz we have
1 Σ
Σ (,,
Σ
my (1)
4 1 IEI
=
th+
SEI JEV+
Uj (WW x+))
jEV
Σ
||r 112 cos(u1y x4+ ||++||
n jEV
|,,,121<<411
(Since / (2)
4n+
JEV,
Σ
|
VpXmmin
(by Lemma
JEV
which suggests Ejere 1|w/)12 < ViXmin" A contradiction.
Therefore, we must have
6 Von X X
(ox 2 V Xmi
2
5t 1144
+ 4nXmaxx()
(39)
v sun Xmin
W2
min
Since the dominant term here is Vun Xmin analysis shows to = = log 4). Therefore t2 = max t t = O log !)
log ('VaX mmn Wa we have t = 0 log 2)). similar
2
Complete the missing pieces We have two claims remaining to be proved. The first is cos(x_ R.)) > when & < t < t Since I+ Liez, and ã, Eier, - Vol(u,, f4(1; W, W))). We simply use the fact that before 15, we have, by Lemma
3
2 < - VVelii W. v)) =5 2
(40)
34

---

Published as a conference paper at ICLR 2024
to show the following
t, Ta) | |all
cos
EnSEZ, (- Vyl(yin (T; w))) (1))
Eijez, (1i4 Eiije, (-Vol(vi, S+ (0;; W, w))) (g)))
Ei.eE, (1i,
EiJez, (1) v EiJEZ, (-V f+ W, w))) (z),
1/8 EEJEE. (1))
2 1
EiJez, (1r, (0t,
VENEEI, (2)2 (2),
since all (o,, ()), ,i, E I+ are non-negative.
The second claim is | X 2 due that
|
Vum+ X
|
Σ
(-Vyllyi f4;;; W; v))) (I,, 2 2 i.ez+
(x)5
2
2
iJJEZ+
where the last inequality is from Lemma
11
D.3 PROOF OF LOW-RANK BIAS
So far have proved the directional convergence the early alignment phase and final (L/tt) convergence of the loss in the later stage. The only thing that remains to be shown is the low-rank bias. The proof is quite straightforward but need some additional notations.
As we proved above, after è,, neurons in S+ (S_) stays in S+ (S.). Thus the loss can be decomposed
L
Σe IEI
1/: Σ Uj (w)
Σ SEI
l Vi, Σ "y (115,
JEV+
jEV_
where 14 { (t) S+} and V_ {i wy(th) E S_}. Therefore, the training after t is decoupled into 1) using neurons in V4 to fit positive data in I4 and 2) using neurons in 12_ fit positive data in I_. We use
W- [W W = [W],
to denote submatrices of W by picking only columns in D_ and v__ respectively, Similarly, we define
= [b] b v= = =-=lvl
for the second layer weight Lastly, we also define
Wasad Wa Udcad
[M] Bans
where
Vicnd = {j (t1) E Saead) ) Given these notations, after t the loss is decomposed as
Σ
(Ur, x. W W+++)++ +
Σ VEI
(y,, x W_).),
and the GF on L equivalent to GF on L4 and f_ separately. suffices to study one of them. For GF on C.. we have the following important invariance Arora al. 2018b) Vt tt:
WI(OW+ (1) + (1) WI (4) (41 -+ (41)0((())) ((1),
35

---

Published as a conference paper at ICLR 2024
from which has
=
s
4eW
Σ (1))( JEV,
vh |+1,
where the last inequality is by Lemma
3
Then one can immediately get
4eW2
(t) - | (t) & ())+(()) - 2+ (0))
max
1u+1,
vh
which is precisely
(W)) (11( + 4eW
max
V+1.
(41)
vh
Similarly, we have
4eW2
(W-(t)] & || (0112 +
max
V-1.
(42)
Vh
Lastly, one has
4eW2
W dead 1/--
Σ
[wy(tz))1
V
(43)
h
JEVen
Adding
together, we have
(t)
[+ | (t)l + [Weendll 1|
(1 + || WvTeW22 max S 22 + 4Vhew22 max
Vh
Finally, since we have shown C + 0 as + oo, then Vi E [n]]. we have E(Ve, f(1; W.e) - 0. This implies
faa; W,,) v) y/
log Elvi, (a;WW, w) o0.
Because we have shown that
f(x;; W. v) - Σ
|| || S || max
jE(hi
f(z; W. + oo enforces | + 00 as t + 0o, thus ||] + co as well. This gets
lim sup t
2.
36

---

Published as a conference paper at ICLR 2024
E EXISTENCE OF CARATHEODORY SOLUTION UNDER FIXED SUBGRADIENT
d((z) 12>0
In this Appendix, we first introduce the notion of solution we are interested in for the GF 2
Caratheodory solutions that satisfy (2) for almost all time f. Next, in Appendix E.2 we show that if
we fix the ReLU subgradient as o' (a) = 12>0- then global Caratheodory solutions exists for (2) under Assumption I cease to exist when ReLU subgradient at zero is chosen to be fixed non-zero value, highlighting the importance of choosing the right subgradient for analysis; 2) Caratheodory solutions are potentially non-unique, the neurons' dynamical behavior could become somewhat irregular certain solutions not excluded, justifying the introduction of regular solutions (Definition
Finally, we use simple examples to illustrate two points: 1) Caratheodory solutions
E. CARATHEODORY SOLUTIONS
Given an differential equation
Ô = F(0), ((0) - 00
(44)
with F potentially be discontinuous, 0(1) is said to be Caratheodory solution of (4) if it satisfies the following integral equation
(() = 00 + / F(O(T))dT.
(45)
for all L E a). where E RRz0 do. In this section, are interested in global Caratheodory solutions: (() that satisfies (45) for all time / 2 0.
E.2 PROOF OF EXISTENCE OF REGULAR CARATHEODORY SOLUTIONS UNDER ASSUMPTION J
In this section, we show the existence of global regular (Definition I Caratheodory solutions to Ø = F()), 8(0) 6.. where 0 == {W. v} and F == V F defined from fixed choice of ReLU subgradient a' d'(x) 12 under Assumption For the sake of clear presentation, we first discuss the case of Saead = 0 where all solutions are regular. then discuss the modifications one needs to make when Saead + 0.
Existence of Caratheodory solutions when Sdead = 0: First of all, notice that V W.,.C is continuous almost everywhere except for a zero measure set A = (W, e 3i E [n], E [h] s.t. (zt, wy) = 0}. since discontinuity only happens when one has to evaluate a' a'((z,, wy)) at (at, wy) 0 for some j. Being a finite union of hyperplanes, A has zero measure.
For points outside A. the existence of local solution is guaranteed by the generalized Caratheodory existence theorem in Persson (1975) (We refer readers to Appendix [E.3] for the construction of such a local solution). The local solution can be extended to global solution, as long it does not encounter any point in A (the set where the flow is discontinuous). Whenever a point A is reached, one requires extra certificates to extend the solution beyond that point. Simply speaking, the existence of a local solution around a point in A requires that the flow around this point does not push trajectories towards A from both sides of the zero measure set, causing an infinite number of crossings of A, called Zeno behavior (van der Schaft & Schumacher 20000; Maennel et al., 2018). See Figure 18 and 19 for an illustration. In Appendix EE.5] we formally show that if there is no Zeno behavior, then solution can be extended until reaching discontinuity in A, and gets extended by leaving A immediately?
One sufficient condition for avoiding Zeno behavior is to show: For each hyperplane Aij ((r,, wyy) = 0), all points in neighborhood around this hyperplane Ay must satisfy that the inner products between the normal vector of A,j and the flow F have the same sign. Formally speaking, we need that there exists > 0, such that for all pair 0x, 0z E {0 = (W, v) : 0 < (2,, w;) < 6}. we have (NAy, F(Ox)) (Na,; F(e:)) > 0, here NA., should be a fixed choice of the normal vector of hyperplane Asj-
Sttrictyy speaking, Appendix P.5 is part of the proof hut discussing the technical part right now disrupts the presentation.
37

---

Published as a conference paper at ICLR 2024
t
F(O)
F(0)
Figure 18: Non-existence of Caratheodory lution around points of discontinuity (Does
Figure 19: Existence of Caratheodory solu-
tion around points of discontinuity (Guaran- teed by Assumption D. When the solution ((1) reaches A, it immediately leaves A since the flow on the opposite side is flowing out-
not happen under Assumption D Since flow
F(0) is continuous The solution f(t) can be extended until it reaches the points of dis-
continuity in A, after which the solution is forced to stay within A (often referred as
ward. This is a valid Caratheodory solution.
Zeno behavior: the solution is crossing A infinitely many times), and Caratheodory so- lution ceases to exist.
This inner product (Nay, F(O))) between the normal vector and the flow exactly computed as (I,, Va,; L). Under Assumption 1 we have much stronger result than what is required in the last paragraph: we can show that the entire parameter space, have (shown in Appendix E.5)
y,sign(v;) (Z,, Vn, C) > 0,
(46)
As such, since (t) does not change sign, Assumption prevents Zeno behavior and ensures the existence local solution around points A.
In summary, from any initialization, the Caratheodory solution be extended (Persson 1975)
until the trajectory encounters points of discontinuity in A. then the existence a local solution is guaranteed by ensuring that the flow forces the solution leave A immediately. Moreover, (46) ensures that A can only be crossed a finite number times (every hyperplane can only be crossed once), after which no discontinuity is encountered and the solution can be extended to t co. Therefore global Caratheodory solution always exists.
Existence of Caratheodory solutions when Saead * 0: Notice that when Saead A 0. A contains boundary of Saead- If the solution gets extended to A where one neuron lands on the boundary of Sdead, then this neuron stays at the boundary of Sdead, i.c. the solution stays at A. Therefore, the previous argument about existence does not apply.
However, one only needs very minor modification: If at time to. the solution enters A by having one neuron (say wy (t)) land on the boundary of Sacad, set (t) (ty) and (t) (to) for t 2 to. then exclude fwy, 2;} from the parameter space and continue constructing and extending local solutions for other parameters via the previous argument. This shows the existence of the Caratheodory solution under non-empty Sacad, and by our construction, the solution is regular.
E.3 NON-EXISTENCE CARATHEODORY SOLUTION UNDER OTHER FIXED SUBGRADIENT
Consider the following simple example: The training data consists of a single data point [1. T y = 11, and the network consists of a single neuron (w, o} initialized at (w(0) = (0, 1] (()) = 1}. See Figure 20 for illustration.
When the ReLU subgradient is chosen to be d'(x) = 1s>0, the Caratheodory solution (w(t) = exists, i.e. the neuron stays at the boundary of Sacad := {w (z. w) 0).
[0, 1]. vtt)
If the ReLU subgradient chosen to be o' e((0) = a > 0, then the Caratheodory solution ceases to exist: the neuron cannot stay the boundary (xr, w)) 0 of Sacad, because the non-zero a' o'(0) pushes
38

---

Published as a conference paper at ICLR 2024
(0)) [0, 1]
w(() = 00, 1
b(0) = 1
Saead
S
[1.0]
r [1
y
y - +1
Figure 20: Non-existence of Caratheodory Figure 21: Non-uniqueness of Caratheodory lution under fixed ReLU subgradient o' (0)
solutions. One solution is that w(l) = wo(0), =\t v(0), i.e. the neuron stays leave
> 0. For any
0, wi(l) cannot stay
at [0, I]" because the subgradient o' o0())
within Sdead. However, the neuron Stead at any time to > then follow the flow outside Saead, which is also valid
is positive, leading to non-zero W.
the same time, it cannot enter the interior of
Salend, because there no flow the interior.
Caratheodory solution.
towards the interior of Stead- However, the neuron cannot enter the interior of Stead because the flow is all zero within the interior of Saead-
To see this formally, suppose {w(t) = w()), vtt) v(0) for t [o, tol (ncuron stay at (w))), })) then by definition of Caratheodory solution, we have
ho ((w(O), =
= 0,
suggesting V (w.x r0)) = 0, which is not true when o(0)) > 0, thus a contradiction. Now suppose w(to) E Int(d)ead) for some to. then it must be that w(t) Int(Seead). VO < to. otherwise leads to the same contradiction as previous paragraph. By definition of Caratheodory solution, we have
w w(to) w()).
w.
The left-hand side is zero because w(t) E Int((de Ww,r (tt) 00 t to. The right-hand side is non-zero because wtoo) E &Int(S ). thus a contradiction. Similarly, w(() cannot enter Slead- Therefore, the Caratheodory solution ((),,(()) does not exist for any > 0.
E.4 NON-UNIQUENESS OF CARATHEODORY SOLUTIONS
Consider the following simple example: The training data consists of a single data point OT. y = 1, and the network consists of a single neuron (w, v) initialized at w(() = [0. 1] v0) 1. See Figure 21 for illustration.
We consider the case when the ReLU subgradient is chosen to be d'xx) = 12 There exists one Caratheodory solution we(t) = [0, 1] v(t) 1, i.e. the neuron stays the boundary of Sdead == {10 (ar, w) & 0}. However, consider (t), i(t) being the solution to the following ode (the one that neuron follows once enters the positive orthant):
= y exexpp-pp -y (x, ūz) = y yexp(-upp yu (z, w)) (x, w) (0)) = w(0), 5(0) = v(0). (47)
Then for any to 0.
wtt) 1, (0) 1,2,, o( to), u(() 1, (0) 1,20, o(t - to)
is a Caratheodory solution. This example shows that the Caratheodory solution could be non-unique.
This is somewhat troublesome for our analysis, one would like that all neurons in Saesd stay within Sdead, but Caratheodory solutions do not have this property, and in fact, long the neuron is
39

---

Published as a conference paper at ICLR 2024
the boundary of Saead, and the flow outside Sdead is pointing away from the boundary, the neuron can leave Sdead at any time and it does not violate the definition of a Caratheodory solution. Therefore, for our main theorem, forcing neurons to stay within Sdead.
added an additional regularity condition (Definition on the solution,
Remark 6. This issue of having irregular solutions is not specific to our choice of the notion of solutions. Even if one considers more generally the Filippov solution Filippov) (977)) of the differential inclusion in (2). the same issue of non-uniqueness persists and needs attention when analyzing neuron dynamics.
Remark 7. Although irregular solutions are not desired for analyzing neuron behaviors, as we see in this example, they are rare cases under very specific initialization of the neurons and thus can be avoided by randomly initializing the weights.
E.5 CONSTRUCTING GLOBAL CARATHEODORY SOLUTION
In this section, we formally show that if there is no Zeno behavior, then a solution can be extended until reaching discontinuity in A. and gets extended by leaving A immediately, leading to a construction of global Caratheodory solution. The only ingredient that is needed is the existence theorem in Persson 1975 Theorem 2.3), showing that if F(0) is continuous and AO
[F(0)] S M(1 + [[0]]=),
(48)
for some M O, then global solution of 0 = F() exists. Obviously, this result cannot be applied directly for two reasons: a) it requires continuity of the flow; b) it requires linear growth of [[(()] w.r.t. ||0||.. The key idea constructing local solution by restricting the flow to a neighborhood of
initial conditions where a) and b) are satisfied, and then extending this solution to global one.
As we discussed in Appendix F.2 we can assume Saead = 0 without loss of generality. Moreover, it suffices to show that starting from an initialization ((0)
(0). v(0)] outside A we can construct either: 1) a global solution without encountering any point in A: or 2) a local solution that lands on A at some to then leave A immediately. Because if 2) happens, we take the end of this local solution as a new initial condition and repeat this argument. Importantly, 2) cannot happen infinitely many times because we have shown in Appendix [..2 that A can only be crossed finitely many times, thus 1) must happen, resulting in a global solution.
Construct local solution from initial condition: Now given an initial condition ((0) (W) v(0)]. define the following two sets (Notation-wise, we drop the dependency on (W) x(0) for simplicity):
O = 10 = (W, v) <(W, S ((W()), (ig)s
(ig Vj e [h]}
+ = (0 = (W, o) vi E [m],] E [M], (zz wy) (z., (0) > 0}:
O₁ is the positive invariant set of (W( rO)): all solutions from (W( [(0}} never leaves O1. so it suffices study the flow within O for the existence of solutions. Moreover, O, is the intersection of a closed set (v. sign(v;) sign(e;(0) (0)) and the pre-image of continuous function L on the range [0, ((W(((W)))))) (0), (()) thus closed. O₂ is the largest connected set that contains (W( (0)) without point of discontinuity.
Consider the following set
b ; ncc(())
(49)
Then Ô is closed. Consider a new flow Fr θ such that FTI = Vw. C for all 0 e Int Int(@
and E2 (0) = 1o F(Ox) for all E \ Int(@r), where O2 E Int(( k = 1, 1, ... is convergent sequence to 0.
Frlle, is continuous by construction, and we can show that (at the end of this section)
(0)le S Cilllle, V0 E θ
(50)
By generalized version of the Tietze extension theorem Erean 1997, there exists continuous F on the entire parameter space, such that
(T) - Fi), ,,
(51)
initial condition within A is taken care of by 2).
40

---

Published as a conference paper at ICLR 2024
and
F (0)l== S CI A8.
(52)
Because now F Fr Fi on l any solution (t) 01 = Fi (01)), = 0(0) (existence guaranteed by Persson 1975)) gives a local solution of = F(0). = (()). for t S to, where to == inf(t: (t) $ 11).
If to oc, one has global solution and the construction is finished. If to < cc, it must be that i, (4x) E A (since 0, must leave θ via the boundary of O2). Now need construct solution that leaves A immediately.
Construct local solution that leaves A: As we discussed, A a union of hyperplanes. For simplicity, let us assume ((6) is not at the intersection of two hyperplanes (the treatment is similar but tedious, we will make remarks in the end).
Now (to) lands on a single hyperplane, let it be {0: (mi wyr) = 0}, we define
2 {0 (W, vi ; * j", (Zi, wz) (z;, wy (0)) > 0, and (Fis, Wy-) (1.-, ()-(())()))) < 0).
and we let
: Oc)
(53)
It is clear that, from the definition of 22 any solution we construct that leaves A immediately after tu must enter ( To construct the solution, we just need to repeat the first part, but now for O₂:
We construct F2 that is continuous on O₂ and agrees with F on the interior, then extends Fgl to F
on the entire parameter space. Consider the solution (2() of
82- = F(02), (0) = (to),
(54)
gives local solution of
0 0= F(),
(to).
(55)
Because we have shown that Zeno behavior does not happen, d2(() leaves immediately and enters Int( In). We just pick any 76 > 0 such that (7O) E Int(@2) then
(t) 1, (t) + 1 toststo 0g to);
(56)
is Caratheodory solution to ê F(6). 8(0) for to + To This is exactly what intended to show.
Remark 8. When ((to) lands the intersection of two (or more) hyperplanes, the only difference is that now there could be more regions to escape to. But under Assumption [ (46) suggests that the solution must cross all hyperplanes after to, leaving one unique region similar θ2. Then one constructs the local solution following previous procedures.
Complete the missing pieces To complete the proof, there are two statements (146) and (50) left to be shown.
To show (46), we start from the derivative
V
((xss W, W), (r)))rrsignn(((,()))) [wy)l
w
2003 expl ys ( wi)sssinn (0) wy
and we have
yisign(1z)) (I; Vwyf) f)
(1x, exp(-xxpppeeeeeee yvk (1ras W)) (yar,, yerk) |my]]
1 (an, ) oa expe--p-wwr f(xi W, ))) | |||=0:||/=||| > 0.
1=1
41

---

Published as a conference paper at ICLR 2024
since there is at least one summand (Séead 0), the summation is always positive.
To show 50
we first consider 0 E [nt((()), and we have
h
wy 2
=
(er.w exp( J W, v) 15
5 1=1
exp( - yu f W, W.
loyl 2.
X
exp( - f W, v))
k=1
h
Σ
10y 12 1.7.(XXma (Xmux ((W,
j=1
$
Σ
l0y 12 ((((O)) X2 (())
similarly, we also have
V 21,
Σ (ox m, >>08 exp(- yk faxs W.v) v)) (arks wyy)
exp( - yk [xrk; W, (
/=1
|
Xmax Σ
expl -o yo f W, w))
k=1
|wsl12 (XmaxC(W, u))²
|w//|| (X (X))) = xx2 X (WW),
Therefore, have v0 Int((11)
h
(0)) |F(0))1 Σ (1IV + [ (13)
5=1
S x2 max (W ((0) + [|
x2 max ((()
which gives (50 with C = Maax C (W(()),
Then for 0 E θ Int(@)), [f) = lim |F(())) < C'ima ||&l]| C| given some
Cauchy sequence 0g E ( ) 1,2, convergent to 0. This finishes proving (50).
WWe show it for exponential loss, the case of logistic loss is similar
42

---

Published as a conference paper at ICLR 2024
F EXTEND MAIN RESULTS TO SOLUTIONS TO DIFFERENTIAL INCLUSION
For Forliippp 1971 solutions (regular according to Definition D to the differential inclusion our Theorem remains the same. The only difference is that the notion of za (w) in no longer singleton, but rather an element from set:
[n(w) E {E2 E. ((22, w) wy)a
(57)
where (1x, u)) ) is subgradient of ReLU activation (z)) at (I,, w). Therefore, the proof of Theorem shall be modified (which can be done) to consider all possible choices of xa (w:)
In the case of d'())) == 0, za (w) become singleton Exter, =0 Ue which simplifies our discussions. This is the main reason opt fix this subgradient of (2) in the main paper.
43


Reviews:
[{"review_id": "snQcIF47cf", "reviewer": "Reviewer_Qhf5", "summary": "This paper studies the problem of training a two-layer ReLU network classifier via gradient flow under small initialization. Training dataset assumes well-separated input vectors. Analysis of the neurons’ directional dynamics establishes an upper bound on the time it takes for all neurons to achieve good alignment with the input data. Numerical experiment on the MNIST dataset validate the theoretical findings.", "questions": "1) Do the results/analysis extend to other activation functions? \n2) How small the initialization shall be to ensure the two-phase convergence? In general, since the training is nonconvex even with correlated/separable data due to the nonlinear relu activation function in the los. SGD/GD converges to a local minimum and indeed, this has been observed and numerically validated in the literature; see also [R1] Brutzkus et al. 2018. SGD learns over-parameterized networks that provably generalize on linearly separable data. ICLR. [R2] Wang et al. 2019. Learning ReLU networks on linearly separable data: Algorithm, optimality, and generalization. IEEE TSP. In [R1], SDG for two-layer ReLU networks under separable data converges to local minimum; yet for leaky ReLU networks, it finds global minimum. In [R2], it also shows that plain SGD on ReLU networks using separable data converges to local minimum numerically. Yet, a bit modification on the SGD helps SGD converge to a global minimum but with random initialization. It would be great if a comparison can be made between the approach in [R2] and the small-initialization SGD for training two-layer ReLU networks on e.g., MINIST. Moreover, is there any transition for such initialization value to go from the two-phase to single-phase gradient flow convergence to local minimum?\n3) It would be great if more numerical tests are provided to demonstrate the two-phase convergence and provide the plots. \n4) If the second-layer weights are initialized not in a balanced manner (although not initialized according to (3)), I understand it would also work and guess that the imbalance between positive v_j and negative v_j values only influences the time it takes for the two-phases. More balanced initalization, faster convergence. It would be interesting to numerically validate if this is the case. \n5) There are some grammar issues and typos. Please correct.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ Interesting alignment behavior of the gradient flow for training two-layer ReLU networks under small initialization and separable data\n+ Rate analysis for the after-alignment-phase convergence", "weaknesses": "- The results only hold for correlated data. \n- Only ReLU activation functions are analyzed.\n- Some of the results have been previously known or observed, e.g., the solution of (stochastic) gradient flow finds in training two-layer ReLU networks on separable data is (almost) low-rank.\n- How small the initialization shall be to ensure the two-phase convergence is not qualitiatively discussed?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uezHKvfizf", "reviewer": "Reviewer_osUg", "summary": "The paper studies the problem of training a two-layer ReLU network for binary\nclassification using gradient flow with small initialization on well-separated input datasets, i.e. datasets $(x_i, y_i)_{i \\in [n]}$ with $ \\min \\frac{\\langle x_iy_i, x_jy_j  \\rangle}{\\Vert x_i \\Vert_2 \\Vert x_j \\Vert_2 } \\geq \\mu$.\nThey show that in time $O(\\frac{\\log(n)}{\\sqrt{\\mu}})$ all neurons are well aligned with the input data, which means that positive neurons show in the same direction as the positively labeled points (or have a negative scalar prouct with all vectors of the dataset) (and an equivalent result for negative neurons).\nFurhter they show that after the early alignment phase, the loss converges to zero at a $O(1/t)$ rate, and the weight\nmatrix on the first layer is approximately low-rank. Numerical experiments are provided.", "questions": "-", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 1, "strengths": "All claims are proven and illustrations are supporting the explanations.", "weaknesses": "The assumption that the dataset is well seperated is very strong. I don't see why one would use a neural network on such a dataset rather than linear regression.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OLj9ffpPqo", "reviewer": "Reviewer_5S56", "summary": "The paper improves on previous theoretical analysis of the early alignement phase of the neurons of a shllow neural network initialized with small weights. This allows them to prove quantitative bounds in terms of the initialization scale, time, number of neurons and number of datapoints required to guarantee convergence.", "questions": "The description of Assumption 2 before the statement of the assumption does not match the statement of the assumption.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is easy to follow and explains well the previous issues and how they are solved. It is nice that the results apply to deterministic initialization of the weights, and do not require a random initialization (though they can of course be applied to this case).", "weaknesses": "The assumption of positively correlated labels and balancedness are very strong, and usually are not true in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HGwImARnts", "reviewer": "Reviewer_Qg2K", "summary": "This paper studies the learning dynamics in the special case of two-layer ReLU neural network with small initialization.", "questions": "In Eq.(5), why the R.H.S. is independent of the network output $f(x_i)$? Or, why there is no such term $y_i-f(x_i)$ which usually appears in the expression of gradients.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The results of this paper extend prior results from infinitesimal initialization to finitely small initialization.\n\nThe relation and difference with prior results are clearly presented. The paper is clearly written and easy to follow. The figures are helpful for understanding the argument.", "weaknesses": "The setting of the neural network is unconventional. It requires the second layer weights to depend on the first layer weights in a way as shown in Eq.(3), instead of independently initialized.  This setting is used neither in practice nor in most theoretical analysis. According to the analysis,  I doubt that the results of this paper hold without this restriction on the second layer weights. \n\n> The paper has some discussion on this setting. However, it does not justify the validity of this setting. That it is commonly assumed in other papers does not directly justify. I would like to see some analysis, or at least some intuition, on why the results would hold on the natural setting.\n\nThe assumption on the data (Assumption 1) is strong, as it is not met by almost all real data. \n\nThe significance of the results is limited, as it is an extension of similar results from $\\epsilon \\to 0$ to the finite but small $\\epsilon$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)