Towards Understanding Learning in Neural Networks with Linear Teachers

Introduction

Neural networks have achieved remarkable performance in many machine learning tasks (Krizhevsky et al., 2012; Silver et al., 2016; Devlin et al., 2019). Although their success has already transformed technology, a theoretical understanding of how this performance is achieved is not complete. Here we focus on one of the simplest learning settings that is still not understood. We consider linearly separable data (i.e., generated by a “linear teacher”) that is being learned by a two layer neural net with leaky ReLU activations and minimization of cross entropy loss using gradient descent or its variants. Two key questions immediately come up in this context:

The Optimization Question: Will the optimization succeed in finding a classifier with zero training error, and arbitrarily low training loss?

The Inductive Bias Question: With a large number of hidden units, the network can find many solutions that will separate the data. Which of these will be found by gradient descent?

Our work addresses these questions as follows. The Optimization Question: We prove that stochastic gradient descent (SGD) will converge to arbitrary low training loss. Concretely, we show that for any $\epsilon>0$, SGD will converge to $\epsilon$ cross-entropy loss in $O\left(\frac{1}{\epsilon^{2}}\right)$ iterations. We consider SGD which performs multiple passes over the data and we devise a novel variant of the perceptron proof to analyze this setting. Our analysis bounds the number of epochs that have high loss examples, and uses this to show convergence to a low loss solution. Importantly, our result holds for any network size and scale of initialization. Therefore, our analysis goes beyond the Neural Tangent Kernel (NTK) analyses which require large network sizes and relatively large initialization scales.

The Inductive Bias Question: We empirically observe that when a small initialization scale is used, the learned network converges to a decision boundary that is very close to linear. See Figure 1(b) for a 2D example.See Section 5.1 and Section 6.3 for more empirical examples. We also observe that all neurons cluster nicely into two sets of vectors (i.e., they form two groups of well-aligned neurons) as in Figure 1(d). To support these empirical findings, we provide the following theoretical results: (1) We prove that an approximate clustering of the neurons implies that the decision boundary of the network is approximately linear. This is a result of a nice property of leaky ReLU networks which we prove in Section 5. (2) We provide a novel sufficient condition on the optimization path of gradient flow which implies convergence to clustered solutions. With the result above, it implies convergence to a linear decision boundary. The condition states that from a certain iteration on, all neurons with the same output sign “agree” on the classification of the data. We observe that this condition holds empirically for several synthetic and real datasets. Finally, we use the latter result to prove that under certain assumptions, the learned network is a solution to an SVM problem with a specific kernel.

a large range of initialization scales, [AB: This is confusing, because the reader may think that this also includes large initialization scales.] the network learns a decision boundary that is very close to linear (note that the network can model highly non-linear rules that separate the data (see Figure 1). Additionally, we observe that all neurons cluster nicely into two sets of vectors (i.e., they form two groups of well-aligned neurons).[AG: add figure for this] [AB: Say that even in certain cases (as in the figure) we see that the decision boundary is exactly linear.]

[AB: I think it would be good to add a figure where the decision boundary is approximately linear. E.g., in the case of a Gaussian distribution with outliers.]

Our results above make significant headway in understanding why optimization is tractable with linear teachers, and why convergence is to approximately linear boundaries. We also provide empirical evaluation that confirms that weight clustering indeed explains why approximate linear decision boundaries are learned.

In this work, we provide an in depth examination of the implicit bias gradient methods yield across all the domains of the optimization, in the case of optimizing a 1 hidden layer Leaky ReLU network with logistic loss. More specifically, we prove that:

Related Work

Since training neural networks is NP-Hard for worst-case datasets (Blum & Rivest, 1992), recent works have analyzed neural networks under certain data assumptions to better understand their performance in practice. One common assumption is to analyze neural networks when the data is linearly separable. Even in this case, the theoretical analysis of optimization and generalization is far from resolved. In a work closely related to ours, Brutzkus et al. (2018) consider this setting and show that SGD converges to zero loss for linearly separable data (which was later extended to ReLU activations using noisy SGD in Wang et al. (2019)). The key difference from our work is that they use the hinge loss instead of the cross entropy loss. The cross entropy loss creates unique challenges for proving convergence as we show in Section 4. Thus, their results cannot be directly applied for the cross entropy loss and we use novel techniques to guarantee convergence of SGD in this case. The second key difference from their work is that we present novel insights on the inductive bias of SGD using results of Lyu & Li (2020) and Ji & Telgarsky (2020) which hold for the cross entropy loss and not for the hinge loss. Finally, our result which shows that a network with clustered neurons has an approximate linear decision boundary is new and holds irrespective of the loss used.

Recently, Phuong & Lampert (2021) analyzed a subclass of linear teachers where data is “orthogonally separable”. In this case they show that training a ReLU network with the cross entropy loss results in a solution where weights are aligned. In terms of our results, this can be viewed as a case of convergence to a particular PAR (see Section 6). Several other works assume that the data is linearly separable but also that the networks are linear (Ji & Telgarsky, 2019a; Moroshko et al., 2020; Gunasekar et al., 2018). We study the more challenging and realistic setting of two layer nonlinear networks with Leaky ReLU activations.

Several works (Lyu & Li, 2020; Ji & Telgarsky, 2020; Nacson et al., 2019) studied the inductive bias of two-layer homogeneous networks and showed connections between gradient methods and margin maximization. Their results hold under the assumption that gradient methods achieve a certain loss value. However, we provide a convergence proof for SGD that shows that it can obtain arbitrary low loss values. Furthermore, we use the results of Lyu & Li (2020) and Ji & Telgarsky (2020) to obtain a more fine grained analysis of the inductive bias of gradient flow for linear teachers. Other works considered the inductive bias of infinite two-layer networks (Chizat & Bach, 2020, 2018; Wei et al., 2019; Mei et al., 2018). Our results hold for networks of any size. An inductive bias towards clustered solutions has been observed in Brutzkus & Globerson (2019) and proved for a simple setup with nonlinear data.

Fully connected networks were also analyzed via the NTK approximation (Du et al., 2019, 2018; Arora et al., 2019; Ji & Telgarsky, 2019b; Cao & Gu, 2019; Jacot et al., 2018; Fiat et al., 2019; Allen-Zhu et al., 2019; Li & Liang, 2018; Daniely et al., 2016). However other works (Yehudai & Shamir, 2019; Daniely & Malach, 2020) have highlighted limitations of the NTK framework, suggesting that it does not accurately model neural networks as they are used in practice. Our convergence analysis in Section 4 holds for any initialization scale and network size and therefore goes beyond the NTK analysis.

Recently, Li et al. (2020) analyzed two-layer networks beyond NTK in the case of Gaussian inputs and squared loss. We assume linearly separable inputs and the cross entropy loss. Allen-Zhu & Li (2019) analyze a three layer ResNet and provide generalization guarantees for sufficiently wide networks in a regression setting. Woodworth et al. (2020) study the inductive bias of gradient methods for a simplified nonlinear model.

Preliminaries

Notations: We use $||\cdot||$ to denote the $L^{2}$ norm on vectors and Frobenius norm on matrices. For a vector ${\bm{v}}$ we denote $\hat{{\bm{v}}}=\frac{{\bm{v}}}{\left\|{\bm{v}}\right\|}$.

It is easy to see that such a network is as expressive as a standard two-layer network where the second layer vector is not fixed (Brutzkus et al., 2018). Furthermore, the assumption that the second layer is fixed is common in previous works (e.g., see Du et al., 2018; Brutzkus et al., 2018; Ji & Telgarsky, 2019b). We denote row $i$ of ${\bm{W}}$ by ${\bm{w}}^{(i)}$ and row $k+i$ by ${\bm{u}}^{(i)}$ for $1\leq i\leq k$. We say that ${\bm{w}}^{(i)}$ are the ${\bm{w}}$ neurons and ${\bm{u}}^{(i)}$ are the ${\bm{u}}$ neurons. Then the network is given by:

Since we use a positive homogeneous activation (Leaky ReLU) the network we consider with $2k$ hidden neurons is as expressive as networks with k hidden neurons and any vector $v$ in the second layer as seen at (brutzkus2017sgd). Hence, we can fix the second layer without limiting the expressive power of the two-layer network. Although it is relatively simpler than the case where the second layer is not fixed, the effect of over-parameterization can be studied in this setting as well.

Optimization Algorithm: The training-loss mimimization optimization problem is to find:

then the update at iteration $t$ is given by

[AB: Define SGD in epochs with sampling without replacement, to fit the setting of Theorem 4.1]

where $||{\bm{W}}||$ is the Frobenius norm of ${\bm{W}}$.

The smoothed margin is defined as:See Remark A.4. in Lyu & Li (2020).

From Lyu & Li (2020) and Ji & Telgarsky (2020) it follows that gradient flow converges to KKT points of the network margin maximization problem (see Supplementary for details). Here we will use this result in Section 6 to characterize the linear decision boundaries of learned networks.

Risk Convergence

We next prove that for any $\varepsilon>0$, SGD converges to $\varepsilon$ empirical-loss (see Eq. (2)) within $O\left(\frac{n^{4}}{\varepsilon^{2}}\right)$ updates.

Define $M(n,\epsilon)=\frac{Cn^{4}}{\varepsilon^{2}}$, where $C$ is a constant that depends polynomially on $R_{x},\alpha,R_{0},k,\eta,v$ and $||{\bm{w}}^{*}||$.In some cases the polynomial dependence is on the inverse of the parameter, e.g., $\frac{1}{\eta}$. See the supplementary for the exact definition of $M(n,\epsilon)$.

The following theorem states that SGD will converge to $\epsilon$ loss within $M(n,\varepsilon)$ updates.

We note that the convergence analysis holds for any $\eta>0$. This is in line with other analyses of learning linearly separable data, which show that convergence holds for any $\eta>0$ (Brutzkus et al., 2018). We next briefly sketch the proof of Theorem 4.1. The full proof is deferred to the supplementary.

Our proof is based on the proof for the hinge loss in Brutzkus et al. (2018) with several novel ideas that enable us to show convergence for the cross entropy loss.

In the case of cross entropy the notation of mistake is not the same as in hinge loss, since every point has a non-zero update, therefor we can’t hope to use a standard perceptron proof.

Next we assume that there exists at least one mistake (as defined above) at every epoch until some time $T=nN_{e}$ where $N_{e}$ is the number of epochs (an epoch based approach as opposed to the online setting of the standard perceptron proof).

Weight Clustering and Linear Separation

As shown in Figure 1, learning with SGD can result in a linear decision boundary, despite the existence of zero-loss solutions that are highly non-linear. In what follows, we provide theoeretical and empirical insights into why an approximately linear boundary is learned.

We next show a nice property of Leaky-ReLU networks that can explain why they converge to linear decision boundaries. Assume that a learned network in Eq. (1) is such that all of its ${\bm{w}}$ neurons form a ball of “small” radius (i.e., they are well clustered) and likewise all the ${\bm{u}}$ neurons (see Figure 1 and Figure 2 for simulations that show such a case). Then, as we show in Theorem 5.1, this implies that the resulting decision boundary will be approximately linear. Later, we give further empirical and theoretical support that learned networks indeed have this clustering structure, and together with Theorem 5.1 this explains the approximate linearity.

Consider the network in Eq. (1). Denote $\overline{{\bm{w}}}=\frac{1}{k}{\sum_{i=1}^{k}{\bm{w}}^{(i)}}$ and $\overline{{\bm{u}}}=\frac{1}{k}{\sum_{i=1}^{k}{\bm{u}}^{(i)}}$. Also, let $r$ denote the maximum radius of the positive and negative weights around their averages. Namely:

The following result says that the decision boundary will be linear except for a region whose size is determined by $r$.

Consider the linear classifier $f({\bm{x}})=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)$. Then $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=f({\bm{x}})$ for all ${\bm{x}}$ such that $|(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}|\geq 2r||{\bm{x}}||$.

The theorem has a simple intuitive implication. The smaller $r$ is, the closer the classifier is to linear. In particular when $r=0$ the classifier is exactly linear.

An alternative interpretation of the theorem comes from rewriting the condition as:

Namely that linearity holds whenever the absolute value of the cosine of the angle between ${\bm{x}}$ and $\overline{{\bm{w}}}-\overline{{\bm{u}}}$ is greater than $\frac{r}{\|\overline{{\bm{w}}}-\overline{{\bm{u}}}\|}$.

We note that the proof strongly relies on two assumptions. The first is that the activation function is Leaky ReLU. The result is not true for ReLU networks (see supplementary for an example). The second is that the clusters correspond to the ${\bm{w}}$ and ${\bm{u}}$ sets of neurons.

Theorem 5.1 states that if neurons cluster, the resulting decision boundary will be approximately linear. But do neurons actually cluster in practice, and what is the resulting $r$? In Figure 2 we show the value of $r$ during training. We use this $r$ to calculate the linear regime in Theorem 5.1 and the fraction of train and test points that fall outside this regime. It can be seen that for small initialization, this fraction converges to zero, implying that the learned classifiers are effectively linear over the data. Additional experiments in the supplementary provide support for the neurons being tightly clustered and $r$ being very small.

Theorem 5.1 shows that a well clustered network leads to a linear decision boundary. However, it does not imply that the network output itself is a linear function of the input. Figure 3 provides a nice illustration of this fact.

In Figure 2 we report results indicating that this is the case.

examine the clusterization for various datasets. The first figure Figure LABEL:fig:Non_Linear_data_points_ratio we measure the percentage of the data points which violate our margin condition. We can see that at convergence the vast majority of data points are in the linear regime. In the second figure Figure LABEL:fig:Minimal_angle_with_separator we see the clusterization of the neurons, the ratio $\frac{r}{||\overline{{\bm{w}}}-\overline{{\bm{u}}}||}$ is approximately zero, and therefore the only directions which are not necessarily in the linear domain are only those which are nearly orthogonal to the separator $\overline{{\bm{w}}}-\overline{{\bm{u}}}$

On Conditions for Convergence to Clustered Solutions

Figure 1 suggests that gradient methods converge to a network with a linear decision boundary when trained on linearly separable data. Understanding when this occurs is important, because a model with a linear decision boundary has good generalization guarantees.For example, standard VC bounds imply $O(\sqrt{d/n})$ sample complexity in this case.

In the previous section we saw that clustering of neurons to two directions implies that the network has an approximate linear decision boundary. Therefore, this reduces the problem of proving that the network has a linear decision boundary to proving that the network neurons are well clustered. It remains to show under which conditions gradient methods converge to clustered solutions.

Providing an end-to-end analysis which shows that gradient methods converge to clustered solutions is a major challenge. In this section we provide initial results for tackling this problem. In Section 6.1 we derive a novel condition on the optimization trajectory which implies that the network converges to a clustered solution and therefore to a linear decision boundary. In Section 6.2 we study a special case where a more fine-grained characterization of the linear decision boundary can be derived using a convex optimization program. Finally, we empirically validate our findings in Section 6.3.

To obtain the results in this section, we apply recent results of Lyu & Li (2020) and (Ji & Telgarsky, 2020) and therefore make the same assumptions presented in these papers. Specifically, we assume that we run gradient flow (GF) as defined in Section A. We further assume that we are in the late phase of training:

We note that by the results in Section 4, SGD can attain the loss value in Assumption 6.1. However, in this section we need this assumption because we consider gradient flow and not SGD.

We first observe that using Theorem 5.1 we can conclude that when the neurons are perfectly clustered around two directions (i.e., $r=0$), the decision boundary is linear. We formally define this below.

A network $N_{{\bm{W}}}({\bm{x}})$ is perfectly clustered if for all $1\leq i,j\leq k$ it holds that: ${\bm{w}}^{(i)}={\bm{w}}^{(j)}$ and ${\bm{u}}^{(i)}={\bm{u}}^{(j)}$.

By applying Theorem 5.1 with $r=0$, we have:

For completeness we provide a proof in the supplementary (this result is easier to prove directly than Theorem 5.1).

Note that the value $c^{{\bm{w}}}_{i}$ determines the agreement of the ${\bm{w}}$ neurons on the point ${\bm{x}}_{i}$. Indeed, if $c^{{\bm{w}}}_{i}=1$, then for $1\leq l\leq k$ it holds that $\hat{{\bm{w}}}^{(l)}\cdot{\bm{x}}_{i}\geq\beta$. Similarly, $c^{{\bm{u}}}_{i}$ determines the agreement of the ${\bm{u}}$ neurons on ${\bm{x}}$.

Importantly, if a network is in an NAR then its neurons can be “far” from being perfectly clustered. Namely, the angles between the normalized weights of different neurons can be relatively large. Next, we show a non-trivial fact: if gradient flow enters an NAR at some time $T_{NAR}$ and stays in it, then it will converge to a perfectly clustered network.

Assume that Assumption 6.1 holds and consider the NAR regime ${\mathcal{N}}$ with parameters $\left(\beta,{\bm{c}}^{{\bm{w}}},{\bm{c}}^{{\bm{u}}}\right)$. Assume that there exists a time $T_{NAR}\geq t_{0}$ such that for all $t\geq T_{NAR}$ it holds that $\overrightarrow{{\bm{W}}}\in{\mathcal{N}}$. Then, gradient flow converges to a solution in ${\mathcal{N}}$ and at convergence the network with normalized parameters $N_{\hat{{\bm{W}}}}({\bm{x}})$ is perfectly clustered.

Theorem 6.1 says that if training is such that the trajectory enters an NAR and never leaves it, then the network will become perfectly clustered. The proof uses results from Lyu & Li (2020) and Ji & Telgarsky (2020) that together guarantee convergence of gradient flow to a KKT point of a minimum norm optimization problem. The theorem then follows from a simple observation that in an NAR, the KKT conditions imply that the network is perfectly clustered. The proof is in the supplementary.

Using Corollary 6.1 we immediately obtain the following.

Under the assumptions in Theorem 6.1, GF converges to a network with a linear decision boundary.

Therefore, we see that if a network is at an NAR from some time $T_{NAR}$, then it will converge to a solution with a linear decision boundary. The question that remains is whether networks indeed converge to an NAR and remain there.

2 The Perfect Agreement Regime

To better understand convergence to NARs, in this section we study a specific NAR for which we provide a more fine-grained analysis. We identify conditions on the training data and optimization trajectory that imply that gradient flow converges to an NAR which we call the Perfect Agreement Regime (PAR). Using Theorem 6.1 and results from Lyu & Li (2020); Ji & Telgarsky (2020), we provide a complete characterization of the weights that gradient flow converges to in this case. Admittedly, the conditions on the data and optimization trajectory are fairly strong. Nonetheless, we show that our theoretical results accurately predict the dynamics that we observe in experiments. Indeed, in Section 6.3 we show empirically that for certain linearly separable datasets, gradient flow converges to a solution in the PAR which is in agreement with our results.

In the PAR, each neuron classifies the data perfectly. Namely, all ${\bm{w}}$ neurons classify like the ground truth ${\bm{w}}^{*}$, and all ${\bm{u}}$ neurons classify like $-{\bm{w}}^{*}$. Formally, let ${\bm{y}}=\left(y_{1},...,y_{n}\right)$. Then PAR is defined as follows.

Note that the fact that a network is in PAR does not mean that ${\bm{w}}_{i}=-{\bm{u}}_{j}$. Indeed, PAR only requires that ${\bm{w}}_{i}$ and $-{\bm{u}}_{j}$ both correctly classify the training set.

Next, we provide conditions under which a network will converge to a PAR. The conditions require a lower bound on the network smoothed margin (Eq. (5)), as well as a separability condition on the data. To define the separability condition we consider the following:

There exists an NAR ${\mathcal{N}}$ and $T_{NAR}\geq t_{0}$ such that for all $t\geq T_{NAR}$ it holds that $\overrightarrow{{\bm{W}}}_{t}\in{\mathcal{N}}$.

Then ${\mathcal{N}}$ is a $\mbox{PAR}(\beta)$ for all $t>T_{Margin}$, and there exists $\delta_{w},\delta_{u}>0$ such that gradient flow converges to a network whose normalized version is perfectly clustered with neuron directions $\hat{{\bm{w}}},\hat{{\bm{u}}}$, where $\left(\delta_{w}\hat{{\bm{w}}},\delta_{u}\hat{{\bm{u}}}\right)$ is the solution to the following convex optimization problem:

We first comment on the assumptions. The first two assumptions are the same assumptions on the optimization trajectory as in Theorem 6.1. Assumption 3 is another assumption on the trajectory that says that sufficiently large smoothed margin is achieved at some stage of the optimization. We note that the lower bound on the smoothed margin can be made small by considering a small $\alpha$.

Assumption 4 refers to the training set. Informally, it corresponds to requiring that the two classes are approximately symmetric with respect to the origin. The next lemma shows that a certain symmetric training set satisfies Assumption 4:

The proof is given in the supplementary. This example suggests that we should observe PAR in symmetric distributions, which produce approximately symmetric training sets. Indeed, we empirically show in Section 6.3 that gradient flow converges to a solution in PAR for a distribution with two symmetric Gaussians. We note that this example shows that Assumption 4 is independent of the maximum margin attainable on the training set. Indeed, by scaling the points, we can obtain any margin and still satisfy the assumption.

We prove Theorem 6.2 in the supplementary, and provide a sketch next. First, we use Theorem 6.1 to show that gradient flow converges to an NAR and the neurons are clustered. Then we show that under Assumption 3 and using the monotonicity of the smoothed margin (Eq. (6)), by Lyu & Li (2020), all ${\bm{w}}$ neurons classify the positive points correctly and all ${\bm{u}}$ neurons classify the negative points correctly for all $t>T_{Margin}$. Then, using Assumption 4 we show that the solution is in PAR. Finally, we use results of Lyu & Li (2020) to show that the network directions solve the convex optimization problem in the theorem.

3 Experiments

In Theorem 6.2 we show that when learning enters the PAR regime the solution will be given by Eq. (8). We performed experiments in several settings that show the above behavior is observed in practice when classes are sampled from Gaussians. Figure 4 shows the decision boundary (Figure 4(a)) and learned weights (Figure 4(b)), for learning from points sampled from two classes corresponding to Gaussians. The figure also shows the PAR predictions for the decision boundary and learned weights, and these show excellent agreement with the empirical results. We have also verified that in this case convergence is indeed to a PAR solution. We performed such experiments also for higher dimensional settings, and the results are in the supplementary. Finally, note that we do not expect learning to always converge to a PAR. In the supplementary we show an example where this does not happen.

Conclusions

Optimization and generalization are closely coupled in deep-learning. Yet both are little understood even for simple models. Here we consider perhaps the simplest “teacher” model where the ground truth is linear. We prove that cross-entropy can be globally minimized by SGD, despite the non-convexity of the loss, and for any initialization scale. We are not aware of any such result for non-linear networks (for example NTK optimization results require large initialization scale, and sufficiently wide networks (Ji & Telgarsky, 2019b)). Our novel proof technique analyzes SGD in an offline setting and uses the notion of loss-violation per epoch, which we believe could be useful elsewhere.

In our setting, small initialization scale leads empirically to approximately linear decision boundaries. We prove that such boundaries are obtained when neurons with same output-weight sign are clustered. Empirically we show that such clustering indeed occurs. Moreover, we provide sufficient conditions for converging to such clustered solutions.

Several open questions remain. The first is reducing the assumptions when proving convergence to a clustered solution. Another interesting direction is extending our results to simple non-linear teachers.

Acknowledgements

This research is supported by the European Research Council (ERC) under the European Unions Horizon 2020 research and innovation programme (grant ERC HOLI 819080) and by the Yandex Initiative in Machine Learning at Tel Aviv University. AB is supported by the Google Doctoral Fellowship in Machine Learning.

References

Appendix A Gradient Flow Definitions

Appendix B Proof of Theorem 4.1

Throughout this proof we will sometimes use the notation $\langle{\bm{x}},{\bm{y}}\rangle$ as the dot product between two vectors ${\bm{x}}$ and ${\bm{y}}$ for readability purposes.

Then, from Cauchy-Schwartz inequality we have:

Recall we define: $N_{{\bm{W}}}({\bm{x}})=v\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})-v\sum\limits_{j=1}^{k}\sigma({\bm{u}}^{(j)}\cdot{\bm{x}})$.

We consider minimizing the objective function:

Optimizing by SGD yields the following update rule:

where ${\bm{W}}_{t}=({\bm{w}}_{t}^{(1)},...,{\bm{w}}_{t}^{(k)},{\bm{u}}_{t}^{(1)},...,{\bm{u}}_{t}^{(k)})$.

For every neuron we get the following updates:

where $p_{t}^{(j)}:=\sigma^{\prime}({\bm{w}}_{t}^{(j)}\cdot{\bm{x}}_{t+1});q_{t}^{(j)}:=\sigma^{\prime}({\bm{u}}_{t}^{(j)}\cdot{\bm{x}}_{t+1})$.

Next we will show recursive upper bounds for $\displaystyle G({\bm{W}}_{t})$ and $\displaystyle F({\bm{W}}_{t})$.

Where we used the inequalities $\displaystyle\langle y_{t}{\bm{x}}_{t},{\bm{w}}^{*}\rangle\geq 1$ and $q_{t}^{(j)},p_{t}^{(j)}\geq\alpha$.

For an upper bound on $G({\bm{W}}_{t})$ we use the following inequalities (which hold for the cross entropy loss):

Using this recursively up until $T=nN_{e}$ we get:

This implies that (recursively using Eq. (18)):

where $N_{e}$ is the number of epochs and $n$ the number of training points, $T=nN_{e}$.

Now, using the Cauchy-Schwartz, Eq. (16) and Eq. (19) we have:

Using $\sqrt{a+b}\leq\sqrt{a}+\sqrt{b}$ the above implies:

Now using $\left|\left|{\bm{w}}_{0}^{(i)}\right|\right|,\left|\left|{\bm{u}}_{0}^{(i)}\right|\right|\leq R_{0}$ we get $G({\bm{W}}_{0})\leq\sqrt{2k}R_{0}$.

Noting that $\left\|\overrightarrow{{\bm{W}}}^{*}\right\|=\sqrt{2k}||{\bm{w}}^{*}||$ and that $N_{e}=\frac{T}{n}$, we get :

Therefore, we have an inequality of the form:

where $\displaystyle a=\frac{2k\eta v\alpha(1-e^{-\varepsilon_{0}})}{n},b=\sqrt{4k^{2}\eta^{2}v^{2}R_{x}^{2}+4k\eta}||{\bm{w}}^{*}||$ and $c=4kR_{0}||{\bm{w}}^{*}||$.

By inspecting the roots of the parabola $P(X)=x^{2}-\frac{b}{a}x-\frac{c}{a}$ we conclude that:

By the inequality $1-e^{-x}>\frac{x}{1+x}$ for $x>0$ (which is equivalent to $\frac{1}{1-e^{-x}}<\frac{x+1}{x}$), with $x=\varepsilon_{0}>0$ we get $\frac{1}{1-e^{-\varepsilon_{0}}}<\frac{\varepsilon_{0}+1}{\varepsilon_{0}}=1+\frac{1}{\varepsilon_{0}}$. Therefore for $\beta>0$ (all arguments are positive):

By using the above inequality we can reach a polynomial bound on $T$:

For any iteration $(i_{e}-1)n+1\leq t\leq i_{e}n$ and $1\leq s\leq n$ we have:

Therefore, if we set $\varepsilon_{0}=\frac{\varepsilon}{1+2v^{2}R_{x}^{2}\eta kn}$ in Eq. (26) we’ll reach our result.

Setting this $\varepsilon_{0}$ at Eq. (B) leads to:

Appendix C Proof of Theorem 5.1

Before we start proving the main theorem we will prove some useful lemmas and corollaries.

if $|(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}|\geq 2r||{\bm{x}}||$ then $|\overline{{\bm{w}}}\cdot{\bm{x}}|\geq r||{\bm{x}}||\lor|\overline{{\bm{u}}}\cdot{\bm{x}}|\geq r||{\bm{x}}||$.

Assume in contradiction that $|\overline{{\bm{w}}}\cdot{\bm{x}}|<r||{\bm{x}}||\land|\overline{{\bm{u}}}\cdot{\bm{x}}|<r||{\bm{x}}||$. then by the triangle inequality and the Cauchy-Shwartz inequality we’ll get:

$|(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}|\leq|\overline{{\bm{w}}}\cdot{\bm{x}}|+|\overline{{\bm{u}}}\cdot{\bm{x}}|<r||{\bm{x}}||+r||{\bm{x}}||=2r||{\bm{x}}||$ in contradiction to the assumption $|(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}|\geq 2r||{\bm{x}}||$. ∎

Next, we prove the following lemma, which will be used throughout the proof of the main theorem. The lemma ties the dot products with the center of the cluster to the dot products with the individual neurons:

Let’s assume that $\overline{{\bm{w}}}\cdot{\bm{x}}\geq r||{\bm{x}}||$, therefore $\forall 1\leq j\leq k:\ \ {\bm{w}}^{(j)}\cdot{\bm{x}}=({\bm{w}}^{(j)}-\overline{{\bm{w}}})\cdot{\bm{x}}+\overline{{\bm{w}}}\cdot{\bm{x}}\geq-||{\bm{w}}^{(j)}-\overline{{\bm{w}}}||\cdot||{\bm{x}}||+r||{\bm{x}}||>-r||{\bm{x}}||+r||{\bm{x}}||=0$ where we had used Cauchy-Shwartz inequality and that $||{\bm{w}}^{(j)}-\overline{{\bm{w}}}||<r$.

If $\overline{{\bm{w}}}\cdot{\bm{x}}\leq-r||{\bm{x}}||$, $\forall 1\leq j\leq k:\ \ {\bm{w}}^{(j)}\cdot{\bm{x}}=({\bm{w}}^{(j)}-\overline{{\bm{w}}})\cdot{\bm{x}}+\overline{{\bm{w}}}\cdot{\bm{x}}<||{\bm{w}}^{(j)}-\overline{{\bm{w}}}||\cdot||{\bm{x}}||-r||{\bm{x}}||<r||{\bm{x}}||-r||{\bm{x}}||=0$ the same derivation would work for ${\bm{u}}$. ∎

We are now ready to move forward with proving the main lemma.

Now we will show that $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)$ in each region, from which the claim follows.

If ${\bm{x}}\in C_{+}^{+}$ then $N_{{\bm{W}}}({\bm{x}})=v\left(\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})-\sigma({\bm{u}}^{(j)}\cdot{\bm{x}})\right)=v\left(\sum\limits_{j=1}^{k}{\bm{w}}^{(j)}-{\bm{u}}^{(j)}\right)\cdot{\bm{x}}$ and therefore $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)$.

If ${\bm{x}}\in C_{-}^{-}$ then $N_{{\bm{W}}}({\bm{x}})=v\left(\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})-\sigma({\bm{u}}^{(j)}\cdot{\bm{x}})\right)=\alpha v\left(\sum\limits_{j=1}^{k}{\bm{w}}^{(j)}-{\bm{u}}^{(j)}\right)\cdot{\bm{x}}$ and therefore $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)$

If ${\bm{x}}\in C_{+}^{-}$ then both $N_{{\bm{W}}}({\bm{x}})=v\left(\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})-\sigma({\bm{u}}^{(j)}\cdot{\bm{x}})\right)=v\left(\sum\limits_{j=1}^{k}{\bm{w}}^{(j)}\cdot{\bm{x}}-\alpha{\bm{u}}^{(j)}\cdot{\bm{x}}\right)>0$ and $\overline{{\bm{w}}}\cdot{\bm{x}}-\overline{{\bm{u}}}\cdot{\bm{x}}>0$. Therefore, $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)$.

If ${\bm{x}}\in C_{-}^{+}$ then both $N_{{\bm{W}}}({\bm{x}})=v\left(\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})-\sigma({\bm{u}}^{(j)}\cdot{\bm{x}})\right)=v\left(\sum\limits_{j=1}^{k}\alpha{\bm{w}}^{(j)}\cdot{\bm{x}}-{\bm{u}}^{(j)}\cdot{\bm{x}}\right)<0$ and $\overline{{\bm{w}}}\cdot{\bm{x}}-\overline{{\bm{u}}}\cdot{\bm{x}}<0$. Therefore, $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)$.

We are left with proving $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)$ holds when exactly one condition holds ,i.e., either $|\overline{{\bm{w}}}\cdot{\bm{x}}|\geq r||{\bm{x}}||$ or $|\overline{{\bm{u}}}\cdot{\bm{x}}|\geq r||{\bm{x}}||$.

and similarly our decision boundary is linear for points in which our condition only holds for $\overline{{\bm{w}}}$:

i.e. our condition only holds for $\overline{{\bm{u}}}$.

There are two cases, and we’ll prove the result for each of them:

If $\overline{{\bm{u}}}\cdot{\bm{x}}\geq r||{\bm{x}}||$:

In this case $N_{{\bm{W}}}({\bm{x}})=v\left(\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})-\sigma({\bm{u}}^{(j)}\cdot{\bm{x}})\right)=v\left(\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})-k\overline{{\bm{u}}}\cdot{\bm{x}}\right)$.

Next, for any ${\bm{x}}$ in the domain, we’ll denote $J_{+}^{w}({\bm{x}})\coloneqq\{j|{\bm{w}}^{(j)}\cdot{\bm{x}}>0\}$ and $k_{+}^{w}({\bm{x}})\coloneqq|J_{+}^{w}({\bm{x}})|$ similarly $J_{-}^{w}({\bm{x}})=\{j|{\bm{w}}^{(j)}\cdot{\bm{x}}<0\}$ and $k_{-}^{w}({\bm{x}})\coloneqq|J_{-}^{w}({\bm{x}})|$. Using these definitions, our network has the following form:

Next, we bound $\forall j\ |{\bm{w}}^{(j)}\cdot{\bm{x}}|=|({\bm{w}}^{(j)}-\overline{{\bm{w}}}+\overline{{\bm{w}}})\cdot{\bm{x}}|\leq||{\bm{w}}^{(j)}-\overline{{\bm{w}}}||\cdot||{\bm{x}}||+|\overline{{\bm{w}}}\cdot{\bm{x}}|<2r||{\bm{x}}||$ where we used $||{\bm{w}}^{(j)}-\overline{{\bm{w}}}||<r$ and $|\overline{{\bm{w}}}\cdot{\bm{x}}|<r||{\bm{x}}||$.

Now, if $(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}\geq 2r||{\bm{x}}||>0$ we get that $N_{{\bm{W}}}({\bm{x}})=v\left(k(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}-(1-\alpha)\sum\limits_{j_{-}\in J_{-}^{w}({\bm{x}})}{\bm{w}}^{(j_{-})}\cdot{\bm{x}}\right)>v\left(2r||{\bm{x}}||k-2r||{\bm{x}}||k_{-}^{w}({\bm{x}})(1-\alpha)\right)>0$ since $(1-\alpha)<1$ and $k_{-}^{w}({\bm{x}})\leq k$ and therefore $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)=1$ for this case.

If $(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}\leq-2r||{\bm{x}}||<0$ we get that $N_{{\bm{W}}}({\bm{x}})=v\left(k(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}-(1-\alpha)\sum\limits_{j_{-}\in J_{-}^{w}({\bm{x}})}{\bm{w}}^{(j_{-})}\cdot{\bm{x}}\right)<v\left(-2r||{\bm{x}}||k+2r||{\bm{x}}||k_{-}^{w}({\bm{x}})(1-\alpha)\right)<0$ since $(1-\alpha)<1$ and $k_{-}^{w}({\bm{x}})\leq k$. Therefore, we get that $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)=-1$ in this case.

If $\overline{{\bm{u}}}\cdot{\bm{x}}\leq-r||{\bm{x}}||$:

First, we notice that $(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}>-r||{\bm{x}}||+r||{\bm{x}}||=0$ so $\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)=1$ again we use Lemma. (C.1) and from our assumption $\overline{{\bm{u}}}\cdot{\bm{x}}\leq-r||{\bm{x}}||$ we have $\forall 1\leq j\leq k\ {\bm{u}}^{(j)}\cdot{\bm{x}}<0$ and we can see that our network takes the form: $N_{{\bm{W}}}({\bm{x}})=v\left(\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})-\sigma({\bm{u}}^{(j)}\cdot{\bm{x}})\right)=v\left(\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})-\alpha\cdot k\overline{{\bm{u}}}\cdot{\bm{x}}\right)\geq v\left(\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})+\alpha kr||{\bm{x}}||\right)$. Next, we prove the following lemma:

If $|\overline{{\bm{w}}}\cdot{\bm{x}}|<r||{\bm{x}}||$ then $\alpha\cdot k\cdot r||{\bm{x}}||>-\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})$.

Let’s assume by contradiction that $-\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})\geq\alpha\cdot k\cdot r||{\bm{x}}||$. We notice that regardless of the sign of the dot product $\forall j:-\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})\leq-\alpha{\bm{w}}^{(j)}\cdot{\bm{x}}$ so we have $-\alpha\sum\limits_{j=1}^{k}{\bm{w}}^{(j)}\cdot{\bm{x}}\geq-\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})\geq\alpha\cdot k\cdot r||{\bm{x}}||$, which leads to $-\alpha k\overline{{\bm{w}}}\cdot{\bm{x}}\geq\alpha\cdot k\cdot r||{\bm{x}}||$ (where we used the definition of $\overline{{\bm{w}}}$) finally we reach $\overline{{\bm{w}}}\cdot{\bm{x}}\leq-r||{\bm{x}}||$. This contradicts $|\overline{{\bm{w}}}\cdot{\bm{x}}|<r||{\bm{x}}||$. ∎

Therefore, we have $-\sum\limits_{j=1}^{k}\sigma({\bm{w}}^{(j)}\cdot{\bm{x}})<\alpha\cdot k\cdot r||{\bm{x}}||$ and $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}})}\right)=1$ as desired.

If $\overline{{\bm{w}}}\cdot{\bm{x}}\geq r||{\bm{x}}||$:

Through a similar derivation for the case of $\overline{{\bm{u}}}\cdot{\bm{x}}\geq r||{\bm{x}}||$, our network has the following form:

where $J_{-}^{u}({\bm{x}})\coloneqq\{j|{\bm{u}}^{(j)}\cdot{\bm{x}}<0\},J_{+}^{u}({\bm{x}})\coloneqq\{j|{\bm{u}}^{(j)}\cdot{\bm{x}}>0\}$ and $k_{-}^{u}({\bm{x}})=|J_{-}^{u}({\bm{x}})|,k_{+}^{u}({\bm{x}})=|J_{+}^{u}({\bm{x}})|$.

If $(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}\geq 2r||{\bm{x}}||>0$ then $N_{{\bm{W}}}({\bm{x}})=v\left(k(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}+(1-\alpha)\sum\limits_{j_{-}\in J_{-}^{u}({\bm{x}})}{\bm{u}}^{(j_{-})}\cdot{\bm{x}}\right)\geq v\left(2kr||{\bm{x}}||-2r||{\bm{x}}||(1-\alpha)k_{-}^{u}({\bm{x}})\right)>0$ (because $(1-\alpha)<1$ and $k_{-}^{u}({\bm{x}})\leq k$) and $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)=1$ (where we used the fact that $\forall j:|{\bm{u}}^{(j)}\cdot{\bm{x}}|<2r||{\bm{x}}||$ which follows from $|\overline{{\bm{u}}}\cdot{\bm{x}}|<r||{\bm{x}}||$ and $||{\bm{u}}^{(j)}-\overline{{\bm{u}}}||<r$).

If $(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}\leq-2r||{\bm{x}}||<0$ we get that $N_{{\bm{W}}}({\bm{x}})\leq v\left(-2r||{\bm{x}}||k+2r||{\bm{x}}||(1-\alpha)k_{-}^{u}({\bm{x}})\right)<0$ and $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)=1$.

If $\overline{{\bm{w}}}\cdot{\bm{x}}\leq-r||{\bm{x}}||$:

We again use Lemma. (C.1) which yields from $\overline{{\bm{w}}}\cdot{\bm{x}}\leq-r||{\bm{x}}||$ that $\forall 1\leq j\leq k\ {\bm{w}}^{(j)}\cdot{\bm{x}}<0$ and we can see that our network takes the form:

If $-\sum\limits_{j=1}^{k}\sigma({\bm{u}}^{(j)}\cdot{\bm{x}})<\alpha kr||{\bm{x}}||$ we have $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)=-1$ as desired.

The same contradiction proof from $\overline{{\bm{u}}}\cdot{\bm{x}}\leq-r||{\bm{x}}||$ segment above (Lemma. (C.2)) would show

$-\sum\limits_{j=1}^{k}\sigma({\bm{u}}^{(j)}\cdot{\bm{x}})<\alpha\cdot k\cdot r||{\bm{x}}||$ (just exchange $w$ and $u$) and we’ll get $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)=-1$.

We can now combine Corollary (C.1), Proposition C.1 and Proposition C.2 and prove Theorem. (5.1):

Therefore, overall for $|(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}|\geq 2r||{\bm{x}}||$ we get $\mbox{sign}\left({N_{{\bm{W}}}({\bm{x}})}\right)=\mbox{sign}\left({(\overline{{\bm{w}}}-\overline{{\bm{u}}})\cdot{\bm{x}}}\right)$ as required.

Since the network is perfectly clustered, the corollary follows by Proposition (C.1) with $r=0$.

Appendix D Additional Experiments - Linear Decision Boundary

In this section we provide additional empirical evaluations of the decision boundary that SGD converges to in our setting.

Theorem. (5.1) addresses the case of Leaky ReLU activation. Here we show that the result is indeed not true for ReLU networks. We compare two perfectly clustered networks (i.e., each with two neurons) one with a Leaky ReLU activation and the other with a ReLU activation. Figure 5 shows a decision boundary for a two neuron network, in the case of Leaky ReLU (Figure 5(a)) and ReLU (Figure 5(b)). It can be seen that the leaky ReLU indeed provides a linear decision boundary, as predicted by Theorem 5.1, whereas the ReLU case is non-linear (we explicitly show the regime where the network output is zero. This can be orange or blue, depending on whether zero is given label positive or negative. In any case the resulting boundary is non-linear).

D.2 MNIST - Linear Regime

In Figure 2 in the main text we saw how for MNIST digit pairs (0,1) and (3,5) the network enters the linear regime at some point in the training process. In Figure 6 we see the robustness of this behavior across the MNIST data-set by showing the above holds for more pairs of digits.

D.3 Clustering of Neurons - Empirical Evidence

In Section 5 in the main text and Figure 6 above, we saw that learning converges to a linear decision boundary on the train and test points. Theorem. (5.1) suggests that this will happen if neurons are well clustered (in the ${\bm{w}}$ and ${\bm{u}}$ groups). Here we show that indeed clustering occurs.

We consider two different measures of clustering. The first is the ratio $\frac{r}{||\overline{{\bm{w}}}-\overline{{\bm{u}}}||}$, and the second is the maximum angle between the neurons of the same type (i.e., the maximal angle between vectors in the same cluster). Figure 7 shows these two measures as a function of the training epochs. They can indeed be seen to converge to zero, which by Theorem. (5.1) implies convergence to a linear decision boundary.

Appendix E Assumptions for Gradient Flow Analysis

In the paper we use results from (Lyu & Li, 2020) and (Ji & Telgarsky, 2020). Here we show that the assumptions required by these theorems are satisfied in our setup.

The assumptions in (Lyu & Li, 2020) and (Ji & Telgarsky, 2020) are:

. (Regularity). For any fixed $\displaystyle{\bm{x}},\Phi(\cdot;{\bm{x}})$ is locally Lipschitz and admits a chain rule;

. (Homogeneity). There exists $L>0$ such that $\displaystyle\forall\alpha>0:\Phi(\alpha{\bm{W}};{\bm{x}})=\alpha^{L}\Phi({\bm{W}};{\bm{x}});$

There exists $b_{f}\geq 0$ such that $\displaystyle f^{\prime}(q)q$ is non-decreasing for $q\in(b_{f},+\infty)$, and $f^{\prime}(q)q\rightarrow+\infty$ as $q\rightarrow+\infty$.

Let $g:[f(b_{f}),+\infty)\rightarrow[b_{f},+\infty)$ be the inverse function of $f$ on the domain $[b_{f},+\infty).$ There exists $b_{g}\geq max\{2f(b_{f}),f(2b_{f})\},K\geq 1$ such that $\displaystyle g^{\prime}(x)\leq Kg^{\prime}(\theta x)$ and $f^{\prime}(y)\leq Kf^{\prime}(\theta y)$ for all $x\in(b_{g},+\infty),y\in(g(b_{g}),+\infty)$ and $\theta\in[1/2,1)$

We next show that these are satisfied in our setup.

And we showed $\Phi(\cdot;{\bm{x}})$ is globally Lipschitz (and therfor locally Lispchitz). Next for the chain rule, as shown in (Davis et al., 2018) (corollary for deep learning therein), any function definable in an o-minimal structure admits a chain rule. Our network is definable because algebraic, composition, inverse, maximum and minimum operations over definable functions are also definable. Leaky ReLUs are definable as maximum operations over two linear functions (linear functions are definable).and because Leaky ReLUs are definable our network is also definable.

(Homogeneity). It is easy to see from the definition that in our case, the trainable parameters are only the first layer weights and the network $\displaystyle\Phi(\cdot;{\bm{x}})$ is $L=1$ homogeneous.

(Separability). This is Assumption 6.1 in the main text. As we mentioned in the main text, this assumption is satisfied with SGD by Theorem. (4.1).

Appendix F Proof of Theorem 6.1

In this proof we will show that the normalized parameters $\hat{{\bm{W}}}_{t}\coloneqq\frac{{\bm{W}}_{t}}{||{\bm{W}}_{t}||}$ under gradient flow optimization, converges to a solution in ${\mathcal{N}}$ and that the network $N_{\hat{{\bm{W}}}}$ at convergence is perfectly clustered. Under our assumption $\forall t\geq T_{NAR}$ $\hat{{\bm{W}}}_{t}\in{\mathcal{N}}$. From the definition of the NAR it’s easy to see that the NAR is a closed domain. Therefore any limit point of $\hat{{\bm{W}}}_{t}$ is also in the NAR. From Ji & Telgarsky (2020) (Theorem 3.1. therein) we have that the normalized parameters flow converges when using gradient flow. To conclude so far, we had shown that $\hat{{\bm{W}}}_{t}$ converges to a point inside the NAR ${\mathcal{N}}$.

We are left with showing that the limit point of $\underset{t\rightarrow\infty}{\lim}\hat{{\bm{W}}}_{t}\coloneqq\hat{{\bm{W}}}_{*}$ has a perfectly clustered form.

Lyu & Li (2020) (Theorem A.8. therein) shows that every limit point of $\hat{{\bm{W}}}_{t}$ is along the direction of a KKT point of the following optimization problem (P):

where $q_{i}({\bm{W}})=y_{i}N_{{\bm{W}}}({\bm{x}}_{i})$ is the network margin on the sample point $(y_{i},{\bm{x}}_{i})$.It is not hard to see that given that the solution is in an NAR, then this optimization problem is convex.

We are left with showing that at convergence the neurons align in two directions. We will use a characterization of the KKT points of (P) and show that they are perfectly clustered. Since every limit point of the normalized parameters flow is along the direction of a KKT point of (P) that would mean $\hat{{\bm{W}}}_{*}$ has a perfectly clustered form.

A feasible point ${\bm{W}}$ of (P) is a KKT point if there exist $\lambda_{1},\dots,\lambda_{n}\geq 0$ such that:

${\bm{W}}-\sum\limits_{i=1}^{n}\lambda_{i}{\bm{h}}_{i}=0$ for some ${\bm{h}}_{1},\dots,{\bm{h}}_{n}$ satisfying ${\bm{h}}_{i}\in\partial^{\circ}q_{i}({\bm{W}})$

$\forall i\in[n]:\lambda_{i}(q_{i}({\bm{W}})-1)=0$

From Lyu & Li (2020) (Theorem A.8. therein) we know $\exists\beta$ s.t. $\beta\hat{{\bm{W}}}_{*}$ is a KKT point of (P). Since our limit point is in an NAR we don’t need to worry about the non differential points of the network because $\forall 1\leq j\leq k,i\in[n]:{\bm{w}}^{(j)}_{*}\cdot{\bm{x}}_{i}\not=0\land{\bm{u}}^{(j)}_{*}\cdot{\bm{x}}_{i}\not=0$. (where ${\bm{w}}^{(j)}_{*}$ and ${\bm{u}}^{(j)}_{*}$ stands for the ${\bm{w}}$ and ${\bm{u}}$ type neurons of ${\bm{W}}_{*}$, respectively). Therefore the Clarke subdifferential coincides with the gradient in our domain, and we can derive it using calculus rules.

By looking at the gradient of the margin for any point $(y_{i},{\bm{x}}_{i})$:

$\displaystyle\frac{\partial q_{i}({\bm{W}})}{\partial{\bm{w}}^{(j)}}=\frac{y_{i}\partial N_{{\bm{W}}}({\bm{x}}_{i})}{\partial{\bm{w}}^{(j)}}=y_{i}v{\bm{x}}_{i}\sigma^{\prime}({{\bm{w}}^{(j)}}\cdot{\bm{x}}_{i})=y_{i}v{\bm{x}}_{i}\sigma^{\prime}({{\bm{w}}^{(j)}}\cdot{\bm{x}}_{i})$

$\displaystyle\frac{\partial q_{i}({\bm{W}})}{\partial{\bm{u}}^{(j)}}=\frac{y_{i}\partial N_{{\bm{W}}}({\bm{x}}_{i})}{\partial{\bm{u}}^{(j)}}=-y_{i}v{\bm{x}}_{i}\sigma^{\prime}({{\bm{u}}^{(j)}}\cdot{\bm{x}}_{i})=-y_{i}v{\bm{x}}_{i}\sigma^{\prime}({{\bm{u}}^{(j)}}\cdot{\bm{x}}_{i})$

Now using the above gradients implies that: $\partial q_{i}({\bm{W}})=y_{i}v{\bm{x}}_{i}(\overbrace{\sigma^{\prime}({{\bm{w}}^{(1)}}\cdot{\bm{x}}_{i}),\dots,\sigma^{\prime}({{\bm{w}}^{(k)}}\cdot{\bm{x}}_{i})}^{k},\overbrace{-\sigma^{\prime}({{\bm{u}}^{(1)}}\cdot{\bm{x}}_{i}),\dots,-\sigma^{\prime}({{\bm{u}}^{(k)}}\cdot{\bm{x}}_{i})}^{k})$

By the definition of the NAR $\mathcal{N}$ with parameters $(\beta,c^{{\bm{w}}}_{i},c^{{\bm{u}}}_{i})$ the dot product of a point ${\bm{x}}_{i}$ with all neurons of the same type is of the same sign, i.e.:

It follows that for ${\bm{W}}\in\mathcal{N}$, $\partial q_{i}({\bm{W}})=y_{i}\cdot v\cdot{\bm{x}}_{i}(\overbrace{c^{{\bm{w}}}_{i},\dots,c^{{\bm{w}}}_{i}}^{k},\overbrace{-c^{{\bm{u}}}_{i},\dots,-c^{{\bm{u}}}_{i}}^{k})$.

Therefore, by the definition of a KKT point we have:

We can see that the first $k$ entries are equal, as well as the next $k$ entries (equal to each other and not to the first $k$ entries).

Therefore the normalized parameters flow $\hat{{\bm{W}}}_{t}$ converges to a perfectly clustered solution.

Appendix G Proof of Theorem 6.2

We divide the proof of Theorem. (6.2) into two parts. First, we show that the NAR is a PAR, and then we show that if a network enters and remains in the PAR the network weights at convergence are proportional to the solutions of the SVM problem we defined in the main text.

To conclude, we have proven so far for all $t>T_{Margin}$:

Next, under the network being in an NAR assumption we have for all $t>T_{Margin}$:

Thus, for all $t>T_{Margin}$, the network is in PAR$(\beta)$.

G.2 PAR alignment direction

Because the solution is in the PAR($\beta$), the network margins are given as follows for positive points:

Using the above notations, the max margin problem in Lyu & Li (2020) (Theorem A.8. therein) takes the form:

Appendix H Proof of Lemma 6.1

Appendix I Entrance to PAR - High Dimensional Gaussians

We will show that the entrance to the PAR indeed happens empirically for two separable Gaussians. We measure the percentage of neurons which are in the PAR of both types. A ${\bm{w}}$ type neuron is considered in the PAR if it classifies like the ground truth ${\bm{w}}^{*}$. A ${\bm{u}}$ type neuron is considered in the PAR if it classifies like $-{\bm{w}}^{*}$.

The percentage of neurons in the PAR throughout the training process is given in Figure 8. We can see that the network enters the PAR.

Appendix J Entrance to NAR which is not a PAR

In this section we show that learning can enter an NAR which is not a PAR. We sample two antipodal Gaussians and add one outlier positive point. Then for each neuron type (${\bm{w}}$ or ${\bm{u}}$) we measure the maximum amount of data points classification disagreements between neurons of the same type denoted $\max(n_{diff})$ and the percentage of neurons which are in the PAR.

In Figure 9(a) we can see that the network yields $100\%$ prediction accuracy. In Figure 9(b) we can see the directions of the neurons (${\bm{w}}$ type in black and ${\bm{u}}$ type in yellow). In Figure 9(c) we can see that the maximal number of points which neurons of the same type classified differently goes to zero, therefore all neurons of the same type agree on the classification of the data points. In Figure 9(d) we can see that the ratio of ${\bm{w}}$ type neurons which perfectly classifies the data does not increase to $1$ so the network does not enter the PAR.

Appendix K Extension - First Layer Bias Term

and extend our neurons to include a bias term:

This reformulation is equivalent to adding a bias term for every neuron in the first layer, and all of the following results would still hold under the above reformulation.

The proofs of Theorem. (4.1) and Theorem. (5.1) follow exactly if we exchange ${\bm{W}}$ with ${\bm{W}}^{\prime}$ while for the proofs of Theorem. (6.1) and Theorem. (6.2) we use results from (Lyu & Li, 2020) and (Ji & Telgarsky, 2020) that require the model to be homogeneous. Note that if we add a bias in the first layer, the model remains homogeneous and the proofs of Theorem. (6.1) and Theorem. (6.2) still hold for those cases as well.