Feature Purification: How Adversarial Training Performs Robust Deep Learning

Introduction

Large scale neural networks have shown great power to learn from a training data set, and generalize to unseen data sampled from similar distributions for applications across different domains . However, recent study has discovered that these trained large models are extremely vulnerable to small “adversarial attacks” . It has been discovered that small perturbations to the input– often small enough to be invisible to humans– can create numerous errors in prediction. Such slightly perturbed inputs are often referred to as “adversarial examples”.

Since the original discovery of “adversarial examples”, a large body of works have been done emphasizing how to improve the robustness of the deep learning models against such perturbations . One seminal approach is called adversarial training , where one iteratively computes adversarial examples from the training examples, and then retrain the model with these adversarial examples instead of the original examples (a.k.a. the clean examples). This approach was reported in as the only approach that can defend against carefully designed adversarial attacks, and many follow-up works are built upon it .

However, despite the great empirical success on improving the robustness of neural networks over various data sets, the theory of the adversarial examples is much less developed. In particular, we found that the following fundamental questions remain largely unaddressed:

To answer these questions, one sequence of theoretical works try to explain the existence of adversarial examples using the high dimensional nature of the input space and the over-fitting behavior due to the sample size and sample noise , and treat adversarial training from the broader view of min-max optimization . However, recent observations indicate that these adversarial examples can also, and arguably often, arise from features (those that do generalize) rather bugs (those that do not generalize due to effect of poor statistical concentration). To the best of our knowledge, all existing works studying adversarial examples either (1) apply generally to the case of arbitrarily unstructured functions $f$ and only consider adversarial examples statistically, or (2) apply to a structured setting but only involving linear learners. These theoretical works, while shedding great lights to the study of adversarial examples, do not yet give concrete mathematical answers to the following questions regarding the specific hidden-layer structure of neural networks:

What are the features (i.e. the hidden weights) learned by the neural network via clean training (i.e., over the original data set)? Why are those features “non-robust”?

What are the differences between the features learned by clean training vs adversarial training (i.e., over a perturbed data set consisting of adversarial examples)?

Why do adversarial examples for a network transfer to other independently-trained networks?

Before going into the above questions regarding robustness, it is inevitable to first study what the features are when learned by a neural network during clean training. Theoretical studies are also limited in this direction. Most of existing works (1) only focus on the case when the training data is spherical Gaussian , and some of them require heavy initialization using tensor decomposition, which might fail to capture the specific structure of the input and the property of a random initialization; or (2) only consider the neural tangent kernel regime, where the neural networks are linearized so the features are not learned (they stay at random initialization) .

In this paper, we present a new routine that enables us to formally study the learned features (i.e. the hidden weights) of a neural network, when the inputs are more naturally structured than being Gaussians. Using this routine, we give, to the best of our knowledge, the first theoretical result towards answering the aforementioned fundamental questions of adversarial examples, for certain neural networks with ReLU activation functions.

Our results. We prove, for certain binary classification data set, when we train a two-layer ReLU neural network using gradient descent,Our theory extends to stochastic gradient descent (SGD) at the expense of complicating notations. starting from random initialization,

Given polynomially manly training examples, in polynomially many iterations, the neural network will learn well-generalizing features for the original data set, and the learned network will have close-to-perfect prediction accuracy for the test data sampled from the same distribution.

However, even with a weight-decay regularizer to avoid over-fitting, even with infinitely many training data, and even when super-polynomially many iterations are used to train the neural network to convergence, the learned network still has near-zero robust accuracy against small-norm adversarial perturbations to the data. In other words, those provably well-generalizing features on the original data set are also provably non-robust to adversarial perturbations to the datas, so they cannot be due to having too few training samples .

Adversarial training, using perturbation algorithms such as Fast Gradient Method (FGM) , can provably and efficiently make the learned neural network achieve near-perfect robust accuracy, against even the worst-case norm-bounded adversarial perturbations, using a principle we refer to as “feature purification”. We illustrate “feature purification” in Figure 1 by an experiment, and explain it in mathematical terms next.

Feature purification: How adversarial training can perform robust deep learning. In this work, we also give precise, mathematical characterizations on the difference between learned features by clean training versus adversarial training in the aforementioned setting, leading to (to our best knowledge) the first theory of how, in certain learning tasks using ReLU neural networks, the provably non-robust features after clean training can be “robustified” via adversarial training.

We emphasize that prior theoretical works mainly study adversarial examples in the context of linear models (such as linear regression, linear regression over prescribed feature mappings, or the neural tangent kernels). In those models, the features are not trained, so adversarial training only changes the weights associated with the linear combination of these features, but not the actual features themselves.

In contrast, this paper develops a theory showing that how, over certain learning tasks, adversarial training can actually change the features of certain neural networks to improve their robustness. We abstract this feature change in our setting into a general principle that we call feature purification, which although we only prove it for two-layer ReLU networks (see Theorem thm:intro:cleansa+Theorem thm:adv1_sc), we empirically observe that it occurs more generally to real-world, deep neural networks on real-world data sets. We sketch its high-level idea as follows.

Mathematically, as a provisional step to measure of change of features in a network, let us use (1) $w_{i}^{(0)}$ to denote the weight vector of the $i$-th neuron at initialization, (2) $w_{i}$ to denote its weight after clean training, and (3) $w^{\prime}_{i}$ to denote its weight after adversarial training (using $w_{i}$ as initialization). The “feature purification” principle, in math, says if we use $\theta(z,z^{\prime}):=\frac{|\langle z,z^{\prime}\rangle|}{\|z\|_{2}\|z^{\prime}\|_{2}}$ as a provisional measure of the correlation between “features”, then (see Figure 2 for real-life experiments):

for most neurons: $\theta(w^{(0)}_{i},w_{i}),\theta(w^{(0)}_{i},w^{\prime}_{i})\leq c$ for a small constant $c$ (such as $0.2$);

for most neurons: $\theta(w_{i},w^{\prime}_{i})\geq C$ for a large constant $C$ (such as $0.8$); and

for most pairs of different neurons: $\theta(w_{i},w_{j})\leq c$ for a small constant $c$ (such as $0.2$).

In words, this says both clean training and adversarial training discover hidden weights $w_{i},w^{\prime}_{i}$ that are fundamentally different from initialization $w^{(0)}_{i}$. However, since $w_{i}$ and $w^{\prime}_{i}$ are close, clean training must have already discovered a big portion of the robust features, and adversarial training merely needs to “purify” some small part of each original feature. In this paper:

we prove this feature purification principle in the case of two-layer ReLU neural networks over certain data sets, with $c=o(1)$ and $C=1-o(1)$ (see Theorem thm:intro:cleansa+Theorem thm:adv1_sc); and

we provide empirical evidence that this feature purification principle holds also for deep neural networks used in real-life datasets (see Figure 2 as well as other experiments in the paper).

Why clean training learns non-robust features? Which part of the features are “purified” during adversarial training? In our setting, we also give mathematical characterizations of where the “non-robust” part of each feature comes from during clean training.

As we shall formally discuss in Section 6.2, training algorithms such as gradient descent will, at every step, add to the current parameters a direction that maximally correlates with the labeling function on average. For two-layer ReLU networks, we prove that such simple correlations will accumulate, in each neuron, a small part of its weight that correlates with the average of the training data, and we refer to it as the dense mixture (see Theorem 5.2). However, under natural assumptions of the data such as the sparse coding model — namely inputs come from sparse combinations of hidden dictionary words/vectors— such dense mixtures cannot have high correlation with any individual, clean example. Thus, even with these “dense mixtures” in the features, the network can still generalize well on the original data set. However, we show that these portions of the features are extremely vulnerable to small, adversarial perturbations along the “dense mixture” directions. As a result, one of the main goals of adversarial training, as we show, is to purify the neurons by removing such dense mixtures. This is the supporting theory behind our feature purification principle, as we also measure and verify it empirically in the experiment section.

We believe our result extends the reach of traditional learning theory, where often statistical properties of the model (such as generalization, etc.) is separate from optimization (i.e., how the models are trained). However, to understand adversarial examples in deep learning, one needs to admit that well-generalizing and adversarially robust neural networks do exist (and can even be found efficiently using adversarial training), thus it is also a global optimal solution of the clean training objective. It is rather a property of traditional clean training process using SGD which biases the network towards learning non-robust networks as another global optimal solution of the training objective.

Moreover, in our setting, these dense mixtures in the hidden weights of the network come from the sparse coding structure of the data and the gradient descent algorithm. It is rather independent of the random initialization of the neural network. Thus, we prove that, at least in our scenario, adversarial examples for one network do transfer to other independently trained ones.

Our contribution to computation complexity. We also prove a lower bound that, for the same sparse coding data model, even when the original data is linearly-separable, any linear classifier, any low-degree polynomial, or even the corresponding neural tangent kernel (NTK) of our studied two-layer neural network, cannot achieve meaningful robust accuracy (although they can easily achieve high clean accuracy). Together with our upper bound, we have shown that using a higher-complexity model (such as a two-layer neural network with ReLU activation, comparing to NTK) can in fact achieve better robustness against adversarial perturbations. Thus, our theory strongly supports the experimental finding in , where experts have noticed that robustness against adversarial examples requires a model with higher complexity. The main intuition is that low-complexity models, including the neural tangent kernel, lacks the power to zero out low magnitude signals to improve model robustness, as illustrated in Figure 3 and Section 3.

Our experimental contributions. We present quite a few experimental results supporting our theory. We find that our sparse coding model can indeed capture real-world data to certain degree, and our principle of feature purification also holds for architectures such as AlexNet and ResNet. We showed during clean training, how the features can emerge from random initialization by wining the “lottery tickets”, as predicted by our theory. Perhaps most importantly, we confirmed through experiments that:

Adversarial training using adversarial examples indeed purify “dense mixtures” in practice.

To gain robust accuracy, it suffices to remove such “dense mixtures” (via a low-rank update).

We present our experiments following each of the theorem statements accordingly. We also include a whole Section 8 for more detailed experiments.

Adversarial examples: Empirical study. Since the seminal paper shows the existence of small adversarial perturbations to change the prediction of the neural networks, many empirical studies have been done to make the trained neural networks robust against perturbations (and we refer to the citations therein). The recent study shows that the seminal approach of adversarial training is the most effective way to make the neural networks robust against adversarial perturbations.

Adversarial examples: Theoretical study. Existing theories mostly explain the existence of adversarial examples as the result of finite-sample data set over-fitting to high-dimensional learning problems . Later, it is discovered by Ilyas et al. that well-generalizing features can also be non-robust. Other theories focus on the Fourier perspective of the robustness , showing that adversarial training might be preventing the network from learning the high frequency signals of the input image. Our theoretical work is fundamentally different from the aspect of poor statistical concentration over finite-sample data set, and our Theorem 5.1 and Theorem 5.3 strongly supports that a well-trained, well-generalizing neural network can still be non-robust to adversarial attacks.

Other theories about adversarial examples focus on how adversarial training might require more training data comparing to clean training , and might decrease clean training accuracy . The works by focus on how adversarial training can be performed efficiently in the neural tangent kernel regime. The purpose of these results are also fundamentally different than ours.

Sparse coding (data) model. We use a data model called sparse coding, which is a popular model to model image, text and speech data . There are many existing theoretical works studying algorithm for sparse coding , however, these algorithms share little similarity to training a neural network.

The seminal work by Arora et al. provides a neurally-plausible algorithm for learning sparse coding along with other works using alternative minimization . However, all of these results require a (carefully picking) warm start, while our theory is for training a neural network starting from random initialization.

Threshold degree and kernel lower bound. We also provide, to the best of our knowledge, the first example when the original classification problem is learnable using a linear classifier but no low-degree polynomial can learn the problem robustly against small adversarial perturbations. Yet, the high-complexity neural networks can provably, efficiently and robustly learn the concept class. The lower bound for the classification accuracy using low-degree polynomials has been widely studied as the (approximate) threshold degree of a function or the sign-rank of a matrix . Our paper give the first example of a function with high (approximate) robust threshold degree, yet efficiently and robustly learnable by training a ReLU neural network using gradient descent.

Other related works prove lower bounds for kernel method in the regression case . Generally speaking, such lower bounds are about the actual (approximate) degree of the function, instead of the (approximate) threshold degree. It is well know that for general functions, the the actual degree can be arbitrary larger than the threshold degree.

Preliminaries

We assume the hidden vector $z$ is “sparse”, in the following sense: for $k\leq d^{0.499}$, we have:

The coordinates of $z$ are independent, symmetric random variables, such that $|z_{i}|\in\{0\}\cup[\frac{1}{\sqrt{k}},1]$. Moreover,

Under Assumption 2.1, w.h.p., $\|z\|_{0}=\Theta(k)$ is a sparse vector.

We study the simplest binary-classification problem, where the labeling function is linear over the hidden vector $z$:

For simplicity, we assume $\forall i\in[D],|w^{\star}_{i}|=\Theta(1)$, so all the coordinates of $z$ have relatively equal contributions. Our theorems extend to other $w^{\star}$ at the expense of complicating notations.

Remark on sparse coding. The sparse coding model is very natural and is widely used to model image, text and speech data . There certainly exist (provable) algorithms for dictionary learning based on sum of squares, linear programming , but they do not shed light on the training process of neural networks. Even the neural algorithm for sparse coding is still far away from training a neural network using SGD or its variants. The main point of this paper is not to show neural networks can do sparse coding. Instead, our main point is to distinguish the adversarial training and the clean training processes of neural networks using the sparse-coding model as a bridging tool.

Noise model. We have allowed the inputs $x=\mathbf{M}z+\xi$ to incorporate a noise vector $\xi$. Our lower bounds hold even when there is no noise ($\xi=0$). Our upper bound theorems not only apply to $\xi=0$, but more generally to “gaussian noise plus spike noise”:

the noise $\xi$ can be of Euclidean norm $O(\sigma_{x})$, larger than the signal $\|\mathbf{M}z\|_{2}\approx 1$; and

the spike noise $\xi^{\prime\prime}_{i}$ can be $\frac{1}{k^{0.501}}$ which is the maximum possible (because $z_{i}$ can be $\frac{1}{k^{0.5}}$).

Clean and robust error. The goal of clean training is to learn a model $f$ so that $\mathsf{sign}(f(x))$ is as close to $y$ as possible. We define the classification error on the original data set as:

Warmup Intuitions

Linear learners are not robust. Given the setting of the data set, one direct approach is to use (the sign of) a linear classifier $f(x)=\langle w^{\star},\mathbf{M}^{\top}x\rangle$ to predict the label of $x$. There are two issues of using such a classifier:

When $\sigma_{x}$ is as large as $\Theta(1)$, such classifier can not even classify $x$ in good clean accuracy. Recall $f(x)=\langle w^{\star},\mathbf{M}^{\top}x\rangle=\langle w^{\star},z\rangle+\langle\mathbf{M}w^{\star},\xi\rangle$. By our assumption, typically $|\langle w^{\star},z\rangle|=O(1)$ and $\langle\mathbf{M}w^{\star},\xi^{\prime}\rangle\sim\mathcal{N}(0,\Theta(\sigma_{x}^{2}))$. Thus, when $\sigma_{x}\geq\Theta(1)$, noise could be much larger than signal, and this linear classifier cannot be used to classify $x$ correctly. In this case, actually no linear classifier (or even constant-degree polynomials One may think that using for example degree-3 polynomial $\sum_{i}w^{\star}_{i}\langle\mathbf{M}_{i},x\rangle^{3}$ can reduce the level of noise, but due to the diversity in the value of $z_{i}$ when $z_{i}\not=0$, one must use something close to linear when $|z_{i}|$ is large. Applying Markov brothers’ inequality, one can show the low-degree polynomial must be close to a linear function.) can give meaningful clean accuracy.

High-complexity models are more robust. Another choice to learn the labeling function is to use a higher-complexity model $f(x)=\sum_{i\in[d]}w^{\star}_{i}\langle\mathbf{M}_{i},x\rangle\mathds{1}_{|\langle\mathbf{M}_{i},x\rangle|\geq\frac{1}{2\sqrt{k}}}$. Here, the “complexity” of $f$ is much higher because an indicator function is used.One concrete measure of “higher complexity” is that $f$ cannot be well-approximated by low degree polynomial. Since $\langle\mathbf{M}_{i},x\rangle=z_{i}+\langle\mathbf{M}_{i},\xi\rangle$, by our noise model, as long as the signal $z_{i}\not=0$ is non-zero, $|\langle\mathbf{M}_{i},x\rangle|\geq\frac{1}{2\sqrt{k}}$ with high probability. Thus, this $f(x)$ is equal to the true labeling function $\langle w^{\star},z\rangle$ w.h.p. over the original data set, so is (much) more robust to noise comparing to linear models.

To sum up, higher-complexity models (such as those using ReLU) have the power to zero out low-magnitude signals to improve adversarial robustness, as illustrated in Figure 3.

Learning robust classifier using neural network. Motivated by the above discussions between linear vs. high-complexity models, our goal is to show that a two-layer neural networks can (after adversarial training) learn a robust function $f(x)$ such as

Here, $\mathsf{ReLU}(y)=\max\{y,0\}$ is the ReLU function and $b\approx\frac{1}{2\sqrt{k}}$. In this paper, we present a theorem stating that adversarial training of a (wlog. symmetric) two-layer neural network can indeed recover a neural network of this form. In other words, after adversarial training, the features learned by the hidden layer of a neural network can indeed form a basis (namely, $\mathbf{M}_{1},\dots,\mathbf{M}_{d}$) of the input $x$ where the coefficients are sparse. We also present a theorem showing why, clean training will not learn this robust function. We also verify experimentally that the features learned by the first layer of AlexNet (after adversarial training) indeed form a sparse basis of the images, see Figure 4.

Learner Network and Adversarial Training

In this paper we consider a simple, two layer (symmetric) We assume the neurons are symmetric (i.e., with $(w_{i},-w_{i})$ pairs) to simplify proofs. neural network with ReLU activation.

To simplify analysis, we fix $a_{i}=1$ throughout the training. We use $w_{i}^{(t)}$ to denote the hidden weights at time $t$, and use $f_{t}(w;x,\rho)$ to denote the network at iteration $t$

Given a training set $\mathcal{Z}=\{x_{j},y_{j}\}_{j\in[N]}$ together with one sample of pre-activation noise $\rho^{(j)}$ for each $(x_{j},y_{j})$, we define

In our case, the (clean, population) classification error at iteration $t$ is

At initialization, we let $w_{i}^{(0)}\sim\mathcal{N}\left(0,\sigma_{0}^{2}\mathbf{I}\right)$ for $\sigma_{0}=\frac{1}{\poly(d)}$ and let $b^{(0)}=\Theta(\sigma_{0}\sqrt{\log d})$. When near initialization, we manually increase the bias $b^{(t+1)}=b^{(t)}+\eta\mathfrak{B}$ where $\mathfrak{B}=\frac{c_{b}}{d}$ for some small constant $c_{b}>0$— this corresponds to the “lottery ticket winning” phase to be discussed later in Section 6.1; and whenever $b^{(t)}$ reaches $\frac{1}{k^{0.5001}}$ we set $\mathfrak{B}=0$— in this phase, the neurons that have won the “lottery ticket” will keep winning and grow significantly, to be discussed in Section 6.2.

We also choose pre-activation noise $\sigma_{\rho}^{(t)}=\frac{b^{(t)}}{\sqrt{\log d}}\cdot\Theta((\log\log\log d)^{3})$ for $t\leq T_{\mathsf{a}}=\frac{1}{\poly(d)\eta}$, and $\sigma_{\rho}^{(t)}=\frac{b^{(t)}}{\log d}\cdot\Theta((\log\log\log d)^{3})$ for $t>T_{\mathsf{a}}$. The explicit choices of $\mathfrak{B}$ and $T_{\mathsf{a}}$ are given in the proofs.

2 Adversarial Training

We state the adversarial training algorithm in Algorithm 2. It takes as input a perturbation algorithm $A$, and repeatedly applies gradient descent over a perturbed data set (that comes from the original data set plus the perturbation given by $A$). Formally,

In contrast, the empirical robust classification error against algorithm $A$ is

Our upper bound theorems apply to all perturbation algorithms under Definition 4.2, and gives small empirical robust error $\widehat{\mathcal{E}^{r}}$. To obtain small (true) robust error $\mathcal{E}^{r}$, as we shall see, one can for instance let $A$ be the fast gradient method (FGM).

Initialized from clean training. In this paper, we assume adversarial training (i.e., Algorithm 2) is initialized from a network that is already clean-trained. In contrast, in practice, adversarial training usually begins directly with random initialization. We remark here that:

First, in practice, adversarial training from a clean-trained initialization performs no worse than from a random initialization, see Table 1 on Page 1. In fact, it is sometimes even beneficial to begin with clean training and gradually switch to adversarial training (see e.g. ).

Second, to prove our main conceptual message— feature purification— it is convenient to start from a clean-trained model, and then try to understand which part of the features are changed after robust training. Since neural nets have lots of equivalent transformations that are not very well-understood (even in two-layer case), if we adversarially train it from random initialization, then it is theoretically very hard to quantify how it is related to another clean-trained model learned from random initialization (since we need to understand all the invariants).Even if one performs clean/adversarial training from the same random initialization, the additional randomness in SGD may quickly make the two models diverge from each other.

Yet, our theory still gives support to what happens in adversarial training from random initialization. As we shall prove, this “$o(1)$ feature change” comes from dense mixtures directions (see Theorem 5.2). Thus, adversarial training from random initialization should directly avoid learning such dense mixtures, as opposed to first learning them (by clean training) and then forgetting (by adversarial training). We illustrate this in Figure 9.

Statements and Explanations of Our Main Results

There exists an absolute constants $C,c>0$ such that for every constant $c_{0}\in(0,c]$, every $d$ and $m$ with $m=d^{1+c_{0}}$, given $N\geq\Omega(d^{C})$ many training data, for every random initialization weight $\sigma_{0}=\frac{1}{\poly(d)}$, for every learning rate \eta\in\big{(}0,\frac{1}{\Omega(d^{C})}\big{]}, if we define $T_{\mathsf{c}}:=\Theta(\frac{d^{1.01}}{\eta})$, then for every $T_{\mathsf{f}}\in[T_{\mathsf{c}},d^{\log d}/\eta]$, the following holds with high probability. The network $f_{t}$ with hidden weights $\{w_{i}^{(t)}\}_{i\in[m]}$ learned by clean training Algorithm 1 satisfies:

Global feature learning: for every $t\in[T_{\mathsf{c}},T_{\mathsf{f}}]$,

Clean training has good clean accuracy: for every $t\in[T_{\mathsf{c}},T_{\mathsf{f}}]$,

Clean training is not robust to small adversarial perturbations: for every $t\in[T_{\mathsf{c}},T_{\mathsf{f}}]$, every $\tau\geq\frac{1}{k^{0.5+10c}}$, using perturbation $\delta=-\tau\frac{y\mathbf{M}w^{\star}}{\|\mathbf{M}w^{\star}\|_{2}}$ (which does not depend on $f_{t}$),

Theorem 5.1 indicates that in our setting, clean training of the neural network has good clean accuracy but terrible robust accuracy. Such terrible robust accuracy is not due to over-fitting, as it holds even when a super-polynomially many iterations and infinitely many training examples are used to train the neural network. In the next theorem, we give a precise characterization of what the hidden weights $\{w_{i}\}$ are after clean training, and why they are not robust.

For every neuron $i\in[m]$, there is a fixed subset $\mathcal{N}_{i}$ of size $|\mathcal{N}_{i}|=O(1)$ such that, for every $t\in[T_{\mathsf{c}},d^{\log d}/\eta)$,

where (1) $|\beta_{i,j}|<\frac{k}{d^{1-c}}$ for some small constant $c\in[0,0.001]$, and (2) for at least $\Omega(d)$ many neurons $i\in[m]$, it satisfies $|\mathcal{N}_{i}|=1$ and $\alpha_{i,j}>d^{-c}$. Moreover,

Theorem 5.2 says that each neuron $w_{i}$ will learn constantly many (allegedly large) components in the directions \big{\{}\mathbf{M}_{j}\colon j\in\mathcal{N}_{j}\big{\}}, and its components in the remaining directions \big{\{}\mathbf{M}_{j}\colon j\not\in\mathcal{N}_{j}\big{\}} are all small. We emphasize that the sets $\mathcal{N}_{i}$ are independent of $t$ but are solely determined by random initialization. In other words, for each neuron $i$, which setset $\mathcal{N}_{i}$ it “wins” is completely determined by the “lottery ticket” (its random initialization). We discuss this in more details in Section 6.1.

Furthermore, Theorem 5.2 shows that instead of learning the pure, robust features $\{\mathbf{M}_{j}\}_{j\in[d]}$, intuitively, ignoring the small $d^{c}$ factors, focusing only on those neurons with $|\mathcal{N}_{i}|=1$, and assuming for simplicity all the $\beta_{i,j}$’s are of similar (positive) magnitude, then, clean training will learn neurons:

Feature purification: mathematical reasoning. Eq. (5.1) says that, after clean training, the neural network will be able to learn a big portion of the robust feature, $\Theta(1)\mathbf{M}_{j}$, plus some small dense mixture v=\sum_{j^{\prime}\not=j}\big{[}\Theta\left(\frac{k}{d}\right)w_{j^{\prime}}^{\star}\mathbf{M}_{j^{\prime}}\big{]}. In our sparse coding model, each $x$ is of form $x=\mathbf{M}z+\xi$, where $z$ is a sparse vector and $\xi$ is the noise. One critical observation is that such dense mixture $v$ has low correlation with almost all inputs $x$ from the original distribution, so it has negligible effect for clean accuracy. However, such dense mixture is extremely vulnerable to small but dense adversarial perturbations of the input along this direction $v$, making the model non-robust.

As we point out, such “dense adversarial perturbation” directions do not exist in the original data.One can try to add these dense mixtures directly to the training data set, which we conjecture to be similar to the approach in Thus, one has to rely on adversarial training to remove dense mixtures to make the model robust. This is the main spirit of our feature purification principle, and we illustrate it in Figure 6.

Where does dense mixture come from? We shall explain in more details in Section 6.2, but at a high level, in each iteration, the gradient $\nabla\mathbf{Obj}$ will bias towards the direction that correlates with the labeling function $y=\mathsf{sign}(\langle w^{\star},z\rangle)$; and since $x=\mathbf{M}z+\xi$ in our model, such direction should be $\mathbf{M}w^{\star}=\sum_{j}w_{j}^{\star}\mathbf{M}_{j}$, so is a dense mixture direction and will be accumulated across time. The accumulation of dense mixture is consistent with the finding (for linearly-separable data).

However, as we have argued in Section 3, in our setting when the noise level $\sigma_{x}\geq\Omega(1)$ is large, such dense mixture direction cannot be used to given even good clean accuracy. Therefore, during clean training, the neural network has the incentive to discover features close to $\{\mathbf{M}_{j}\}_{j\in[d]}$ because they can “de-noise” $\xi$ better (see discussions in Section 3). Yet, our critical observation is that, even for well-trained neural network which aims to de-noise $\xi$, even when the neurons are close to being pure features $\{\mathbf{M}_{j}\}_{j\in[d]}$, the “dense direction” still locally correlates with the labeling function $y$, and thus can still be accumulated during the course of a local training algorithm such as gradient descent, leading to the small, non-robust part of each feature.

Next, we state the theorem for adversarial training (recall Algorithm 2). It shows that adversarial training indeed purifies the small dense mixtures, leading to local changes of the weights.

Empirical robust accuracy: for $t=T_{\mathsf{f}}+T_{\mathsf{g}}$

Provable robust accuracy: when $A$ is the fast gradient method (FGM), for $t=T_{\mathsf{f}}+T_{\mathsf{g}}$

Feature (local) purification: for every $t\in[T_{\mathsf{f}},T_{\mathsf{f}}+T_{\mathsf{g}}-1]$,

We emphasize that Theorem thm:adv1_sa holds for any perturbation algorithm $A$ satisfying Definition 4.2, when we only concern the robustness of the network against $A$. Meaning that the local feature purification happens regardless of which adversarial perturbation algorithm is used to find the adversarial examples. More surprisingly, Theorem thm:adv1_sb says when a good perturbation algorithm such as FGM is used, then not only the robustness generalizes to unseen examples, it also generalizes to any worst-case perturbation algorithm.We point out that since our results hold for any perturbation algorithm $A$, it is impossible to characterize exactly what are the learned features after training (for example showing that they corresponds to the actual dictionary) since $A$ might be a bad adversarial perturbation finding algorithm and the network does not even need to remove all the dense mixture in order to fool $A$. However, the true robustness given by Theorem thm:adv1_sb does imply that the dense mixture should be removed at least in terms of functionality, if one uses a good adversarial perturbation finding algorithm.

Given that this “dense mixture” direction $\mathbf{M}w^{*}$ is common across neurons, one may think that during adversarial training, instead of training all the parameters, it may be sufficient to train a low-rank update on top of the clean-trained model. In practice, this indeed works very well, see Section 8.5.

Density of adversarial perturbation. Our previous theorem suggests that one of the main goals of adversarial training is to remove dense mixtures to make the network more robust. Therefore, before adversarial training, the adversarial perturbations are dense in the basis of $\{\mathbf{M}_{j}\}_{j\in[d]}$; and after adversarial training, the adversarial perturbations ought to be more sparse and aligned with inputs from the original data set. Figure 7 has confirmed this theoretical finding using real-life data sets. Later in Section 8.3, we also present concrete measurements of the sparsity of these adversarial perturbations, and compare them in Figure 12.

This gives a gap because $c_{0},c_{1}$ can be made arbitrarily small. (see Theorem E.1 and Theorem F.4)

We also show a lower bound that no low-degree polynomial, or even the corresponding neural tangent kernel (NTK), can robustly learn the concept class. Recall for our two-layer ReLU network,

The feature mapping of the neural tangent kernel for our two-layer network $f$ is

Therefore, given weights $\{v_{i}\}_{i\in[m]}$, the NTK function $p(x)$ is given as

Without loss of generality, assume each $w_{i}\sim\mathcal{N}(0,\mathbf{I})$.

In this paper, we consider a wide range of NTK parameters: $\rho_{i}\sim\mathcal{N}(0,\sigma_{\rho_{i}}^{2})$ for arbitrary $\sigma_{\rho_{i}}\in[0,d^{o(1)}]$ and $|b_{i}|\leq d^{o(1)}$.

Our lower bound holds even for a most simple case $\mathbf{M}=\mathbf{I}$ and $\sigma_{x}=0$, so the original concept class is linearly separable. We prove the following:

(In contrast, Theorem 5.5 says adversarial training of neural network gives robust radius $\tau=\frac{1}{k^{1.76}}$.)

Since a poly-sized NTK kernel is known to be powerful enough to incorporate any low complexity functions (such as constant-degree polynomials) , we have the following corollary.

In the same setting as Theorem 5.7, if $q(x)$ is a constant degree polynomial, then we also have the robust error $\mathcal{E}^{r}(q)\geq\frac{1-o(1)}{2}$.

Overview of the Training Process

In this section, we present an overview of the proof for the training process, using gradient descent starting from random initialization. The complete proof is deferred to the Appendix.

Our proof begins by showing how the features in the neural network are emerged from random initialization. In this phase, the loss function is not sufficiently minimized yet, so the classification accuracy remains around $50\%$. However, we prove in this phase, gradient descent can already drive the neural network to learn a rich set of interesting features out of the random initialization. We call this process “lottery ticket winning” near random initialization, which is related to the study of .

This “lottery ticket winning” process is fundamentally different from the neural tangent kernel analysis (e.g. ). In this phase, although the loss is not sufficiently minimized, the activation patterns of the ReLU activations have changed dramatically, so that they have little correlations with the random initialization. Yet, we develop a new theoretical technique that allows us to control the change of the weights of the neurons, as we summarize below.

We derive the following property at random initialization. At iteration $t=0$, the hidden weights are initialized as w_{i}^{(0)}\sim\mathcal{N}\big{(}0,\sigma_{0}^{2}\mathbf{I}_{d\times d}\big{)}. Using standard properties of Gaussians, we show the following critical property: as long as $m\geq d^{1.01}$, there exists small constants $c_{3}>c_{4}>0$ such that

For most of the neurons $i\in[m]$, $\max_{j\in[d]}\{\langle\mathbf{M}_{j},w_{i}^{(0)}\rangle^{2}\}\leq 2\sigma_{0}^{2}\log d$.

For at most $\frac{1}{d^{c_{4}}}$ fraction of of the neurons $i\in[m]$, there is a dimension $j\in[d]$ with $\langle\mathbf{M}_{j},w_{i}^{(0)}\rangle^{2}\geq 2.01\sigma_{0}^{2}\log d$.

For at least $\frac{1}{d^{c_{3}}}$ fraction of of the neurons $i\in[m]$, there is one and only one $j\in[d]$ such that $\langle\mathbf{M}_{j},w_{i}^{(0)}\rangle^{2}\geq 2.02\sigma_{0}^{2}\log d$, and all the other $j^{\prime}\in[d]$ satisfies $\langle\mathbf{M}_{j^{\prime}},w_{i}^{(0)}\rangle^{2}\leq 2.01\sigma_{0}^{2}\log d$.

In other words, even with very mild over-parameterization $m\geq d^{1.01}$, by the property of random gaussian initialization, there will be some “potentially lucky neurons” in (ii), where the maximum correlation to one of the features $\mathbf{M}_{j}$ is slightly higher than usual. Moreover, there will be some “surely lucky neurons” in (iii), where such “slightly higher correlation” appears in one and only one of the target features $\mathbf{M}_{j}$.

In our proof, we denote the set of the neurons in (iii) whose correlation with $\mathbf{M}_{j}$ is slightly higher than usual as the set $\mathcal{S}_{j,sure}^{(0)}$, and denote those in (ii) as $\mathcal{S}_{j,pot}^{(0)}$. We will identify the following process during the training, as given in Theorem C.1:

In other words, if neuron $i$ wins the lottery ticket at random initialization, then eventually, it will deviate from random initialization and grow to a feature that is more close to (a scaling of) $\mathbf{M}_{j}$. Our other main observation is that if we slightly over-parameterize the network with $m\geq d^{1.001}$, then for each $j\in[d]$, $|\mathcal{S}_{j,sure}^{(0)}|\geq 1$ and $|\mathcal{S}_{j,pot}^{(0)}|\leq d^{0.01}$. Or in words, for each dimension $j\in[d]$, the number of lottery tickets across all neurons is at most $d^{0.01}$, but at least one neuron will win a lottery ticket (see Lemma B.2). We also illustrate the lottery ticket winning process experimentally in Figure 8.

2 The Formation of “Dense Mixtures” During Training

The next phase of our analysis begins when all the neurons already won their lottery tickets near random initialization. After that, the loss starts to decrease significantly, so the (clean) classification error starts to drop. We shall prove that in this phase, gradient descent will also accumulate, in each neuron, a small “dense mixture” that is extremely vulnerable to small but adversarial perturbations. To show this, we maintain the following critical property as given in Theorem C.2:

If a neuron $i$ wins the lottery ticket for feature $\mathbf{M}_{j}$ near random initialization, then it will keep this “lottery ticket” throughout the training.

Or in math words, for each neuron $i$, after $\langle\mathbf{M}_{j},w_{i}^{(t)}\rangle^{2}$ becomes sufficiently larger than all the other $\langle\mathbf{M}_{j^{\prime}},w_{i}^{(t)}\rangle^{2}$ at the first stage, it will stay much larger than other $\langle\mathbf{M}_{j^{\prime}},w_{i}^{(t)}\rangle^{2}$ for the remaining of the training process. To prove this, we introduce a careful coupling between the (directional) gradient of the neuron, and the (directional) Lipschitz continuity of the network $f_{t}$, this is given in Section C.4.2.

The vulnerable dense mixtures. The most critical observation in this phase is the formation of “dense mixtures”, where we show that even for the “lucky neuron” that wins the lottery ticket, the hidden weight of this neuron will look like (see Theorem 5.2)

In other words, up to scaling, these neurons will look like $w_{i}\approx\mathbf{M}_{j}+v_{i}$, where $v_{i}$ is a “dense mixture” $v_{i}=\Theta\left(\frac{k}{d}\right)\sum_{j^{\prime}\not=j}w_{j^{\prime}}^{\star}\mathbf{M}_{j^{\prime}}$.

The key observation is that $v_{i}$ is small and dense, in the sense that it is a mixture of all the other features $\{\mathbf{M}_{j^{\prime}}\}_{j^{\prime}\in[d]}$, but each of the feature has a much smaller contribution comparing to the leading term $\mathbf{M}_{j}$. Recall in our sparse coding model, each input $x=\mathbf{M}z+\xi$; so with high probability:

This value is even smaller than $\frac{1}{k}$ when $k\leq\sqrt{d}$. Thus, this dense mixture will not be correlated with any particular natural input, and thus the existence of these mixtures will have negligible contribution to the output of $f_{t}$ on clean data.

However, if we perturb input $x$ along the dense direction $\delta\propto\sum_{j^{\prime}\in[d]}\mathbf{M}_{j^{\prime}}$, we can observe that:

Comparing this with Eq (6.2), such “dense perturbation” can change the output of the neural network $f_{t}$ by a lot, using a small $\delta$ whose norm is much smaller than that of $x$. Thus, at this phase, even when the network has a good clean accuracy, it is still non-robust to these small yet dense adversarial perturbations. Moreover, this perturbation direction is “universal”, in the sense that it does not depend on the randomness of the model at initialization, or the randomness we use during the training. This explains transfer attacks in practice: that is, the adversarial perturbation found in one model can also attack other models that are independently trained.

Feature purification. Since Eq. (6.2) suggests most original inputs have negligible correlations with each dense mixture, during clean training, gradient descent will have no incentive to remove those mixtures. Thus, we have to rely on adversarial training to purify those dense mixtures by introducing adversarial examples. Those examples have correlation with $v_{i}$’s that are higher than usual. As we prove in Theorem 5.1 and illustrate in Figure 1, such “purifications”, albeit imposing only a small change to each neuron, will greatly improve the robustness of the neural network.

The formation of the dense mixtures. To further help the readers understand how those “dense mixtures” are formed, we sketch the proof of Theorem 5.2, which shows why clean training is provably non-robust. The main observation is that when the dense mixtures are small, the negative gradient of the (say, population) loss with respect to each neuron $w_{i}$ is approximately given by (recall $x=\mathbf{M}z+\xi$):

We emphasize that this is indeed a special property of gradient descent. Consider again the case $\sigma_{x}=\Omega(1)$ discussed in Section 3, where $x=\mathbf{M}z+\xi$ with $\|\xi\|_{2}=\Omega(1)=\Omega(\|\mathbf{M}z\|_{2})$. With high probability, a linear classifier using direction $\mathbf{M}w^{\star}$ cannot be used to classify $x$ correctly. Yet, this direction $\mathbf{M}w^{\star}$ is still locally positively correlated with the labeling function $y$, especially for well-trained, well-generalizing neural networks when the $\xi$ can be “de-noised”. (Stochastic) gradient descent, as a local update algorithm, only exams the local correlation between the update direction and the labeling function, and it does not exam whether this direction can be used in the final result. Thus, this dense direction $\mathbf{M}w^{\star}$ will be accumulated step by step, leading to a non-robust part of each of the features during clean training. In fact, even if we use $w_{i}=\mathbf{M}_{i}$ as initialization as opposed to random initialization, continuing clean training will still accumulate these small but dense mixtures. We illustrate this in Figure 9.

Conclusion

In this paper, we made a first step towards understanding how, in principle, the features in a neural network are learned during the training process, and why after clean training, these provably well-generalizing features are still provably non-robust. Our main conclusion is that during the clean training process using (stochastic) gradient descent, neural network will accumulate, in all features, some “dense mixture directions” that have low correlations with any natural input, but are extremely vulnerable to (dense) adversarial perturbations. During adversarial training, such dense mixtures are purified to make the model more robust. Our results suggest that the non-robustness of clean training is mainly due to two reasons:

the inductive bias of (stochastic) gradient descent, and

the “sparse coding” structure of the data.

Both reasons are necessary in some sense. First, a robust model is also a global minimizer of the clean training objective (at least in our setting); but even with proper regularization and infinite training examples to avoid over-fitting, gradient descent still has inductive bias towards finding a non-robust model. Second, it is easy to come up with data sets— such as linear-classifier labels over well-conditioned mixture-of-Gaussians like inputs— where clean training using gradient descent directly achieves the best robust accuracy. Thus, to understand the non-robustness of neural networks, we more or less have to take into account the gradient descent algorithm and the structure of the inputs.

Indeed, our step is still very provisional. We immediately see a plethora of extensions from our work. First of all, natural images have much richer structures than sparsity; hence, those “non-robust mixtures” accumulated by clean training might also carry structural properties other than density. Moreover, we would like to extend our work to the clean and robust training of multi-layer neural networks, possibly with hierarchical feature purification processes. (Our experiments in Figure 10 have confirmed on such hierarchical feature purification phenomenon.) Indeed, understanding the whole picture of adversarial examples and adversarial training might require a complete understanding of deep learning.

Experiment Details

We perform experiments using three standard architectures, AlexNet, ResNet-16, and ResNet-34 with basic blocks, and tested on the CIFAR-10 dataset.We used the implementations from https://github.com/bearpaw/pytorch-classification. We used their default random crop and random flip as data augmentation.

We discover that learning rate $0.1$ for good for ResNet and $0.02$ is good for AlexNet; while weight decay $0.0001$ is good for ResNet and $0.0005$ is good for AlexNet (this was also recommended by the git repo authors). We use standard SGD with 0.9 momentum as the training algorithm. During adversarial training, we have implemented:

In Table 1, we present robust/clean accuracies against such attackers after vanilla clean training / vanilla adversarial training. We emphasize here that, in practice, one can also first perform clean training and then apply adversarial traing using the clean-trained weights as initialization (like we have theoretically studied in this paper). This does not affect the overall performance of both robust and clean accuracies.

Visualizing the first layer of any trained architecture is trivial: for instance, for AlexNet, the weight tensor of the first layer is $3\times 11\times 11$ which gives the RGB color of $11\times 11$ patches (and this was precisely what we presented in Figure 1). However, such visualization can be less meaningful for ResNet because the tensors are of dimension $3\times 3\times 3$.

Visualizing the features presented by deeper convolutional layers is an active research area, dating back at least to . Perhaps the most naive approach is to start from a randomly initialize image (of size $3\times 32\times 32$), then take a specific neuron $n$ at some layer, and repeatedly take its gradient with respect to the image. If we keep adding this gradient to the input image, then ideally this gives us the image which “excites” $n$ the most. Unfortunately, it is a common knowledge in this area that this naive approach does not lead to “visually meaningful” images as we go (even slightly) deeper into a network (see e.g. the left column of Figure 10).

In existing literature, researchers have tried to various ways to resolve this issue (see e.g. an extensive survey by Olah et al. and the references therein). At a high level, some penalizes the image to remove high-frequency noise ; some searches for images that can still excite the given neuron after jittering ; and some searches only in the space of “real data” by building a model (e.g. using GAN) to capture the prior .

We observe that, if the model is robustly trained, then one can directly apply the naive approach to visualize features of the deep layers, and the resulting images can be “visually very meaningful.” See Figure 10.This should not be surprising given that the “jittering” technique is known to work in practice on visualizing clean models. Our theory in fact explains this phenomenon: the dense mixtures accumulated during clean training are extremely harmful to the visualization effect, since they are “visually meaningless.” After robust training, such dense mixtures are removed so the visualization starts to align better with human concepts.

Throughout this paper we stick to this naive approach for visualizing features of deep layers.Specifically, starting from a random input image, we take 2000 gradient steps to update the image so that the given neuron at a specific layer is excited the most. We added a weight decay factor to incentivize the image to go to RGB (128,128,128) — except in Figure 6 we incentivize the image to go to RGB (0,0,0).

2 Feature Purification at Deeper Layers

(stage 1) perform $T$ epochs of adversarial training;

3 Sparse Reconstruction of Input Data and of Adversarial Perturbation

Perhaps more importantly, our theory suggests that for clean-trained models, adversarial perturbations (we refer to as clean_delta) have “dense mixtures”; while for robust-trained models, adversarial perturbations (we refer to as robust_delta) are “more pure.” This was visually illustrated in Figure 7. Now, to better quantify this observation, we compare how sparse clean_delta and robust_delta can be reconstructed from robust features. See the second row of Figure 12.In fact, we have also re-scaled the perturbations so that they have similar mean and standard deviations comparing to real input images. This allows one to also compare the two rows of Figure 12. From this experiment, we confirm that in practice, adversarial perturbations on robust models are more “pure” and closer to real input images.

We point out when comparing how sparse clean_delta and robust_delta can be reconstructed from robust features, we did not cheat. For instance, in principle clean_delta may not lie in the span of robust features and if so, it cannot be (sparsely) reconstructed from them. In our experiments (namely, the second row of Figure 12), we noticed that clean_delta almost lies in the span of robust features (with regression error $<0.00005$ for AlexNet and $<10^{-9}$ for ResNet).

4 Comparing Different Attackers

We also demonstrate in Figure 13 that feature purification occurs against several different attackers.

5 Feature Purification is a Low-Rank Update

Recall from Theorem 5.2 and illustrated in Section 6.2 that the non-robustness of neurons in a clean-trained model, only comes from a common dense mixture direction $\mathbf{M}w^{\star}$. This suggests, during the robust (e.g., adversarial) training, we do not need to re-train all the parameters; it suffices to search only for a hidden mixture direction. We use experiments to support this finding.

Before we do so, please note we derived this “common dense direction” theory using a two-layered, binary classification setting. When there are multiple classes, one would expect there to be more dense mixture directions. In such a case, “low-rank update” is a more suitable choice.Specifically, consider a Conv2D unit of $k\times k$ kernel size from $n$ in-channels to $m$ out-channels, its weight matrix (tensor) $\mathbf{W}$ is of dimension $m\times n\times k\times k$. When performing a rank-$r$ update to it, we can construct two other Conv2D units, $\mathbf{A}$ of dimension $r\times n\times k\times k$, and $\mathbf{B}$ of dimension $m\times r\times 1\times 1$. Now, during adversarial training, we perform low-rank update $\mathbf{W}+\mathbf{B}\mathbf{A}$ by fixing $\mathbf{W}$ to be the clean-train parameters, and only letting $\mathbf{A},\mathbf{B}$ be trainable. For simplicity, we use zero initialization for $\mathbf{A}$ and Gaussian initialization for $\mathbf{B}$, to ensure a smooth transition between clean and adversarial training. As for parameters, for this simple illustration we did not tune much, and simply set $0.05$ for learning rate and 5e-5 for weight decay. Both of them are just slightly smaller than the commonly used choices (for training ResNet) as we now have much fewer number of trainable parameters.

Now on CIFAR-10, we compare clean training, (traditional) adversarial training, as well as first conducting clean training and then performing an adversarially trained, low-rank update on all the convolutional parameters. We use the (pre-activation) ResNet-28 architecture as well as its widened versions ResNet-28-3/5/10 . Recall ResNet-28-$x$ has three groups of basic convolutional layers, each of $16x,32x$ and $64x$ channels respectively.

We use “rank $r$” to denote a rank-$r$ update to all of the convolutional parameters.

We use “rank $r^{*}$” to denote rank-$r$, $2r$, $4r$ updates respectively to the convolutional parameters in the three groups.

We present our findings in Figure 14. For example, for the ResNet-28-10 model:

Using the clean-trained model weights alone, the robust accuracies are zeros.

Now, on top of such clean-trained weights, barely training a rank-1 (or rank $1^{*}$) update for each convolutional matrix, or equivalently only about $\frac{1}{160}$ (less than 1 percent) of the parameters, one can already recover more than 80% of the maximum robust accuracy.

Similarly, in the “rank $8^{*}$” case we train only about $5\%=\frac{8}{160}$ of the parameters, one can recover more than 90% of the maximum robust accuracy.

We believe this preliminary experiment can be useful in supporting our theory, and may be of independent interests for other applications. (Indeed, in a follow-up work we found an application of such low-rank update to language model fine-tuning .)

We give a quick overview of the structure of our appendix sections.

In Section A, we warm up the readers by calculating the gradient of the objective, and demonstrating that polynomially many samples are sufficient for the training.

In Section B, we formally introduce $\mathcal{S}_{j,pot}^{(t)}$, the set of “potentially lucky neurons” and $\mathcal{S}_{j,sure}^{(t)}$, the set of “surely lucky neurons” at iteration $t$. In particular, we shall emphasize on how those notions evolve as $t$ increases.

In Section C, we formally prove how “lucky neurons” continue to be lucky, and more importantly, for every neuron $i$ that is lucky in direction $j$, why it grows faster than other unlucky directions $j^{\prime}$, and how much faster. Specifically, Theorem C.1 corresponds to the initial “lottery-winning” phase where the accuracy remains around $50\%$; and Theorem C.2 corresponds to the later phase where large signals become even larger and eventually most neurons become “pure + dense mix” of the form (6.1). This is the most difficult section of this paper.

In Section D, we prove that why clean training gives good clean (testing) accuracy. It is based on the structural theorem given by Theorem C.2, and requires some non-trivial manipulations of probability theory results (such as introducing a high-probability, Bernstein form of the McDiarmid’s inequality).

In Section E, we prove that why the model obtained from clean training is non-robust. It formally shows how the “dense mixtures” become accumulated step by step during clean training.

In Section G, we prove lower bounds for the neural tangent kernel model given by two-layer networks.

In Section H, we give missing details of some probability theory lemmas.

Appendix A Notations and Warmups

We find it perhaps a good exercise to do some simple calculations to warmup the readers with our notations, before going into the proofs.

Global Assumptions. Throughout the proof,

We choose $m=d^{1+c_{0}}$ for a very small constant $c_{0}\in(0,1)$.

(One should think of $c_{0}=0.0001$ for a simple reading. Our proof generalizers to larger $m=\poly(d)$ since having more neurons does not hurt performance, but we ignore the analysis so as to provide the simplest notations.)

We choose $\lambda=\frac{\log\log\log d}{d}$ for simplicity.

(The purpose of $\log\log\log d$ factor is to simplify notations, and it can be tightened to constant.)

Whenever we write “for random $x$”, “for random $z$” or “for random $\xi$”, we mean that the come from the distributions introduced in Section 2 with $x=\mathbf{M}z+\xi$.

Suppose $\mathcal{Z}=\{x^{(1)},\dots,x^{(N)}\}$ are i.i.d. samples from $\mathcal{D}$ and $y^{(i)}=y(x^{(i)})$, and suppose $N\geq\poly(d)$ for some sufficiently large polynomial. Let $f=f_{t}$ and suppose $b^{(t)}\leq\poly(d)$ and $\sigma_{\rho}\geq\frac{1}{\poly(d)}$. Then, for every $w_{1},\dots,w_{N}$ that may depend on the randomness of $\mathcal{Z}$ and satisfies $\|w_{i}\|\leq\poly(d)$, it satisfies

In addition, suppose for every $i\in[N]$, we have an i.i.d. random sample $\rho^{(i)}\sim\mathcal{N}(0,\sigma_{\rho}^{2}\mathbf{I})$ that is independent of $\mathcal{Z}$ and $w$. Then, with probability at least $1-e^{-\Omega(\log^{2}d)}$ over $\rho$, we have

The proof of the second part can be done by trivial Hoeffding bounds. ∎

Appendix B Neuron Structure and Initialization Properties

We consider $m=d^{1+c_{0}}$ for a very small constant $c_{0}\in(0,1)$, and consider constants $c_{1}>c_{2}$ to be chosen shortly. Let us define a few notations to characterize each neuron’s behavior.

Recall $w_{i}^{(t)}$ is the weight for the $i$-th neuron at iteration $t$. We shall choose a parameter $\sigma_{w}^{(t)}$ at each iteration $t\geq 0$ and define the following notions. Consider any dimension $j\in[d]$.

Let $\mathcal{S}_{j,sure}^{(t)}\subseteq[m]$ be those neurons $i\in[m]$ satisfying

$\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle^{2}\geq(c_{1}+c_{2})(\sigma_{w}^{(t)})^{2}\log d$,

$\langle w_{i}^{(t)},\mathbf{M}_{j^{\prime}}\rangle^{2}<(c_{1}-c_{2})(\sigma_{w}^{(t)})^{2}\log d$ for every $j^{\prime}\neq j$,

$\mathsf{sign}(\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle)=\mathsf{sign}(w^{\star}_{j})$.

Let $\mathcal{S}_{j,pot}^{(t)}\subseteq[m]$ be those neurons $i\in[m]$ satisfying

$\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle^{2}\geq(c_{1}-c_{2})(\sigma_{w}^{(t)})^{2}\log d$

Let $\mathcal{S}_{ept}^{(t)}\subseteq[m]$ be the set of neurons $i\in[m]$ satisfying

$\|w_{i}^{(t)}\|_{2}^{2}\leq 2(\sigma_{w}^{(t)})^{2}d$

$\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle^{2}\geq(c_{1}-c_{2})(\sigma_{w}^{(t)})^{2}\log d$ for at most $O(1)$ many $j\in[d]$.

$\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle^{2}\geq 2(\sigma_{w}^{(t)})^{2}\sqrt{\log d}$ for at most $2^{-\sqrt{\log d}}d$ many $j\in[d]$.

$|\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle|\leq\frac{\sigma_{w}^{(t)}}{\log d}$ for at least $\Omega(\frac{d}{\log d})$ many $j\in[d]$.

Suppose each $w_{i}^{(0)}\sim\mathcal{N}(0,\sigma_{0}^{2}\mathbf{I})$ and suppose $\sigma_{w}^{(0)}=\sigma_{0}$. For every constants $c_{0}\in(0,1)$ and $\gamma\in(0,0.1)$, by choosing $c_{1}=2+2(1-\gamma)c_{0}$ and $c_{2}=\gamma c_{0}$, we have with probability $\geq 1-o(1/d^{3})$ over the random initialization, for all $j\in[d]$:

Recall $b^{(t)}$ is the bias at iteration $t$, and let us introduce more notions.

Let $\mathcal{S}_{ept+}^{(t)}\subseteq[m]$ be the set of neurons $i\in[m]$ satisfying

$\|w_{i}^{(t)}\|_{2}^{2}\leq\frac{(\sigma_{w}^{(t)})^{2}d}{\log^{2}d}$,

$\left|\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle\right|\geq\frac{\sigma_{w}^{(t)}}{\log d}$ for at most $O(1)$ many $j\in[d]$.

Let $\mathcal{S}_{ept++}^{(t)}\subseteq[m]$ be the set of neurons $i\in[m]$ satisfying

Let $\mathcal{S}_{j,pot+}^{(t)}\subseteq[m]$ be the set of neurons $i\in[m]$ satisfying

$|\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle|\geq\frac{k}{d\beta}b^{(t)}$.

Let $\mathcal{S}_{j,sure+}^{(t)}\subseteq[m]$ be the set of neurons $i\in[m]$ satisfying

$\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle^{2}\geq 4k(b^{(t)})^{2}$,

$\mathsf{sign}(\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle)=\mathsf{sign}(w^{\star}_{j})$.

Note that we do not have good properties on $\mathcal{S}_{ept+}^{(t)}$, $\mathcal{S}_{ept++}^{(t)}$, $\mathcal{S}_{j,pot+}^{(t)}$ or $\mathcal{S}_{j,sure+}^{(t)}$ at initialization $t=0$; however, they will gradually begin to satisfy certain properties as the training process goes. See Section C for details.

Recall if $g$ is standard Gaussian, then for every $t>0$,

Therefore, for every $i\in[m]$ and $j\in[d]$,

$p_{1}=\operatornamewithlimits{\mathbf{Pr}}[\langle w_{i}^{(0)},\mathbf{M}_{j}\rangle^{2}\geq(c_{1}+c_{2})\sigma_{0}^{2}\log d]=\Theta(\frac{1}{\log d})\cdot\frac{1}{d^{(c_{1}+c_{2})/2}}=\Theta(\frac{1}{\sqrt{\log d}})\cdot\frac{1}{d\cdot d^{(1-\gamma/2)c_{0}}}$

$p_{2}=\operatornamewithlimits{\mathbf{Pr}}[\langle w_{i}^{(0)},\mathbf{M}_{j}\rangle^{2}\geq(c_{1}-c_{2})\sigma_{0}^{2}\log d]=\Theta(\frac{1}{\log d})\cdot\frac{1}{d^{(c_{1}-c_{2})/2}}=\Theta(\frac{1}{\sqrt{\log d}})\cdot\frac{1}{d\cdot d^{(1-3\gamma/2)c_{0}}}$

We first lower bound $|\mathcal{S}_{j,sure}^{(0)}|$. For every $i\in[m]$, with probability at least $p_{1}/2\cdot(1-p_{2})^{d-1}\geq\Omega(\frac{1}{\sqrt{\log d}})\cdot\frac{d^{\frac{\gamma}{2}c_{0}}}{m}$ it satisfies

By concentration with respect to all $m$ choices of $i\in[m]$, we know with probability at least $1-o(\frac{1}{d^{3}})$ it satisfies $|\mathcal{S}_{j,sure}^{(0)}|=\Omega\left(d^{\frac{\gamma}{4}c_{0}}\right)$.

We next upper bound $|\mathcal{S}_{j,pot}^{(0)}|$. For every $i\in[m]$, with probability at most $p_{2}<O(\frac{1}{\sqrt{\log d}})\cdot\frac{d^{\frac{3\gamma}{2}c_{0}}}{m}$ it satisfies

By concentration with respect to all $m$ choices of $i$, we know with probability at least $1-o(\frac{1}{d^{3}})$ it satisfies $|\mathcal{S}_{j,pot}^{(0)}|=O(d^{2\gamma c_{0}})$.

As for $\mathcal{S}_{ept}^{(0)}$, we first note that for every $i\in[m]$, by chi-square distribution’s tail bound, with probability at least $1-o(1/d^{3})$ it satisfies $\|w_{i}^{(0)}\|_{2}^{2}\in\left[\frac{\sigma_{0}^{2}d}{2},2\sigma_{0}^{2}d\right]$.

For every $i\in[m]$, the probability of existing $q=20/c_{0}$ different

is at most $d^{q}\cdot(p_{2})^{q}\leq d^{-q\cdot\frac{c_{0}}{2}}\leq\frac{1}{d^{4}m}$. Union bounding over all possible $i\in[m]$ gives the proof that, with probability at least $1-1/d^{4}$, for all but at most $q=O(1)$ values of $j\in[d]$, it satisfies $\langle w_{i}^{(0)},\mathbf{M}_{j}\rangle^{2}<(c_{1}-c_{2})(\sigma_{w}^{(t)})^{2}\log d$.

For every $i\in[m]$ and $j\in[d]$, with probability at least $1-e^{-\sqrt{\log d}}$ it satisfies $\langle w_{i}^{(0)},\mathbf{M}_{j}\rangle^{2}\leq 2\sigma_{0}^{2}\sqrt{\log d}$. Therefore, with probability at least $1-o(1/d^{3})$, there are $\geq d(1-2^{-\sqrt{\log d}})$ indices $j\in[d]$ satisfying $\langle w_{i}^{(0)},\mathbf{M}_{j}\rangle^{2}\leq 2\sigma_{0}^{2}\sqrt{\log d}$.

For every $i\in[m]$ and $j\in[d]$, with probability at least $\frac{1}{40000\sqrt{\log d}}$ it satisfies $|\langle w_{i}^{(0)},\mathbf{M}_{j}\rangle|\leq\frac{\sigma_{0}}{10000\sqrt{\log d}}$. Therefore, with probability at least $1-o(1/d^{3})$, there are $\frac{1}{100000\sqrt{\log d}}$ indices $j\in[d]$ satisfying $|\langle w_{i}^{(0)},\mathbf{M}_{j}\rangle|\leq\frac{\sigma_{0}}{10000\sqrt{\log d}}$. ∎

Appendix C Neuron Structure Change During Training

For analysis purpose, we consider two phases during training. In Phase I, the neurons have moved so little so that the accuracy remains $50\%$ for binary classification; however, some neurons shall start to win lottery and form “singleton” structures. We summarize this as the following theorem.

$\mathcal{S}_{j,sure}^{(0)}\subseteq\mathcal{S}_{j,sure}^{(t)}$ for every $j\in[d]$.

$\mathcal{S}_{j,pot}^{(0)}\supseteq\mathcal{S}_{j,pot}^{(t)}$ for every $j\in[d]$.

$\mathcal{S}_{ept++}^{(t)}=[m]$ and $\mathcal{S}_{j,pot}^{(0)}\supseteq\mathcal{S}_{j,pot+}^{(t)}$ for every $j\in[d]$ at this iteration $t=T_{\mathsf{b}}$.

(Recall according to Definition B.1 we have $\mathcal{S}_{j,sure}^{(t)}\subseteq\mathcal{S}_{j,pot}^{(t)}$.)

In Phase II, the neurons start to move much more so that the network output becomes more meaningful; in phase II, the “singleton” neurons become even more singleton.

In the same setting as Theorem C.1, with probability at least $1-e^{-\Omega(\log^{2}d)}$, the following holds for all t\in\big{[}T_{\mathsf{b}},\,d^{O(\log d)}/\eta\big{]}.

$\mathcal{S}_{ept++}^{(t)}=\mathcal{S}_{ept+}^{(t)}=[m]$.

$\mathcal{S}_{j,sure}^{(0)}\subseteq\mathcal{S}_{j,sure}^{(t)}$ for every $j\in[d]$.

$\mathcal{S}_{j,pot}^{(0)}\supseteq\mathcal{S}_{j,pot+}^{(t)}\supseteq\mathcal{S}_{j,pot}^{(t)}$ for every $j\in[d]$.

(Recall according to Definition B.3 we have $\mathcal{S}_{j,sure+}^{(t)}\subseteq\mathcal{S}_{j,pot}^{(t)}$.)

Theorem C.2 immediately implies the first claim of Theorem 5.1 and the first claim of Theorem 5.2, after plugging in the definitions of those neuron structure sets introduced in Section B. For instance, we can write

We make several observations, when $t\geq T_{\mathsf{e}}$:

$\mathcal{S}_{ept}^{(t)}=[m]$ implies the cardinality of $\{j\in[d]\colon i\in\mathcal{S}_{j,pot}^{(0)}\}$ is $\leq O(1)$, so we can define it as $\mathcal{N}_{i}$.

For every $i\not\in\mathcal{S}_{j,pot}^{(0)}$, it also satisfies $i\not\in\mathcal{S}_{j,pot+}^{(t)}$, so we have $|\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle|<\frac{k}{d\beta}b^{(t)}\leq\frac{k\Xi_{2}^{2}}{d}$ (using $b^{(t)}\leq\beta\Xi_{2}^{2}$ from Definition C.12).

For every $i\in\mathcal{S}_{j,sure+}^{(t)}\subseteq\mathcal{S}_{j,pot}^{(0)}$, we have $\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle\cdot\mathsf{sign}(w_{i}^{*})\geq 2\sqrt{k}b^{(t)}>\frac{1}{\Xi_{2}^{8}}$ (using $b^{(t)}=\beta\Xi_{2}^{2}$ from Definition C.12 and our choice of $\beta$).

They together imply the first claim of Theorem 5.2. One can similarly derive the first claim of Theorem 5.1.

We present a lemma to bound the size of the pre-activation signal.

For every $t\geq 0$, every $i\in[m]$, every $\lambda\geq 0$, every $j\in[d]$:

Let $\mathcal{F}$ be the event where there exists $j^{\prime}\in[d]$ with $z_{j^{\prime}}\not=0$ and $\langle w_{i}^{(t)},\mathbf{M}_{j^{\prime}}\rangle^{2}\geq\Omega((\sigma_{w}^{(t)})^{2}\log d)$. Again, by the definition of $\mathcal{S}_{ept}^{(t)}$, we know that

Thus, when neither $\mathcal{E}$ or $\mathcal{F}$ happens, we have for every $j^{\prime}\in[d]$:

Apply Bernstein concentration bound we complete the proof that

Finally, for the $\xi$ part, let us recall $\langle w_{i}^{(t)},\xi\rangle$ variable with variance at most $O(\frac{\|w_{i}^{(t)}\|^{2}\sigma_{x}^{2}}{d})\leq O((\sigma_{w}^{(t)})^{2})$ and each $|\langle w_{i}^{(t)},\mathbf{M}_{j}\rangle\langle\mathbf{M}_{j},\xi\rangle|\leq\frac{\sigma_{w}^{(t)}}{\log^{2}d}$ w.h.p. Using Bernstein concentration of random variables, we finish the proof. ∎

Let $\mathcal{F}$ be the event where there exists $j^{\prime}\in[d]$ with $z_{j^{\prime}}\not=0$ and $|\langle w_{i}^{(t)},\mathbf{M}_{j^{\prime}}\rangle|\geq\Omega(\frac{\sigma_{w}^{(t)}}{\log d})$. By the definition of $\mathcal{S}_{ept+}^{(t)}$, we know that

When $\mathcal{F}$ does not happen, we have for every $j^{\prime}\in[d]$: $|\langle w_{i}^{(t)},\mathbf{M}_{j^{\prime}}z_{j^{\prime}}\rangle|\leq O(\frac{\sigma_{w}^{(t)}}{\log d})$ and at the same time

Apply Bernstein concentration bound we complete the proof that

Finally, for the $\xi$ part, let us recall $\langle w_{i}^{(t)},\xi\rangle$ is a random variable with variance at most $O(\frac{\|w_{i}^{(t)}\|^{2}\sigma_{x}^{2}}{d})\leq O(\frac{(\sigma_{w}^{(t)})^{2}}{\log^{2}d})$. Using the Bernstein concentration bound, we finish the proof. ∎

C.2 Auxiliary Lemma 2: A Critical Lemma for Gradient Bound

In this section we present a critical lemma that shall be used multiple times to bound the gradient in many of the following sections. Recall $\frac{c_{2}}{c_{1}}\in(0,0.1)$ is a constant from Lemma B.2.

For every $\alpha>0$, suppose $\rho\sim\mathcal{N}(0,\sigma_{\rho}^{2})$, define quantity

it always satisfies $|\Delta|\leq O\left(\frac{\sqrt{V}}{\sigma_{\rho}}+L\right)$

Furthermore, suppose we can write $S_{1}=(S_{1}^{\prime},S_{1}^{\prime\prime})$ and $S_{2}=S_{2}^{\prime}+S_{2}^{\prime\prime}$ for $(S_{1}^{\prime},S_{2}^{\prime})$ and $(S_{1}^{\prime\prime},S_{2}^{\prime\prime})$ being independent (although $S_{1}^{\prime},S_{2}^{\prime}$ may be dependent, and $S_{1}^{\prime\prime},S_{2}^{\prime\prime}$ may be dependent). Then, we have