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

Introduction

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 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 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].

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 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 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 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 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 ${\bm{x}}_{i}$ can serve as a perfect linear separator, Phuong and Lampert 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. , Sarussi et al. 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.

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 ${\bm{x}}$ is in the training set, then so is $-{\bm{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 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. . 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.

Related Works

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

Some works studied the training dynamics of (nonlinear) neural networks on linearly separable data (labels are generated by a linear teacher). Brutzkus et al. 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. 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. 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 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. 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 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.

Kalimeris et al. 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. 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. can be seen as another theoretical justification for online SGD on aribitrary data distribution. Shah et al. pointed out that extreme simplicity bias can lead to suboptimal generalization and negative effects on adversarial robustness.

Several theoretical works studying neural network training with small initialization can be connected to simplicity bias. Maennel et al. 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. 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. , Moroshko et al. , Mehta et al. studied feature learning when the initialization scale transitions from small to large scale.

Preliminaries

Throughout this paper, we restrict our attention to $L$-homogeneous neural nets with $f_{{\bm{\theta}}}({\bm{x}})$ definable with respect to ${\bm{\theta}}$ in an o-minimal structure for all ${\bm{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 $\mathcal{S}=\{({\bm{x}}_{1},y_{1}),\dots,({\bm{x}}_{n},y_{n})\}$, we define $q_{i}({\bm{\theta}}):=y_{i}f_{{\bm{\theta}}}({\bm{x}}_{i})$ to be the output margin on the data point $({\bm{x}}_{i},y_{i})$, and $q_{\min}({\bm{\theta}}):=\min_{i\in[n]}q_{i}({\bm{\theta}})$ to be the output margin on the dataset $\mathcal{S}$ (or margin for short). It is easy to see that $q_{1}({\bm{\theta}}),\dots,q_{n}({\bm{\theta}})$ are $L$-homogeneous functions, and so is $q_{\min}({\bm{\theta}})$. We define the normalized margin $\gamma({\bm{\theta}}):=q_{\min}\left(\frac{{\bm{\theta}}}{\|{\bm{\theta}}\|_{2}}\right)=\frac{q_{\min}({\bm{\theta}})}{\|{\bm{\theta}}\|_{2}^{L}}$ to be the output margin (on the dataset) for the normalized parameter $\frac{{\bm{\theta}}}{\|{\bm{\theta}}\|_{2}}$.

Alternatively, we can also constrain the margin to have $q_{\min}\geq 1$ and minimize the norm:

One can easily show that ${\bm{\theta}}^{*}$ is a global maximizer of (M) if and only if $\frac{{\bm{\theta}}^{*}}{(q_{\min}({\bm{\theta}}^{*}))^{1/L}}$ is a global minimizer of (P). For convenience, we make the following convention: if $\frac{{\bm{\theta}}}{\|{\bm{\theta}}\|_{2}}$ is a local/global maximizer of (M), then we say ${\bm{\theta}}$ is along a local-max-margin direction/global-max-margin direction; if $\frac{{\bm{\theta}}}{(q_{\min}({\bm{\theta}}))^{1/L}}$ satisfies the KKT conditions of (P), then we say ${\bm{\theta}}$ is along a KKT-margin direction.

Gradient flow with logistic loss is defined by the following differential inclusion,

For homogeneous neural networks, if $\mathcal{L}({\bm{\theta}}(0))<\frac{\ln 2}{n}$, then $\mathcal{L}({\bm{\theta}}(t))\to 0$, $\|{\bm{\theta}}(t)\|_{2}\to+\infty$, and $\frac{{\bm{\theta}}(t)}{\|{\bm{\theta}}(t)\|_{2}}$ converges to a KKT-margin direction as $t\to+\infty$.

2 Two-Layer Leaky ReLU Networks on Linearly Separable Data

Let $\mathcal{S}:=\{({\bm{x}}_{1},y_{1}),\dots,({\bm{x}}_{n},y_{n})\}$ be the training set. For simplicity, we assume that $\|{\bm{x}}_{i}\|_{2}\leq 1$. We focus on linearly separable data, thus we assume that $\mathcal{S}$ is linearly separable throughout the paper.

Training on Linearly Separable and Symmetric Data

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 ${\bm{x}}$ is present in the training set, the input $-{\bm{x}}$ is also present. By linear separability, ${\bm{x}}$ and $-{\bm{x}}$ must have different labels because $\left<{\bm{w}}^{*},{\bm{x}}\right>=-\left<{\bm{w}}^{*},-{\bm{x}}\right>$, where ${\bm{w}}^{*}$ is the max-margin linear separator. The formal statement for this assumption is given below.

$n$ is even and ${\bm{x}}_{i}=-{\bm{x}}_{i+n/2},y_{i}=1,y_{i+n/2}=-1$ for $1\leq i\leq n/2$.

This symmetry can be ensured via data augmentation. Given a dataset, if it is known that the ground-truth labels are produced by an unknown linear classifier, then one can augment each data point $({\bm{x}},y)$ by flipping the sign, i.e., replace it with two data points $({\bm{x}},y)$, $(-{\bm{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.

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 max-margin linear classifier (i.e., $f_{{\bm{\theta}}}({\bm{x}})=c\left<{\bm{w}}^{*},{\bm{x}}\right>$ for some $c>0$) can simultaneously achieve the goals of being simple and maximizing the margin.

The result of Theorem 4.2 is based on the observation that replacing each neuron $(a_{k},{\bm{w}}_{k})$ in a network with two neurons of oppositing parameters $(a_{k},{\bm{w}}_{k})$ and $(-a_{k},-{\bm{w}}_{k})$ 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 ${\bm{w}}_{k}$ of this linear classifier must be in the direction of ${\bm{w}}^{*}$ or $-{\bm{w}}^{*}$. Then the original classifier must also be linear in the same direction.

2 Convergence to Global-Max-Margin Directions

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.

Combining Theorem 4.2 and Theorem 4.3, we can conclude that gradient flow achieves the global max margin in our case.

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)}$.

3 Additional Notations and Assumptions

We make the following technical assumption, which holds if we are allowed to add a slight perturbation to the training set.

For all $i\in[n]$, $\left<{\bm{\mu}},{\bm{x}}_{i}\right>\neq 0$.

Another technical issue we face is that the gradient flow may not be unique due to non-smoothness. It is possible that $\varphi({\bm{\theta}}_{0},t)$ is not well-defined as the solution of (1) may not be unique. See Section I.2 for more discussions. In this case, we assign $\varphi({\bm{\theta}}_{0},\,\cdot\,)$ to be an arbitrary gradient flow trajectory starting from ${\bm{\theta}}_{0}$. In the case where $\varphi({\bm{\theta}}_{0},t)$ has only one possible value for all $t\geq 0$, we say that ${\bm{\theta}}_{0}$ is a non-branching starting point. We assume the following technical assumption.

Proof Sketch for the Symmetric Case

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.

Expanding $f_{{\bm{\theta}}}({\bm{x}}_{i})$ and reorganizing the terms, we have

where $G$-function [Maennel et al., 2018] is defined below:

This means gradient flow optimizes each $-a_{k}G({\bm{w}}_{k})$ separately near origin.

2 Phase II: Near-Two-Neuron Dynamics

It is easy to check that $f_{\hat{{\bm{\theta}}}}({\bm{x}})=f_{\pi_{{\bm{b}}}(\hat{{\bm{\theta}}})}({\bm{x}})$ by the homogeneity of the activation ($\phi(cz)=c\phi(z)$ for $c>0$):

Moreover, by taking the chain rule, we can obtain the following lemma showing that the trajectories starting from $\hat{{\bm{\theta}}}$ and $\pi_{{\bm{b}}}(\hat{{\bm{\theta}}})$ are essentially the same.

Given $\hat{{\bm{\theta}}}:=(\hat{{\bm{w}}}_{1},\hat{{\bm{w}}}_{2},\hat{a}_{1},\hat{a}_{2})$ with $\hat{a}_{1}>0$ and $\hat{a}_{2}<0$, if both $\hat{{\bm{\theta}}}$ and $\pi_{{\bm{b}}}(\hat{{\bm{\theta}}})$ are non-branching starting points, then $\varphi(\pi_{{\bm{b}}}(\hat{{\bm{\theta}}}),t)=\pi_{{\bm{b}}}(\varphi(\hat{{\bm{\theta}}},t))$ for all $t\geq 0$.

and moreover, for the $m$-neuron dynamics of ${\bm{\theta}}(t)$, the following holds for all $t$,

3 Phase III: Dynamics near Global-Max-Margin Direction

With some efforts, we have the following characterization for the two-neuron dynamics.

For $m=2$, if initially $a_{1}=\|{\bm{w}}_{1}\|_{2}$, $a_{2}=-\|{\bm{w}}_{2}\|_{2}$, $\left<{\bm{w}}_{1},{\bm{w}}^{*}\right>>0$ and $\left<{\bm{w}}_{2},{\bm{w}}^{*}\right><0$, then ${\bm{\theta}}(t)$ directionally converges to the following global-max-margin direction,

where ${\bm{w}}^{*}$ is the max-margin linear separator.

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.

Non-symmetric Data Complicates the Picture

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 global-max-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.Here linear classifier refers to a classifier whose decision boundary is linear. See Definition 6.1 for the construction and Figure 1 (right) for an example.

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 separator ${\bm{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.

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

References

Appendix A Theorem Statements for the Non-symmetric Case

Theorem 6.2 is indeed a simple corollary of Theorem A.7 below which holds for a broader class of datasets. Now we illustrate the assumptions one by one.

There exists a unit-norm vector ${\bm{w}}^{\diamond}$ such that $\gamma^{\diamond}:=\min_{i\in[n]}y_{i}\left<{\bm{w}}^{\diamond},{\bm{x}}_{i}\right>>0$ and

where ${\bm{P}}^{\diamond}:={\bm{I}}-{\bm{w}}^{\diamond}{{\bm{w}}^{\diamond}}^{\top}$ is the projection matrix onto the space perpendicular to ${\bm{w}}^{\diamond}$, and ${\bm{\mu}}:=\frac{1}{n}\sum_{i\in[n]}y_{i}{\bm{x}}_{i}$ is the mean vector of $y_{i}{\bm{x}}_{i}$.

Indeed, our main theorem is based on a weaker assumption than Assumption A.1, which is Assumption A.2 below, but the geometric meaning of Assumption A.2 is not as clear as Assumption A.1. We will show in Lemma G.1 that Assumption A.1 implies Assumption A.2.

In general, the norms $\|{\bm{\mu}}^{+}\|_{2}$ and $\|{\bm{\mu}}^{-}\|_{2}$ should not be equal: for any given dataset $\mathcal{S}$, we can make $\|{\bm{\mu}}^{+}\|_{2}\neq\|{\bm{\mu}}^{-}\|_{2}$ by adding arbitrarily small perturbations to the data points. This motivates us to assume that $\|{\bm{\mu}}^{+}\|_{2}\neq\|{\bm{\mu}}^{-}\|_{2}$. Without loss of generality, we can assume that $\|{\bm{\mu}}^{+}\|_{2}>\|{\bm{\mu}}^{-}\|_{2}$ for convenience (Assumption A.3). When the reverse is true, i.e., $\|{\bm{\mu}}^{+}\|_{2}<\|{\bm{\mu}}^{-}\|_{2}$, we can change the direction of the inequality by flipping all the labels in the dataset so that our theorems can apply. We include the theorem statements for this reversed case in Section A.3.

The norm of ${\bm{\mu}}^{+}$ is strictly larger than ${\bm{\mu}}^{-}$, i.e., $\|{\bm{\mu}}^{+}\|_{2}>\|{\bm{\mu}}^{-}\|_{2}$.

Now we define ${\bm{w}}^{+}$ to be the max-margin linear separator of the dataset consisting of $({\bm{x}}^{+}_{i},y_{i})$, where $i\in[n]$, and define $\gamma^{+}$ to be this max margin. That is,

The reason that we care about ${\bm{w}}^{+}$ and $\gamma^{+}$ is because that it can be related to margin maximization on one-neuron Leaky ReLU nets. The following lemma is easy to prove.

The third assumption we made is that this margin cannot be obtained when all $a_{i}$ are negative, regardless of the width. This assumption holds when all the support vectors ${\bm{x}}_{i}^{+}$ have positive labels, i.e., $y_{i}=1$. Conceptually, this assumption is about whether nearly all the support vectors have positive labels (or negative labels in the reversed case where $\|{\bm{\mu}}^{+}\|_{2}<\|{\bm{\mu}}^{-}\|_{2}$).

Similar to Assumption 4.6 in the symmetric case, we need Assumption A.6 on non-branching starting point due to the technical difficulty for the potential non-uniqueness of gradient flow trajectory.

Now we are ready to state our theorem, and we defer the proofs to Appendix G.

A.2 Applying Theorem A.7 to prove Theorem 6.2

We give a proof of Theorem 6.2 here given the result of Theorem A.7.

With a ($H$, $K$ $\epsilon$, ${\bm{w}}_{\perp}$)-Hinted Dataset (Definition 6.1) with proper $H,K,\epsilon$, we only need to show that Assumptions A.2, A.3 and A.5 hold for Theorem 6.2. Specifically, we choose the parameters such that

Notice that $H_{0}$ is indepenent of $H$ as the data point ${\bm{x}}_{1}$ has projection $\|{\bm{P}}^{*}x_{1}\|_{2}=0$. For Assumption A.1, ${\bm{w}}^{\diamond}={\bm{w}}^{*}$ is a valid principal direction in this case, as

Then Assumption A.2 follows from Assumption A.1 by Lemma G.1. Since $H>n\left\|{\bm{\mu}}^{-}\right\|_{2}+\|\sum_{j>1}y_{j}{\bm{x}}^{+}_{j}\|_{2}$,

A.3 Results in the Reversed Case

In a reversed case where $\|{\bm{\mu}}^{+}\|_{2}<\|{\bm{\mu}}^{-}\|_{2}$, we can apply Theorem A.7 by flipping the labels in the dataset. Below we state the assumptions and the theorem in the reversed case.

$\|{\bm{\mu}}^{+}\|_{2}<\|{\bm{\mu}}^{-}\|_{2}$.

Now similarly we define ${\bm{w}}^{-}$ and $\gamma^{-}$.

Appendix B Additional Preliminaries and Lemmas

In this section, we will introduce additional notations and give some preliminary results for the dynamics of the two-layer Leaky ReLU network. The only assumption we will use for the results in the section is that the input norm is bounded $\max_{i\in[n]}\left\|{\bm{x}}_{i}\right\|_{2}\leq 1$ and we do not assume other properties of the dataset (such as symmetry) except we assume it explicitly.

For notational convenience for calculation with subgradients, we generalize the following notations for vectors to vector sets. More specifically, we define

Furthermore, we use the following notations to denote the radial and spherical components of $\bar{\partial}^{\circ}\mathcal{L}({\bm{\theta}})$ (which will be used in analyzing Phase III):

Let $\Omega_{\mathcal{S}}$ be the set of parameter vectors ${\bm{\theta}}=({\bm{w}}_{1},\dots,{\bm{w}}_{m},a_{1},\dots,a_{m})$ so that $\langle{\bm{w}}_{k},{\bm{x}}_{i}\rangle\neq 0$ for all $i\in[n],k\in[m]$, i.e., no activation function has zero input. For any ${\bm{\theta}}\in\Omega_{\mathcal{S}}$, $f_{{\bm{\theta}}}({\bm{x}}_{i})$ and $\mathcal{L}({\bm{\theta}})$ are continuously differentiable at ${\bm{\theta}}$, and the gradients are given by

Then the Clarke’s subdifferential for any ${\bm{\theta}}$ can be computed from (8) with $\Omega=\Omega_{\mathcal{S}}$ if needed.

Recall that $G$-function (Section 5.1) is defined by

B.2 Grönwall’s Inequality

We frequently use Grönwall’s inequality in our analysis.

Let $\alpha,\beta,u$ be real-valued functions defined on $[a,b)$. Suppose that $\beta,u$ are continuous and $\min\{\alpha,0\}$ is integrable on every compact subinterval of $[a,b)$. If $\beta\geq 0$ and $u$ satisfies the following inequality for all $t\in[a,b)$:

Furthermore, if $\alpha$ is non-decreasing, then for all $t\in[a,b]$,

B.3 Homogeneous Functions

The following is a direct corollary of Lemma B.4.

B.4 Karush-Kuhn-Tucker Conditions for Margin Maximization

We say that ${\bm{\theta}}$ is a feasible point if $g_{i}({\bm{\theta}})\leq 0$ for all $i\in[n]$. A feasible point ${\bm{\theta}}$ is a KKT point if it satisfies Karush-Kuhn-Tucker Conditions: there exist $\lambda_{1},\dots,\lambda_{n}\geq 0$ such that

${\bm{0}}\in\partial^{\circ}f({\bm{\theta}})+\sum_{i\in[n]}\lambda_{i}\partial^{\circ}g_{i}({\bm{\theta}})$;

$\forall i\in[n]:\lambda_{i}g_{i}({\bm{\theta}})=0$.

${\bm{\theta}}\in\sum_{i\in[n]}\lambda_{i}\partial^{\circ}q_{i}({\bm{\theta}})$;

For all $i\in[n]$, if $q_{i}({\bm{\theta}})\neq q_{\min}({\bm{\theta}})$ then $\lambda_{i}=0$.

For two-layer Leaky ReLU network, $q_{i}({\bm{\theta}}):=y_{i}\sum_{k\in[m]}a_{k}\phi({\bm{w}}_{k}^{\top}{\bm{x}}_{i})$. Then the KKT-margin direction is defined as follows.

For all $k\in[m]$, ${\bm{w}}_{k}\in\sum_{i\in[n]}\lambda_{i}y_{i}a_{k}\phi^{\circ}({\bm{w}}_{k}^{\top}{\bm{x}}_{i}){\bm{x}}_{i}$;

For all $k\in[m]$, $a_{k}=\sum_{i\in[n]}\lambda_{i}y_{i}\phi({\bm{w}}_{k}^{\top}{\bm{x}}_{i})$;

For all $i\in[n]$, if $q_{i}({\bm{\theta}})\neq q_{\min}({\bm{\theta}})$ then $\lambda_{i}=0$.

For ${\bm{\theta}}$ along a KKT-margin direction of two-layer Leaky ReLU network, Lemma B.9 below shows that $\lvert a_{k}\rvert=\|{\bm{w}}_{k}\|_{2}$ for all $k\in[m]$.

By Definition B.8 and Theorem B.3, we have

Therefore $\|{\bm{w}}_{k}\|_{2}^{2}=\lvert a_{k}\rvert^{2}$. ∎

B.5 Lemmas for Perturbation Bounds

Writing the formula with respect to ${\bm{w}}_{k},a_{k}$, we have

which completes the proof for ${\bm{\theta}}\in\Omega_{\mathcal{S}}$ and thus the same bounds hold for the general case. ∎

Lemma B.11 is a lemma for bounding the partial subderivatives. For the full subgradient, we have the following lemma.

If ${\bm{\theta}}\in\Omega_{\mathcal{S}}$, by (13), there exists $\delta_{i}\in[-\lvert f_{{\bm{\theta}}}({\bm{x}}_{i})\rvert,\lvert f_{{\bm{\theta}}}({\bm{x}}_{i})\rvert]$ for all $i\in[n]$ such that

Writing it with respect to ${\bm{w}}_{k}$, we have

We conclude the proof by noticing that $[-2\lvert f_{{\bm{\theta}}}({\bm{x}}_{i})\rvert,2\lvert f_{{\bm{\theta}}}({\bm{x}}_{i})\rvert]\subseteq[-{\epsilon},{\epsilon}]$ by Lemma B.10. ∎

B.6 Basic Properties of Gradient Flow

The following lemma is a simple corollary from Davis et al. .

For gradient flow ${\bm{\theta}}(t)$ on a two-layer Leaky ReLU network with logistic loss, we have

The following lemma is from Du et al. . We provide a simple proof here for completeness.

For gradient flow ${\bm{\theta}}(t)=({\bm{w}}_{1}(t),\dots,{\bm{w}}_{m}(t),a_{1}(t),\dots,a_{m}(t))$ on a two-layer Leaky ReLU network with logistic loss, the following holds for all $t\geq 0$,

where $q_{i}({\bm{\theta}}):=y_{i}f_{{\bm{\theta}}}({\bm{x}}_{i})$. Therefore, $\frac{\textup{{d}}}{\textup{{d}}t}(\|{\bm{w}}_{k}\|_{2}^{2}-\lvert a_{k}\rvert^{2})=0$ for all $t\geq 0$.

By (9), we have the following for any ${\bm{\theta}}\in\Omega_{\mathcal{S}}$,

By $1$-homogeneity of $\phi$ and Theorem B.3, we have $\phi^{\prime}({\bm{w}}_{k}^{\top}{\bm{x}}_{i}){\bm{w}}_{k}^{\top}{\bm{x}}_{i}=\phi({\bm{w}}_{k}^{\top}{\bm{x}}_{i})$, which implies that $\left<{\bm{w}}_{k},\frac{\partial f_{{\bm{\theta}}}({\bm{x}})}{\partial{\bm{w}}_{k}}\right>=a_{k}\phi({\bm{w}}_{k}^{\top}{\bm{x}}_{i})$.

By chain rule, for a.e. $t\geq 0$ we have

The following lemma shows that if a neuron has zero weights, then it stays with zero weights forever. Conversely, this also implies that the weights stay non-zero if they are initially non-zero.

If $a_{k}(t_{0})=0$ and ${\bm{w}}_{k}(t_{0})={\bm{0}}$ at some time $t_{0}\geq 0$, then $a_{k}(t)=0$ and ${\bm{w}}_{k}(t)={\bm{0}}$ for all $t\geq 0$.

By Lemma B.16, we know that $\|{\bm{w}}_{k}\|_{2}=\lvert a_{k}\rvert$ hold for all $t\geq 0$. Also, we have $\frac{1}{2}\left\lvert\frac{\textup{{d}}\|{\bm{w}}_{k}\|_{2}^{2}}{\textup{{d}}t}\right\rvert=\frac{1}{2}\left\lvert\frac{\textup{{d}}\lvert a_{k}\rvert^{2}}{\textup{{d}}t}\right\rvert\leq C\cdot\lvert a_{k}\rvert\|{\bm{w}}_{k}\|_{2}=C\|{\bm{w}}_{k}\|_{2}^{2}$, where $C>0$ is some constant. Then

By Grönwall’s inequality (12) this implies that $\|{\bm{w}}_{k}(t)\|_{2}=0$ for all $t\geq t_{0}$. Similarly,

By Grönwall’s inequality (12) again, $\|{\bm{w}}_{k}(t)\|_{2}=0$ for all $t\leq t_{0}$, which completes the proof. ∎

A direct corollary of Lemma B.16 and Lemma B.17 is the following characterization in the case where the weights are initially balanced.

If $\left\lvert a_{k}\right\rvert=\|{\bm{w}}_{k}\|_{2}$ initially for $t=0$, then this equation holds for all $t\geq 0$. Moreover,

If $a_{k}(0)=\|{\bm{w}}_{k}(0)\|_{2}$, then $a_{k}(t)=\|{\bm{w}}_{k}(t)\|_{2}$ for all $t\geq 0$;

If $a_{k}(0)=-\|{\bm{w}}_{k}(0)\|_{2}$, then $a_{k}(t)=-\|{\bm{w}}_{k}(t)\|_{2}$ for all $t\geq 0$.

B.7 A Useful Theorem for Loss Convergence

In this section we prove a useful theorem for loss convergence, which will be used later in our analysis for both symmetric and non-symmetric datasets.

Under Assumption 3.2, for any linear seprator ${\bm{w}}^{*}$ of the data with positive linear margin (e.g. $y_{i}\left<{\bm{w}}^{*},x_{i}\right>\geq\gamma^{*}>0$ for all $i\in[n]$), if initially there exists $k\in[m]$ such that

then $a_{k}(t)\neq 0$ for all $t>0$, and $\mathcal{L}({\bm{\theta}}(t))\to 0$ and $\|{\bm{\theta}}(t)\|_{2}\to+\infty$ as $t\to+\infty$.

Before proving Theorem B.19, we first prove a lemma on gradient lower bounds.

We only need to show that there exists $t_{0}$ such that $\mathcal{L}({\bm{\theta}}(t_{0}))<\frac{\ln 2}{n}$, then we can apply Theorem 3.1 to show that $\mathcal{L}({\bm{\theta}}(t))\to 0$. Assume to the contrary that $\mathcal{L}({\bm{\theta}}(t))\geq\frac{\ln 2}{n}$ for all $t\geq 0$. By Lemma B.20,

Lemma B.15 ensures that $-\frac{\textup{{d}}\mathcal{L}}{\textup{{d}}t}=\left\|\frac{\textup{{d}}{\bm{\theta}}}{\textup{{d}}t}\right\|_{2}^{2}$ for a.e. $t\geq 0$. Then we have

Integrating on $t$ from to $+\infty$, we can see that the LHS is upper bounded by $\mathcal{L}({\bm{\theta}}(0))-\frac{\ln 2}{n}$ while the RHS is unbounded, which leads to a contradiction. Therefore, there exist time $t_{0}$ such that $\mathcal{L}({\bm{\theta}}(t_{0}))<\frac{\ln 2}{n}$, and thus $\mathcal{L}({\bm{\theta}}(t))\to 0$ as $t\to+\infty$. ∎

Appendix C Proofs for Linear Maximality for the Symmetric Case

For linearly separable and symmetric data, we show that all global-max-margin directions represent linear functions in Theorem 4.2. We give a proof here.

Now we define $A:=\sqrt{\sum_{k\in[m]}a_{k}^{2}}$ and let ${\bm{\theta}}^{\prime}=({\bm{w}}^{\prime}_{1},\dots,{\bm{w}}^{\prime}_{m},a^{\prime}_{1},\dots,a^{\prime}_{m})$ where

Meanwhile, by the Cauchy-Schwarz inequality,

Thus $\gamma({\bm{\theta}}^{\prime})=\frac{q_{\min}({\bm{\theta}}^{\prime})}{\|{\bm{\theta}}^{\prime}\|_{2}^{2}}\geq\gamma({\bm{\theta}}^{*})$. As ${\bm{\theta}}^{*}$ is already a global-max-margin direction, equalities should hold in all the inequalities above, so

There is $j\in[n]$ that $y_{j}f_{{\bm{\theta}}^{*}}({\bm{x}}_{j})=-y_{j}f_{{\bm{\theta}}^{*}}(-{\bm{x}}_{j})=\gamma({\bm{\theta}}^{*})$.

Certainly ${\bm{c}}^{\top}{\bm{x}}_{j}\neq 0$ as otherwise the margin would be zero. Then $\sum_{k\in[m]}a_{k}|a_{k}|=0$, which means $\sum_{k:a_{k}\geq 0}a^{2}_{k}=\sum_{k:a_{k}<0}a^{2}_{k}=\frac{1}{2}A^{2}$, and therefore

$\min_{i\in[n]}y_{i}{\bm{c}}^{\top}{\bm{x}}_{i}=\left\|{\bm{c}}\right\|_{2}\gamma_{{\bm{w}}^{*}}$;

$\frac{\left\|{\bm{c}}\right\|_{2}}{1+\left\|{\bm{c}}\right\|_{2}^{2}}=\frac{1}{2}$.

Appendix D Proofs for Phase I

In the subsequent sections we first show the proofs for the symmetric datasets under Assumption 4.1. Additional proofs for the non-symmetric counterparts are provided in Appendix G.

By definition and Cauchy-Schwartz inequality,

For initial point ${\bm{\theta}}_{0}\neq{\bm{0}}$, we have

Appendix E Proofs for Phase II

To prove Lemma 5.3, we start from the following lemma.

Given $\hat{{\bm{\theta}}}_{0}:=(\hat{{\bm{w}}}_{1},\hat{{\bm{w}}}_{2},\hat{a}_{1},\hat{a}_{2})$ with $\hat{a}_{1}>0$ and $\hat{a}_{2}<0$, then ${\bm{\theta}}(t)=\pi_{{\bm{b}}}(\varphi(\hat{{\bm{\theta}}}_{0},t))$ is a gradient flow trajectory on $\mathcal{L}({\bm{\theta}})$ starting from ${\bm{\theta}}(0)=\pi_{{\bm{b}}}(\hat{{\bm{\theta}}}_{0})$.

For any $\hat{{\bm{\theta}}}$ and ${\bm{g}}\in\partial^{\circ}\mathcal{L}(\hat{{\bm{\theta}}})$, $\pi_{{\bm{b}}}({\bm{g}})\in\partial^{\circ}\mathcal{L}(\pi_{{\bm{b}}}(\hat{{\bm{\theta}}}))$.

Below we use $\pi_{{\bm{b}}}(S)=\{\pi_{{\bm{b}}}(s):{\bm{s}}\in S\}$ to denote the embedding of a parameter set.

For every $\hat{{\bm{\theta}}}=(\hat{{\bm{w}}}_{1},\hat{{\bm{w}}}_{2},\hat{a}_{1},\hat{a}_{2})\in\Omega_{\mathcal{S}}$ (i.e., no activation function has zero input), let ${\bm{\theta}}=\pi_{{\bm{b}}}(\hat{{\bm{\theta}}})=({\bm{w}}_{1},\dots,{\bm{w}}_{m},a_{1},\dots,a_{m})$, and clearly ${\bm{\theta}}\in\Omega_{\mathcal{S}}$. Then $\partial^{\circ}\mathcal{L}(\hat{{\bm{\theta}}})=\{\nabla\mathcal{L}(\hat{{\bm{\theta}}})\}$ and $\partial^{\circ}\mathcal{L}({\bm{\theta}})=\{\nabla\mathcal{L}({\bm{\theta}})\}$ are the usual differentials. In this case, we can apply the chain rule as

Notice that the embedding preserves the function value,

so $\frac{\partial f_{{\bm{\theta}}}({\bm{x}}_{i})}{\partial{\bm{\theta}}}=\pi_{{\bm{b}}}\left(\frac{\partial f_{\hat{{\bm{\theta}}}}({\bm{x}}_{i})}{\partial\hat{{\bm{\theta}}}}\right)$. Then from the chain rule above we can see $\nabla\mathcal{L}({\bm{\theta}})=\pi_{{\bm{b}}}(\nabla\mathcal{L}(\hat{{\bm{\theta}}}))$, and we proved the lemma in this case.

In the general case, by the definition of Clarke’s subdifferential,

For any $\hat{{\bm{\theta}}}_{n}\to\hat{{\bm{\theta}}}$ with $\hat{{\bm{\theta}}}_{n}\in\Omega_{\mathcal{S}}$, $\pi_{{\bm{b}}}(\hat{{\bm{\theta}}}_{n})\to\pi_{{\bm{b}}}(\hat{{\bm{\theta}}})$, and

Taking the convex hull, it follows that $\pi_{{\bm{b}}}(\partial^{\circ}\mathcal{L}(\hat{{\bm{\theta}}}))\subseteq\partial^{\circ}\mathcal{L}(\pi_{{\bm{b}}}(\hat{{\bm{\theta}}}))$, and we finished the proof. ∎

For notations we write $\hat{{\bm{\theta}}}(t):=\varphi(\hat{{\bm{\theta}}}_{0},t)$ and ${\bm{\theta}}(t)=\pi_{{\bm{b}}}(\hat{{\bm{\theta}}}(t))$. Then $\frac{\textup{{d}}}{\textup{{d}}t}\hat{{\bm{\theta}}}(t)\in-\partial^{\circ}\mathcal{L}(\hat{{\bm{\theta}}}(t))$ for a.e. $t$. At these $t$, $\frac{\textup{{d}}}{\textup{{d}}t}{\bm{\theta}}(t)=\pi_{{\bm{b}}}(\frac{\textup{{d}}}{\textup{{d}}t}\hat{{\bm{\theta}}}(t))\in\pi_{{\bm{b}}}(-\partial^{\circ}\mathcal{L}(\hat{{\bm{\theta}}}(t)))$. From Lemma E.2 we know $\pi_{{\bm{b}}}(\partial^{\circ}\mathcal{L}(\hat{{\bm{\theta}}}(t)))\subseteq\partial^{\circ}\mathcal{L}({\bm{\theta}}(t))$. Then $\frac{\textup{{d}}}{\textup{{d}}t}{\bm{\theta}}(t)\in-\partial^{\circ}\mathcal{L}({\bm{\theta}}(t))$ for a.e. $t$, and therefore ${\bm{\theta}}(t)$ is indeed a gradient flow trajectory. ∎

By Lemma E.1, $\pi_{{\bm{b}}}(\varphi(\hat{{\bm{\theta}}}_{0},t))$ is indeed a gradient flow trajectory. Then, as $\pi_{{\bm{b}}}(\varphi(\hat{{\bm{\theta}}}_{0},0))=\pi_{{\bm{b}}}(\hat{{\bm{\theta}}}_{0})$, as well as the fact that $\hat{{\bm{\theta}}}_{0}$ and $\pi_{{\bm{b}}}(\hat{{\bm{\theta}}}_{0})$ are non-branching starting points, the gradient flow trajectory is unique and therefore $\pi_{{\bm{b}}}(\varphi(\hat{{\bm{\theta}}}_{0},t))=\varphi(\pi_{{\bm{b}}}(\hat{{\bm{\theta}}}_{0}),t)$ for all $t\geq 0$. ∎

E.2 A General Theorem for Limiting Trajectory Near Zero

We say that $\hat{{\bm{\theta}}}:=(\hat{{\bm{w}}}_{1},\dots,\hat{{\bm{w}}}_{m},\hat{a}_{1},\dots,\hat{a}_{m})$ is a well-aligned parameter vector if it satisfies the following for some $1\leq p\leq m$:

For $1\leq k\leq p$, $\left<\hat{{\bm{w}}}_{k},{\bm{x}}_{i}\right>\neq 0$ for all $i\in[n]$;

For $p+1\leq k\leq m$, $\hat{{\bm{w}}}_{k}={\bm{0}}$, $\hat{a}_{k}=0$.

Our analysis for Phase I shows that weight vectors approximately align to either of $\bar{{\bm{\mu}}}$ or $-\bar{{\bm{\mu}}}$, and both of them are maximizers of $\lvert G({\bm{w}})\rvert$. Therefore, gradient flow goes near a well-aligned parameter vector (with $p=m$) at the end of Phase I.

The following is the main theorem of this subsection.

Then for all $t\in(-\infty,t_{0}]$, the following is true:

$\lim_{r\to 0}\varphi(r\hat{{\bm{\theta}}},T_{2}(r)+t)$ exists. This limit is independent of the choice of $\varphi$ when the gradient flow may not be unique.

$\lim_{r\to 0}\varphi(r\hat{{\bm{\theta}}},T_{2}(r)+t)$ lies near $e^{\lambda t}\hat{{\bm{\theta}}}$:

Let ${\bm{\theta}}_{1},{\bm{\theta}}_{2},\dots$ be a series of parameters converging to ${\bm{0}}$, $r_{1},r_{2},\dots$ be a series of positive real numbers converging to . If $\|{\bm{\theta}}_{s}-r_{s}\hat{{\bm{\theta}}}\|_{2}\leq Cr_{s}^{1+\kappa}$ for some $C>0,\kappa>0$, then

Now we prove Theorem E.4. Throughout this subsection, we fix a well-aligned parameter vector $\hat{{\bm{\theta}}}:=(\hat{{\bm{w}}}_{1},\dots,\hat{{\bm{w}}}_{m},\hat{a}_{1},\dots,\hat{a}_{m})$ with constant $p\in[m]$. We also use $t_{0}$ and $T_{2}(r)$ to denote the same constant $t_{0}$ defined by (14) and the same function $T_{2}(r):=\frac{1}{\lambda}\ln\frac{1}{r}$ as in Theorem E.4.

For all $k\in[p]$, ${\bm{w}}_{k}(T_{2}(r)+t)\in\widehat{\mathcal{W}}_{k}$;

within the time interval $[0,T_{2}(r)+t_{\max})$ (and thus it also holds for $[0,T_{2}(r)+t_{\max}]$ by continuity), and to show that $t_{0}$ is actually equal to $t_{\max}$, i.e., $t_{0}$ is the minimum among $t_{0},t_{1},t_{2}$. It is easy to see that proving these suffice to deduce the original lemma statement, given the translation of time $e^{\lambda T_{2}(r)}=\frac{1}{r}$.

For $k\in[p]$, $\partial^{\circ}G({\bm{w}}_{k})=\{\nabla G({\bm{w}}_{k})\}=\{\nabla G(\hat{{\bm{w}}}_{k})\}$. Also note that $\nabla G(\hat{{\bm{w}}}_{k})=\nabla G(\frac{\hat{{\bm{w}}}_{k}}{\|\hat{{\bm{w}}}_{k}\|_{2}})=\lambda\frac{\hat{{\bm{w}}}_{k}}{\|\hat{{\bm{w}}}_{k}\|_{2}}$ by Lemma B.5. Then $a_{k}\partial^{\circ}G({\bm{w}}_{k})=\{\lambda a_{k}\frac{\hat{{\bm{w}}}_{k}}{\|\hat{{\bm{w}}}_{k}\|_{2}}\}$ and $G({\bm{w}}_{k})=\lambda\left<\frac{\hat{{\bm{w}}}_{k}}{\|\hat{{\bm{w}}}_{k}\|_{2}},{\bm{w}}_{k}\right>$. Combining these with (15) gives

For $p<k\leq m$, we can combine Theorem B.3 and (15) to give the following bound for the norm growth:

To prove the lemma, now we only need to show that $t_{\max}=t_{0}$. Combining (19) and (21), we have for $t\leq T_{2}(r)+t_{\max}$,

For all time $0\leq t<T_{2}(r)+t_{\max}$, we can use (22) to deduce

For norm growth, we can again use (22) to deduce

Now we have $t_{1}>t_{\max},t_{2}>t_{\max}$. Recall that $t_{\max}:=\min\{t_{0},t_{1},t_{2}\}$ by definition. Then $t_{\max}=t_{0}$ must hold, which completes the proof. ∎

For $\Lambda(t)$, by triangle inequality and Lemma B.5 we have

For $\Delta(t)$, we use triangle inequality again to give the following bound:

At time $T_{2}(r)+t\in[0,T_{2}(r)+t_{0}]$, this bound can be rewritten as

First we show that $\lim_{r\to 0}\varphi(r\hat{{\bm{\theta}}},T_{2}(r)+t)$ exists. We consider the case of $r\leq r_{\max}$, where $r_{\max}$ is chosen to be small enough so that the properties in Lemma E.5 hold. For any $r^{\prime}<r$, by Lemma E.5 we have

Note that $T_{2}(r^{\prime})+\frac{1}{\lambda}\ln r+T_{2}(r)+t=T_{2}(r^{\prime})+t$. So this proves

For any fixed $t\leq t_{0}$, the RHS converges to as $r\to 0$, which implies Cauchy convergence of the limit $\lim_{r\to 0}\varphi(r\hat{{\bm{\theta}}},T_{2}(r)+t)$ and thus the limit exists. By the 1st property in Lemma E.5, we know that there is no activation pattern switch in the time interval $t\in[0,T_{2}(r)+t_{0}]$ if $r$ is small enough. This means $\mathcal{L}$ is locally smooth near the trajectory of $\varphi(r\hat{{\bm{\theta}}},T_{2}(r)+t)$ and thus the trajectory is unique. Therefore, the limit $\lim_{r\to 0}\varphi(r\hat{{\bm{\theta}}},T_{2}(r)+t)$ is uniquely defined.

Taking $r\to 0$ on both sides gives the range of the limit $\lim_{r\to 0}\varphi(r\hat{{\bm{\theta}}},T_{2}(r)+t)$:

So $\lim_{s\to\infty}\varphi({\bm{\theta}}_{s},T_{2}(r_{s})+t)=\lim_{r\to 0}\varphi(r\hat{{\bm{\theta}}},T_{2}(r)+t)$ is proved. ∎

E.3 Proof for Approximate Embedding

To analyze Phase II, we need to deal with approximate embedding instead of the exact one. For this, we further divide Phase II into Phase II.1 and II.2 and analyze them in order. At the end of this subsection we will prove Lemma 5.4.

Given the discussions in the previous sections, we are ready to present proofs for the phase II dynamics (Lemma 5.4) here.

For the $m$-neuron dynamics ${\bm{\theta}}(t)$, the following holds for all $t\in(-\infty,t_{0}]$,

Applying Theorem E.4 proves the following for all $t\in(-\infty,t_{0}]$:

E.4 Proofs for Phase II.2

Next, at the end of Phase II.1, ${\bm{\theta}}(T_{12}+t_{0})$ has a constant norm. Then we show the trajectory convergence with respect to the initialization scale in Phase II.2.

We first start with a simple lemma on gradient upper bounds, and then show that the trajectory of gradient flow is Lipschitz with time.

Finally we show that ${\bm{q}}(t)$ is indeed a valid gradient flow trajectory. Notice that ${\bm{q}}(t)$ is $(Ce^{T-t_{0}})$-Lipschitz, then by Rademacher theorem for ${\bm{q}}(t)$ is differentiable for a.e. $t\in[t_{0},T]$. We are left to show ${\bm{q}}^{\prime}(t)\in\partial^{\circ}\mathcal{L}({\bm{q}}(t))$ whenever ${\bm{q}}$ is differentiable at $t$.

For any $\epsilon>0$ that $[t,t+\epsilon]\subseteq[t_{0},T]$, we investigate the behaviour of ${\bm{q}}(t)$ in the $\epsilon$-neighborhood of $t$. Let $\Omega_{j}$ be the set of $\tau\in[t_{0},T]$ so that $\frac{\textup{{d}}}{\textup{{d}}\tau}{\bm{\theta}}_{i_{k}}(T_{12}+\tau)\in-\partial^{\circ}\mathcal{L}({\bm{\theta}}_{i_{k}}(T_{12}+\tau))$. By definition of differential inclusion, $\Omega_{j}$ has full measure in $[t_{0},T]$. Define $B_{j,\epsilon}$ be the following closed convex hull:

It is easy to see that $B_{j,\epsilon}$ is monotonic with respect to $j$. Then we know that for any $j$,

Then taking the limits $j\to\infty$, as all $B_{j,\epsilon}$ are closed, we know $\frac{{\bm{q}}(t+\epsilon)-{\bm{q}}(t)}{\epsilon}\in\lim_{j\to\infty}B_{j,\epsilon}$.

Now let $C_{j,\epsilon},C_{\epsilon}$ be the following closed convex hull of subgradients:

Then we know $B_{j,\epsilon}\subseteq-C_{j,\epsilon}$ for all $j\geq 1$ and $\epsilon>0$. Notice that $C_{j,\epsilon}$ and $C_{\epsilon}$ are also monotonic with respect to $j$ and $\epsilon$ respectively so we can take the respective limit. As for $\tau\in[t,t+\epsilon]$, $\lim_{j\to\infty}{\bm{\theta}}_{i_{j}}(T_{12}+\tau)={\bm{q}}(T_{12}+\tau)$, by the upper-semicontinuity of $\partial^{\circ}\mathcal{L}$, $\lim_{j\to\infty}C_{j,\epsilon}\subseteq C_{\epsilon}$. Then $\frac{{\bm{q}}(t+\epsilon)-{\bm{q}}(t)}{\epsilon}\in\lim\limits_{j\to\infty}B_{j,\epsilon}\subseteq\lim\limits_{j\to\infty}C_{j,\epsilon}\subseteq C_{\epsilon}$.

When $t\in[t_{0},T)$ and ${\bm{q}}(t)$ is differential at $t$, we can take the limit $\epsilon\to 0$, and by the upper-semicontinuity of $\partial^{\circ}\mathcal{L}$ again, we have

as $\partial^{\circ}\mathcal{L}({\bm{q}}(T_{12}+t))$ is closed convex for any $t$. Therefore ${\bm{q}}(t)$ is indeed a gradient flow trajectory. ∎

Appendix F Proofs for Phase III

In this subsection we prove Theorem 5.5 for the symmetric datasets. By Theorem B.19 and Theorem 3.1, we know that gradient flow must converge in a KKT-margin direction of width-$2$ two-layer Leaky ReLU network (Definition B.8). Thus we first give some characterizations for KKT-margin directions by proving Lemma F.1 and Lemma F.2.

for all Clarke’s sub-differentials $h^{(1)}_{i}\in\phi^{\circ}(\langle{\bm{u}}_{1},{\bm{x}}_{i}\rangle),h^{(2)}_{i}\in\phi^{\circ}(-\langle{\bm{u}}_{2},{\bm{x}}_{i}\rangle)$.

We prove by cases for any fixed $i\in[n]$. By Assumption 4.1 we have

Otherwise, $h^{(1)}_{i}\neq h^{(2)}_{i}$, then we have $h^{(1)}_{i}\langle{\bm{u}}_{2},{\bm{x}}_{i}\rangle=\phi(\langle{\bm{u}}_{2},{\bm{x}}_{i}\rangle)$, $h^{(2)}_{i}\langle{\bm{u}}_{1},{\bm{x}}_{i}\rangle=-\phi(-\langle{\bm{u}}_{1},{\bm{x}}_{i}\rangle)$, and thus

Suppose that $\langle{\bm{u}}_{1},{\bm{x}}_{i}\rangle=0$ or $\langle{\bm{u}}_{2},{\bm{x}}_{i}\rangle=0$. WLOG we assume that $\langle{\bm{u}}_{1},{\bm{x}}_{i}\rangle=0$ (the case of $\langle{\bm{u}}_{2},{\bm{x}}_{i}\rangle=0$ can be proved similarly). Then we have

If $({\bm{w}}_{1},{\bm{w}}_{2},a_{1},a_{2})$ is along a KKT-margin direction of width-$2$ two-layer Leaky ReLU network and $a_{1}>0,a_{2}<0$, then ${\bm{w}}_{1}=-{\bm{w}}_{2}$, $a_{1}=-a_{2}=\|{\bm{w}}_{1}\|_{2}$.

${\bm{w}}_{1}=a_{1}\sum_{i\in[n]}\lambda_{i}y_{i}h^{(1)}_{i}{\bm{x}}_{i}$, ${\bm{w}}_{2}=a_{2}\sum_{i\in[n]}\lambda_{i}y_{i}h^{(2)}_{i}{\bm{x}}_{i}$;

$a_{1}=\|{\bm{w}}_{1}\|_{2}$, $a_{2}=-\|{\bm{w}}_{2}\|_{2}$;

For all $i\in[n]$, if $q_{i}({\bm{\theta}})\neq 1$ then $\lambda_{i}=0$.

Let ${\bm{u}}_{1}=a_{1}{\bm{w}}_{1}$ and ${\bm{u}}_{2}=-a_{2}{\bm{w}}_{2}$. Let $\bar{{\bm{u}}}_{1}:=\frac{{\bm{u}}_{1}}{\|{\bm{u}}_{1}\|_{2}},\bar{{\bm{u}}}_{2}:=-\frac{{\bm{u}}_{2}}{\|{\bm{u}}_{2}\|_{2}}$. Then the following conditions hold for all $i\in[n]$:

By homogeneity, $h^{(1)}_{i}\cdot\langle{\bm{u}}_{1},{\bm{x}}_{i}\rangle=\phi(\langle{\bm{u}}_{1},{\bm{x}}_{i}\rangle),h^{(2)}_{i}\cdot\langle{\bm{u}}_{2},{\bm{x}}_{i}\rangle=-\phi(-\langle{\bm{u}}_{2},{\bm{x}}_{i}\rangle)$. Left-multiplying $({\bm{u}}_{1})^{\top}$ or $({\bm{u}}_{2})^{\top}$ on both sides of (25), we have

where the last inequality is due to Lemma F.1. Since we have deduced that $\|{\bm{u}}_{1}\|_{2}+\|{\bm{u}}_{2}\|_{2}=\sum_{i=1}^{n}\lambda_{i}$, we further have

Combining this with $\langle\bar{{\bm{u}}}_{1},\bar{{\bm{u}}}_{2}\rangle\leq\|\bar{{\bm{u}}}_{1}\|_{2}\|\bar{{\bm{u}}}_{2}\|_{2}\leq 1$, we have $1\leq\langle\bar{{\bm{u}}}_{1},\bar{{\bm{u}}}_{2}\rangle\leq 1$. So all the inequalities become equalities, and thus $\bar{{\bm{u}}}_{1}=\bar{{\bm{u}}}_{2}$. (36) also equals to (37), so

By (27), we have $y_{i}\left(h^{(1)}_{i}\langle{\bm{u}}_{1},{\bm{x}}_{i}\rangle+h^{(2)}_{i}\langle{\bm{u}}_{2},{\bm{x}}_{i}\rangle\right)=1$ whenever $\lambda_{i}\neq 0$. Combining this with (38), we have

Then we prove that $\langle{\bm{u}}_{1},{\bm{x}}_{i}\rangle=\langle{\bm{u}}_{2},{\bm{x}}_{i}\rangle$ by discussing two cases:

If $\langle{\bm{u}}_{1},{\bm{x}}_{i}\rangle=0$, then $\langle{\bm{u}}_{2},{\bm{x}}_{i}\rangle=0$ since $\bar{{\bm{u}}}_{1}=\bar{{\bm{u}}}_{2}$;

This means ${\bm{u}}_{1}$ and ${\bm{u}}_{2}$ have the same projection onto the linear space spanned by $\{{\bm{x}}_{i}:\lambda_{i}\neq 0\}$. By (25) and (26), ${\bm{u}}_{1}$ and ${\bm{u}}_{2}$ are in the span of $\{{\bm{x}}_{i}:i\in[n],\lambda_{i}\neq 0\}$. Therefore, ${\bm{u}}_{1}={\bm{u}}_{2}$ and we can easily deduce that ${\bm{w}}_{1}=-{\bm{w}}_{2}$, $a_{1}=-a_{2}=\|{\bm{w}}_{1}\|_{2}$. ∎