Towards Understanding Clean Generalization and Robust Overfitting in Adversarial Training

Throughout this work, we use letters for scalars and bold letters for vectors. We will use $[k]$ to indicate the index set $\{1,2,\cdots,k\}$. The indicator of an event is defined as $\mathbb{I}\{\cdot\}=1$ if the event $\cdot$ holds and $\mathbb{I}\{\cdot\}=0$ otherwise. We use $\operatorname{sgn}(\cdot)$ to denote the sign function of the real number $\cdot$, and we use $\operatorname{span}(\bm{v}_{1},\bm{v}_{2},\cdots,\bm{v}_{n})$ to denote the linear span of the vectors $\bm{v}_{1},\bm{v}_{2},\cdots,\bm{v}_{n}\in\mathbb{R}^{D}$.

For any given two sequences $\left\{A_{n}\right\}_{n=0}^{\infty}$ and $\left\{B_{n}\right\}_{n=0}^{\infty}$, we denote $A_{n}=O\left(B_{n}\right)$ if there exist some absolute constant $C_{1}>0$ and $N_{1}>0$ such that $\left|A_{n}\right|\leq C_{1}\left|B_{n}\right|$ for all $n\geq N_{1}$. Similarly, we denote $A_{n}=\Omega\left(B_{n}\right)$ if there exist $C_{2}>0$ and $N_{2}>0$ such that $\left|A_{n}\right|\geq C_{2}\left|B_{n}\right|$ for all $n>N_{2}$. We say $A_{n}=\Theta\left(B_{n}\right)$ if $A_{n}=O\left(B_{n}\right)$ and $A_{n}=\Omega\left(B_{n}\right)$ both holds. We use $\widetilde{O}(\cdot),\widetilde{\Omega}(\cdot)$, and $\widetilde{\Theta}(\cdot)$ to hide logarithmic factors in these notations respectively. Moreover, we denote $A_{n}=\operatorname{poly}\left(B_{n}\right)$ if $A_{n}=O\left(B_{n}^{K}\right)$ for some positive constant $K$, and $A_{n}=\operatorname{polylog}\left(B_{n}\right)$ if $B_{n}=\operatorname{poly}\left(\log\left(B_{n}\right)\right)$. We also say $A_{n}=o(B_{n})$ if for arbitrary positive constant $C_{3}>0$, there exists $N_{3}>0$ such that $|A_{n}|<C_{3}|B_{n}|$ for all $n>N_{3}$. An event is said that it happens with high probability (or w.h.p. for short) if it happens with probability at least $1-o(1)$.

We use the notation $\left\|\cdot\right\|_{p},p\in[1,+\infty]$ to denote the $\ell_{p}$ norm in the vector space $\mathbb{R}^{D}$. For two sets $\mathcal{A},\mathcal{B}\subset\mathbb{R}^{D}$, we can define the $\ell_{p}$-distance between $\mathcal{A}$ and $\mathcal{B}$ as $\operatorname{dist}_{p}(\mathcal{A},\mathcal{B}):=\operatorname{inf}\left\{\|\bm{X}-\bm{Y}\|_{p}:\bm{X}\in\mathcal{A},\bm{Y}\in\mathcal{B}\right\}$. For $r>0$, $\mathbb{B}_{p}(\bm{X},r):=\left\{\bm{Z}\in\mathbb{R}^{D}:\left\|\bm{Z}-\bm{X}\right\|_{p}\leq r\right\}$ is defined as the $\ell_{p}$-ball with radius $r$ centered at $\bm{X}$.

A multi-layer neural network is a function from input in $\mathbb{R}^{D}$ to output in $\mathbb{R}^{m}$, which is defined as follows:

and

where the weight $\bm{W}=[\bm{W}_{0},\bm{W}_{1},\cdots,\bm{W}_{L}]$ and the bias $\bm{B}=[\bm{B}_{0},\bm{B}_{1},\cdots,\bm{B}_{L}]$ are learnable parameters, and $\max\{D,m_{0},m_{1},\cdots,m_{L-1},m\}$ and $L+1$ are the width and depth of the neural network, respectively. We use $\sigma(\cdot)$ to denote the (non-linear) entry-wise activation function, and ReLU activation function is defined as $\sigma(\cdot):=\max(\cdot,0)$.

We consider a binary classification setting, where we use $\bm{X}\in\mathcal{X}\subset\mathbb{R}^{D}$ to denote the data input and the binary label $y$ is in $\mathcal{Y}=\{-1,1\}$. Given a data distribution $\mathcal{D}$ that is a joint distribution over the supporting set $\mathcal{X}\times\mathcal{Y}$, and a function $f:\mathcal{X}\rightarrow\mathbb{R}$ as the classifier, we can define the following measurements to describe the clean (robust) classification performance of the classifier $f$ on data $\mathcal{D}$.

In our work, we mainly focus on the cases when $p=2$ and $p=\infty$, which can be extended to the general $p$ case as well.

In adversarial training, with access to the training dataset $\mathcal{S}=\{(\bm{X}_{1},y_{1}),(\bm{X}_{2},y_{2}),\dots,(\bm{X}_{N},y_{N})\}$ randomly sampled from the data distribution $\mathcal{D}$, we aim to minimize the following robust training error to derive the robust classifier.

Now, we present the concept CGRO classifier that we mainly study in this paper as the following definition.

In this section, we consider the case that data input $\bm{X}\in\mathbb{R}^{D}$ has a patch structure as follow, which is similar to a list of theoretical works about feature learning (Allen-Zhu & Li, 2020b; Chen et al., 2021; Jelassi & Li, 2022; Jelassi et al., 2022; Kou et al., 2023; Chidambaram et al., 2023).

To learn our structured data, we use two-layer convolutional neural network (CNN) (LeCun et al., 1998; Krizhevsky et al., 2012) with non-linear activation as the learner network.

$\operatorname{ReLU}^{q}$ activation functions are useful to address the non-smoothness of ReLU function at zero, which are widely applied in literatures of deep learning theory (Kileel et al., 2019; Allen-Zhu & Li, 2020a, b; Chen et al., 2021; Jelassi & Li, 2022; Jelassi et al., 2022). To simplify our analysis, we fix the second layer weights as $a_{r}^{+}=a_{r}^{-}=\frac{1}{2}$. We also assume that $q$ is odd and it holds that $w_{r}:=w_{r}^{+}=-w_{r}^{-}$ during optimization process. At initialization, we sample $\bm{w}_{r}$ i.i.d from $\mathcal{N}(0,\sigma_{0}^{2}\mathbf{I}_{d})$. Our results can be extended to the case when $q$ is even and $w_{r}^{+},-w_{r}^{-}$ are differently.

The Role of Non-linearity. Indeed, a series of recent theoretical works (Li et al., 2019; Chen et al., 2021; Javanmard & Soltanolkotabi, 2022) show that linear model can achieve robust generalization for adversarial training under certain settings, which but fails to explain the CGRO phenomenon observed in practice. To mitigate this gap, we improve the expressive power of model by using non-linear activation that can characterize the data structure and learning process more precisely.

In adversarial training, we aim to minimize the following adversarial training loss, which is a trade-off between natural risk and robust regularization defined as follow.

Here, we use the logistic loss defined as $\mathcal{L}(\bm{W};\bm{X},y):=\log\left(1+\exp\{-yf_{W}\left(\bm{X}\right)\}\right)$. And we apply the perturbed range defined as $\mathbb{A}_{p}\left(\bm{X},\delta\right):=\mathbb{B}_{p}\left(\bm{X},\delta\right)\cap\bm{X}+\bm{\Delta}\left(\bm{X}\right)$, where $\bm{\Delta}\left(\bm{X}\right)\subset\mathbb{R}^{Pd}$ satisfies that

This perturbed range ensures that adversarial perturbations used to generate adversarial examples are always aligned with the meaningful signal vector $\bm{w}^{*}$ during adversarial training, which exactly simplifies the optimization analysis.

Next, we make the following assumptions about hyperparameters we introduced in our structured-data setting.

Discussion of Hyperparameter Choices. Actually, the choices of these hyperparameters are not unique. We make concrete choices above for the sake of calculations in our proofs, but only the relationships between them are important. Namely, since the norm of signal patch is $\alpha$ and the norm of noise patch w.h.p. is $\Theta(\sigma_{p}\sqrt{d})$, our choices ensure that meaningful signal is stronger than noise. Without this assumption, in other word, if the signal-to-noise ratio is very low, there even exists no clean generalization, which has been theoretically shown under the similar patch-structured data setting (Cao et al., 2021; Chen et al., 2021; Frei et al., 2022). The width of network learner is $\tilde{O}(ND)$ to achieve mildly over-parameterization we mentioned in Theorem 4.4. The separation in Assumption 4.2 also holds due to $\delta<\alpha$.

First, we introduce the concept called feature learning to characterize what the model learns from training dynamics.

Thus, it holds that the model correctly classify the clean data in two cases. One is that the model learns the true feature and ignores the spurious features, i.e. $\mathcal{U}=\Omega(1)\gg|\mathcal{V}|$. Another is that the model doesn’t learn the true feature but memorizes the spurious features, i.e. $\mathcal{U}=o(1)$ and $\mathcal{V}=\Omega(1)\gg 0$. We can analyze the perturbed data similarly.

Now, we give our main result about feature learning process.

Next, we provide a proof sketch of Theorem 5.9. To obtain a detailed analysis of learning process, we consider the following objects that can be viewed as weight-wise version of $\mathcal{U}^{(t)}$ and $\mathcal{V}_{i}^{(t)}$. We define $u^{(t)}$ and $v_{i,j}^{(t)}$ as

for each $r\in[m]$, $i\in[N]$ and $j\in[P]\setminus\operatorname{signal}(\bm{X}_{i})$.

We then propose a three-stage analysis technique to decouple the complicated feature learning process as follows.

Lemma 5.11 manifests that the signal component increases quadratically at initialization. Therefore, we know that, after $T_{0}=\Theta\left(\eta^{-1}\alpha^{-q}\sigma_{0}^{-1}\right)$ iterations, the signal component $u^{(T_{0})}$ attains the order $\tilde{\Omega}(\alpha^{-1})$, which implies the model learns partial true feature at this point.

Lemma 5.12 shows that, after partial true feature learning, the increment of signal component is mostly dominated by the noise component $\mathcal{V}_{i}^{(t)}$, which implies that the growth of signal component will converge when $\mathcal{V}_{i}^{(t)}=\Omega(1)$.

The practical implication of Lemma 5.13 is two-fold. First, by the quadratic increment of noise component, we derive that, after $T_{1}=\Theta\left(N\eta^{-1}\sigma_{0}^{-1}\sigma_{p}^{-3}d^{-\frac{3}{2}}\right)$ iterations, the total noise memorization $\mathcal{V}_{i}^{(T)}$ attains the order $\Omega(1)$, which suggests that the model is able to robustly classify adversarial examples by memorizing the data-wise noise. Second, combined with Lemma 5.12, the signal component $u^{(t)}$ will maintain the order $\Theta(\alpha^{-1})$, which immediately implies the main conclusion of Theorem 5.9. And the full detailed proof of Theorem 5.9 can be see in Appendix E.

Experiment Settings. For the MNIST dataset, we consider a simple convolutional network, LeNet-5 (LeCun et al., 1998), and study how its performance changes as we increases the size of network (i.e. we expand the number of convolutional filters and the size of the fully connected layer to the size number multiple of the initial, where the size numbers are $1,4,8,12,16$). The original convolutional network has a convolutional layer with $1$ filters, followed by another convolutional layer with $2$ filters, and a fully connected hidden layer with 32 units. Convolutional layers are followed by $2\times 2$ max-pooling layers. For the CIFAR10 dataset, we apply WideResNet-34 (Zagoruyko & Komodakis, 2016) with different widen factors $1,2,5,10$. We use the standared projected gradient descent (PGD) (Madry et al., 2017) to train the network by adversarial training. We choose the classical $\ell_{\infty}$-perturbation with radius $0.3$ for MNIST and $8/255$ for CIFAR10. All models are trained for 100 epoches on MNIST and 200 epoches on CIFAR10 by using a single NVIDIA RTX 4090 GPU.

Experiment Results. We report our results about the performance (including the clean test accuracy, robust test accuracy and robust training accuracy) of models with different sizes in Table 1 and Figure 2 (a)(b). It shows that when the model size becomes larger, first the robust training loss decreases but the robust generalization gap remains large, and then when the model gets even larger, the robust generalization gap gradually decreases, and we also find that, in the small size case (see LeNet with the size number $1,4,8$), adversarial training converges to a trivia solution that always predicts a fixed class, while it can learn an accurate clean classifier through the standard training, which corresponds to the theoretical results in Theorem 4.4 and Theorem 4.7.

Experiment Settings. Here we generate synthetic data exactly following Definition 5.1 and apply the two-layer convolutional network in Definition 5.3. We train the network by the adversarial training algorithm we mentioned in Definition 5.6. We choose the hyperparameters that we need as: $d=100,P=2,\alpha=10,\sigma_{p}=1,\sigma_{0}=0.01,q=3,N=20,m=10,p=2,\delta=10,\lambda=0.9,\eta=0.1,T=100$, which is a feasible one under Assumption 5.7. We characterize the true feature learning and noise memorization via calculating $\mathcal{U}^{(t)}$ and the smallest/largest/average of $\{\mathcal{V}_{i}^{(t)}\}_{i\in[N]}$. We calculate the robust train and test accuracy of the model by using the standard PGD attack.

Experiment Results. We plot the dynamics of true feature learning and noise memorization in Figure 2 (c). It is clear that a three-stage phase transitions happen during adversarial training , which is consistent with our theoretical analysis of learning process (Lemma 5.11, Lemma 5.12 and Lemma 5.13), and finally the model partially learns true feature but exactly memorizes all training data (Theorem 5.9). The performance of the model is presented in Table 2, which shows that the network converges to a CGRO classifier eventually.