Random Feature Amplification: Feature Learning and Generalization in Neural Networks

Introduction

A number of recent works have developed optimization and generalization guarantees for neural networks in the ‘neural tangent kernel regime’, namely, where the behavior of the neural network can be well-approximated by the linearization of the network around its random initialization [JGH18, ALS19, Zou+19, Du+19, Aro+19, SJL19]. Although these works provide a deep understanding of the behavior of neural networks in the early stages of training—where the network parameters are close to their initial values—they fail to capture a number of meaningful characteristics of practical neural networks such as the ability to learn features that differ significantly from those found at random initialization [For+20, Lon21]. This points to the need for analyses of neural network training that can characterize how gradient descent is able to learn meaningful features.

A remarkable feature of neural networks is that despite their capacity to overfit, when trained by gradient descent they are capable of feature-learning even when there is significant label noise in the training data. Label noise is a common feature in modern machine learning datasets like ImageNet [Sha+20], and moreover, some of the most interesting behaviors of neural networks have been observed when they are trained on datasets with artificially introduced random label noise [Zha+17]. This points to the importance of theoretically understanding the effect of noisy labels on the neural network training process. A handful of recent works have sought to understand the training dynamics of neural networks in the presence of noisy labels, but were either restricted to neural networks in the neural tangent kernel (NTK) regime, where feature learning is impossible [HLY20, JLT21]; failed to provide generalization guarantees for the resulting network [LSO19]; or only applied in settings where linear classifiers perform well [FCG21].

Our proof follows by characterizing the types of features that individual neurons learn throughout the training process. We show that at random initialization, provided the width of the network is a sufficiently large constant, most neurons are ‘weak’ random features: they have a normalized correlation of order $O(1/\sqrt{d})$, where $d$ is the input dimension, with at least one of the cluster means $\{\pm\mu_{1},\pm\mu_{2}\}$. After initialization, provided the learning rate is sufficiently large, a single step of gradient descent amplifies these neurons from ‘weak’ random features to ‘strong’, learned features: the normalized correlations with the cluster means improve from order $O(1/\sqrt{d})$ to order $O(1)$. In the later part of the training process, we show that the gradient descent dynamics ensure that if a neuron is highly correlated with a given cluster center $\mu_{s}$ after the first step, then (1) its norm increases throughout training, so that the network relies more upon this neuron to determine the network output, and (2) the neuron becomes orthogonal to the opposing cluster center $\mu_{s^{\prime}}$, $s^{\prime}\neq s$, so that the neuron is useful only for samples from the cluster center $\mu_{s}$. We show that having properties (1) and (2) is sufficient for producing a network that classifies all of the clean samples correctly and noisy samples incorrectly. A key difficulty in showing each of these facts is the presence of noisy training labels, which could in principle prevent the network from learning useful features; a careful analysis shows that this barrier is surmountable provided the fraction of noisy labels is smaller than an absolute constant.

As mentioned in the previous subsection, a number of works have highlighted the need to develop analyses of neural network training that go ‘beyond’ the NTK, or equivalently, neural networks that lie in the ‘feature learning regime’. One collection of works has focused on developing separations between what hypothesis classes can be learned efficiently using neural networks in the feature learning regime versus what can be learned using approaches based on kernels or random features [YS19, AL19, Gho+19, Wei+19, DM20, AL21, Mal+21, Abb+21]. One example of such a hypothesis class includes single neurons $x\mapsto\phi(\langle w,x\rangle)$, which can be efficiently learned using gradient descent on neural networks beyond the kernel regime [FCG20, YS20] but cannot be efficiently learned using random features or kernel-based methods [YS19, KMS20]. For a more detailed comparison of recent work on separations between what is learnable using kernel methods versus what is learnable using neural networks in the feature learning regime, we refer the reader to Table 2 and Appendix A of [Mal+21]. We note that two concurrent works have shown that a single step of gradient descent suffices for feature-learning behavior in neural networks [Ba+22, DLS22]. We also show that a single step of gradient descent suffices for learning data-dependent features, but our analysis also requires training for more than one step so that the learned features become more ‘refined’ (see Conditions 4.7 and 4.8 as well as Lemma 4.9 below).

Another line of work utilizes the mean field approximation to connect the training dynamics of infinitely wide neural networks to that of the solution to a partial differential equation [MMN18, CB18, Wei+19, Che+20, Fan+21]. This approach allows for the network weights to traverse far from the initialization and learn features. These works provide a useful characterization of the limiting behavior of neural networks as they become infinitely wide. By contrast, in this work we provide a guarantee for neural network optimization and generalization for networks of constant width (for a constant level of failure probability).

A handful of other works have explored the behavior of neural networks trained by gradient descent for variants of the XOR distribution we consider in this work. [Wei+19] used the mean field approximation to show that infinite-width two-layer networks trained by gradient flow will generalize well. [BL20] considered two-layer neural networks with smooth activations trained with additional ‘random sign’ and $\|W\|_{2,4}^{8}$ penalty regularization. They showed that when training with a large random initialization and a very large network, the second-order term of the Taylor expansion of the network around its initialization dominates the training dynamics and has a good optimization landscape provided the weights are close enough to initialization. They used this to derive a generalization guarantee for the resulting network. Although the work [BL20] is a strict improvement over standard NTK-based approaches, their analysis is more similar to the kernel-based analysis than the feature-learning approach we take here. Finally, [DM20] provided a characterization of learning a noiseless parity over the binary cube when performing gradient descent on the population risk (i.e., assuming infinite samples). Their analysis relies upon a neuron-by-neuron characterization of the learning process, similar to ours, but it is unclear how their analysis would proceed without access to infinite samples or if there are noisy labels. Indeed, much of the difficulty in characterizing feature-learning in neural networks comes from the possibility that neural networks could simply memorize the sampled training data rather than learn useful representations that enable generalization to unseen test data. In contrast to all of the above works, our work provides a novel characterization of how feature-learning occurs in finite-width neural networks that are trained in the finite-sample setting and when a substantial portion of the training labels are adversarially corrupted.

Finally, since our analysis shows that early-stopped gradient descent with a small initialization variance produces neural networks with rather simple decision boundaries which essentially ignore the noisy labels (see Fig. 1), our work is related to a series of works on the simplicity bias of gradient descent [PL21, Lyu+21, BPF22, Fre+23]. The aforementioned works all rely upon data that is either nearly-orthogonal or exactly orthogonal, while we make no such assumption. On the other hand, these other works characterize the behavior of gradient descent throughout the entire training trajectory, while we require early-stopping.

Preliminaries

We assume the training data $S$ is generated as i.i.d. samples from $\mathsf{P}$,

We will consider the regime where the noise rate $\eta\approx|{\mathcal{N}}|/n$ is smaller than a constant. In Figure 1, we illustrate what samples from this distribution look like.

We will use the convention that $\phi$ is applied entry-wise, so that $\phi(Wx)$ has $j$-th component $\phi(\langle w_{j},x\rangle)$. For simplicity, we assume that $m$ is an even number and that half of the second layer weights $a_{j}$ are initialized at the value of $+1/\sqrt{m}$, and the other half are initialized at the value $-1/\sqrt{m}$. (Our results hold for odd $m$ by setting $a_{m}=0$.) We assume the second layer weights are fixed at their initialized values throughout training. This assumption allows for a more simplified analysis as it allows for a static partition of the neurons into ‘positive’ neurons (those for which $a_{j}>0$) and ‘negative’ neurons ($a_{j}<0$) throughout training. We believe it is possible to extend our analysis to the setting where both layers are trained but we do not pursue this question in this work.

The population risk under the logistic loss is defined as

Note that since the ReLU activation $\phi(q)=\max(0,q)$ is not differentiable at 0, we use any subgradient value $\phi^{\prime}(0)\in$ when performing gradient descent. (Our results do not depend on the value chosen for the subgradient.)

We let $C>1$ denote a positive absolute constant that is large enough. Given a failure probability $\delta\in(0,1/2)$ we make the following assumptions going forward:

The dimension $d\geq C\max\left\{\log^{2}(n/\delta),\log(m/\delta)\right\}$;

The in-cluster variance $\sigma^{2}\leq 1/(C^{2}d)$;

The number of hidden nodes satisfies $m\geq C\log(1/\delta)$;

The step-size $\alpha$ satisfies $1/(2\sqrt{C})\leq\alpha\leq 1/\sqrt{C}$.

The first four assumptions above concern the distribution and the relationship between the number of samples, dimension, and number of neurons in the network. These assumptions are relatively mild as they only require that the dimension and number of samples are logarithmically large. These assumptions ensure that the signal-to-noise ratio in the model is quite high, and that in the setting with no label noise $\eta=0$, the optimal test error achievable is $o_{n}(1)$ (see Appendix D for more details). The final three assumptions concern the hyperparameters for the model and the optimization algorithm. Assumption (A5) ensures that the network is wide enough to ensure there are enough random features at initialization for gradient descent to “amplify”. It is noteworthy that assumption (A6) permits arbitrarily small (but nonzero) initialization variance. The assumption (A7) ensures that the step-size is large enough so that significant features can be learned after a single step of gradient descent but small enough so that optimization is stable.

Main results

Our main contribution is summarized in the following theorem.

Let $\delta\in(0,1/2)$. For all $C>1$ sufficiently large, under the assumptions (A1) through (A7), by running gradient descent with step-size $\alpha$ for $T=1+1/(4\alpha)$ iterations, with probability at least $1-4\delta$ over the random initialization and the draws of the samples we have,

Theorem 3.1 shows that at time $T$, gradient descent learns a network that accurately classifies every clean sample, and incorrectly classifies every noisy sample, and achieves population risk close to the noise rate $\eta$. In Figure 1, we plot the decision boundary for a neural network trained by gradient descent when 15% of the training labels are flipped and we observe that indeed every noisy sample is incorrectly classified and every clean sample is correctly classified.

It is worth noting that the decision boundary displayed in Figure 1 is rather simple. Our proof below will show that this simplicity is due to the fact that nearly every neuron in the neural network will become highly correlated to one of the four cluster means $\{\pm\mu_{1},\pm\mu_{2}\}$ so that the neural network essentially acts as the low-complexity classifier $x\mapsto\operatorname{sgn}(|\langle\mu_{1},x\rangle|-|\langle\mu_{2},x\rangle|)$. The main technical contribution of our work is the characterization of this feature-learning process and an examination of how it proceeds in the presence of noisy labels.

Let us remark that previous works on the generalization of neural networks in the feature-learning regime for variants of the XOR problem we study (without label noise) have sample complexities of order $O(\sqrt{d/n})$, which is an improvement over kernel-based methods which have sample complexity $\Omega(\sqrt{d^{2}/n})$ [Wei+19, BL20]. By contrast, Theorem 3.1 provides a dimension-independent rate of $O(\sqrt{1/n})$. This difference is due to the fact that they consider an XOR problem with a lower signal-to-noise ratio than the one we consider. In particular, they assume the features are uniform on the hypercube $\{\pm 1\}^{d}$ with labels given by $y=\operatorname{sgn}(x_{i}x_{j})$ for distinct coordinates $i\neq j$. Since the variance in every direction is the same, the signal-to-noise ratio is thus of order $\Theta(1/d)$. In our setting, the variance in the signal directions is larger: the variance in the direction of $\mu_{1}$ and $\mu_{2}$ is equal to $1+\sigma^{2}$ while the variance in the direction of any vector orthogonal to $\mu_{1}$ and $\mu_{2}$ is $\sigma^{2}$. Thus, the signal-to-noise ratio in our setting is of order $\Theta\left(\frac{1+\sigma^{2}}{d\sigma^{2}}\right)=\Omega(1)$ by Assumption (A2).

We note that our analysis does not rely upon the neural tangent kernel approximation. One way to see this is to observe that the assumption on the width of the network given in Assumption (A5) only requires the width to be larger than a fixed constant for a constant level of failure probability. Moreover, we show explicitly in the following proposition that for each sample, the feature maps given by the hidden layer activations change significantly from their values at random initialization, an essential characteristic of neural networks in the feature-learning regime [YH21].

Under the settings of Theorem 3.1, with probability at least $1-4\delta$ over the random initialization and draws of the samples, the feature maps of the neural network at time $T=1+1/(4\alpha)$ satisfy, for all $i\in[n]$,

The proof of Proposition 3.2 is given in Appendix C.

In the next section, we provide the proof of Theorem 3.1. The proof follows by concretely characterizing the type of features that different neurons learn throughout the training process.

Proofs

In this section, we provide an overview of the proof of Theorem 3.1. The detailed proofs are collected below in Appendix A.

We note that there exists a natural neural network consisting of four ReLU neurons that can classify the (clean) data with high accuracy:

This ideal low-complexity classifier is suggestive of the following possibility: for positive neurons, corresponding to second-layer weights satisfying $a_{j}>0$, the neurons become adapted to either the $+\mu_{1}$ cluster or the $-\mu_{1}$ cluster, depending upon the sign of $\langle w^{(0)}_{j},\mu_{1}\rangle$ at initialization. For negative neurons, corresponding to neurons with $a_{j}<0$, the neurons become adapted to either the $+\mu_{2}$ cluster or the $-\mu_{2}$ cluster depending on the sign of $\langle w^{(0)}_{j},\mu_{2}\rangle$ at initialization. This is at a high-level the argument that we show below.

In the remainder of this section assume that Assumptions (A1) through (A7) are in force.

We begin with an analysis of the properties of the random initialization. In the lemma below, we derive concentration results on the norm of the random weights, as well as a count for the number of neurons that are correlated with a fixed vector at a given threshold level. The correlation part of the lemma will be the basis of a ‘random feature amplification’ phenomenon, whereby the relatively small (random) correlations of the neurons with different cluster means at initialization will be amplified into strong correlations by gradient descent.

We show below that this ‘capturing’ phenomenon can be shown through a two-step process: first, at initialization, most of the positive neurons will have a normalized correlation with $\mu_{1}$ of order $\Theta(1/\sqrt{d})$, and similarly most of the negative neurons will have a normalized correlation with $\mu_{2}$ of order $\Theta(1/\sqrt{d})$. This is Lemma 4.2 below. Next, we show that by taking a single gradient step with a sufficiently large step-size, the normalized correlations for these neurons will improve from order $\Theta(1/\sqrt{d})$ to order $\Theta(1)$. This result, shown later in Lemma 4.10, is what we refer to as the ‘random feature amplification’ phenomenon, whereby the random features at initialization are amplified into useful features by gradient descent. Towards this end, we characterize the correlations of the neurons with the cluster means at initialization in the following lemma.

Let $\delta\in(0,1/2)$. For any absolute constant $C_{0}>1$, if $C$ is sufficiently large, with probability at least $1-\delta$ over the random initialization, there exist sets of neurons $J_{+\mu_{1}},J_{-\mu_{1}},J_{+\mu_{2}},J_{-\mu_{2}}\subset[m]$ satisfying the following:

In particular, $J:=J_{\pm\mu_{1}}\cup J_{\pm\mu_{2}}$ satisfies $|J|\geq m(1-1/C_{0})^{2}$.

Lemma 4.2 identifies a set of candidate neurons that are partially correlated with the cluster means $\{\pm\mu_{1},\pm\mu_{2}\}$. We would like to translate this result into a statement about the data, and to do so, we first need to provide some basic facts about samples from the distribution. The reader may find it helpful to refer back to the beginning of Section 4 where we introduce the $I_{\pm\mu_{i}}$ notation.

There is a universal constant $C_{1}\geq 2$ such that the following holds. For any $\delta\in(0,1/2)$, for all $C>1$ large enough, with probability at least $1-\delta$ over $S\sim\mathsf{P}^{n}$, the following holds.

For each $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$ and $\mu^{\perp}$ orthogonal to $\mu$,

For all $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$, for any $i\in I_{\mu}$, $\|x_{i}-\mu\|^{2}\leq C_{1}\sigma^{2}d\leq 1/C_{1}.$

The fraction of noisy points $\frac{|{\mathcal{N}}|}{n}\leq\eta+C_{1}\sqrt{\log(1/\delta)/n}\leq\eta+1/C_{1}$.

For any cluster $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$ and any $0\leq t\leq T-1$, we have

Now, recall that Lemma 4.2 shows that a large fraction of the neurons will ‘capture’ at least one of the four cluster centers with a normalized correlation of $\langle w_{j}^{(0)}/\|w^{(0)}_{j}\|,\mu_{s}\rangle\geq\Omega(1/\sqrt{d})$. Since the within-cluster variance is of order $\sigma=O(1/\sqrt{d})$, there is not enough signal for these neurons to capture all samples within each cluster. However, the following lemma demonstrates that capturing the cluster mean with a normalized correlation threshold of order $1/\sqrt{d}$ suffices to guarantee that a strictly larger portion of the samples from that cluster will be captured than not. This technical lemma will be key to our subsequent analysis.

There exists a universal constant $C_{2}>1$ such that for any $\delta\in(0,1/2)$, for all $C>1$ large enough, with probability at least $1-2\delta$, both Lemma 4.3 and the following event holds. For any $j\in[m]$ satisfying $\langle w^{(0)}_{j}/\|w^{(0)}_{j}\|,\mu\rangle\geq 1/(3C_{0}\sqrt{d})$ for some $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$, it holds that

In light of the above, we introduce the following definition.

We define the event where all parts of Lemma 4.1, Lemma 4.2 (with $C_{0}=4^{5}\cdot 1024^{2}\exp(4)$), Lemma 4.3, and Lemma 4.4 hold a good run.

By the above lemmas, we know that for any $\delta\in(0,1/2)$, for all $C>1$ large enough, a good run occurs with probability at least $1-4\delta$. In the remainder of this section, we will assume that a good run occurs.

2 Sufficient Conditions for a Large Margin Classifier via a Good Subnetwork

The below lemma demonstrates that in order to show that the neural network produces a good margin, it suffices to show that there exists a large subnetwork that produces a good margin provided that the weights of the network are bounded.

Lemma 4.6 demonstrates that in order to show the neural network classifies an example correctly, it suffices to identify a large subnetwork that does so. The rest of our proof is dedicated to showing that this happens. The subnetwork that performs well is defined in terms of the neurons $j\in J_{\pm\mu_{1}}\cup J_{\pm\mu_{2}}$, where the index sets $J_{\pm\mu_{1}}\cup J_{\pm\mu_{2}}$ are defined in Lemma 4.2 and are shown to constitute a large fraction of all of the neurons: for each $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$, the set $J_{\mu}$ has cardinality at least $|J_{\mu}|\geq\frac{m}{4}(1-1/C_{0})^{2}$, where $C_{0}>1$ is a large constant. We next define two conditions that we will show suffice for showing this subnetwork classifies examples correctly, which we refer to as the neuron alignment condition and the almost-orthogonality condition. We describe the first of these below.

We say that the neuron alignment condition holds at time $t$ if the subsets of neurons $J_{\pm\mu_{1}}$ and $J_{\pm\mu_{2}}$ defined in Lemma 4.2 satisfy the following: for every $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$, and for all $j\in J_{\mu}$,

The neuron alignment condition loosely states that there is a substantial number of neurons (the neurons in the sets $J_{\pm\mu_{1}}\cup J_{\pm\mu_{2}}$) that completely capture each of the clusters in the sense that all samples within each cluster have the same ReLU activation, which is “on” on one of the clusters and “off” on the opposing cluster. By Lemma 4.2, we know that there is a large fraction of neurons that catch the cluster means $\{\pm\mu_{1},\pm\mu_{2}\}$ at initialization. However, as we argued prior to Lemma 4.4, because the normalized correlation between the neurons at initialization and the cluster means is of order $1/\sqrt{d}$ while the variance within each cluster is also of order $1/\sqrt{d}$, a substantial portion of the examples within each cluster will not be captured by a neuron at initialization. We briefly note here that in the next section, we will show that a single step of gradient descent suffices to address this problem.

We say almost-orthogonality holds up to time $\tau$ if for all $t\leq\tau$,

The almost-orthogonality condition ensures that the projection of the $J_{\pm\mu_{1}}$ (resp. $J_{\pm\mu_{2}}$) neurons onto the space spanned by $\mu_{2}$ (resp. $\mu_{1}$) remains small for all iterates of gradient descent up to time $\tau$.

In the next lemma, we show that the combination of neuron alignment and almost-orthogonality suffices to produce a good subnetwork margin. Note that we consider times $t\geq 1$ with foresight, as we shall eventually show that neuron alignment and almost-orthogonality hold for all $t\geq 1$.

Let $J=J_{\pm\mu_{1}}\cup J_{\pm\mu_{2}}$, where the sets $J_{\pm\mu_{1}}$ and $J_{\pm\mu_{2}}$ are defined in Lemma 4.2. Suppose that neuron alignment (Condition 4.7) and almost-orthogonality (Condition 4.8) hold at times $\tau=1,\dots,T-1=1/(4\alpha)$. Then, on a good run, for all $C>1$ large enough, at time $T=1+1/(4\alpha)$, we have $\|W^{(T)}\|_{F}\leq 1$, and that

where $C_{3}=4096\exp(2)/(1-1/C_{0})^{2}$ and $C_{0}>1$ is the constant from Lemma 4.2.

Lemma 4.9 demonstrates that in order to show that a given subnetwork $f^{J}(x;W)$ accurately classifies all of the clean data, it suffices to show that neuron alignment and almost-orthogonality hold for a sufficiently large (but constant) number of steps. By Lemma 4.6, this translates to a guarantee for the entire network $f(x;W)$ if we can show that the subnetwork is sufficiently large, which we can ensure by taking $J$ as the union of the sets $J_{\pm\mu_{1}},J_{\pm\mu_{2}}\subset[m]$ as in Lemma 4.2 and by taking the constant $C_{0}>1$ from that lemma to be sufficiently large. Thus, to complete the proof, we need only verify that neuron alignment and almost-orthogonality hold for a sufficiently large but constant number of steps. This is what we show in the next subsection.

We note that both neuron alignment and almost-orthogonality are needed in order to ensure that the subnetwork $f^{J}(x;W)$ behaves like the simple classifier $|\langle\mu_{1},x\rangle|-|\langle\mu_{2},x\rangle|$. For instance, consider what happens if half of the positive neurons (corresponding to $j\in[m]$ with $a_{j}>0$) are proportional to $\mu_{1}+100\mu_{2}$ and the other half are proportional to $-\mu_{1}-100\mu_{2}$, and likewise half of the negative neurons are proportional to $\mu_{2}$ and the other half are proportional to $-\mu_{2}$. Then the neuron alignment condition would hold, but almost-orthogonality would not hold, and the network would behave like the predictor $|\langle\mu_{1}+100\mu_{2},x\rangle|-|\langle\mu_{2},x\rangle|$ and not generalize well. Thus, in addition to showing that the neurons are highly correlated with the cluster means from a given class, we must also show that they are nearly orthogonal to the cluster means from the opposite class.

3 Gradient Descent Produces a Large Margin Classifier

As mentioned previously, we cannot expect neuron alignment to hold at initialization, as the random features that define the subnetwork $f^{J}$ have per-neuron normalized correlations $\langle w^{(0)}_{j}/\|w^{(0)}_{j}\|,\mu_{1}\rangle$ of order $O(1/\sqrt{d})$, while the fluctuations within each cluster $\langle w^{(0)}_{j}/\|w^{(0)}_{j}\|,\mu_{s}-x_{i}\rangle$ are also of order $\sigma=O(1/\sqrt{d})$. This means that many samples $x_{i}$ belonging to a cluster $\mu_{s}$ will satisfy $\operatorname{sgn}(\langle w^{(0)}_{j},x_{i}\rangle)\neq\operatorname{sgn}(\langle w^{(0)}_{j},\mu_{s}\rangle)$, preventing the satisfaction of the neuron alignment condition. This is where Lemma 4.4 will play a role: although the random features have normalized correlations of order $\Theta(1/\sqrt{d})$ with the cluster means, this signal provides an ‘edge’ in terms of the ReLU activations of samples within the cluster. That is, having $\langle w^{(0)}_{j}/\|w^{(0)}_{j}\|,\mu_{s}\rangle\geq c/\sqrt{d}$ is sufficient to guarantee that the fraction of samples within the $\mu_{s}$ cluster sharing the same sign as $\langle w^{(0)}_{j},\mu_{s}\rangle$ is at least $\nicefrac{{1}}{{2}}+\Delta$ for some absolute constant $\Delta>0$. This provides enough signal for gradient descent to latch onto and ‘amplify’ the normalized per-neuron correlations from $\langle w^{(0)}_{j}/\|w^{(0)}_{j}\|,\mu_{s}\rangle\geq c/\sqrt{d}$ to $\langle w^{(1)}_{j}/\|w^{(1)}_{j}\|,\mu_{2}\rangle\geq c^{\prime}$ after one sufficiently large step. Since now the normalized correlations are of order $1$ while the within-cluster fluctuations are of order $1/\sqrt{d}$, this allows for neuron alignment to hold after a single step of gradient descent.

For $C>1$ sufficiently large, on a good run Condition 4.7 holds at time $t=1$. Moreover, letting $C_{2}>1$ denote the constant from Lemma 4.4, the per-neuron normalized correlations satisfy

We now know that neuron alignment holds at time $t=1$, and that the number of neurons that are characterized by the alignment condition is quite large (precisely, $m(1-1/C_{0})^{2}$ for a large constant $C_{0}$). By Lemma 4.9, if we can show that $(i)$ neuron alignment continues to hold for a certain number of steps, $(ii)$ almost-orthogonality holds throughout these steps, and $(iii)$ we early-stop so that the hidden layer weights are not too large, then there will be a large subnetwork that classifies clean examples with a positive margin. In the next lemma, we inductively argue that this is the case.

For $C>1$ sufficiently large, on a good run, for every time $t=1,\dots,1/(4\alpha)$, neuron alignment (Condition 4.7) holds at time $t$ and almost-orthogonality (Condition 4.8) holds up to time $t$.

We emphasize that although Lemma 4.10 shows that neuron alignment holds at time $t=1$, this is not sufficient to guarantee generalization since we must ensure that the positive (respectively negative) neurons are not highly correlated to $\pm\mu_{2}$ (respectively $\pm\mu_{1}$) since this could result in inaccurate predictions as outlined at the end of Section 4.2. This potential problem is precisely what almost-orthogonality (Condition 4.8) prevents, and Lemma 4.11 shows that by running gradient descent for a large (but constant) number of steps, we can guarantee that both neuron alignment and almost-orthogonality hold up to time $T-1=1/(4\alpha)$. By Lemma 4.9, this implies that at time $T$ the subnetwork $f^{J}(x;W^{(T)})$ classifies all of the clean examples correctly, and by Lemma 4.6 this implies that the full network $f(x;W^{(T)})$ classifies all of the clean examples correctly with small $\|W^{(T)}\|_{F}$. From here the proof of Theorem 3.1 is a straightforward Rademacher-complexity based argument; the details are provided in Appendix A.4.

Discussion

We have shown that two-layer neural networks with ReLU activations trained by gradient descent can achieve small test error on a distribution for which linear classifiers perform no better than random guessing. We developed a novel proof technique that detailed how using a random initialization provides a collection of random features that gradient descent is able to amplify into stronger, useful features for prediction. Importantly, our analysis holds when a constant fraction of the training labels are arbitrarily corrupted.

In Figure 2, we examine the behavior of two-layer ReLU networks trained by gradient descent on the logistic loss for the 2-XOR distribution we consider when 15% of the labels are flipped (for full experimental details, see Appendix E). We consider two distinct settings: a low-dimensional setting where $n\gg d$ and a high-dimensional setting where $d\gg n$. In the low-dimensional setting, the test accuracy decreases after the network overfits to the noisy training data, while in the high-dimensional setting the test accuracy remains at the optimal 85% level even after reaching the point of interpolation. Since our assumptions only require that the number of samples and dimension are not super-exponential in the other, this suggests that we would need to introduce new techniques, separately tailored to the low-dimensional and high-dimensional settings, in order to characterize the generalization behavior of the network after the point of interpolation.

Another natural direction for future research is to understand whether or not the random feature amplification phenomenon that we identified in two-layer networks has an analogue in deeper networks. Yet another direction is to understand whether this analysis technique can be generalized to settings with more cluster centers.

Acknowledgements

We thank the anonymous reviewers for their numerous suggestions which helped improve the presentation of the paper. We thank Hongren Yan and Yutong Wang for pointing out issues in a previous version of this paper. We gratefully acknowledge the support of the NSF and the Simons Foundation for the Collaboration on the Theoretical Foundations of Deep Learning through awards DMS-2023505, DMS-2031883, and #814639.

Appendix A Omitted Proofs from Section 4

In this appendix, we provide the proofs for all of the lemmas in Section 4. In Section A.1, we provide the proofs for the lemmas that involve concentration inequalities: Lemmas 4.1, 4.2, 4.3, and 4.4. Next, we prove Lemmas 4.6 and 4.9, which show that producing a good subnetwork suffices for the neural network to classify the clean examples correctly and provide a sufficient condition for producing a good subnetwork. In Section A.3, we show that gradient descent produces a good subnetwork. Finally, in Section C, we provide a proof of Proposition 3.2, which emphasizes that the feature maps produced by gradient descent differ significantly from those found at initialization.

We remind the reader that throughout this section we assume that Assumptions (A1)-(A7) are in effect. We also note that $C>1$ is always used to denote the constant used in these assumptions.

In this subsection we provide the proofs for Lemmas 4.1, 4.2, 4.3, and 4.4.

We restate the lemma here for the reader’s convenience. See 4.1

We first prove the first part of the lemma. Note that for fixed $j\in[m]$, there are i.i.d. $z_{i}\sim\mathsf{N}(0,1)$ such that

By concentration of the $\chi^{2}$ distribution [Wai19, see,], for any $t\in(0,1)$,

In particular, by taking $t=\sqrt{8\log(4m/\delta)/d}$, we have that for fixed $j\in[m]$, with probability at least $1-\delta/2m$,

where we have used Assumption (A1), that is, $d\geq C\log(m/\delta)$for a sufficiently large constant $C>1$ implies $t\leq 1/2$. Applying a union bound over $j\in[m]$ shows that the bound on the norms at initialization holds over all $j$ with probability at least $1-\delta/2$.

For the neuron-counting argument, let $z\sim\mathsf{N}(0,1)$ denote a standard normal random variable. Denote by $p$ the probability

By anti-concentration of the Gaussian, we have

Thus, with probability at least $1-\delta/2$, we have

Taking a union bound over the first and second parts of the proof shows that both claims hold simultaneously with probability at least $1-\delta$. ∎

A.1.2 Proof of Lemma 4.2

We restate and prove Lemma 4.2 below. See 4.2

Fix $C_{0}>1$. Apply Lemma 4.1 to the positive neurons $j$ satisfying $a_{j}>0$ with $\mu_{1}$. This tells us that with probability at least $1-\delta/16$,

Taking a union bound, we see that with probability at least $1-\delta$,

A.1.3 Proof of Lemma 4.3

We restate and prove Lemma 4.3 below. See 4.3

We shall show that each part of the lemma holds with a large enough probability and then take a union bound to establish our claim.

Proof of parts (a) and (b): We consider the case $i\in I_{+\mu_{1}}$. The cases of $i\in I_{\mu}$ for $\mu\in\{-\mu_{1},\pm\mu_{2}\}$ follow using identical arguments.

Let $i\in I_{\mu_{1}}$. We begin by noting that, since $\|\mu\|=1$, we have by Cauchy–Schwarz,

Therefore, to derive a lower bound on $\langle x_{i},\mu_{1}\rangle$ when $i\in I_{+\mu_{1}}$, it suffices to derive an upper bound on $\|x_{i}-\mu\|$ for each $i$, so that we will first prove part (b).

Since $(x_{i}-\mu)/\sigma$ is isotropic and log-concave, by concentration of the Euclidean norm of isotropic log-concave random vectors [Ada+14, Theorem 1], there is a universal constant $c>0$ such that,

In particular, using Assumption (A1), we can take $d\geq\log^{2}(32n/\delta)/c^{2}$ so that $\exp(-cu\sqrt{d})\leq\delta/(32n)$ and thus we have with probability at least $1-\delta/32$, for all $i\in I_{+\mu_{1}}$,

This, along with Assumption (A2) proves part (b).

Using (4) and Assumption (A2) so that $c\sigma\sqrt{d}<1$, we have

which proves the first half of part (a) of the lemma when $i\in I_{+\mu_{1}}$. When $i\in I_{+\mu_{1}}$, the cluster mean $\mu^{\perp}$ orthogonal to $\mu_{1}$ is $\mu_{2}$, and so we have,

which completes the proof of the second part of (a) when $\mu=\mu_{1}$. Taking a union bound over $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$ shows that parts (a) and (b) hold with probability at least $1-\delta/8$.

In particular, for $u=\sqrt{\frac{2c\log(2/\delta)}{n}}$, by Assumption (A3) we have with probability at least $1-\delta/2$, $|{\mathcal{N}}|/n\leq\eta+\sqrt{2c\log(2/\delta)/n}\leq\eta+1/C_{1}$ by Assumption (A4).

Since $\delta\in(0,1/2)$, there is a larger constant $c^{\prime}>0$ such that $c\sqrt{\frac{\log(32/\delta)}{n}}\leq c^{\prime}\sqrt{\frac{\log(1/\delta)}{n}}$. Taking a union bound over the four clusters shows that part (d) holds with probability at least $1-\delta/4$.

Thus all four parts (a), (b), (c), (d) hold with probability at least $1-\delta$. ∎

A.1.4 Proof of Lemma 4.4

We restate and prove Lemma 4.4 below. See 4.4

We shall prove this lemma in two parts. First, we shall define a “good event” $\mathcal{E}$ that occurs with probability at least $1-2\delta$. Then via a deterministic argument, we shall show that the lemma holds whenever this good event occurs.

This motivates deriving upper and lower bounds for the cardinality of the sets $I_{\mu}^{\mathcal{C}}$ and $I_{-\mu}^{\mathcal{C}}$. To do so, we first note that with probability at least $1-\delta$, all of the events in Lemma 4.3 hold. In particular, by Part (d) of that lemma, we have with probability at least $1-\delta$, for any $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$,

We thus have with probability at least $1-\delta$, all of the events in Lemma 4.3 hold, and, for all $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$,

where the inequality $(i)$ uses Assumptions (A3) and (A4).

Now, since $N_{\mu}(j)\sim\mathsf{Binomial}(|I_{\mu}^{\mathcal{C}}|,p)$, by Hoeffding’s inequality and a union bound (over the neurons and over the clusters), there is some $c>0$ such that with probability at least $1-\delta$, for all $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$, for all $j\in[m]$,

Let us define $\mathcal{E}$ to be the event where the events in Lemma 4.3, and inequalities (7) and (8) all simultaneously hold. By a union bound this happens with probability at least $1-2\delta$. This shall determine the success probability of the lemma.

In the remainder of the proof let us assume that this event $\mathcal{E}$ occurs; we will show that the lemma holds as a deterministic consequence of these events.

Since the event $\mathcal{E}$ occurs, for all $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$ and all $j\in[m]$ we have,

Above, $(i)$ uses Eq. (8), while $(ii)$ uses Eq. (7). Inequality $(iii)$ uses Assumption (A3) so that $n\geq\log(64m/\delta)$ and by taking $c$ to be a larger constant. Further, for all $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$ and all $j\in[m]$, we have,

Above, $(i)$ uses (7) and $(ii)$ uses the assumption $n\geq C\log(1/\delta)$ given by (A3). Thus, we have shown that, when the good event $\mathcal{E}$ occurs, then inequalities (9) and (10) hold. In the remainder of the proof, we will show that the lemma follows as a consequence of the inequalities (9) and (10).

In order to show $N_{\mu}(j)\gg N_{-\mu}(j)$, it suffices to show that $p$ is large enough so that there is sufficient ‘edge’ for more samples to be captured by $w_{j}$ than not. To this end, we have for any $j$ such that $\langle w^{(0)}_{j},\mu\rangle>0$,

Recall that we are considering neurons $j\in[m]$ such that $\langle w^{(0)}_{j}/\|w^{(0)}_{j}\|,\mu\rangle\geq 1/(3C_{0}\sqrt{d})$. By assumption (A2), for $C$ sufficiently large we have $\sigma\sqrt{d}\leq 1/C\leq 3/C_{0}$ so that the inclusion $[-1/9,0]\subset[-1/(3C_{0}\sigma\sqrt{d}),0]$ holds. Thus, we have,

Note that $\langle z,w^{(0)}_{j}/\|w^{(0)}_{j}\|\rangle$ is the projection of a log-concave isotropic random vector onto the one dimensional subspace spanned by $w^{(0)}_{j}/\|w^{(0)}_{j}\|$, and thus by [Dia+20, Definition 1.2, Fact A.4 ] there exists an absolute constant $c_{1}>0$ such that

and continuing from the previous display we thus have

We can thus use the inequalities given in events (9) and (10) to see that for any $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$ and for any $j\in[m]$ such that $\langle w^{(0)}_{j}/\|w^{(0)}_{j}\|,\mu\rangle\geq 1/(3C_{0}\sqrt{d})$,

In $(i)$, we have used (13). Inequality $(ii)$ follows by a direct calculation. Finally, $(iii)$ uses that Assumption (A3) ensures $n\geq 4\cdot 100c^{2}c_{1}^{-2}\log(64m/\delta)$, as well as Lemma 4.3(c) and Assumption (A4).

This shows that there exists a universal constant $C_{2}>1$ such that whenever event $\mathcal{E}$ occurs, for all $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$ and

Recall that above we argued that the event $\mathcal{E}$ (which is when the events in Lemma 4.3 and Equations (7) and (8) all hold simultaneously) occurs with probability at least $1-2\delta$, and since this implies the claim in Equation (15) holds, this completes the proof. ∎

A.2 Sufficient Conditions for a Large Margin Classifier via a Good Subnetwork

In this subsection, we prove Lemmas 4.6 and 4.9, which demonstrate that in order to show the neural network correctly classifies all clean samples, it suffices to show that there exists a large subnetwork that classifies the points correctly. Before we prove this, we introduce the following auxiliary lemma, which bounds the growth of the weights of the network over time. This lemma will be used in a number of places in the remaining proofs.

For $C>1$ large enough, on a good run we have the following bound on the norms of the weights for times $t\geq 1$:

For all $j\in[m]$, $\|w^{(t)}_{j}\|\leq 2|a_{j}|\alpha t$;

First, note that since a good run occurs, Lemma 4.3 and Assumption (A2) imply that for any sample $i\in[n]$, we have $\|x_{i}-\mu_{s}\|^{2}\leq C_{1}\sigma^{2}d+\frac{1}{C_{1}}\leq 1/3$, where $\mu_{s}$ is the cluster mean corresponding to $x_{i}$. Therefore, we have for any $i\in[n]$,

A.2.2 Proof of Lemma 4.6

With the above lemma in hand, we now restate and prove Lemma 4.6. See 4.6

In $(i)$ we use the Cauchy–Schwarz inequality, and that $\phi$ is 1-Lipschitz with $\phi(0)=0$. The final claim follows as $\|W_{J^{c}}\|_{F}\leq\|W\|_{F}\leq 1$, so that

A.2.3 Proof of Lemma 4.9

In this section we restate and prove Lemma 4.9.

By Lemma A.1, we have that for all $\tau\in\{1,\ldots,T\}$

since $\tau\leq T=1/(4\alpha)+1$. This shows the claimed guarantee for the norm.

Using Lemma 4.3(b), we can therefore bound the neural network output at time $\tau$ by

We will now show that for sufficiently large $t$, the network produces a positive margin on the $+\mu_{1}$. This shall be crucial in showing that the network produces a positive margin on the clean points associated with this cluster, and a negative margin on the noisy points in the cluster.

Summing (20) from $\tau=1,\dots,T-1$ and using that $j\in J_{+\mu_{1}}$ implies $\langle w^{(1)}_{j},\mu_{1}\rangle>0$, we get that

Thus, we have the following lower bound on the network output at $\mu_{1}$:

In $(i)$ we use the neuron alignment condition. In $(ii)$ we use that $\phi$ is 1-Lipschitz. In $(iii)$ we use the almost-orthogonality (Condition 4.8) and that $|a_{j}|=1/\sqrt{m}$. In $(iv)$ we use Eq. (21) and again use the fact that $|a_{j}|=1/\sqrt{m}$. Finally, $(v)$ uses Lemma 4.2, so that we have $|J_{+\mu_{1}}|/m\geq\frac{1}{4}(1-1/C_{0})^{2}$, as well as the fact that $|J_{\pm\mu_{2}}|\leq m$. In particular, we see that for $T-1=1/(4\alpha)$, we have

In the last inequality, we use the Assumption (A7) and take $C>1$ large enough so that $\alpha\leq\exp(-2)(1-1/C_{0})^{2}/(6\cdot 1024)$. With a lower bound on the margin for the cluster center $\mu_{1}$ established, we can translate this result to one for samples using Lemma 4.3. To do so, note that the sub-network $f^{J}(\cdot;W)$ is $\|W\|_{F}$-Lipschitz in the network input, i.e., we have

where the first inequality follows by Cauchy–Schwarz inequality and the last inequality follows since $\lVert a\rVert=\sum_{j=1}^{m}a_{j}^{2}=1$ and $\|W(x-x^{\prime})\|\leq\|W\|_{F}\|x-x^{\prime}\|$. Therefore we can use Lemma 4.3 (b) to translate (22) into a guarantee for the samples. For any $i\in I_{+\mu_{1}}^{\mathcal{C}}$, so that $y_{i}=+1$,

The second inequality uses that $\|W^{(T)}\|_{F}\leq 1$ and Lemma 4.3, while the last inequality uses Assumption (A2) so that $C_{1}\sigma\sqrt{d}$ can be taken smaller than any absolute constant for $C>1$ sufficiently large.

This completes the proof for samples $i\in I_{+\mu_{1}}^{\mathcal{C}}$. To see that the network also incorrectly classifies noisy samples, take $i\in I_{+\mu_{1}}^{\mathcal{N}}$, so that $y_{i}=-1$. Then, again using Lemma 4.3(b),

where the last inequality follows since $\lVert W^{(T)}_{F}\rVert\leq 1$ as we proved above.

For the other clusters, an identical argument to (20) yields

We can utilize the identities (23) and similar arguments to show that the desired margin condition holds for other clusters $I_{-\mu_{1}}^{\mathcal{C}},I_{\pm\mu_{2}}^{\mathcal{C}}$ so the result holds for all $i\in{\mathcal{C}}$. ∎

A.3 Gradient Descent Produces a Large Margin Classifier

In this section, we show that the sufficient conditions necessary for producing a good subnetwork described in Lemma 4.9 hold. The first step for this is to show that neuron alignment holds at time $t=1$.

We restate and prove Lemma 4.10 below. See 4.10

Since a good run occurs, all of the events in Lemma 4.1, Lemma 4.2, Lemma 4.3, and Lemma 4.4 hold. Recall that the sets $J_{\pm\mu_{1}}$ and $J_{\pm\mu_{2}}$ were defined in Lemma 4.2. We will now show that Condition 4.7 holds for these sets at time $t=1$. We will demonstrate the first claim in the condition statement (regarding $\mu_{1}$), that is, for all $j\in J_{+\mu_{1}}$:

The remaining parts of the neuron alignment condition concerning $j\in J_{-\mu_{1}}\cup J_{\pm\mu_{2}}$ shall follow by using an identical argument.

There are two parts to the neuron alignment condition, let us begin by proving that the first part holds.

Let us begin by showing that for all $j\in J_{+\mu_{1}}$:

Recall that by the definition of the set $J_{+\mu_{1}}$, we have that for all $j\in J_{+\mu_{1}}$,

To show that the first part of NAC holds for the subset $J_{+\mu_{1}}$, we need to show that a step of gradient descent takes ensures that all of the samples from this cluster are captured by the neurons in $J_{+\mu_{1}}$. We shall prove this in stages.

First, we shall establish a relation between the parameters after one the first step of gradient descent $w_{j}^{(1)}$ and those at initialization $w_{j}^{(0)}$.

Then, we shall leverage this relation to show that the angle between $w_{j}^{(1)}$ and $\mu_{1}$ is small.

This, along with the fact that the samples from this cluster are close to its center, shall be sufficient to ensure that (24) is satisfied.

With this in place, let us analyze the gradient update for a neuron in the set $J_{+\mu_{1}}$. Recall that for such nodes, $a_{j}=1/\sqrt{m}>0$ and therefore,

where we have defined the second “error vector” $\zeta_{2}$ as,

Applying the triangle inequality and Equation (25),

Substituting the above definition into (27), we then have

Continuing with the plan outlined above, we will now show that $\langle w^{(1)}_{j}/\|w^{(1)}_{j}\|,\mu_{1}\rangle\geq c$ for a universal constant $c$. We have,

In the last line we have applied the inequalities (26) and (28). Thus, it suffices to derive a lower bound for $N_{+\mu_{1}}(j)-N_{-\mu_{1}}(j)$, which is precisely the result that Lemma 4.4 provides. We have,

In $(i)$ we use Lemma 4.4, while in $(ii)$ we use Assumption (A4) so that $|{\mathcal{N}}|/n\leq 2\eta\leq 1/2C_{2}$. Thus, plugging this in to (31) we get that

This completes the proof for the normalized margin claim.

To show that the first part of the neuron alignment holds, we want to show that $\langle w^{(1)}_{j},x_{i}\rangle>0$. We have,

Above, $(i)$ uses (35) and the Cauchy–Schwarz inequality. Inequality $(ii)$ uses Lemma 4.3. The final inequality $(iii)$ uses Assumption (A2), so that $C_{1}\sigma\sqrt{d}\leq 1/64C_{2}$. This completes the part of neuron alignment concerning neurons $J_{+\mu_{1}}$ and for samples in cluster $I_{+\mu_{1}}$.

To show the part of neuron alignment concerning samples in cluster $I_{-\mu_{1}}$, note that we still have the identity (35). But for samples $i\in I_{-\mu_{1}}$, we have

where $\|x_{i}+\mu_{1}\|$ is small, and so the inequality $\langle w^{(1)}_{j},x_{i}\rangle<0$ follows using the same argument as above. Hence, we have shown that $\phi^{\prime}(\langle w^{(1)}_{j},x_{i}\rangle)$ for all $i\in I_{-\mu_{1}}$.

This completes the proof of neuron alignment for the neurons in $J_{+\mu_{1}}$. An analogous argument can also be used to establish the claim for the neurons in $J_{-\mu_{1}}\cup J_{\pm\mu_{2}}$. ∎

A.3.2 Proof of Lemma 4.11

We now show that the neuron alignment condition and almost-orthogonality condition hold for a sufficiently large amount of time.

The proof is by induction. To see the base case $t=1$, first, note that neuron alignment holds at time $t=1$ by Lemma 4.10. Further, almost-orthogonality holds at time $t=1$ since by Lemma A.1 we have $|\langle w^{(1)}_{j},\mu\rangle|\leq\|w^{(1)}_{j}\|\leq 2|a_{j}|\alpha t$ for any $\mu\in\{\pm\mu_{1},\pm\mu_{2}\}$. So let us now assume that neuron alignment and almost-orthogonality hold at every time step until time $t$, and consider the case $t+1\leq 1/(4\alpha)$. By Lemma A.1, since $t+1\leq 1/(4\alpha)$, we have $\|W^{(\tau)}\|_{F}\leq 1$ for every $\tau\leq t+1$. Using an identical argument to (19), this implies for all $i\in[n]$ and $\tau\leq\left\{1,\ldots,1/(4\alpha)\right\}$,

This key property will allow us to show that neuron alignment holds at time $t+1$.

𝑡1t+1. We will first show the result for neurons $j\in J_{+\mu_{1}}$; the result for neurons in $J_{-\mu_{1}}\cup J_{\pm\mu_{2}}$ will follow similarly.

Let $j\in J_{+\mu_{1}}$, so $a_{j}=|a_{j}|=1/\sqrt{m}$. It suffices to show that for $k\in I_{+\mu_{1}}$, we have $\langle w^{(t+1)}_{j},x_{k}\rangle>0$, and for $k\in I_{-\mu_{1}}$, we have $\langle w^{(t+1)}_{j},x_{k}\rangle<0$. To show this, we will utilize an argument similar to that we used in the proof of Lemma 4.10 (see eqs. (35) and (36)), in that we will first show that $\langle w^{(t+1)}_{j}/\|w^{(t+1)}_{j}\|,+\mu_{1}\rangle\geq c$ for some constant $c>0$, and then use that the within-cluster variance is of order $\sigma^{2}d$ and that $\sigma^{2}\ll 1/d$. Towards this end, we first derive a consequence of neuron alignment. Let $\tau$ be a time satisfying $1\leq\tau\leq t$. Then neuron alignment holds at time $\tau$ by the induction hypothesis, so that,

In $(i)$ we have used that the neuron alignment condition holds at time $\tau$, and thus $\phi^{\prime}(\langle w^{(\tau)}_{j},x_{i}\rangle)=1$ for $i\in I_{+\mu_{1}}$ and $\phi^{\prime}(\langle w^{(\tau)}_{j},x_{i}\rangle)=0$ for $i\in I_{-\mu_{1}}$. We can bound the terms $\langle x_{i},+\mu_{1}\rangle$ appearing above with Lemma 4.3, so that

where inequality $(i)$ uses Lemma 4.10. By Lemma A.1, we have $\|w^{(t+1)}_{j}\|\leq 2\alpha|a_{j}|(t+1)$, so that,

Using an identical argument to (36), since by Lemma 4.3(b) and Assumption (A2) we have the inequalities $\|x_{k}-\mu_{1}\|\leq C_{1}\sigma\sqrt{d}\leq C_{1}/C$, by taking $C>512C_{1}\exp(2)$ we have $\langle w^{(t+1)}_{j},x_{k}\rangle>0$ for $k\in I_{+\mu_{1}}$. A symmetric argument shows that $\langle w^{(t+1)}_{j},x_{k}\rangle<0$ for $k\in I_{-\mu_{1}}$. This completes the proof that neuron alignment holds for neurons $j\in J_{+\mu_{1}}$. We can show that neuron alignment holds for neurons in $J_{-\mu_{1}}\cup J_{\pm\mu_{2}}$ using an analogous argument.

𝑡1t+1. We now show that almost-orthogonality continues to hold at time $t+1$ given it holds at time $t$. We will prove the result for neurons $j\in J_{+\mu_{1}}$ with an analogous argument holding for the neurons in $J_{-\mu_{1}}\cup J_{\pm\mu_{2}}$.

We want to show that, for any neuron $j\in J_{+\mu_{1}}$ satisfying

we have that $|\langle w^{(t+1)}_{j},\mu_{2}\rangle|\leq 3\alpha|a_{j}|$ as well. We will show this by demonstrating that if at time $t$ we have $|\langle w^{(t)}_{j},\mu_{2}\rangle|\geq\alpha|a_{j}|$, then $\langle w^{(t+1)}_{j},\mu_{2}\rangle$ will either change sign or will decrease in magnitude at the next iteration; since the order of norm changes for a single neuron in one step is $O(\alpha|a_{j}|)$, this will complete the proof.

Consider the case that $\langle w^{(t)}_{j},\mu_{2}\rangle\geq\alpha|a_{j}|$; the negative case will follow using a symmetric argument. Since neuron alignment holds, an identical argument used to derive Equations (38) through (39) implies that

From the above, one can see that if $\sum_{i\in I_{+\mu_{2}}^{\mathcal{C}}}\phi^{\prime}(\langle w^{(t)}_{j},x_{i}\rangle)\gg\sum_{i\in I_{-\mu_{2}}^{\mathcal{C}}}\phi^{\prime}(\langle w^{(t)}_{j},x_{i}\rangle)$, then we will have that the above quantity is negative, showing that $\langle w^{(t)}_{j},\mu_{2}\rangle$ will decrease. When $\langle w^{(t)}_{j},\mu_{2}\rangle$ is large, then this is likely to occur; this is precisely the second part of Lemma 4.4. In particular, since $\langle w^{(t)}_{j},\mu_{2}\rangle\geq\alpha|a_{j}|$ by assumption, we have

where $(i)$ follows by Lemma A.1 and the fact that we are considering the case when $\langle w^{(t)}_{j},\mu_{2}\rangle\geq\alpha|a_{j}|$; inequality $(ii)$ uses that $t\leq 1/(4\alpha)$; and $(iii)$ uses Assumption (A7), so that $\alpha\geq 1/(2\sqrt{C})$. Since the correlation with the cluster mean is of constant order, we can repeat the argument used in (36) to show that the sign of $\langle w^{(t)}_{j},x_{i}\rangle$ is the same as the sign of $\langle w^{(t)}_{j},\mu_{2}\rangle$ for $i\in I_{+\mu_{2}}^{\mathcal{C}}$:

Inequality $(i)$ uses the lower bound in (42) as well as Lemma 4.3. Inequality $(ii)$ uses assumption (A2), so that $C_{1}\sigma\sqrt{d}\leq C_{1}/C<1/(2\sqrt{C})$ for $C$ sufficiently large relative to $C_{1}$. Using a symmetric argument, we thus have for positive neurons satisfying $\langle w^{(t)}_{j},\mu_{2}\rangle\geq\alpha|a_{j}|$,

Substituting the above into (41), we get,

The inequality $(i)$ uses eq. (43). Inequality $(ii)$ uses that $|I_{+\mu_{2}}^{\mathcal{C}}|\geq|I_{+\mu_{2}}|-|I_{+\mu_{2}}^{\mathcal{N}}|\geq|I_{+\mu_{2}}|-|{\mathcal{N}}|$. Inequality $(iii)$ uses the lower bound on the number of points in cluster $\mu_{2}$ given in Lemma 4.3 together with Assumption (A3). The final inequality follows by using Assumption (A2) and Lemma 4.3, which allow for us to take $\sigma\sqrt{d}$ and $|{\mathcal{N}}|/n$ smaller than an absolute constant. This shows that, in the case that $\langle w^{(t)}_{j},\mu_{2}\rangle\geq\alpha|a_{j}|$, the value of $\langle w^{(t+1)}_{j},\mu_{2}\rangle$ is strictly less than $\langle w^{(t)}_{j},\mu_{2}\rangle$. Since by Lemma 4.3 we have $\|x_{i}\|\leq\sqrt{2}$, we have,

As we have shown $\langle w^{(t)}_{j},\mu_{2}\rangle\geq\alpha|a_{j}|$, this implies $\langle w^{(t+1)}_{j},\mu_{2}\rangle\in\left[(1-2)\alpha|a_{j}|,\alpha|a_{j}|\right)$, and thus the inequality $|\langle w^{(t+1)}_{j},\mu_{2}\rangle|\leq 3\alpha|a_{j}|$ holds as desired. This completes the induction in the case that $\langle w^{(t)}_{j},\mu_{2}\rangle\geq\alpha|a_{j}|$.

For the case $\langle w^{(t)}_{j},\mu_{2}\rangle\leq-\alpha|a_{j}|$, we can use a nearly identical argument as above to show that $\newline \langle w^{(t+1)}_{j}~{}-~{}w^{(t)}_{j},\mu_{2}\rangle>0$ so that $\langle w^{(t+1)}_{j},\mu_{2}\rangle\in(-\alpha|a_{j}|,(-1+2)\alpha|a_{j}|]$. This again gives $|\langle w^{(t+1)},\mu_{2}\rangle|\leq 3\alpha|a_{j}|$.

The only remaining case is when $|\langle w^{(t)}_{j},\mu_{2}\rangle|\leq\alpha|a_{j}|$. In this case, (45) implies that we have the inequality $|\langle w^{(t+1)}_{j},\mu_{2}\rangle|\leq(1+2)\alpha|a_{j}|=3\alpha|a_{j}|$, completing the induction for the $J_{+\mu_{1}}$ neurons. The proof that almost-orthogonality holds for neurons $j\in J_{-\mu_{1}}\cup J_{\pm\mu_{2}}$ holds using an analogous argument.

A.4 Proof of Theorem 3.1

For the reader’s convenience, we restate the theorem below before completing its proof. See 3.1

First, note that with probability at least $1-4\delta$, a good run occurs, so that the results of Lemma 4.1, Lemma 4.2, Lemma 4.3, and Lemma 4.4 all hold for the absolute constant $C_{0}=4^{5}\cdot 1024^{2}\exp(4)$. We thus can apply Lemma 4.11 so that neuron alignment and almost-orthogonality hold for times $t=1,\dots,1/4\alpha$. Since neuron alignment and almost-orthogonality hold, by Lemma 4.9, we have,

In order to apply Lemma 4.6, which relates the prediction on the subnetwork to the entire network, we need to ensure that for $C_{f}=4\cdot 1024\exp(2)/(1-1/C_{0})^{2}$ we have $|J^{c}|/m\leq 1/16C_{f}^{2}$. If we denote by $J^{c}=[m]\setminus(J_{\pm\mu_{1}}\cup J_{\pm\mu_{2}})$, then $|J^{c}|/m\leq 1-(1-1/C_{0})^{2}\leq 2/C_{0}$, so that,

The equality $(i)$ follows since $C_{0}=4^{5}\cdot 1024^{2}\exp(4)$. Thus we may apply Lemma 4.6. Since $\|W^{(T)}\|_{F}\leq 1$ by Lemma 4.9, and since $\|x_{i}\|\leq 2$ by Lemma 4.3, the lower bound for clean samples given in (46) can be used in Lemma 4.6 to get,

Using a symmetric argument, we have that noisy samples satisfy

This shows that the neural network accurately classifiers all of the clean samples correctly at a margin of $\gamma>0$, and misclassifies all noisy samples incorrectly. Since we have the Frobenius norm bound $\|W^{(T)}\|_{F}\leq 1$, we can therefore use a simple Rademacher complexity-based argument to derive a generalization bound for the neural network. In particular, let us define the ramp loss

Then $r_{\gamma}$ is $\gamma^{-1}$-Lipschitz, and if we denote by