Gradient Descent on Two-layer Nets: Margin Maximization and Simplicity Bias
Authors: Kaifeng Lyu
Zhiyuan Li
Runzhe Wang
Sanjeev Arora
Venue: NeurIPS 2021 Poster (2021)
Keywords: margin maximization
gradient flow
gradient descent
implicit bias
linearly separable
two-layer neural networks

Abstract: The generalization mystery of overparametrized deep nets has motivated efforts to understand how gradient descent (GD) converges to low-loss solutions that generalize well. Real-life neural networks are initialized from small random values and trained with cross-entropy loss for classification (unlike the "lazy" or "NTK" regime of training where analysis was more successful), and a recent sequence of results (Lyu and Li, 2020; Chizat and Bach, 2020; Ji and Telgarsky, 2020) provide theoretical evidence that GD may converge to the "max-margin" solution with zero loss, which presumably generalizes well. However, the global optimality of margin is proved only in some settings where neural nets are infinitely or exponentially wide. The current paper is able to establish this global optimality for two-layer Leaky ReLU nets trained with gradient flow on linearly separable and symmetric data, regardless of the width. The analysis also gives some theoretical justification for recent empirical findings (Kalimeris et al., 2019) on the so-called simplicity bias of GD towards linear or other "simple" classes of solutions, especially early in training. On the pessimistic side, the paper suggests that such results are fragile. A simple data manipulation can make gradient flow converge to a linear classifier with suboptimal margin.

Gradient Descent on Two-layer Nets: Margin Maximization and Simplicity Bias
Kaifeng Lyu"t
Zhiyuan Li"
Princeton University klyudes, wwtuuuss.princce. pwwnceton.ee.ee edu
Princeton University
zhiyuanli@cs wewecs.ppinncee pwwnceton.eeeeee edu
Runzhe Wang"! Princeton University runzheweprinceton.. edu
Sanjeev Arora Princeton University arorašcs, arora@cs.princettnn wrwnceton. edu
Abstract
The generalization mystery of overparametrized deep nets has motivated efforts to understand how gradient descent (GD) converges to low-loss solutions that generalize well. Real-life neural networks are initialized from small random values and trained with cross-entropy loss for classification (unlike the "lazy" or "NTK" more successful), and recent sequence of results (Lyu and Li, 2020; Chizat and Bach, 2020; Ji and Telgarsky, 2020a) provide theoretical evidence that GD may converge the maxx-arginn solution with zero loss, which presumably generalizes well. However, the global optimality of margin is proved only in
regime of training where analysis
settings where neural nets infinitely exponentially wide. The current paper is able to establish this global optimality for two-layer Leaky ReLU nets trained with gradient flow linearly separable and symmetric data, regardless of the width. The analysis also gives some theoretical justification for recent empirical findings (Kalimeris al., 2019) the so-called simplicity bias of GD towards linear or other "simple" classes of solutions, especially early in training. On the pessimistic side, the paper suggests that such results are fragile. A simple data manipulation can make gradient flow converge to a linear classifier with suboptimal margin.
Introduction
One major mystery in deep learning is why deep neural networks generalize despite overparameteri- zation (Zhang et al., 2017). To tackle this issue, many recent works turn to study the implicit bias of gradient descent (GD) - what kind of theoretical characterization can we give for the low-loss solution found by GD?
The seminal works by Soudry et al. (201 8a,b) revealed an interesting connection between GD and margin maximization: for linear logistic regression on linearly separable data, there can be multiple linear classifiers that perfectly fit the data, but GD with any initialization always converges to the max- margin (hard-margin SVM) solution, even when there is no explicit regularization. Thus the solution found by GD has the same margin-based generalization bounds hard-margin SVM. Subsequent works on linear models have extended this theoretical understanding of GD to SGD (Nacson et al., 2019b), other gradient-based methods (Gunasekar et al., 2018a), other loss functions with certain poly-exponential tails (Nacson et al., 2019a), linearly non-separable data (Ji and Telgarsky, 2018, 20190 2019b), deep linear nets (Ji and Telgarsky, 2019a; Gunasekar et al., 2018b).
Equaal contribution
Most of the work is done when Kaifeng Lyu and Runzhe Wang were at Tsinghua University.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Given the above results, a natural question ask is whether GD has the same implicit bias towards max-margin solutions for machine learning models in general. Lyu and Li (2020) studied the relationship between GD and margin maximization on deep homogeneous neural network, i.c. neural network whose output function is (positively) homogeneous with respect to its parameters. For homogeneous neural networks, only the direction of parameter matters for classification tasks. For logistic and exponential loss, Lyu and Li (2020) assumed that GD decreases the loss to a small value and achieves full training accuracy at some time point, and then provided an analysis for the training dynamics after this time point (Theorem 3. 1. which we refer to as late phase analysis. It is shown that GD decreases the loss to 0 in the end and converges to direction satisfying the Karush- -aaushh-Khn----e Kuhn-Tucker (KKT)) conditions of constrained optimization problem (P) margin maximization.
However, given the non-convex nature of neural networks, KKT conditions do not imply global optimality for margins. Several attempts are made to prove the global optimality specifically for two-layer nets, Chizat and Bach (2020) provided a mean-field analysis for infinitely wide two-layer Squared ReLU nets showing that gradient flow converges to the solution with global max margin, which also corresponds to the max-margin classifier in some non-Hilbertian space of functions. Ji and Telgarsky (2020a) extended the proof to finite- witth neural nets, but the width needs to be exponential in the input dimension (due to the use of covering condition). Both works build upon late phase analyses. Under a restrictive assumption that the data is orthogonally separable, i.e., any data point a, can serve as a perfect linear separator, Phuong and Lampert (2021) analyzed the full trajectory of gradient flow on two-layer ReLU nets with small initialization, and established the convergence to a piecewise linear classifier that maximizes the margin, irrespective of network width.
In this paper, we study the implicit bias of gradient flow on two-layer neural nets with Leaky ReLU activation (Maas et al., 2013) and logistic loss. To avoid the lazy or Neural Tangent Kernel (NTK) regime where the weights are initialized to large random values and do not change much during training (Jacot et al., 2018: Chizat et al., 2019; Du et al., 2019b,a; Allen-Zhu al., 2018, 2019; Zou et al., 2018; Arora et al., 2019b), we use small initialization to encourage the model to learn features actively, which is closer to real-life neural network training.
When analyzing convergence behavior of training on neural networks, one can simplify the prob- lem and gain insights by assuming that the data distribution has simple structure. Many works particularly study the case where the labels are generated by an unknown teacher network that is much smaller/simpler than the (student) neural network to be trained. Following Brutzkus al. (2018); Sarussi et al. (2021) and many other works, we consider the case where the dataset is linearly separable, namely the labels are generated by a linear teacher, and study the training dynamics of two-layer Leaky ReLU nets on such dataset.
1.1 Our Contribution
Among all the classifiers that can be represented by the two-layer Leaky ReLUU nets, we show any maaallmm--marrginn classifier is exactly linear under one more data assumption: the dataset is symmetric, i..., if in the training set, then so is x. Note that such symmetry can be ensured by simple data augmentation.
Still, little is known about what kind of classifiers neural network trained by GD learns. Though Lyu and Li (2020) showed that gradient flow converges to a classifier along KKT-margin direction, we note that this result is not sufficient to guarantee the global optimality since such classifier can have nonlinear decision boundaries. See Figure (left) for an example.
In this paper, we provide multi-phase analysis for the full trajectory of gradient flow, in contrast with previous late phase analyses which only analyzes the trajectory after achieving 100% training accuracy. We show that gradient flow with small initialization converges to global-max-margh linear classifier (Theorem 4.2). The proof leverages power iteration to show that neuron weights align in two directions in an early phase of training, inspired by Li et al. (2021). We further show the alignment at any constant training time by associating the dynamics of wide neural net with that of two-neuron neural net, and finally, extend the alignment to the infinite time limit by applying Lojasiewicz (KL) inquality in a similar way as Ji and Telgarsky (2020a). The alignment at convergence implies that the convergent classifier is linear.
The above results also justify a recent line of works studying the so-called simplicity bias: GD first learns linear functions in the carly phase of training, and the complexity of the solution increases
2

---

as training goes on (Kalimeris al., 2019; Hu et al., 2020; Shah al., 2020). Indeed, our result establishes a form of extreme simplicity bias of GI: if the dataset can be fitted by a linear classifter, then GD learns a linear classifier not only in the beginning but also at convergence.
On the pessimistic side, this paper suggests that such global margin maximization result could be fragile. Even for linearly separable data, maaaal- classifiers may be nonlinear without the symmetry assumption. In particular, we show that for any linearly separable dataset, gradient flow can be led to converge to linear classifier with suboptimal margin by adding only 3 extra data points (Theorem 6.2). See Figure (right) for an example.
2 Related Works
Generalization Aspect of Margin Maxmization. Margin often appears in the generalization bounds for neural networks (Bartlett et al., 2017; Neyshabur al., 2018), and larger margin leads to smaller bounds. Jiang et al. (2020) conducted an empirical study for the causal relationships between complexity measures and generalization errors, and showed positive results for normalized margin, which is defined by the output margin divided by the product (or powers of the sum) of Frobenius norms of weight matrices from each layer. On the pessimistic side, negative results are also shown if Frobenius norm is replaced by spectral norm. In this paper, we do use the normalized margin with Frobenius norm (see Section 3).
Learning on Linearly Separable Data. Some works studied the training dynamics of (nonlincar) neural networks on linearly separable data (labels are generated by linear teacher). Brutzkus et al. (2018) showed that SGD on two-layer Leaky ReLU nets with hinge loss fits the training set in finite steps and generalizes well. Frei et al. (2021) studied online SGD (taking a fresh sample from the population in each step) on the two-layer Leaky ReLU nets with logistic loss. For any data distribution, they proved that there exists a time step in the early phase such that the net has test error competitive with that of the best linear classifier over the distribution, and hence generalizes well on linearly separable data. Both two papers reveal that the weight vectors in the first layer have positive correlations with the weight of the linear teacher, but their analyses do not imply that the learned classifier is linear. In the NTK regime, Ji and Telgarsky (2020b); Chen al. (2021) showed that GD on shallow/deep neural nets learns a kernel predictor with good generalization on linearly separable data, and it suffices have width polylogarithmic in the number of training samples. Still, they do not imply that the learned classifier is linear. Pellegrini and Biroli (2020) provided a mean-field analysis for two-layer ReLU net showing that training with hinge loss and infinite data leads to a linear classifier, but their analysis requires the data distribution to be spherically symmetric (i.e., the probability density only depends on the distance to origin), which is a more restrictive assumption than ours. Sarussi et al. (2021) provided a late phase analysis for gradient flow on two-layer Leaky ReL.U nets with logistic loss, which establishes the convergence to linear classifier based on an assumption called Neural Agreement Regime (NAR): starting from some time point, for any training sample, the outputs of all the neurons have the same sign. However, it is unclear why this can happen priori. Comparing with our work, we analyze the full trajectory of gradient flow and establish the convergence to linear classifier without assuming NAR. Phuong and Lampert (2021) analyzed the full trajectory for gradient flow on orthogonally separable data, but every KKT-margin direction attains the global max margin (see Appendix H) in their setting, which is not necessarily true in general. In our setting, KKT-margin direction with suboptimal margin does exist.
Simplicity Bias. Kalimeris et al. (2019) empirically observed that neural networks in the early phase of training are learning linear classifiers, and provided evidence that SGD learns functions of increasing complexity. Hu ct al. (2020) justified this view by proving that the lcarning dynamics of two-layer neural nets and simple linear classifiers are close to each other in the early phase, for dataset drawn from a data distribution where input coordinates are independent after some linear transformation. The aforementioned work by Frei et al. (2021) can be seen another theoretical justification for online SGD on aribitrary data distribution. Shah et al. (2020) pointed out that extreme simplicity bias can lead to suboptimal generalization and negative effects on adversarial robustness.
Small Initialization. Several theoretical works studying neural network training with small initial- ization can be connected to simplicity bias. Maennel et al. (2018) uncovered a weight quantization effect in training two-layer nets with small initialization: gradient flow biases the weight vectors to a certain number of directions determined by the input data (independent of neural network width). It
3

---

Symmetric Data pwfthout Blas)
Non-sym-mebric Data (with BlaS)
Non- Non-symmetric hinted Data (without blas)
1.88
0.00
x
-1.08
sebal--as-mangg boundary KKT-margin
Gobalmss-marf
Mac-marg miliner boundary
Mes-marpnsimer boundary
Figure I: Two-layer Leaky ReLU nets (Mesky = 1/2) with KKT margin and global max margin on linearly separable data. See Appendix 1.1 for detailed discussions. Left: Theorem 4.3 is not vacuous: a symmetric dataset can have KKT directions with suboptimal margin, but our theory shows that gradicnt flow from small initialization goes to global max margin. Middle: The linear classifier (orange) is along a KKT-margin direction with a much smaller margin comparing to the (nonlinear) maobal-max-margeee classifier (black), but theory suggests that gradient flow converges to the linear classifier. Right: Adding three extra data points (marked as "x"; see Definition 6.1) 1) to a linearly separable dataset makes the linear classifier (orange) has suboptimal margin but causes the neural to be biased to it.
is hence argued that gradient flow has bias towards "simple" functions, but their proof is not entirely rigorous and no clear definition of simplicity is given. This weight quantization effect has also been studied under the names of weight clustering (Brutzkus and Globerson, 2019), condensation (Luo et al., 2021; Xu et al., 2021). Williams et al. (2019) studied univariate regression and showed that two-layer ReLU nets with small initialization tend to learn linear splines. For the matrix factorization problem, which can be related to training neural networks with linear or quadratic activations, we can measure the complexity of the learned solution by rank. A line of works showed that gradient descent learns solutions with gradually increasing rank (Li et al., 2018; Arora et al., 2019a; Gidel et al., 2019; Gissin al., 2020; Li et al., 2021). Such results have been generalized to tensor factorization where the complexity measure is replaced by tensor rank (Razin et al., 2021). Beyond small initialization of our interest and large initialization in the lazy or NTK regime, Woodworth et al. (2020); Moroshko et al. (2020); Mehta et al. (2021) studied feature learning when the initialization scale transitions from small large scale.
3 Preliminaries
We denote the set ..) by [n] and the unit sphere {x E IRd ||0||| = 1} by gd-1 We call function h : RD + R L-homogeneous if h(cl) c' (0) for all 0 E and > For
S C R conv(S) denotes the convex hull of S. For locally Lipschitz function f R D + R. we
define Clarke's subdifferential (Clarke, 1975; Clarke et al.. 2008; Davis al., 2020) be 0 f())
conv aoima V (0,)) differentiable at On, lim O, = 0} (see also Appendix B.
3.1 Logistic Loss Minimization and Margin Maximization
For neural net,
So (x)) E R denote the output logit input æ E IRd when the parameter is E D We say that the neural net is T-hhmmggeneous if fo fo(a)) (x) is I--homogeneous with respect to 0,
i... (z)(x) ch fø(x) for all R and e O. VGG-like CNNs can be made homogeneous if remove all the bias terms expect those in the first layer (Lyu and Li, 2020).
Throughout this paper, restrict attention to L-homogeneous neural nets with fa definable with respect to an o-minimal structure for æ. (See Coste 2000 for reference for o-minimal structures.) This is technical condition needed by Theorem 3.1, and it is mild regularity condition as almost all modern neural networks satisfy this condition, including the two-layer Leaky ReLU networks studied in this paper.
For dataset ((11,91)..... (In,, y. )) we define qi (0) = vaJo((r;) to the output margin on the data point (æ y.). and gmin (0) = min.e(n) (0) be the output margin on the dataset S (or

---

margin for short). It is easy that 91 (0),....,..... T-homogeneous functions, and so is
Imin 0). We define the normalized margin r()) :- Jmin
0
Amin (0)
to be the output margin
[0lla
0
(on the dataset) for the normalized parameter Toll2"
We refer the problem of finding 0 that maximizes (() as margin maximization. Note that once we have found optimal solution 0* E IR=, c0 is also optimal for all c >> 0. We can put the norm constraint on 0 to eliminate this freedom on rescaling:
max
(0).
(M)
BESD-1
Alternatively, we can also constrain the margin to have Quii 2 and minimize the norm:
min 100 S.L. (0) 2 E [m].
One can easily show that 0" is a global maximizer of (M) if and only if [Qmin (0 Juuz is global minimizer of (P). For convenience, we make the following convention: if O is a local/global
0
TTell
maximizer of (M), then say 9 is along local-max-mane direction| global-max-margin direction;
if 0 (umin (Qmin(0))) (0) satisfies the KKT conditions of (P), then we say 0 is along a KKT-margin direction.
Gradient flow with logistic loss is defined by the following differential inclusion,
dl E - with L()) = (Q)())).
n
(1)
dt
where k(g) == In(1+ + is the logistic loss. Lyu and Li (2020); Ji and Telgarsky (2020a) showed that ((t)/][((t]]l always converges to KKT-margin direction. We restate the results below.
Theorem 3.1 (Lyu and Li 2020; Ji and Telgarsky 2020a). For homogeneous neural networks, if ((O(0)) < ln2 then [(O(t)) - 0, [[0(t)]la + +oG, and [0(1)]|2 converges to a KKT-margin n
0(t)
direction as t +o0.
3.2 Two-Layer Leaky ReLU Networks on Linearly Separable Data
Let (bx)) max (xr (leaky be Leaky ReLU, where Alaaky E (0, 1). Throughout the following sections, we consider two-layer neural net defined as below,
fo(
S a)).
where wi..... Wm Rdd are the weights in the first layer, 0l,..., E R are the weights in the second layer, and O = (wi..... Wm, al...., am) E R is the concatenation of all trainable parameters, where = md + m. We can verify that fo(x) ) is 2-homogeneous with respect to O.
Let S = (<11,1)).... vi).... (æn, }n)} be the training set. For simplicity, we assume that ||æ 2| S 1. We focus on linearly separable data, thus we assume that linearly separable throughout the paper.
Assumption 3.2 (Linear Separable). There exists E Rd such that yi (w, an:) 2 1 for all i [n]] Definition 3.3 (Max-margin Linear Separator). For the linearly separable dataset S, we say that w² gd-1 is the max-margin linear separator wn maximizes min4e(n) Be (w, æx)) over ur $d-1
Training on Linearly Separable and Symmetric Data
In this section, study the implicit bias of gradient flow assuming the training data is linearly separable and symmetric. We say a dataset is symmetric whenever & is present in the training set, the input
also present. By linear separability, and must have different labels because (w=, xx) = (w", =x), where wn is the max-margin linear separator. The formal statement for this assumption is given below.
Assumption 4.1 (Symmetric). is even and T; 22tnn22 Vi Mitn/2 -1 for s n/2.
5

---

This symmetry be ensured via data augmentation. Given dataset, if it is known that the ground- truth lahels are produced by an unknown linear classifier, then one can augment each data point (æz, y) by flipping the sign, i.c., replace it with two data points (z.y). y). (-x, ----yyl y) (and thus the dataset size is doubled).
Our results show that gradient flow directionally converges to a global-max-marghn direction for two-layer Leaky ReLU networks, when the dataset is linearly separable and symmetric. To achieve such result, the key insight is that any global-mas max-mmaa-aaaaaaa direction represents linear classifier, which we will see in Section 4.1. Then we will present our main convergence results in Section 4.2.
4.1 Global-Max-Margin Classifiers are Linear
Theorem 4.2 below characterizes the global-max-manee margin direction in our case by showing that margin maximization and simplicity bias coincide with each other: a network that representing the max- margin linear classifier (i.e., fo) (ww", ..) for some > 0) can simultaneously achieve the goals of being simple and maximizing the margin.
Theorem 4.2. Under Assumptions 3.2 and 4.1, for the two-layer Leaky ReLU network with width m 2 2, any global-max-marg direction 0 SD-1, for represents a linear classifier. Moreover, we have fe- (x)
140lakky (w², for all a E R where w" is the max max-margin linear separator.
4
The result of Theorem 4.2 is based on the observation that replacing each neuron (ax, wx) in
network with two neurons of oppositing parameters (ag, w/) and (-ak. (-ak:-wy) does not decrease the normalized margin on the symmetric dataset, while making the classifier linear in function space. Thus if any direction attains the global max margin, we can construct a new global-max-man max-magiinn direction which corresponds to a linear classifier. We can show that every weight vector wr of this linear classifier must be in the direction of w" or ur.. Then the original classifier must also be linear in the same direction.
4.2 Convergence to Global-Max-Marye Directions
Though Theorem 3.1 guarantees that gradient flow directionally converges a KKT-margin direction the loss is optimized successfully, we note that KKT--mrgiin directions can be non-linear and have complicated decision boundaries. See Figure 1 (left) for example. Therefore, establish the convergence to linear classifiers, Theorem 3.1 is not enough and we need a new analysis for the trajectory of gradient flow.
lid. We use initialization wy ~ NO0 i I). ag ~ N00 caiit Fiit). where Cainit is fixed constant throughout this paper and Fnit controls the initialization scale. We call this distribution as Og Dinit (onit() An alternative way to generate this distribution is to first draw G, ~ Dinit (1), and then set 0y Oni With small initialization, we can establish the following convergence result.
iid.
Theorem 4.3. Under Assumptions 3.2 and 4.1 and certain regularity conditions (see Assumptions 4.5 and 4.6 below), consider gradient flow on Leaky ReLU network with width m Z 2 and initialization 0, Oi where 0, Dini(1), With probability 1 (m over the random draw of Oo. if the initialization scale is sufficiently small, then gradient flow directionally converges and poo (x) ;= limm S0(0//10(0)1a (x) represents the max-margin linear classifier. That is,
Pr (0-Dmntt1)
Boxx > S.s.I. Vonnit < mmnt Vxx E Rdd ( =C =C0" (,00,xx)) ;a)]] 2 1 - 2 2-(mm11)
where C := 1+ is scaling factor.
Combining Theorem 4.2 and Theorem 4.3, we can conclude that gradient flow achieves the global max margin in our case.
Corollary 4.4. the settings of Theorem 4.3, gradient flow on linearly separable and symmetric
data directionally converges to the global maxball-xx-arggin direction with probability 1 - 2 (m (m-1))
4.3 Additional Notations and Assumptions
Let μ 1, Em y,sr, which non-zero since (n,, w.) n Eie(n) V.w, xx 1. Let j == TA
We use (00,, d to the value of 0 at time for (()) On.

---

We make the following technical assumption, which holds if we are allowed to add a slight perturbation to the training set. Assumption 4.5. For all i [n]. (ja, (x)) A
Another technical issue we face is that the gradient flow may not be unique due to non-smoothness.
is possible that p(0u, t) is not well-defined as the solution of (1) may not be unique. See Appendix 1.2 for more discussions. In this case, we assign ço(On, -) to be an arbitrary gradient flow trajectory starting from Og. In the case where (0(,; 1) has only one possible value for all / 2 0, we say that Oo a non-branching starting point. We assume the following technical assumption.
Assumption 4.6. For any 2 2, there exist ree > 0 such that O is non-branching starting point if its neurons can be partitioned into two groups: in the first group, ag || E (0,r) and all wy point to the same direction w+ E gá-1 with || ااl2 e; in the second group, -ag
| (0, r) and all wE point to the same direction w E s with || + øll2 S E.
Proof Sketch for the Symmetric Case
In this section, we provide proof sketch for Theorem 4.3. Our proof uses multi-phase analysis, which divides the training process into phases, from small initialization to the final convergence. We will now elaborate the analyses for them one by one.
5.1 Phase I: Dynamics Near Zero
Gradient flow starts with small initialization. In Phase I, analyze the dynamics when gradient flow does not go far away from zero. Inspired by Li et al. (2021), we relate such dynamics to power iterations and show that every weight vector wy in the first layer moves towards the directions of either j or -p. To see this, the first step is to note that fa x 0 when 0 is close to O. Applying Taylor expansion on ((yifo(x;))
((0)
Σ ((y:f((x))) =
Σ (E(0) (0)
n iE[n]
n 3E/m]
Expanding (x;;;) and rcorganizing the terms, we have
1 Σ ((0) C(0)=
Σ (0)
Σ
sandwwp æ.) E(0) + (0())
Σ
yaan pww xs))
kE|M]
sE(M]
kE[m]
kE[m) (E(m)
= =((0) Σ ag GxG((u):
km(ml
where G-function (Maennel et al., 2018) defined below:
G(w) = (0)
Σ voo(w T xx))
Σ vd mx))
= 2n
n
mE|m
iE|m]
This means gradient flow optimizes each -a, G(w)) separately near origin.
dw
daş G(w)).
dt a G(wx))
(3)
In the case where Assumption 4.1 holds, we can pair each x, with - -x; and use the identity d(2) d( ((=z) max (z, (leakyz)) max (2,, Menkk 2} (1 Meaky )= show that G(w) is linear:
1
G(w) 2n
Σ
(ó(w To:)=
x;) - d( w x:)) 2n vE(n/2)
Σ
(1+ Cleaky w xx (w, pt).
4E/n/2]
where ş = 1+0tlsky
Eie(n) y.s. Substituting this formula for G into (3) reveals that the dynamics of two-layer neural nets near zero has close relationship power iteration (or matrix exponentiation) of matrix Ma. E R(d+1)x((+1) that only depends data.
1+atesky 2n
d
Wk
Wk
0
j4
z M
where
M =
df
0
ak
0
7

---

1) RR++1 are the unique top eigenvalue and eigenvector of Ma. which suggests that (ux ax(()) Rd+1 aligns to this top eigenvector direction the approximation (3) holds for a sufficiently long time. With small initialization, this can indeed be true and we obtain the following lemma.
Simple linear algebra shows that Ao | 1/2 (1,
Definition 5.1 (M-norm). For parameter vector (W), Wm; 01 Am), we define the M-norm to be ||0|| = mm) (|axxl// |ax 1]}.
Lemma 5.2. Let
be small value. With probability 1 over the random draw of 0o ,
(wi...., :.. 01,..., am) Dinit (1), we we take mit I <m /m time (i(r)= Xo In
then any neuron (wk, 0x) ar
Vimorais 1 can be decomposed into
wx (x(((rr)) rb&f + Awkk
ax(Ti(F)) rbs + Dak:
where by := (10g, D)+ and the error term A = (Aw),.... Awm; Aui.... Amm) is bounded
2m/m 0o
Cr by | JM S vm for some universal constant C.
5.2 Phase II: Near-Two-Neuron Dynamics
By Lemma 5.2, we know that at time (r) we have wT = rbeft and og (T)()) = rbe. where b E I is some fixed vector. This motivates to couple the training dynamies of 0(1) (wn (l)....,,u) Wm (4), (0,a1(()) (1))....., am (4)) after the time Ti (r) with another gradient flow starting from the point (rdjn. rbm//, rbi,.... ..., r..)). Interestingly, the latter dynamic can be seen a dynamic of two neurons "embedded" into the m-neuron neural net, and we will show that ((t) is close to this "embedded" two-neuron dynamic for long time. Now we first introduce our idea of embedding two-ncuron network into an m-neuron network.
Embedding. For any b R say that b is good embedding vector if it has least one positive entry and one negative entry, and all the entries are non-zero. For good embedding vector b we use b+ == v EzE(m Dy and b_ == - Eyetmm bo to denote the root-sum-squaree of the positive entries and the negative root-sum sum-suuared of the negative entries. For parameter Ô = (w,, U22 ài, å2) of a two-neuron neural with ài > 0 and å2 from two-neuron into m-neuron neural nets as Tb(#0) (Ū), 102, âi, a2) = (W),.... Wm, al..... am) where
define the embedding
by be by 0_ àz, by
if be > 0
bi wi if bb > 0 ba L22 if by < 0
as
wk
is easy to check that (a( = S (2r) by the homogeneity of the activation (Ø(cz) = cop(2) for 0):
(x)
Σ
and(wf x)) + Σ
and(w x:)
be >0
12
Σ
dn s u x)
Σ by
à x) án x)) 02 x)) fo(a)).
Moreover, by taking the chain rule, we can obtain the following lemma showing that the trajectories starting from i and To are essentially the same.
Lemma 5.3. Given Ô = (wo,, 22, âi, â2) with ài > and å2 < 0, if both Ô and (0) are non-branching starting points, then (0), = To(p( f)) for all t 2 0.
Approximate Embedding. Back to our analysis for Phase II, b is good embedding vector with high probability (see lemma below). Let θ := 164,644, b_ b_ By Lemma 5.2, TT(r (rb,j,.... .... rbm,k rbı,..., rbm,) R (T))))), which means r - ((Ti(r)) is approximately an embedding. Suppose that the approximation happens to be exact, namely T5(rÔ) = 0(Ti then ((Ti(r) (r) + 1) = m ()) by Lemma 5.3. Inspired by this, we consider the case where Finit + 0, r - so that the approximate embedding is infinitely close to the exact one, and prove the following lemma. We shift the training time by (2(r) to avoid trivial limits (such as 0).
8

---

Lemma 5.4. Follow the notations in Lemma 5.2 and take Vinit.
Let T2(r) = 50
// /o
then Ti2 == (r) + T2(r) = + In Vmmmmi Oul1x
regardless the choice of For width m 2 2.
with probability 2-(m-1) over the random draw of Oo Dinit (1), the vector b E Rn is good embedding vector, and for the two-neuron dynamics starting with rescaled initialization the direction of 0 (b+, by fa, b_. b- p). the the following limit exists for all t,
((t) = lim vp r-ù
(0 (2(r) #0. 0
(4)
and moreover, for the m-neuron dynamics of ((), the following holds for all L,
lim (T12 + t) = (T()t)))
(5)
Cinat +00
5.3 Phase III: Dynamics near Global-Max-Margee Direction
With some efforts, we have the following characterization for the two-ncuron dynamics.
Theorem 5.S. Form 2, if initially an = || ||2. 02 =11/wal12 (wi w"), > 0 and (w2, w²) then (t) directionally converges to the following tataalmmx-maraa direction,
((t)
1
lim +0+000
1 (w0 - -w", ,,, 11,,,
[0(1)]22
where w is the max- max-margin linear separator.
It is not hard to verify that (t) satisfies the conditions required by Theorem 5.5. Given this result, first attempt to establish the convergence of (tt) to global-max max-margin direction is to take t - too on both sides of (5). However, this only proves that 0 (T'2 f) directionally converges to the global- max-margin direction if we take the limit Jii - first then take t - +00, while we are interested in the convergent solution when t - to0 first then Jiit - 0 (i.e., solution gradient flow converges to with infinite training time, it starts from sufficiently small initialization). These two double limits not equivalent because the order of limits cannot be exchanged without extra conditions.
To overcome this issue, we follow similar proof strategy as and Telgarsky (2020a) to prove local convergence near a local-max- marx-nn direction, as formally stated below. Theorem 5.6 holds for L-homogeneous neural networks in general and we believe of independent interest.
Theorem 5.6. Consider any T--homoeeeous neural networks with logistic loss. Given local-max- margin direction 0" E D-1 and any > 0, there exists Eo > O and po 2 such that for any Og with norm ||0gl|2 2 Po and direction converges to some direction 0 with the same normalized margin as 0", and 0 0 S 5.
| 0-1 || S <b. gradient flow starting with 0g directionally
Using Theorem 5.6, neuron mlobal-max-marginn direction [ww", -w", 1, 11) after embedding is global-max-margii direction for m-neurons, and constant 8 > 0 is still global-max-marge direction. Then we can take i to be large enough so that T6 ((())) satisfies the conditions in Theorem 5.6. According to (5), we can also make the conditions hold for 0 (Ti2 1) by taking iint and to be sufficiently small. Finally, applying Theorem 5.6 finishes the proof.
finish the proof for Theorem 4.3 as follows. First note that the two-
prove that any direction with distance more than small
Non-symmetrie Data Complicates the Picture
Now we turn to study the case without assuming symmetry and the question whether the implicit bias to dlobal-max-man max-margin classifier very fragile for any linearly separable dataset, we can add 3 extra data points so that every linear classifier has suboptimal margin but still gradient flow with small initialization converges to a linear classifier." See Definition 6.1 for the construction and Figure (right) for example.
solution still holds. Unfortunately, it turns out the convergence to global-
è Here linear classifier refers classifier whose decision boundary is linear.

---

Unlike the symmetric case, we balanced Gaussian initialization instead of purely random Gaussian
initialization: ur ~ (0, o J), ag = 8% [tx ll, where 8/: ~ unif(+1). ). We call this distribution as 0g ~ Dinie ( This adaptation can greatly simplify our analysis since it ensures that ax (t) = $y [t)(()))z (t)llz for all 2 0 (Corollary B. 1.18). Similar as the symmetric case, an alternative way to generate this distribution is first draw O.- ~ iniic and then set 0,- = In
Definition 6.1 ((H, K c, w Hi--teeeed Dataset). Given a linearly separable dataset S with max margin linear separator w", for constants H, K, e > 0 and unit vector w1 E gd-1 perpendicular to w", we define the (H, K, e, )-hinted dataset S' by the dataset containing all the data points in S and the following 3 data points (numbered by 1, 2,3) 3) that can serve as hints to the max-margin linear separator w":
(21, 91) (Hw", .1),
(x22: y2) = (cw" + 1.),
(x3. Us) (cw" Kw 1).
Theorem 6.2. Given linearly separable dataset S and unit vector w E g perpendicular to the max-margin linear separator w², for any sufficiently large H > 0, K > and sufficiently small >> 0, the following statement holds for the (H, K, 6, W-)-Hinted Dataset S'. Under regularity assumption for gradient flow (see Assumption A.6), consider gradient flow on a Leaky ReLU network with width m 2 and initialization 0n init where Og ~ (1)nt()). With probability 1 2- over the draw of Oo. if the initialization scale is sufficiently small, then gradient flow directionally converges and (x) := -00 J0(0//1000))))) (x)) represents the one- Leaky- ReLU classifier 2 ( ('0""))))))"
x)) with linear decision boundary. That is,
Pr
max > 0 s.I. Vainit <oin , Va E R { = _0( m²)) 2 1 6.
max
Moreover, the convergent classifier only attains suboptimal margin.
Theorem 6.2 is actually simple corollary general theorem under data assumptions that hold for broader class of linearly separable data. From a high-level perspective, we only require two assumptions: (1). There is a direction such that data points have large inner products with this direction on average; (2). The support vectors for the max-margin linear separator w" have nearly the same labels. The first hint data point is for the first condition and the second and third data point is for the second condition. We defer formal statements of the assumptions and theorems to Appendix A.
7 Conclusions and Future Works
We study the implicit bias of gradient flow in training two-layer Leaky ReLU networks on linearly separable datasets. When the dataset is symmetric, we show any saoballmax-margaee classifier is exactly linear and gradient flow converges to a global-max-margee direction. On the pessimistic side, we show such margin maximization result is fragile - for any linearly separable dataset, we can lead gradient flow to converge to linear classifier with suboptimal margin by adding only 3 extra data points. A critical assumption for convergence analysis is the linear separability of data. We left it as a future work to study simplicity bias and global margin maximization without assuming lincar separability.
Acknowledgments and Disclosure of Funding
The authors acknowledge support from NSF, ONR, Simons Foundation, DARPA and SRC. ZL is also supported by Microsoft Research PhD Fellowship.
References
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers, 2018.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 242-252, Long Beach, Californi, USA, 09-15 Jun 2019. PMLR
10

---

Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 7411-7422. Curran Associates, Inc., 2019a.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Reseurch, pages 322-332.
PMLR, 09-15 Jun 2019b.
Peter L Bartlett, Dylan J Foster, and Matus J Telgarsky. Spectrally-normaaizee normalized margin bounds for neural networks. In I. Guyon, U. V. Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30, pages 6240 2240. Curran Associates, Inc., 2017.
Jérôme Bolte, Aris Daniliidis, Olivier Ley, and Laurent Mazet. Characterizations of tojasiewicz inequalities: subgradient flows, talweg, convexity. Transactions of the American Mathematical Society, 362(6):3319-3366 2010.
Alon Brutzkus and Amir Globerson. Why do larger models generalize better? A theoretical perspective via the XOR problem. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 822-830. PMLR, 09-15 Jun 2019.
Alon Brutzkus, Amir Globerson, Eran Malach, and Shai Shalev-Shwartz. SGD learns over- parameterized networks that provably generalize on linearly separable data. In International Conference on Learning Representations, 2018.
Zixiang Chen, Yuan Cao, Difan Zou, and Quanquan Gu. How much over-parameterizztionn is sufficient learn deep ReLU networks? In International Conference on Learning Representations, 2021.
Lénaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 1338-1338. PMLR, 09-12 Jul 2020.
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 2937-2947. Curran Associates, Inc., 2019.
Francis H. Clarke, Yuri S. Ledyaev, Ronald J. Stern, and Peter R. Wolenski. Nonsmooth analysis and control theory, volume 178. Springer Science & Business Media, 2008
Frank H. Clarke. Generalized gradients and applications. Transactions of the American Mathematical Society, 205:247-262, 1975.
Michel Coste. An introduction to o-minimal geometry. Istituti editoriali poligrafici internazionali Pisa, 2000.
Damek Davis, Dmitriy Drusvyatskiy, Sham Kakade, and Jason D. Lee. Stochastic subgradient method converges on tame functions. Foundations of Computational Mathematics, 20(1):119-155) Feb 2020.
Simon Du Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai Gradient descent finds global minima of deep neural networks. Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 1675-1685, Long Beach, California, USA, 09-15 Jun 201 9a. PMLR.
11

---

Simon S. Du, Wei Hu, and Jason D. Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and Garnett, editors, Advances in Neural Information Processing Systems 31, pages 382-393, Curran Associates, Inc., 2018.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes wverpparameterized neural networks. In International Conference Learning Representations, 2019b.
Joydeep Dutta, Kalyanmoy Deb, Rupesh Tulshyan, and Ramnik Arora. Approximate KKT points and a proximity measure for termination. Journal of Global Optimization, 56(4): 2013.
Spencer Frei, Yuan Cao, and Quanquan Gu. Provable generalization of sgd-trained neural networks any width in the presence of adversarial label noise. Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 3427-3438. PMLR, 18-24 Jul 2021.
Gauthier Gidel, Francis Bach, and Simon Lacoste-Julies Implicit regularization of discrete gradient dynamics in linear neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché-Bue E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 3196-3206. Curran Associates, Inc., 2019.
Daniel Gissin, Shai Shalev-Shwartz, and Amit Daniely. The implicit bias of depth: How incremental learning drives generalization. International Conference on Learning Representations, 2020.
Suriya Gunasckar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in
terms of optimization geometry. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1832-1888 1841, Stockholmsmässann Stockholm Sweden, 10-15 Jul 201 8a. PMLR.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett, editors, Advances Neural Information Processing Systems 31, pages 9482-9491. Curran Associates, Inc., 2018b.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification. In The IEEE International Conference on Computer Vision (ICCV), December 2015.
Wei Hu, Lechao Xiao, Ben Adlam, and Jeffrey Pennington. The surprising simplicity of the early- time learning dynamics of neural networks. H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 17116-17128. Curran Associates, Inc., 2020.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett, editors, Advances Neural Information Processing Systems 31, pages 8571-8580. Curran Associates, Inc., 2018.
Ziwei Ji and Matus Telgarsky. Risk and parameter convergence of logistic regression. arXiv preprint /tt3:/880.073500 2018.
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations, 2019a.
Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Alina Beygelzimer and Daniel Hsu, editors, Proceedings of the Thirty-Secom Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pages 1772-1798, Phoenix, USA, 25-28 Jun 2019b. PMLR
12

---

Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neu- ral Information Processing Systems, volume 33, pages 17176-17186. Curran Associates, Inc., 2020a.
Ziwei and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. In International Conference on Learning Representations, 2020b.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In International Conference on Learning Representations, 2020.
Dimitris Kalimeris, Gal Kaplun, Preetum Nakkiran, Benjamin Edelman, Tristan Yang, Boaz Barak, and Haofeng Zhang. SGD on neural networks learns functions of increasing complexity. In H. Wallach. H. Larochelle, A. Beygelzimer. F. d'Alché- Buché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over parameterrzeed matrix sensing and neural networks with quadratic activations. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet, editors, Proceedings of the 31st Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pages 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.
Tao Luo, Zhi-Que thi-Qin John Xu, Zheng Ma, and Yaoyu Zhang. Phase diagram for two-layer relu neural networks at width limit. Journal of Machine Learning Research, (2(71)):-47, 2021.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2020.
Andrew L. Maas, Awni Y. Hannun, and Andrew Y. Ng. Rectifier nonlinearities improve neural network acoustic models. In ICML Workshop on Deep Learning for Audio, Speech and Language Processing. 2013.
Hartmut Maennel, Olivier Bousquet, and Sylvain Gelly, Gradient descent quantizes relu network features. arXiv preprint arXiv: 2018
Harsh Mehta, Ashok Cutkosky, and Behnam Neyshabur. Extreme memorization via scale of initial- ization. In International Conference on Learning Representations, 2021.
Edward Moroshko, Blake E Woodworth, Suriya Gunasckar, Jason D Lee, Nati Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale vs training accuracy. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 22182-22199 Curran Associates, Inc., 2020.
Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent on separable data. In Kamalika Chaudhuri and Masashi Sugiyama, editors, Proceedings of Machine Learning Research, volume 89 of Proceedings of Machine Learning Research, pages 1420-3423 PMLR, 16-18 Apr 2019a.
Mor Shpigel Nacson, Nathan Srebro, and Daniel Soudry. Stochastic gradient descent on separable data: Exact convergence with fixed learning rate. In Kamalika Chaudhuri and Masashi Sugiyama, editors, Proceedings of Machine Learning Research, volume 89 of Proceedings of Machine Learning Research, pages 3051-3059. PMLR, 16- 18 Apr 2019b.
Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A PAC-bayesian approach to spectrally-normalized margin bounds for neural networks. In International Conference on Learning Representations, 2018.
13

---

Franco Pellegrini and Giulio Biroli. An analytic theory of shallow networks dynamics for hinge loss classification. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 5356-5367, Curran Associates, Inc., 2020.
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 Nadav Cohen. Implicit regularization in tensor factorization. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8913-8924. PMLR, 18-24 Jul 2021.
Roei Sarussi, Alon Brutzkus, and Amir Globerson. Towards understanding learning in neural networks with linear teachers. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9313-9322. PMLR, 18-24 Jul 2021.
Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balean, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 9573-9585. Curran Associates, Inc., 2020.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70):1-57, 20.
Daniel Soudry, Elad Hoffer, and Nathan Srebro. The implicit bias of gradient descent on separable data. In International Conference Learning Representations, 2018b.
Francis Williams, Matthew Trager, Daniele Panozzo, Claudio Silva, Denis Zorin, and Joan Bruna. Gradient dynamics of shallow univariate relu networks. In H. Wallach, H. Larochelle, A. Beygelz- imer, F. d'Alché-Buc, Buch E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 3635-3673 PMLR, 09-12 Jul 2020.
Zhi-Qin John Xu, Hanxu Zhou, Tao Luo, and Yaoyu Zhang. Towards understanding the condensation of two-layer neural networks at initial training. arXiv preprint 1rXiv:2105.11688 2021.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes over-parameterrzzed deep relu networks. arXiv preprint 1811::8811.08888 2018.
14


Reviews:
[{"review_id": "di_quutxNc6", "reviewer": "Reviewer_zJ7E", "summary": "The paper studies the implicit bias for finite two layers Leaky ReLU networks. For linearly separable symmetric data, the authors show that the global max-margin direction corresponds to a linear classifier. In addition, under additional regularity conditions and for a sufficiently small initialization scale, the authors show gradient flow achieves the global max-margin. For non-symmetric yet linearly separable data, the authors show that gradient flow can be led to converge to a linear classifier with suboptimal margin by adding 3 data points.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "The paper studies the implicit bias for finite two layers Leaky ReLU networks and logistic loss.\n\u001e\nThe paper main contributions are:\n- For linearly separable *symmetric* datasets, the authors show that the global max-margin separator is linear.\n- Moreover, the authors show that in this setting and under additional regularity conditions and sufficiently small initialization scale, gradient flow converges to a global max-margin *linear* classifier. \nIn particular, the authors note that previous works only showed convergence to a KKT-margin direction which in general may be non-linear.\n- These results connect current implicit bias max-margin results with simplicity bias works.\n- However, for non-symmetric yet linearly separable data, the authors show that gradient flow can be led to converge to a linear classifier with suboptimal margin by adding 3 data points.\n\u001e\nMy main questions\\concerns:\n- How small should the initialization scale be for the result stated in theorem 4.3 to hold?\nThe authors mention that small initialization is needed to avoid the NTK regime and yet, to get a better sense of the results, it’s important to understand if they only hold for infinitesimal initialization or does they also hold for intermediate initialization scales (where the training accuracy might also have some effect, e.g. [1]).\nSome experimental results on a synthetic dataset that demonstrate the paper’s main findings might help clarify this point.\n- Symmetric dataset: The authors justified the symmetric dataset assumption by mentioning that this assumption can be easily satisfied with data augmentation. They also give a negative result for what happens when the dataset is not symmetric. However, since the mentioned data augmentation is not common practice, I worry that the symmetry assumption is too restrictive and weakens the paper’s positive results.\n\u001e\nClarity/writing: \nIn general, the paper is well written, and I did not see any critical issues. Two suggestions that I think can improve the clarity:\n- The o-minimal structure should be briefly explained\\given some intuition.\n- Some explanation should be given for why assumption A.5 is necessary.\n\u001e\u001e\nMinor comments:\n- Line 31: nature-->natural\n- Equation between lines 159 and 160: ‘.’-->’,’\n- Line 344: corollary general theorem--> corollary of the general theorem\n\u001e\nCitations:\n\u001e\n[1] E. Moroshko, S. Gunasekar, B. Woodworth, J. D. Lee, N. Srebro, D. Soudry, \"Implicit Bias in Deep Linear Classification: Initialization Scale vs Training Accuracy\", NeurIPS 2020.\n\u001e\n---\n\u001e\nUpdate: Thank you for your response and for explaining the assumption.\n\u001e", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors mentioned and justified most of the limitations of their results. However, some of the limitations should be more thoroughly discussed, as mentioned in the main review.", "review_text": "", "extra_scores": {"ethical_concerns": "None.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "12"}, "extra_text": {}}, {"review_id": "c4JnmpMQGHv", "reviewer": "Reviewer_Eg4Q", "summary": "This paper considers margin maximization and simplicity bias properties of gradient flow on a two-layer leaky-ReLU network with linearly separable data. If the data is also symmetric (i.e., whenever (x,y) is in the dataset, (-x,-y) is also in the dataset), this paper proves that the max-margin network basically represents the max-margin linear classifier. Moreover, while prior results ensure convergence to KKT points which may not give the maximum margin, this paper shows that under a few technical conditions gradient flow can indeed maximize the margin. On the other hand, if the data is not symmetric, this paper provides an example where gradient flow cannot globally maximize the margin. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "I think this paper provides a few interesting results and techniques:\n\n(1) Theorem 4.2 uses the symmetry condition to ensure the max-margin network is basically the max-margin linear classifier; I think this is interesting, and the proof of Theorem 4.2 is also nice. \n\n(2) Section 5 gives a multi-phase analysis: in the first phase, even though random initialization is used, as long as the initialization is small enough, this paper shows gradient flow will approximately move the weight vectors on the first layer to two opposite directions. Then in the second phase, the gradient flow dynamic is approximated by a dynamic of two neurons, and finally in the third phase, this paper proves margin maximization. I think this multi-phase analysis is novel and interesting.\n\n(3) The negative result given by Theorem 6.2 is also interesting, confirming that margin maximization may not hold in general.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "As mentioned in Section 7, the linear separability assumption is restrictive. However, in my opinion the symmetry condition is more restrictive: under this condition, if the data is not linearly separable, then it cannot be separated by a network either (this can be proved using the proof of Theorem 4.2). Moreover, one common sufficient condition for a positive definite NTK is that no feature vectors are parallel, which do not hold under the symmetry condition. The symmetry condition is also used in phase I of the gradient flow analysis. I think these points should be discussed more clearly in the body.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "R1wZTjPjyYl", "reviewer": "Reviewer_6TUB", "summary": "This paper proves that gradient descent may converge to the “max-margin” linear solution that attains the global optimum of the loss on finite two-layer Leaky ReLU nets trained on linearly separable and symmetric data. \n\nThe contribution of this paper is mainly reflected in the strict assumption of data distribution (linearly separable and symmetric data). It is proved that GD has good convergence property in relatively simple nonconvex optimization problems (finite two-layer Leaky ReLU nets), and tends to get relatively “simple” solutions (the model complexity of “max-margin” linear solution is low, and it is resistant to overfitting risk).\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Originality: This paper studies the properties of the convergence solution of GD under the specific data distribution and simplified network model, which is not a new problem. The originality of this paper is mainly reflected in the new assumption of data distribution (linearly separable and symmetric data), which leads to stronger conclusions than previous work (from KKT margin direction to global Max margin direction).\n\nQuality & Clarity: This paper is theoretical research, and the process of proof is rigorous. In writing, the main results and the secondary lemmas are properly arranged and expressed clearly. Although there are some types in the appendix, it does not affect the understanding.\n\nIn the appendix, line 551 and line 560  have the missing references.\n\nSignificance: It is difficult to extend the assumptions to the deep neural network structure. Its significance only lies in helping us understand some simple structures in DNNs. Considering that the assumptions of the model and data distribution in this paper are not difficult to be realized through experiments, I hope that the authors can verify their theoretical results through experimental design by synthesizing the dataset.\n\nIf the author can provide some experimental results to verify the theoretical conclusion, I would be willing to raise the score to 7-8.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "In this paper, some non-symmetric results are given in the appendix.\n\nTheoretical research has no social impact.", "review_text": "", "extra_scores": {"ethical_concerns": "NAN", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours"}, "extra_text": {}}, {"review_id": "243NZZXM7uN", "reviewer": "Reviewer_ZQSg", "summary": "This paper studies the gradient descent for two-layer leaky-ReLU neural networks on binary classification problems with the logistic loss and shows the simplicity bias where the neural network converges to the linear function that maximizes the margin under appropriate assumptions. ", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "[Contributions]\n\nThe line of research that characterizes the implicit bias of the optimization method is important to explain the generalization capability without explicit regularization, especially for the nonconvex models such as neural networks. In this sense, the paper makes certain contributions summarized below:\n\n- Under linear separable and symmetric assumptions, the study shows that the two-layer leaky-ReLU network with the parameter maximizing the normalized margin takes the linear function.\n\n- The study shows that the gradient flow of the two-layer leaky ReLU converges to the max-margin linear classifier in direction.\n\nAdvantages over related studies are that (i) the study shows the convergence to the max-margin linear function rather than KKT-margin direction which yields a somewhat complicated function, and (ii) the study does not require an assumption called Neural Agreement Regime (NAR) which is used to conclude the similar result in [Sarussi et al. (2021)].\n\n[Weaknesses]\n\nI think this study is interesting, but I'm not sure when Assumption 4.6 is satisfied. Since one advantage is the validity of the required assumption, it is better to elaborate on the description of this assumption. Could you explain the situation where this assumption is satisfied? \nI would like to increase the score if the authors well address this concern.\n\n[Questions]\n\n- There is no explanation about the dependency on the constant $c_{ainit}$ which is introduced for the initialization of $a$. I would like to know how this constant will affect the results. The following paper shows that the training dynamics fall into the kernel regime when the initialization scale of $a$ is small depending on the network width, although the problem setting and the way of initialization is different from this submission.\n\nE Weinan, Chao Ma, and Lei Wu. A comparative analysis of optimization and generalization properties of two-layer neural network and random feature models under gradient descent dynamics, 2020.\n\n- I would like to know how important small-scale initialization is. I acknowledge its importance in the context of feature learning. I want to know the importance in the context of the margin. Why does it fail to converge to a linear function when the initialization scale is large?\n\n[Minor comments]\n\nThere are some typos in Appendix. For instance, see lines 551, 556, and 561. Moreover in Lemma B.5, $\\ell$ should be replaced with $\\ell'$\n\n-----\nAfter reading the rebuttal:\nThe authors have addressed my concern about Assumption 4.6. They well explained the motivation of this assumption and provided an example. Although a given example is rather limited, it is useful in finding a new research direction. I would like to increase the score to 6. I hope some comments on this assumption will be added in the revision.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations and potential negative social impact", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": ["I don’t know"], "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper refines the implicit bias story for two-layer networks, analyzing settings where it is guaranteed to prefer certain simple linear solutions.  While the reviewers had some concerns, overall reviews and discussion were positive.  I urge the authors to carefully address reviewer comments during their final revisions.
Decision: Accept (Poster)