Gradient Descent on Two-layer Nets: Margin Maximization and Simplicity Bias
Year: 2021

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, 2020a) 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.

heading: 1 Introduction
text: One major mystery in deep learning is why deep neural networks generalize despite overparameterization (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. (2018a,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 maxmargin (hard-margin SVM) solution, even when there is no explicit regularization. Thus the solution found by GD has the same margin-based generalization bounds as 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, 2019b), deep linear nets (Ji and Telgarsky, 2019a; Gunasekar et al., 2018b). ∗Equal contribution †Most of the work is done when Kaifeng Lyu and Runzhe Wang were at Tsinghua University.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Given the above results, a natural question to 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.e., 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 a direction satisfying the Karush-Kuhn-Tucker (KKT) conditions of a constrained optimization problem (P) on 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-width neural nets, but the width needs to be exponential in the input dimension (due to the use of a covering condition). Both works build upon late phase analyses. Under a restrictive assumption that the data is orthogonally separable, i.e., any data point xi 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 et 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 problem and gain insights by assuming that the data distribution has a 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 et 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.

heading: 1.1 Our Contribution
text: Among all the classifiers that can be represented by the two-layer Leaky ReLU nets, we show any global-max-margin classifier is exactly linear under one more data assumption: the dataset is symmetric, i.e., if x is 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 1 (left) for an example.
In this paper, we provide a 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 a global-max-margin 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 Kurdyka-Łojasiewicz (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 early phase of training, and the complexity of the solution increases
as training goes on (Kalimeris et al., 2019; Hu et al., 2020; Shah et al., 2020). Indeed, our result establishes a form of extreme simplicity bias of GD: if the dataset can be fitted by a linear classifier, 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, global-max-margin 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 a linear classifier with suboptimal margin by adding only 3 extra data points (Theorem 6.2). See Figure 1 (right) for an example.

heading: 2 Related Works
text: Generalization Aspect of Margin Maxmization. Margin often appears in the generalization bounds for neural networks (Bartlett et al., 2017; Neyshabur et 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 (nonlinear) neural networks on linearly separable data (labels are generated by a 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 a 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 et al. (2021) showed that GD on shallow/deep neural nets learns a kernel predictor with good generalization on linearly separable data, and it suffices to 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 ReLU 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 a 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 it 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 et al. (2020) justified this view by proving that the learning 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 as 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 initialization 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
is hence argued that gradient flow has a 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 et 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 to large scale.

heading: 3 Preliminaries
text: We denote the set {1, . . . , n} by [n] and the unit sphere {x ∈ Rd : ‖x‖2 = 1} by Sd−1. We call a function h : RD → R L-homogeneous if h(cθ) = cLh(θ) for all θ ∈ RD and c > 0. For S ⊆ RD, conv(S) denotes the convex hull of S. For locally Lipschitz function f : RD → R, we define Clarke’s subdifferential (Clarke, 1975; Clarke et al., 2008; Davis et al., 2020) to be ∂◦f(θ) := conv {limn→∞∇f(θn) : f differentiable at θn, limn→∞ θn = θ} (see also Appendix B.1).

heading: 3.1 Logistic Loss Minimization and Margin Maximization
text: For a neural net, we use fθ(x) ∈ R to denote the output logit on input x ∈ Rd when the parameter is θ ∈ RD. We say that the neural net is L-homogeneous if fθ(x) is L-homogeneous with respect to θ, i.e., fcθ(x) = cLfθ(x) for all θ ∈ RD and c > 0. VGG-like CNNs can be made homogeneous if we remove all the bias terms expect those in the first layer (Lyu and Li, 2020).
Throughout this paper, we restrict our attention to L-homogeneous neural nets with fθ(x) definable with respect to θ in an o-minimal structure for all x. (See Coste 2000 for reference for o-minimal structures.) This is a technical condition needed by Theorem 3.1, and it is a mild regularity condition as almost all modern neural networks satisfy this condition, including the two-layer Leaky ReLU networks studied in this paper.
For a dataset S = {(x1, y1), . . . , (xn, yn)}, we define qi(θ) := yifθ(xi) to be the output margin on the data point (xi, yi), and qmin(θ) := mini∈[n] qi(θ) to be the output margin on the dataset S (or
margin for short). It is easy to see that q1(θ), . . . , qn(θ) are L-homogeneous functions, and so is qmin(θ). We define the normalized margin γ(θ) := qmin ( θ ‖θ‖2 ) = qmin(θ)‖θ‖L2 to be the output margin (on the dataset) for the normalized parameter θ‖θ‖2 .
We refer the problem of finding θ that maximizes γ(θ) as margin maximization. Note that once we have found an optimal solution θ∗ ∈ RD, cθ∗ is also optimal for all c > 0. We can put the norm constraint on θ to eliminate this freedom on rescaling:
max θ∈SD−1 γ(θ). (M)
Alternatively, we can also constrain the margin to have qmin ≥ 1 and minimize the norm:
min 1
2 ‖θ‖22 s.t. qi(θ) ≥ 1, ∀i ∈ [n]. (P)
One can easily show that θ∗ is a global maximizer of (M) if and only if θ ∗
(qmin(θ∗))1/L is a global
minimizer of (P). For convenience, we make the following convention: if θ‖θ‖2 is a local/global maximizer of (M), then we say θ is along a local-max-margin direction/global-max-margin direction; if θ
(qmin(θ))1/L satisfies the KKT conditions of (P), then we say θ is along a KKT-margin direction.
Gradient flow with logistic loss is defined by the following differential inclusion,
dθ dt ∈ −∂◦L(θ), with L(θ) := 1 n n∑ i=1 `(qi(θ)), (1)
where `(q) := ln(1 + e−q) is the logistic loss. Lyu and Li (2020); Ji and Telgarsky (2020a) showed that θ(t)/‖θ(t)‖2 always converges to a KKT-margin direction. We restate the results below. Theorem 3.1 (Lyu and Li 2020; Ji and Telgarsky 2020a). For homogeneous neural networks, if L(θ(0)) < ln 2n , then L(θ(t)) → 0, ‖θ(t)‖2 → +∞, and θ(t) ‖θ(t)‖2 converges to a KKT-margin direction as t→ +∞.

heading: 3.2 Two-Layer Leaky ReLU Networks on Linearly Separable Data
text: Let φ(x) = max{x, αleakyx} be Leaky ReLU, where αleaky ∈ (0, 1). Throughout the following sections, we consider a two-layer neural net defined as below,
fθ(x) = m∑ k=1 akφ(w > k x).
where w1, . . . ,wm ∈ Rd are the weights in the first layer, a1, . . . , am ∈ R are the weights in the second layer, and θ = (w1, . . . ,wm, a1, . . . , am) ∈ RD is the concatenation of all trainable parameters, where D = md+m. We can verify that fθ(x) is 2-homogeneous with respect to θ.
Let S := {(x1, y1), . . . , (xn, yn)} be the training set. For simplicity, we assume that ‖xi‖2 ≤ 1. We focus on linearly separable data, thus we assume that S is linearly separable throughout the paper. Assumption 3.2 (Linear Separable). There exists a w ∈ Rd such that yi 〈w,xi〉 ≥ 1 for all i ∈ [n]. Definition 3.3 (Max-margin Linear Separator). For the linearly separable dataset S, we say that w∗ ∈ Sd−1 is the max-margin linear separator if w∗ maximizes mini∈[n] yi 〈w,xi〉 over w ∈ Sd−1.

heading: 4 Training on Linearly Separable and Symmetric Data
text: In this section, we study the implicit bias of gradient flow assuming the training data is linearly separable and symmetric. We say a dataset is symmetric if whenever x is present in the training set, the input −x is also present. By linear separability, x and −x must have different labels because 〈w∗,x〉 = −〈w∗,−x〉, where w∗ is the max-margin linear separator. The formal statement for this assumption is given below. Assumption 4.1 (Symmetric). n is even and xi = −xi+n/2, yi = 1, yi+n/2 = −1 for 1 ≤ i ≤ n/2.
This symmetry can be ensured via data augmentation. Given a dataset, if it is known that the groundtruth labels are produced by an unknown linear classifier, then one can augment each data point (x, y) by flipping the sign, i.e., replace it with two data points (x, y), (−x,−y) (and thus the dataset size is doubled).
Our results show that gradient flow directionally converges to a global-max-margin 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-max-margin direction represents a linear classifier, which we will see in Section 4.1. Then we will present our main convergence results in Section 4.2.

heading: 4.1 Global-Max-Margin Classifiers are Linear
text: Theorem 4.2 below characterizes the global-max-margin direction in our case by showing that margin maximization and simplicity bias coincide with each other: a network that representing the maxmargin linear classifier (i.e., fθ(x) = c 〈w∗,x〉 for some c > 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, any global-max-margin direction θ∗ ∈ SD−1, fθ∗ represents a linear classifier. Moreover, we have fθ∗(x) = 1+αleaky 4 〈w ∗,x〉 for all x ∈ Rd, where w∗ is the max-margin linear separator.
The result of Theorem 4.2 is based on the observation that replacing each neuron (ak,wk) in a network with two neurons of oppositing parameters (ak,wk) and (−ak,−wk) 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-margin direction which corresponds to a linear classifier. We can show that every weight vector wk of this linear classifier must be in the direction of w∗ or −w∗. Then the original classifier must also be linear in the same direction.

heading: 4.2 Convergence to Global-Max-Margin Directions
text: Though Theorem 3.1 guarantees that gradient flow directionally converges to a KKT-margin direction if the loss is optimized successfully, we note that KKT-margin directions can be non-linear and have complicated decision boundaries. See Figure 1 (left) for an example. Therefore, to establish the convergence to linear classifiers, Theorem 3.1 is not enough and we need a new analysis for the trajectory of gradient flow.
We use initialization wk i.i.d.∼ N (0, σ2initI), ak i.i.d.∼ N (0, c2ainitσ2init), where cainit is a fixed constant throughout this paper and σinit controls the initialization scale. We call this distribution as θ0 ∼ Dinit(σinit). An alternative way to generate this distribution is to first draw θ̄0 ∼ Dinit(1), and then set θ0 = σinitθ̄0. With small initialization, we can establish the following convergence result. 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 a Leaky ReLU network with width m ≥ 2 and initialization θ0 = σinitθ̄0 where θ̄0 ∼ Dinit(1). With probability 1 − 2−(m−1) over the random draw of θ̄0, if the initialization scale is sufficiently small, then gradient flow directionally converges and f∞(x) := limt→+∞ fθ(t)/‖θ(t)‖2(x) represents the max-margin linear classifier. That is,
Pr θ̄0∼Dinit(1)
[ ∃σmaxinit > 0 s.t. ∀σinit < σmaxinit ,∀x ∈ Rd, f∞(x) = C 〈w∗,x〉 ] ≥ 1− 2−(m−1),
where C := 1+αleaky4 is a 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. In the settings of Theorem 4.3, gradient flow on linearly separable and symmetric data directionally converges to the global-max-margin direction with probability 1− 2−(m−1).

heading: 4.3 Additional Notations and Assumptions
text: Let µ := 1n ∑n i=1 yixi, which is non-zero since 〈µ,w∗〉 = 1 n ∑ i∈[n] yiw > ∗ xi ≥ 1. Let µ̄ := µ ‖µ‖2 . We use ϕ(θ0, t) ∈ Rd to the value of θ at time t for θ(0) = θ0.
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], 〈µ,xi〉 6= 0.
Another technical issue we face is that the gradient flow may not be unique due to non-smoothness. It is possible that ϕ(θ0, t) is not well-defined as the solution of (1) may not be unique. See Appendix I.2 for more discussions. In this case, we assign ϕ(θ0, · ) to be an arbitrary gradient flow trajectory starting from θ0. In the case where ϕ(θ0, t) has only one possible value for all t ≥ 0, we say that θ0 is a non-branching starting point. We assume the following technical assumption. Assumption 4.6. For any m ≥ 2, there exist r, > 0 such that θ is a non-branching starting point if its neurons can be partitioned into two groups: in the first group, ak = ‖wk‖2 ∈ (0, r) and all wk point to the same direction w+ ∈ Sd−1 with ‖w+ − µ̄‖2 ≤ ; in the second group, −ak = ‖wk‖2 ∈ (0, r) and all wk point to the same direction w− ∈ Sd−1 with ‖w− + µ̄‖2 ≤ .

heading: 5 Proof Sketch for the Symmetric Case
text: In this section, we provide a proof sketch for Theorem 4.3. Our proof uses a multi-phase analysis, which divides the training process into 3 phases, from small initialization to the final convergence. We will now elaborate the analyses for them one by one.

heading: 5.1 Phase I: Dynamics Near Zero
text: Gradient flow starts with small initialization. In Phase I, we 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 wk in the first layer moves towards the directions of either µ̄ or −µ̄. To see this, the first step is to note that fθ(xi) ≈ 0 when θ is close to 0. Applying Taylor expansion on `(yifθ(xi)),
L(θ) = 1 n ∑ i∈[n] `(yifθ(xi)) ≈ 1 n ∑ i∈[n] (`(0) + `′(0)yifθ(xi)) . (2)
Expanding fθ(xi) and reorganizing the terms, we have
L(θ) ≈ 1 n ∑ i∈[n] `(0) + 1 n ∑ i∈[n] `′(0) ∑ k∈[m] yiakφ(w > k xi) = `(0) + `′(0) n ∑ k∈[m] ∑ i∈[n] yiakφ(w > k xi)
= `(0)− ∑ k∈[m] akG(wk),
where G-function (Maennel et al., 2018) is defined below:
G(w) := −`′(0) n ∑ i∈[n] yiφ(w >xi) = 1 2n ∑ i∈[n] yiφ(w >xi).
This means gradient flow optimizes each −akG(wk) separately near origin. dwk dt ≈ ak∂◦G(wk), dak dt ≈ G(wk). (3)
In the case where Assumption 4.1 holds, we can pair each xi with −xi and use the identity φ(z)− φ(−z) = max{z, αleakyz} −max{−z,−αleakyz} = (1 + αleaky)z to show that G(w) is linear:
G(w) = 1
2n ∑ i∈[n/2] ( φ(w>xi)− φ(−w>xi) ) = 1 2n ∑ i∈[n/2] (1 + αleaky)w >xi = 〈w, µ̃〉 ,
where µ̃ := 1+αleaky2 µ = 1+αleaky 2n ∑ i∈[n] yixi. Substituting this formula for G into (3) reveals that the dynamics of two-layer neural nets near zero has a close relationship to power iteration (or matrix exponentiation) of a matrixMµ̃ ∈ R(d+1)×(d+1) that only depends on data.
d dt [ wk ak ] ≈Mµ̃ [ wk ak ] , where Mµ̃ := [ 0 µ̃ µ̃> 0 ] .
Simple linear algebra shows that λ0 := ‖µ̃‖2, 1√2 (µ̄, 1) ∈ R d+1 are the unique top eigenvalue and eigenvector of Mµ̃, which suggests that (wk(t), ak(t)) ∈ Rd+1 aligns to this top eigenvector direction if the approximation (3) holds for a sufficiently long time. With small initialization, this can indeed be true and we obtain the following lemma. Definition 5.1 (M-norm). For parameter vector θ = (w1, . . . ,wm, a1, . . . , am), we define the M-norm to be ‖θ‖M = maxk∈[m] {max{‖wk‖2, |ak|}}. Lemma 5.2. Let r > 0 be a small value. With probability 1 over the random draw of θ̄0 = (w̄1, . . . , w̄m, ā1, . . . , ām) ∼ Dinit(1), if we take σinit ≤ r 3
√ m‖θ̄0‖
M
, then any neuron (wk, ak) at
time T1(r) := 1λ0 ln r√
mσinit‖θ̄0‖ M
can be decomposed into
wk(T1(r)) = rb̄kµ̄+ ∆wk, ak(T1(r)) = rb̄k + ∆ak,
where b̄k := 〈w̄k,µ̄〉+āk 2 √ m‖θ̄0‖
M
and the error term ∆θ := (∆w1, . . . ,∆wm,∆a1, . . . ,∆am) is bounded
by ‖∆θ‖M ≤ Cr 3
√ m for some universal constant C.

heading: 5.2 Phase II: Near-Two-Neuron Dynamics
text: By Lemma 5.2, we know that at time T1(r) we have wk(T1(r)) ≈ rb̄kµ̄ and ak(T1(r)) ≈ rb̄k, where b̄ ∈ Rd is some fixed vector. This motivates us to couple the training dynamics of θ(t) = (w1(t), . . . ,wm(t), a1(t), . . . , am(t)) after the time T1(r) with another gradient flow starting from the point (rb̄1µ̄, . . . , rb̄mµ̄, rb̄1, . . . , rb̄m). Interestingly, the latter dynamic can be seen as 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 a long time. Now we first introduce our idea of embedding a two-neuron network into an m-neuron network.
Embedding. For any b ∈ Rm, we say that b is a good embedding vector if it has at least one positive entry and one negative entry, and all the entries are non-zero. For a good embedding vector b, we use b+ := √∑ j∈[m] 1[bj>0]b 2 j and b− := − √∑ j∈[m] 1[bj<0]b 2 j to denote the root-sum-squared
of the positive entries and the negative root-sum-squared of the negative entries. For parameter θ̂ := (ŵ1, ŵ2, â1, â2) of a two-neuron neural net with â1 > 0 and â2 < 0, we define the embedding from two-neuron into m-neuron neural nets as πb(ŵ1, ŵ2, â1, â2) = (w1, . . . ,wm, a1, . . . , am), where
ak =
{ bk b+ â1, if bk > 0
bk b− â2, if bk < 0
, wk =
{ bk b+ ŵ1, if bk > 0
bk b− ŵ2, if bk < 0
.
It is easy to check that fθ̂(x) = fπb(θ̂)(x) by the homogeneity of the activation (φ(cz) = cφ(z) for c > 0):
fπb(θ̂)(x) = ∑ bk>0 akφ(w > k x) + ∑ bk<0 akφ(w > k x)
= ∑ bk>0 b2k b2+ â1φ(ŵ > 1 x) + ∑ bk<0 b2k b2− â2φ(ŵ > 2 x) = â1φ(ŵ > 1 x) + â2φ(ŵ > 2 x) = fθ̂(x).
Moreover, by taking the chain rule, we can obtain the following lemma showing that the trajectories starting from θ̂ and πb(θ̂) are essentially the same.
Lemma 5.3. Given θ̂ := (ŵ1, ŵ2, â1, â2) with â1 > 0 and â2 < 0, if both θ̂ and πb(θ̂) are non-branching starting points, then ϕ(πb(θ̂), t) = πb(ϕ(θ̂, t)) for all t ≥ 0.
Approximate Embedding. Back to our analysis for Phase II, b̄ is a good embedding vector with high probability (see lemma below). Let θ̂ := (b̄+, b̄+µ̄, b̄−, b̄−µ̄). By Lemma 5.2, πb̄(rθ̂) = (rb̄1µ̄, . . . , rb̄mµ̄, rb̄1, . . . , rb̄m) ≈ θ(T1(r)), which means rθ̂ → θ(T1(r)) is approximately an embedding. Suppose that the approximation happens to be exact, namely πb̄(rθ̂) = θ(T1(r)), then θ(T1(r) + t) = πb̄(ϕ(rθ̂, t)) by Lemma 5.3. Inspired by this, we consider the case where σinit → 0, r → 0 so that the approximate embedding is infinitely close to the exact one, and prove the following lemma. We shift the training time by T2(r) to avoid trivial limits (such as 0).
Lemma 5.4. Follow the notations in Lemma 5.2 and take σinit ≤ r 3
√ m‖θ̄0‖
M
. Let T2(r) := 1λ0 ln 1 r ,
then T12 := T1(r) + T2(r) = 1λ0 ln 1√
mσinit‖θ̄0‖ M
regardless the choice of r. For width m ≥ 2,
with probability 1 − 2−(m−1) over the random draw of θ̄0 ∼ Dinit(1), the vector b̄ ∈ Rm is a good embedding vector, and for the two-neuron dynamics starting with rescaled initialization in the direction of θ̂ := (b̄+, b̄+µ̄, b̄−, b̄−µ̄), the following limit exists for all t,
θ̃(t) := lim r→0
ϕ ( rθ̂, T2(r) + t ) 6= 0, (4)
and moreover, for the m-neuron dynamics of θ(t), the following holds for all t,
lim σinit→0 θ (T12 + t) = πb̄(θ̃(t)). (5)

heading: 5.3 Phase III: Dynamics near Global-Max-Margin Direction
text: With some efforts, we have the following characterization for the two-neuron dynamics. Theorem 5.5. Form = 2, if initially a1 = ‖w1‖2, a2 = −‖w2‖2, 〈w1,w∗〉 > 0 and 〈w2,w∗〉 < 0, then θ(t) directionally converges to the following global-max-margin direction,
lim t→+∞
θ(t)
‖θ(t)‖2 =
1 4 (w∗,−w∗, 1,−1),
where w∗ is the max-margin linear separator.
It is not hard to verify that θ̃(t) satisfies the conditions required by Theorem 5.5. Given this result, a first attempt to establish the convergence of θ(t) to global-max-margin direction is to take t→ +∞ on both sides of (5). However, this only proves that θ (T12 + t) directionally converges to the globalmax-margin direction if we take the limit σinit → 0 first then take t→ +∞, while we are interested in the convergent solution when t→ +∞ first then σinit → 0 (i.e., solution gradient flow converges to with infinite training time, if it starts from sufficiently small initialization). These two double limits are not equivalent because the order of limits cannot be exchanged without extra conditions.
To overcome this issue, we follow a similar proof strategy as Ji and Telgarsky (2020a) to prove local convergence near a local-max-margin direction, as formally stated below. Theorem 5.6 holds for L-homogeneous neural networks in general and we believe is of independent interest. Theorem 5.6. Consider any L-homogeneous neural networks with logistic loss. Given a local-maxmargin direction θ̄∗ ∈ SD−1 and any δ > 0, there exists 0 > 0 and ρ0 ≥ 1 such that for any θ0 with norm ‖θ0‖2 ≥ ρ0 and direction
∥∥∥ θ0‖θ0‖2 − θ̄∗∥∥∥2 ≤ 0, gradient flow starting with θ0 directionally converges to some direction θ̄ with the same normalized margin γ as θ̄∗, and ‖θ̄ − θ̄∗‖2 ≤ δ.
Using Theorem 5.6, we can finish the proof for Theorem 4.3 as follows. First we note that the twoneuron global-max-margin direction 14 (w
∗,−w∗, 1,−1) after embedding is a global-max-margin direction for m-neurons, and we can prove that any direction with distance no more than a small constant δ > 0 is still a global-max-margin direction. Then we can take t to be large enough so that πb̄(θ̃(t)) satisfies the conditions in Theorem 5.6. According to (5), we can also make the conditions hold for θ (T12 + t) by taking σinit and r to be sufficiently small. Finally, applying Theorem 5.6 finishes the proof.

heading: 6 Non-symmetric Data Complicates the Picture
text: Now we turn to study the case without assuming symmetry and the question is whether the implicit bias to global-max-margin solution still holds. Unfortunately, it turns out the convergence to globalmax-margin classifier is 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.3 See Definition 6.1 for the construction and Figure 1 (right) for an example.
3Here linear classifier refers to a classifier whose decision boundary is linear.
Unlike the symmetric case, we use balanced Gaussian initialization instead of purely random Gaussian initialization: wk ∼ N (0, σ2initI), ak = sk‖wk‖2, where sk ∼ unif{±1}. We call this distribution as θ0 ∼ D̃init(σinit). This adaptation can greatly simplify our analysis since it ensures that ak(t) = sk‖wk(t)‖2 for all t ≥ 0 (Corollary B.18). Similar as the symmetric case, an alternative way to generate this distribution is to first draw θ̄0 ∼ D̃init(1), and then set θ0 = σinitθ̄0. Definition 6.1 ((H ,K ,w⊥)-Hinted Dataset). Given a linearly separable dataset S with max-margin linear separator w∗, for constants H,K, > 0 and unit vector w⊥ ∈ Sd−1 perpendicular to w∗, we define the (H , K, , w⊥)-hinted dataset S ′ by the dataset containing all the data points in S and the following 3 data points (numbered by 1, 2, 3) that can serve as hints to the max-margin linear separator w∗:
(x1, y1) = (Hw ∗, 1), (x2, y2) = ( w ∗ +Kw⊥, 1), (x3, y3) = ( w ∗ −Kw⊥, 1).
Theorem 6.2. Given a linearly separable dataset S and a unit vector w⊥ ∈ Sd−1 perpendicular to the max-margin linear separator w∗, for any sufficiently large H > 0,K > 0 and sufficiently small > 0, the following statement holds for the (H , K, , w⊥)-Hinted Dataset S ′. Under a regularity assumption for gradient flow (see Assumption A.6), consider gradient flow on a Leaky ReLU network with width m ≥ 1 and initialization θ0 = σinitθ̄0 where θ̄0 ∼ D̃init(1). With probability 1− 2−m over the draw of θ̄0, if the initialization scale is sufficiently small, then gradient flow directionally converges and f∞(x) := limt→+∞ fθ(t)/‖θ(t)‖2(x) represents the one-LeakyReLU classifier 12φ(〈w ∗,x〉) with linear decision boundary. That is,
Pr θ̄0∼Dinit(1)
[ ∃σmaxinit > 0 s.t. ∀σinit < σmaxinit ,∀x ∈ Rd, f∞(x) = 1
2 φ(〈w∗,x〉)
] ≥ 1− δ.
Moreover, the convergent classifier only attains a suboptimal margin.
Theorem 6.2 is actually a simple corollary general theorem under data assumptions that hold for a 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 separatorw∗ 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.

heading: 7 Conclusions and Future Works
text: 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 global-max-margin classifier is exactly linear and gradient flow converges to a global-max-margin 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 a linear classifier with suboptimal margin by adding only 3 extra data points. A critical assumption for our 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 linear 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.

Review summaries:
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.

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. 

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.   The 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).