Feature Averaging: An Implicit Bias of Gradient Descent Leading to Non-Robustness in Neural Networks

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.

[W] The assumptions on the data (orthonormal equinorm and multi-cluster) is restrictive.

[E] To address the concern about the restrictive assumptions on the data, we have conducted several ablation studies to examine the effects of sample size, learning rate, initialization magnitude, signal-to-noise ratio, and orthogonal conditions. These studies demonstrate that our results can be extended to more relaxed settings. In particular, regarding the equal-norm assumption, we have included an ablation study in Appendix B in the revised version of our paper.

Besides, we believe that real-world data can exhibit a multi-cluster structure. To support this, we employed a CLIP model for feature extraction on the CIFAR-10 dataset. Our experiment result shows that the features extracted by CLIP exhibit a multi-cluster structure, which aligns with our data assumption.

[Q1] In the fine-grained supervision case of the multi-cluster data, I understand that the binary classification using the multi-class network results in a robust solution. But, I wonder if the learned network would still be non-robust for the multi-class classification.

[A1] The network remains robust for multi-class classification. Almost the same as the one for binary classification, we can prove that the network achieves O(\sqrt{d})-robustness, which is optimal for our setting.

[Q2] It is not an ask to prove the following. I am only curious to understand and know the authors thought. The fine-grained supervision setting can also be seen as learning a more expressive/complex classifier (learn k classes to do binary classification). Then, can the authors comment on the adversarial min-max training? Would similar feature decoupling, ie the weight of each neuron is aligned with one cluster feature, happen naturally if one does adversarial training of the binary class data?

[A2] Thank you for this excellent question, which provides an important perspective for understanding and extending our work. We conducted experiments on synthetic data (see the revised vision of our paper), and the results indicate that networks trained with adversarial training exhibit a tendency to learn feature decoupling solutions. However, the degree of decoupling is less pronounced compared to networks trained with fine-grained supervision. In terms of robustness, adversarial training does provide significant improvements over standard GD, but it remains slightly inferior to the robustness achieved through fine-grained supervision.

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 the strength of our theory and experiments, 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.

[W] Furthermore, in line 110, the authors mention the work of Min & Vidal (2024) and their conjecture for the first time. It would be more appropriate to introduce this work and its conjecture earlier in the introduction, making it easier for readers to appreciate the contribution.

[E] We thank the reviewer for the suggestion. We have highlighted the work of [1] and their conjecture in line 70-77 of the revision.

[Q1] Can the fact that the network reaches an "averaging" solution be deduced from the KKT equations in Frei et al. (2022)?

[A1] As we mentioned in the paper, the convergence of the network to KKT points is notably slow. Both Theorem A.9 in the appendix of [2] and Theorem 5 in [3] demonstrate that reaching KKT points requires exponentially many iterations, far exceeding the duration of practical training. In contrast, our analysis is applicable within polynomial time.

In fact, we also attempted to derive our feature averaging conclusion solely from the KKT conditions presented by [4] but were unsuccessful. The primary challenges we encountered are as follows:

1. The difficulty in analyzing the specific magnitude of the dual variables in the KKT conditions.
2. We were unable to characterize the activation regions induced by the ReLU activations (like our Lemma C6) only from KKT conditions.
3. We suspect that there exists certain KKT solution that does not correspond to the uniform average of features (with nearly uniform weights). But this is hard to verify even empirically since reaching KKT points is quite slow.

[Q2] The implicit bias of GD for this problem in 2-layer neural networks is undesirable for robustness. What about perturbations? It seems that this bias would help for perturbations, in contrast to the feature decoupling solution . I would appreciate your thoughts on this.

[A2] If we consider the case where the features consist of k standard orthogonal basis ($e_1,e_2,\cdots,e_k$), the robustness under the $\ell_{\infty}$ norm also leads to similar conclusions. Specifically, we can simply prove that:

1. The network trained by 2-classification labels is non-robust to $\ell_{\infty}$-perturbation radius $\Omega(\sqrt(d)/k)$;

2. The network trained by fine-grained(k-classification) labels is robust to $\ell_{\infty}$-perturbation radius $O(\sqrt(d))$.

For the general case, the theoretical analysis for $\ell_{\infty}$-​ perturbations may require additional assumptions, as the orthogonality condition holds less significance in this context. Therefore, we leave the study of the $\ell_{\infty}$ scenario for future work.

[Q3] Relatedly, would it be possible to measure robustness in the experiments on CIFAR10/MNIST against perturbations? If there is evidence for larger robustness when training with additional supervision, then this would make the authors' claims much stronger.

[A3] We conducted the experiments and presented the results in Appendix H in the revised version of our paper. The results show that fine-grained supervision also improves $\ell_{\infty}$ robustness.

[Q4] Figure 3: How do you measure perturbation radius? If, for example, you normalise MNIST images to be between 0 and 1, then a perturbation radius of 2.0 does not make any sense. Can you mention the ranges of the input for MNIST and CIFAR10?

[A4] We follow the standard practice to normalize an image. (transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) for CIFAR-10 dataset and transforms.Normalize((0.5,), (0.5,)) for MNIST dataset) Given the high dimensionality of the data, the $\ell_2$ norm of the normalized input is significantly larger than 2. Therefore, an $l_2$ perturbation radius of 2.0 is reasonable.

[Q5] In line 882, you mention that "In image classification tasks, Ilyas et al. (2019) visualized both robust and nonrobust features.". However, this is not true to the best of my knowledge. I believe there are 2 works that have provided such visualizations [1, 2].

[A5] The "visualization" referred to in this sentence corresponds to Figure 1 in [5]. We appreciate you pointing out these additional sources and have referenced them in the revised version of the paper.

Reference

[1] Min, H., & Vidal, R. (2024). Can Implicit Bias Imply Adversarial Robustness?. arXiv preprint arXiv:2405.15942.

[2] Lyu, K., & Li, J. (2019). Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv:1906.05890.

[3] Soudry, D., Hoffer, E., Nacson, M. S., Gunasekar, S., & Srebro, N. (2018). The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70), 1-57.

[4] Frei, S., Vardi, G., Bartlett, P., & Srebro, N. (2024). The double-edged sword of implicit bias: Generalization vs. robustness in relu networks. Advances in Neural Information Processing Systems, 36.

[5] Ilyas, A., Santurkar, S., Tsipras, D., Engstrom, L., Tran, B., & 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 positive feedback! We greatly appreciate the recognition of the novelty and significance of our contribution to the topic of adversarial robustness in the deep learning community, as well as the positive remarks on 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-1] First, the number of neurons assumed to be finite as the number of clusters of features in assumption 4.3. It is against the purpose of using the neural tangent kernel (NTK) theory to approximate infinite wide neural networks. Thus, the result in this paper is limited since “training deep neural network is a highly non-convex and over-parametrized optimization problem”.

[E1-1] We would like to clarify that we did not use the NTK assumption or NTK analysis in this paper. Unlike the infinite width (excessive over-parameterization) requirement of the NTK regime, we focus on a more realistic setup with moderately finite width (mild over-parameterization) and small initialization. This setup leads the neural network to the feature learning regime (i.e., more discriminative features are learnt during the training process, unlike in the NKT (also called lazy training) regime where the features barely change during training. Within this feature learning regime, we aim to understand the type of features (averaged features or decoupled features) learned by the GD algorithm based on the multi-cluster data distribution.

Indeed, due to the non-linearity of the ReLU activation function, training two-layer ReLU networks with gradient descent method under this feature learning regime becomes a highly non-convex optimization problem. To address this challenge, we develop novel techniques for analyzing the full training dynamics by leveraging margin auto-balance (Lemma C.4 and Lemma E.7), which enables us to bound the loss derivative ratio (equation (6)) and analyze ReLU activation regions (Lemma C.6). See our detailed proof sketch in Appendix C.

[W1-2] Second, as the solution to adversarial robustness, “fine-grained supervision” suggests including the original cluster labels for all data points, which is extremely unrealistic, because the response is simply assigned from the cluster labels. It is not surprising that there is an improvement in adverbial robustness as the exact latent information is provided. Feature decoupling is a simple result as the network can learn each cluster’s mean given the cluster labels.

[E1-2] While it might intuitively appear that annotation refined to the cluster level could ensure that a multi-class neural network eventually learns a feature-decoupling solution, we emphasize that convergence to robust networks is nontrivial and requires implicit bias. In particular, in Proposition 4.8, we provide a specific construction of multi-class network that achieves perfect clean accuracy but is non-robust. Indeed, due to non-linearity of ReLU activation function, it is also highly non-trivial to rigorously prove that the multi-class network converges to feature-decoupling solution, where we apply a margin auto-balance technique (Lemma F.10) similar to that in proof of feature-averaging regime to show it.

[W2] The evidence of feature averaging is shown Figure 2 as the strong correlations/cosine value between cluster’s mean and network weights, but the scale of correlations is not consistent. In particular, the results from CIFAR-10 are not convincing, where the average cosine value can be lower than 0.1.

[E2] In real-world datasets, such as the CIFAR-10 dataset, the data distribution is considerably more complex than the theoretical distributions assumed in our analysis. Consequently, even though we theoretically prove that the features exhibit exact uniform averaging, networks trained on real-world datasets often show feature averaging with moderately non-uniform coefficients.

[Q1] Minor issues in notations: Big notation on line 185 is different from the others; notation is not defined.

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

[Q2] Why does strong correlations indicate the existence of feature averaging? Is there a rigorous proof? How do you define the threshold of having a strong correlation?

[A2] We would like to clarify that we consider the following decomposition: $\boldsymbol{w} \approx c_1 \boldsymbol{\mu}_1 + c_2 \boldsymbol{\mu}_2 + \cdots + c_k \boldsymbol{\mu}_k$, where $\boldsymbol{w}$ represents the weight vector, and $\boldsymbol{\mu}_i$ represents the features of the $i$-th cluster (the first half corresponds to positive clusters, and the remaining part corresponds to negative clusters). Assuming that the cluster features are nearly orthogonal, the feature averaging regime means that:

1. For positive neurons, $c_1 \approx c_2 \approx \cdots \approx c_{k/2} \gg c_{k/2 + 1} \approx \cdots \approx c_k \approx 0$.

2. For negative neurons, $c_{k/2 + 1} \approx c_{k/2 + 2} \approx \cdots \approx c_k \gg c_1 \approx \cdots \approx c_{k/2} \approx 0$.

In our experiments, the coefficients $c_i$ are approximately computed based on the correlation between the weight and the cluster features. The term "strong correlation" here is a relative concept, indicating that the correlation between weight vectors and features of the corresponding classes is significantly stronger than the correlation with features of the other classes.

[Q3] Proposition 4.8 is used as an example of “adding more fine-grained supervision signals does not trivially lead to decoupled features and robustness”, but it does not follow the setup of your analysis, since . How is this multi-class network obtained? Why is its parameter assumed to be the combination of ?

[A3] We would like to clarify that we manually construct the multi-class network to show the existence of a multi-class network that achieves perfect clean accuracy but is non-robust, which shows that adding more fine-grained supervision signals does not trivially lead to decoupled features and robustness. We design this network, and it aligns with the setup. If $h=1$ represents an important distinction, then it is possible to replicate $h$ identical neurons.

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 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] The feature average is clearly defined in two-layer networks - the linear input combinations. However, in the deeper layer, such a definition does not seem reasonable. How do we analyze features in this case?

[A1] We thank the reviewer for the insightful suggestion about feature learning in deep networks. Indeed, we also believe a similar feature-averaging bias may emerge in more complex scenarios beyond our theoretical settings. Conceptually, although there may not be a predefined set of features in deeper networks, gradient descent could still have a tendency to combine many localized, semantically meaningful (and thus, more robust ([1][2])) features into a single discriminative but non-robust feature.

Specifically, deep neural networks can be viewed as consisting of two parts: a feature extractor that is a mapping from the input space to the latent space, and a shallow classifier that predicts the classification results based on the extracted features (in the latent space). Feature averaging means that the shallow classifier mixes extracted features corresponding to different classes in the latent space. To verify this, we provide a CLIP-based experiment in the revised version of the paper, where we first apply the CLIP model to extract latent representation of the input image and then we train a two-layer network based on the latent representation. See the CLIP experiment results in Appendix H in the revised version of the paper for details.

Indeed, if all layers (including the feature extractor layers) are trainable, the feature averaging bias may further influence the feature extracted in a more complicated way, and we leave the study for deeper networks as a challenging future direction.

Finally, we would like to mention that similar “averaging” or “superposition” effects also observed in ref [3]. In particular, the authors observe that in CLIP models each neuron may encodes multiple, often unrelated concepts (e.g., ships and cars), and they leveraged such effect to construct "semantic" adversarial examples. See Figure 1 in [3] for details.

[Q2] In image tasks, reasonable and robust features are often local and isolated. However, there are other scenarios where the features are required to have global or long-term dependency. Such as time series, and natural language. In those cases, it seems like the real interpretable and robust features are some "mixing" of the original data.

[A2] Thank the reviewer for giving the valuable suggestion, and we clarify that our study focuses on adversarial examples in the computer vision field. While extending to natural language processing or time series is very important, it falls beyond the scope of this work.

Indeed, as the reviewer pointed out, many interpretable features in NLP have global or long-term dependency. However, only local attacks can be very harmful for NLP models. For instance, the Greedy Coordinate Gradient (GCG) [4] attack only appends a suffix to the original prompt (which is a form of local attack) and it is sufficient to induce jailbreaks in large language models (LLMs). So far, the GCG attack continues to be challenging to defend against effectively. Hence, there seems to be a different mechanism for the non-robustness in NLP models and we leave it as an exciting and challenging direction for future research.

[Q3] There are researches that have out that, the neural network tends to learn low-frequency features (which are nonlocal and varied smoothly) more effectively than high-frequency features. Some connect such behaviour to the generalization capacity of over-parameterized neural networks. Do you think there is any relation between this theory and the ones developed in this paper?

[A3] We thank the reviewer for pointing out the related works. The papers mentioned by the reviewer refers to the phenomena that deep neural networks generally learn lower-frequency features first, and then the higher-frequency ones. This can be seen as a particular form of simplicity bias. The feature averaging bias studied in this paper is also a form of simplicity bias: under our theoretical setup, the simplicity refers to the linear combination of cluster features, and it is closely related to the approximate linearity of the decision boundary, as discussed in Appendix A. Hence, both studies assert that neural networks tend to favor simplicity during the initial stages of training, sharing a similar underlying spirit.

Reference

[1] 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.

[2] Tsilivis, N. and Kempe, J. (2022). What can the neural tangent kernel tell us about adversarial robustness? Advances in Neural Information Processing Systems, 35 18116–18130.

[3] Gandelsman, Y., Efros, A. A., & Steinhardt, J. (2024). Interpreting the Second-Order Effects of Neurons in CLIP. arXiv preprint arXiv:2406.04341.

[4] Zou, A., Wang, Z., Kolter, J. Z., & Fredrikson, M. (2023). Universal and transferable adversarial attacks on aligned language models. arXiv preprint arXiv:2307.15043.