Adversarial Training Can Provably Improve Robustness: Theoretical Analysis of Feature Learning Process Under Structured Data

We sincerely thank the reviewer for the thoughtful review and positive feedback, and for highlighting the strength of our theoretical contributions and the clarity of our writing. We are very glad to address the questions and suggestions raised by the reviewer, which we believe will help further refine our work. Below are our responses to the questions and suggestions raised by the reviewer.

[W1] The adversarial training is with just one step gradient ascent which is usually not the case in practice.

[E1] While adversarial training against one-step attacks (such as FGSM) does not produce robust models in practice, we would first like to point out that, under our theoretical framework, adversarial training based on one-step attacks is sufficient to improve network robustness against arbitrary attacks (see details in Theorem 4.4). Furthermore, we believe that, under our synthetic structured data setup, multi-step adversarial training performs similarly to the one-step method.

We also emphasize that analyzing one-step adversarial training for two-layer neural networks is highly challenging due to the inherently non-convex and non-concave nature of the min-max optimization problem in adversarial training. To address this challenge, we developed a novel and highly non-trivial analysis technique to track the dynamics of the training process. To the best of our knowledge, we are the first to provide a refined analysis of the feature learning process in standard one-step adversarial training for non-linear two-layer networks.

Additionally, we empirically verify our theoretical findings on real-world image datasets (MNIST, CIFAR10, and SVHN) using multi-step adversarial training and multi-step attacks. The numerical results are consistent with our theoretical insights derived from analyzing one-step adversarial training in our toy theoretical model.

[W2] The definition of robust and non-robust features require parallel to the coordinate axis along with orthonormal condition which is restrictive.

[E2] We would like to clarify that the one-hot assumption is adopted for simplicity in mathematical calculations. Indeed, since our study focuses on $\ell_{\infty}$-robustness, considering only the orthogonality of class signal vectors is insufficient to fully characterize the properties of robustness under the $\ell_{\infty}$ norm. For example, two vectors with equal $\ell_2$ norms can have $\ell_{\infty}$ norms that differ by a factor of $d$, where $d$ is the dimensionality of the vectors. Therefore, it is necessary to explicitly describe the $\ell_{\infty}$ properties of the class signal vectors. For simplicity, we assume them to be one-hot vectors in this context.

[Q1] From the analysis, there seems to be no tradeoff between standard and robust accuracies in adversarial training contrary to the empirical observations [P1]. Can the analysis show this tradeoff? Perhaps, in theorem 4.4, it might be worthwhile to include what happens to the standard accuracy $\mathbb{P} _{(\boldsymbol{X},y)\sim \mathcal{D}} [\operatorname{argmax}F_i^{(T)}(\boldsymbol{X})\ne y]$ and also what happens when $(\boldsymbol{X}_f, y)\sim\mathcal{D} _{NR}$.

[A1] We would like to clarify that, while there is a tradeoff between standard and robust accuracies in adversarial training, contrary to the empirical observations [1], perfect robust accuracy can imply perfect standard accuracy under our synthetic data setup. We believe that extending our theoretical framework to study the tradeoff between robustness and accuracy represents a valuable and interesting direction for future research.

[Q2] The phase transition is also not really observed in real data and neural networks with activations other than smoothed ReLU. Can the authors comment on why there is a discrepancy in the results? Adding a discussion contrasting to these results would be helpful.

[A2] Indeed, the time when phase transition happens is related to the one-step gradient ascent learning rate $\tilde{\eta}$ (roughly speaking, the threshold time is nearly $\Theta(\tilde{\eta}^{-1})$). Thus, in practice, due to the relatively large learning rate (i.e., step size) of PGD for finding adversarial examples (e.g., for CIFAR10, $\tilde{\eta}$ may be chosen as $2/255$, whereas the gradient descent learning rate is typically just $10^{-3}$), the duration of the first stage is very short, which may lead to the non-observation of phase transition in real-world experiments.

[Q3] The analysis considers smoothed ReLU activation following the feature learning analysis from Allen-Zhu & Li 2023a. Why is this particular choice made? What happens with other activations?

[A3] We would like to point out that we use the smoothed ReLU to replace the original ReLU due to the discontinuity of ReLU at zero, which is a common approach in the feature learning theory literature, as seen in [2][3] in the related work section. However, we believe that the theoretical insights derived from analyzing smoothed ReLU networks can also be applied to other non-linear activation functions commonly used in practice, such as the standard ReLU and GeLU activation functions.

[Q4] Minor comment: in line 179, 'where u,v are the corresponding' - u,v should be boldfaced.

[A4] Thank the reviewer for pointing out this typo. We have fixed it in the revision of our paper.

Reference

[1] Tsipras et al. "Robustness May Be at Odds with Accuracy." ICLR 2019.

[2] Allen-Zhu, Z. and Li, Y. (2023b). Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. In The Eleventh International Conference on Learning Representations.

[3] Chidambaram, M., Wang, X., Wu, C. and Ge, R. (2023). Provably learning diverse features in multi view data with midpoint mixup. In International Conference on Machine Learning. PMLR.

We sincerely thank the reviewer for the thoughtful feedback, and for highlighting the strength of our theoretical contributions and the clarity of our writing. We are very glad to address the questions and suggestions raised by the reviewer, which we believe will help further refine our work. Below are our responses to the questions and suggestions raised by the reviewer.

[Q1] In Proposition 3 and Theorem 4.3, what if the adversarial perturbation is designed as $\boldsymbol{\delta}_p = -\beta_p \boldsymbol{v}_y + \epsilon \boldsymbol{v}_t$, where $v_t$ is orthogonal to all $v_i$ for $(i∈[k])$? For example, could an adversarial perturbation be constructed without using any known non-robust features from the set (JNR)? Given the high dimensionality of (d), and the potential infinitude of non-robust features, this scenario might be feasible.

[A1] We thank the reviewer for the valuable suggestion. We would like to point out that the perturbation proposed by the reviewer is a type of untargeted attack. For instance, assuming the network learns only the non-robust features $\boldsymbol{v}_y$ , the output class prediction of the network would be random due to the orthogonal property and the random initialization of network weights. However, this construction still includes a known non-robust feature $\boldsymbol{v}_y$.

[Q2] In the definition of “Small Perturbation Radius” (lines 281-284), two data points $\boldsymbol{X}$ and $\boldsymbol{X}'$ with distinct labels may have significant differences with high probability. While this may hold for real data, what if $\boldsymbol{X}'$ is also an adversarial example?

[A2] We would like to clarify that this data assumption is directly applied to real data. As mentioned in the paper, we emphasize that we require $\boldsymbol{X}$ and $\boldsymbol{X}'$ to both be sampled from the real data distribution $\mathcal{D}$ (in line 282 of our paper), so it does not include cases where $\boldsymbol{X}'$ is an adversarial example. This assumption is consistent with the empirical observation that the typical perturbation radius is often much smaller than the separation distance between different classes [1].

[Q3] In line 285, the authors assume that $\boldsymbol{F}$ is an accurate classifier on $(\boldsymbol{X}, y)$. However, since $\boldsymbol{F}$ is only a two-layer neural network, this assumption may not be entirely valid. There could be a substantial portion of the dataset on which $\boldsymbol{F}$ misclassifies. Can the authors discuss which parts of the theory would be impacted by relaxing this assumption? For instance, in Propositions 3.1 and 3.2, $\operatorname{argmax}F_i(\boldsymbol{X} + \boldsymbol{\Delta})$ could hold even with if $F_i(\boldsymbol{X}) \ne y$ .

[A3] We thank the reviewer for the insightful suggestion. In fact, we would like to clarify that this assumption can be relaxed by defining adversarial examples without requiring correct classification on the clean data. All our analysis techniques remain valid under this relaxed assumption. We have revised the assumption in the updated version of our paper. Thank you again for the valuable suggestion.

Reference

[1] Yang, Y.-Y., Rashtchian, C., Zhang, H., Salakhutdinov, R. R. and Chaudhuri, K. (2020). A closer look at accuracy vs. robustness. Advances in neural information processing systems, 33 8588–8601.

We sincerely thank the reviewer for the encouraging and insightful feedback, and for highlighting the strength of our theoretical contributions and the clarity of our writing. We are very glad to address the question raised by the reviewer, which we believe will help further refine our work. Below is our response to the question raised by the reviewer.

[Q] How would you motivate the assumption that the class signal vectors are one-hot? How important it is from a technical point of view?

[A] We would like to clarify that the one-hot assumption is adopted for simplicity in mathematical calculations. Indeed, since our study focuses on $\ell_{\infty}$-robustness, considering only the orthogonality of class signal vectors is insufficient to fully characterize the properties of robustness under the $\ell_{\infty}$ norm. For example, two vectors with equal $\ell_2$ norms can have $\ell_{\infty}$ norms that differ by a factor of $d$, where $d$ is the dimensionality of the vectors. Therefore, it is necessary to explicitly describe the $\ell_{\infty}$ properties of the class signal vectors. For simplicity, we assume them to be one-hot vectors in this context.

We sincerely thank the reviewer for the positive support and valuable feedback! We greatly appreciate the insightful review, and the recognition of highlighting the significance of our contribution and solidity of our theory, as well as the clarity of our writing. We are very glad to address the questions and suggestions raised by the reviewer, which we believe will help further refine our work. Below are our responses to the questions and suggestions raised by the reviewer.

[Q1] Can the authors clarify whether their theoretical framework can be applied to standard ReLU-activated neural nets and multi-step PGD?

[A1] We would first like to clarify the reasons for our choice of smoothed ReLU and one-step PGD in our theoretical model. Specifically, we use the smoothed ReLU to replace the original ReLU due to the discontinuous property of ReLU at zero, which is a common approach in the literature of feature learning theory papers, such as [1][2] in the related work section.

Additionally, we point out that, under our theoretical framework, adversarial training based on one-step attacks is sufficient to improve network robustness against arbitrary attacks (see details in Theorem 4.4). We also emphasize that analyzing one-step adversarial training for two-layer neural networks is highly challenging due to the inherently non-convex and non-concave nature of the min-max optimization problem in adversarial training. To address this, we developed a novel and highly non-trivial analysis technique to track the dynamics of the training process.

We believe that extending our results to standard ReLU-activated neural networks and multi-step PGD is an important and interesting future direction for our work.

[Q2] It is hard to tell from my perspective which data model is more realistic since they are all synthetic data distributions. Can the authors provide some explanations or references?

[A2] As mentioned in the related work section (lines 117–130), we note that several previous theoretical studies [3][4][5] have considered the mixture-of-Gaussian data setup to analyze adversarial robustness with linear classifiers. However, standard training does not explicitly converge to non-robust solutions under these conditions. This result contrasts with empirical findings, where networks trained using standard methods exhibit poor robustness performance (e.g., [6][7][8]). For instance, as highlighted in [5], adversarial training, similar to standard training, also directionally converges to the maximum $\ell_2$-margin solution when applied to a Gaussian-mixture data model with $\ell_2$ perturbations. This implies that, under their assumptions, standard training alone can achieve adversarial robustness due to the maximum-margin implicit bias, even though neural networks trained with standard methods generally lack robustness in practice.

In this paper, we aim to bridge the gap between theory and practice by adopting a more structured data assumption and employing a non-linear two-layer CNN as the learner. This setup ensures the existence of both robust and non-robust global minima due to the non-linearity of the data model and the non-convexity of the learner model (see the detailed discussion in Section 3). Specifically, we follow the patch-structured data framework proposed in [2], providing a synthetic data setup based on the empirical observations of [9]. Our Assumption 2.3 can be linked to a downsized version of convolutional networks applied to image classification data. With a small kernel size, high-magnitude features that are easily perceivable by humans in an image typically appear only in a few patches (e.g., the ears of a cat or the nose of an elephant), while most other patches resemble random noise to human observers (e.g., the textures of cats and elephants blending into a random background). Illustrations of real images and our patch data are provided in Figures 1 and 2.

[Q3] The relationship between robust/non-robust features and the input data is sort of linear, up to some noise. Is it possible to consider a non-linear relationship (e.g., using a ground-truth feature extractor to map the input space to some latent space)?

[A3] We thank the reviewer for the thoughtful suggestion. Indeed, due to the non-linearity of the smoothed ReLU activation function, the one-step adversarial training method results in a highly non-convex and non-concave min-max optimization problem. To address this, we developed a novel and highly non-trivial analysis technique to track the dynamics of the training process. Adding a non-linear feature extractor to our data structure would make the analysis of the feature learning process even more challenging, which we believe represents an important and interesting direction for future research.

[Q4] The robust/non-robust features are assumed to be well-separated with respect to particular input dimensions. Generally speaking, there could be cases where robust and non-robust features overlap with each other within the input, and for each input, the robust/non-robust feature dimensions may vary. Is it possible to use the proposed theoretical framework to analyze the feature learning process for such a more generalized setting?

[A4] We thank the reviewer for the insightful suggestion. We would like to point out that our theoretical framework for analyzing the feature learning process can be applied to the more generalized setting mentioned by the reviewer. Indeed, we can consider the case where robust and non-robust features overlap with each other within the input patch. Formally, the data input can be written as $\boldsymbol{X} = (\boldsymbol{x}_1, \boldsymbol{x}_2, \dots, \boldsymbol{x}_P)$ and $\boldsymbol{x}_p = \alpha_p \boldsymbol{u}_y + \beta_p \boldsymbol{v} + \sum _{ \boldsymbol{f} \in \mathcal{F}\setminus \{\boldsymbol{u}_y, \boldsymbol{v}_y\} } \gamma _{p, f}\boldsymbol{f} + \boldsymbol{\xi}_p$, where the first two terms are class signals, the third term is feature noise, and the last term is random noise. Then, the modified Assumption 2.3 can be presented as follows:

1. There exists at least one robust patch $p$ such that for any arbitrary patch $p’$ (possibly $p’=p$), we have $\alpha_p \gg \beta_{p’}$.
2. There exists a constant $\gamma \geq 0$ such that $\sum _{p \in [P]} \alpha_p^{\gamma} \ll \sum _{p \in [P]} \beta_p^{\gamma}$.
3. For each feature noise, the coefficient satisfies $\gamma _{p,f} = \Theta(\sigma_n / \sqrt{d})$, where random noise $\boldsymbol{\xi}_p \sim \mathcal{N}(\boldsymbol{0}, \sigma_n^2 \mathcal{I}_d)$.

Under our generalized data setup and Assumption 2.3, we can apply a similar analysis technique to prove results analogous to Theorem 4.3 and Theorem 4.4.

Reference

[1] Chidambaram, M., Wang, X., Wu, C. and Ge, R. (2023). Provably learning diverse features in multi view data with midpoint mixup. In International Conference on Machine Learning. PMLR.

[2] Allen-Zhu, Z. and Li, Y. (2023b). Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. In The Eleventh International Conference on Learning Representations.

[3] Li, Y., X.Fang, E., Xu, H. and Zhao, T. (2020). Implicit bias of gradient descent based adversarial training on separable data. In International Conference on Learning Representations.

[4] Javanmard, A. and Soltanolkotabi, M. (2022). Precise statistical analysis of classification accuracies for adversarial training. The Annals of Statistics, 50 2127–2156.

[5] Chen, J., Cao, Y. and Gu, Q. (2023). Benign overfitting in adversarially robust linear classification. In Uncertainty in Artificial Intelligence. PMLR.

[6] Biggio, B., Corona, I., Maiorca, D., Nelson, B., Šrndi´ c, N., Laskov, P., Giacinto, G. and Roli, F. (2013). Evasion attacks against machine learning at test time. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2013, Prague, Czech Republic, September 23-27, 2013, Proceedings, Part III 13. Springer.

[7] Szegedy, C., Zaremba, W., Sutskever, I., Bruna, J., Erhan, D., Goodfellow, I. and Fergus, R. (2013). Intriguing properties of neural networks. arXiv preprint arXiv:1312.6199.

[8] Goodfellow, I. J., Shlens, J. and Szegedy, C. (2014). Explaining and harnessing adversarial examples. arXiv preprint arXiv:1412.6572.

[9] Ilyas, A., Santurkar, S., Tsipras, D., Engstrom, L., Tran, B. and Madry, A. (2019). Adversarial examples are not bugs, they are features. Advances in neural information processing systems, 32.

We sincerely thank the reviewer for the thoughtful feedback, and for highlighting our solid theoretical contributions and the clarity of our writing. We are very glad to address the questions and suggestions raised by the reviewer, which we believe will help further refine our work. Below are our responses to the questions and suggestions raised by the reviewer.

[W1] The title is a bit misleading. Usually, 'provable robustness' means 'robustness verification', the conclusion of this work does not involve robustness verification.

[E1] We would like to clarify that we did not claim anything about “provable robustness” or “robustness verification” in our paper. Specifically, the title of our paper indicates that, under our patch-structured data assumption, we rigorously prove (i.e., provably demonstrate) that adversarial training ensures the network learns robust features, thereby improving model robustness. The term “provably” is widely used in the literature of learning theory, as seen in [1][2][3] in the related work section, where it uniformly refers to the rigorous mathematical demonstration of certain important properties of an algorithm.

[W2] The conclusions of theoretical analyses should involve the magnitude of perturbation $\epsilon$. As we can see when $\epsilon = 0$, adversarial training will be degraded to normal training. Can the authors extend your results in this manner? In the current manuscript, Lemma 5.2 cannot degrade to Lemma 5.1 by setting $\epsilon = 0$.

[E2] We would like to clarify that Lemma 5.2 can exactly reduce to Lemma 5.1 by setting $\epsilon = 0$ in the current manuscript. Notably, Lemma 5.1 and Lemma 5.2 apply to different time ranges: Lemma 5.1 is valid for any training time $t \geq 0$, whereas Lemma 5.2 focuses specifically on the early stage (the polynomial part) of adversarial training, i.e., it applies for time $0 \leq t \leq T_0$. Within the overlapping time range of Lemma 5.1 and Lemma 5.2, Lemma 5.2 reduces to Lemma 5.1 when $\epsilon = 0$, as the smoothed ReLU lies in the polynomial regime and $1 - \operatorname{logit} = \Theta(1)$ (see details in Lemma D.5 and Lemma D.19).

[W3] The authors assume to use one-step attack in adversarial training (Remark 2.5). However, adversarial training against one-step attacks (like FGSM) does not actually produce robust models. The models trained in this way may be easily broken by a stronger attack.

[E3] While adversarial training against one-step attacks (such as FGSM) does not produce robust models in practice, we would first like to point out that, under our theoretical framework, adversarial training based on one-step attacks is sufficient to improve network robustness against arbitrary attacks (see details in Theorem 4.4). Furthermore, we believe that, under our synthetic structured data setup, multi-step adversarial training performs similarly to the one-step method.

We also emphasize that analyzing one-step adversarial training for two-layer neural networks is highly challenging due to the inherently non-convex and non-concave nature of the min-max optimization problem in adversarial training. To address this, we developed a novel and highly non-trivial analysis technique to track the dynamics of the training process. To the best of our knowledge, we are the first to provide a refined analysis of the feature learning process in standard one-step adversarial training for non-linear two-layer networks.

Additionally, we empirically verify our theoretical findings on real-world image datasets (MNIST, CIFAR10, and SVHN) using multi-step adversarial training and multi-step attacks. The numerical results align with our theoretical insights derived from analyzing one-step adversarial training in our toy theoretical model.

[Q1] In line 247, the author does not provide any intuition or reference about the design of the smoothed ReLU. Why don't the authors use softplus or GeLU? They are popular smooth alternatives of ReLU. In addition, the function proposed by authors here seems discontinuous at the point $z=\rho$?

[A1] We would like to clarify that the smoothed ReLU is continuous at the point $z = \rho$, as we can calculate the values and derivatives of two parts at $z = \rho$:

1. At the polynomial part, $\operatorname{ReLU}(z) = \rho / q$, $\operatorname{ReLU}'(z) = z^{q-1} / \rho^{q-1} = 1$ when $z \to \rho^-$.
2. At the linear part, $\operatorname{ReLU}(z) = \rho / q$, $\operatorname{ReLU}'(z) = 1$ when $z \to \rho^+$.

Here, we use the smoothed ReLU to replace the original ReLU due to the discontinuous property of ReLU at zero, which is widely applied in the literature of feature learning theory papers, such as [3][4] in the related work section.

[Q2] In line 281, why $\epsilon = \Theta(\sigma_n \sqrt{d})$? In practice, we typically use smaller ϵ for l∞ bounded perturbations when the dimension d is larger. For example, we usually use $\epsilon = 0.3$ for MNIST, $\epsilon = 8/255$ for CIFAR10, and $\epsilon = 2/255$ for ImageNet, these are the datasets with increasing input dimensions.

[A2] We would like to point out that, under our theoretical framework, we use a relatively small perturbation, even though the perturbation is a function of the data dimension $d$. As mentioned in lines 281–284 of our paper, under our setting, for two data points $(\boldsymbol{X}, y), (\boldsymbol{X}', y') \sim \mathcal{D}$ with distinct labels $y \ne y' \in [k]$, it can be checked that w.h.p. $\|\|\boldsymbol{X}-\boldsymbol{X}'\|\|_{\infty} \geq \Omega(\mathbb{E}[\alpha_p]) \gg \Theta(\sigma_n \sqrt{d}) = \Theta(\epsilon)$, which is consistent with the empirical observation that the typical perturbation radius is often much smaller than the separation distance between different classes [5].

Reference

[1] Du, S. S., Zhai, X., Poczos, B. and Singh, A. (2019b). Gradient descent provably optimizes over parameterized neural networks. In International Conference on Learning Representations.

[2] Jelassi, S., Sander, M. and Li, Y. (2022). Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35 37822–37836.

[3] Chidambaram, M., Wang, X., Wu, C. and Ge, R. (2023). Provably learning diverse features in multi view data with midpoint mixup. In International Conference on Machine Learning. PMLR.

[4] Allen-Zhu, Z. and Li, Y. (2023b). Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. In The Eleventh International Conference on Learning Representations.

[5] Yang, Y.-Y., Rashtchian, C., Zhang, H., Salakhutdinov, R. R. and Chaudhuri, K. (2020). A closer look at accuracy vs. robustness. Advances in neural information processing systems, 33 8588–8601.

Dear reviewer wbRQ,

Thanks very much for your update. We will address the questions raised by the review as below.

[Q1] However, why does the smoothed ReLU function exhibit the polynomial part in the early stage of training?

[A1] As we mentioned in line 457 of our paper, due to our small initialization technique, the values of $A_{i,r}^{(t)}$ and $B_{i,r}^{(t)}$ will be sufficiently small such that the smoothed ReLU function exhibits the polynomial part in the early stage of training. This small initialization technique is widely used in literature of feature learning paper with smoothed ReLU activation, such as [1][2][3].

[Q2] In addition, I believe a comprehensive investigate on the effect of $\epsilon$ on feature learning can make the manuscript stronger.

[A2] We sincerely thank the reviewer for the valuable suggestion. We will add a comprehensive investigate on the effect of $\epsilon$ on feature learning in the revision of our paper.

Thanks for the positive feedback again, and we would greatly appreciate it if you would consider raising your score!

Reference

[1] Allen-Zhu, Z. and Li, Y. (2023b). Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. In The Eleventh International Conference on Learning Representations.

[2] Chidambaram, M., Wang, X., Wu, C. and Ge, R. (2023). Provably learning diverse features in multi view data with midpoint mixup. In International Conference on Machine Learning. PMLR.

[3] Li, J., Pan, J., Tan, V. Y., Toh, K. C., & Zhou, P. (2024). Towards Understanding Why FixMatch Generalizes Better Than Supervised Learning. arXiv preprint arXiv:2410.11206.