Boosting Adversarial Robustness with CLAT: Criticality Leveraged Adversarial Training

Authors: Bhavna Gopal, Huanrui Yang, Jingyang Zhang, Mark Horton, Yiran Chen
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Adversarial robustness, criticality, Computer vision, Adversarial attacks, defense

Abstract:
Adversarial training (AT) is a common technique for enhancing neural network robustness. Typically, AT updates all trainable parameters, but such comprehensive adjustments can lead to overfitting and increased generalization errors on clean data. Research suggests that fine-tuning specific parameters may be more effective; however, methods for identifying these essential parameters and establishing effective optimization objectives remain unclear and inadequately addressed. We present CLAT, an innovative adversarial fine-tuning algorithm that mitigates adversarial overfitting by integrating "criticality" into the training process. Instead of tuning the entire model, CLAT identifies and fine-tunes fewer parameters in robustness-critical layers—those predominantly learning non-robust features—while keeping the rest of the model fixed. Additionally, CLAT employs a dynamic layer selection process that adapts to changes in layer criticality during training. Empirical results demonstrate that CLAT can be seamlessly integrated with existing adversarial training methods, enhancing clean accuracy and adversarial robustness by over 2% compared to baseline approaches.

Decision: reject
Reviewer scores: 6, 3, 6, 5, 6

Reviews and discussion:
<thread>
creation date:
1730195058194

thread_type:
official_review

<thread_object>
creator: reviewer_od5s

<content>
summary:
This work studies the adversarial robustness of image classification models. It focuses on improving adversarial robustness through fine-tuning the parameters of selected layers. Specifically, it first proposes to identify robustness-critical layers by the (approximate) curvature of loss landscape. During fine-tuning, it only updates the parameters in the robustness-critical layers, while keeping the rest fixed.

strengths:
1. The idea of identifying and fine-tuning robustness-critical layers is interesting and less studied before. It has the potential of improving the trade-off between accuracy and robustness through reducing the over-regularization of adversarial training.

weaknesses:
1. The paper uses the term "adversarial overfitting" without a clear definition. It is confusing because there exists phenomena such as "robust overfitting" [1] and "catastrophic overfitting" [2]. The authors should clarify if they are proposing something new or just referring to either of them. If latter, I would suggest the authors to avoid the creation of new term and adopt the existing. 
2. The motivation of the proposed method is weak. The paper could benefit from a more in-depth discussion of how updating critical layers, instead of all, mitigates adversarial overfitting. 
3. Technically, the fine-tuning loss proposed in section 3.2 is similar to the loss function of TRADES.
4. 10-step PGD without restart is too weak so the conducted adversarial evaluation using it is unreliable. It is therefore unable for me to judge the soundness of the results e.g. in the Tab. 1. 
5. It has been demonstrated that data augmentation methods such as IDBH [3] and generative models [4] are effective in mitigating overfitting in adversarial training. It could enhance the strength of the proposed method if it improves the performance when combined with them.

[1]  Rice, Leslie, Eric Wong, and Zico Kolter. "Overfitting in adversarially robust deep learning." International conference on machine learning. PMLR, 2020.

[2] Wong, Eric, Leslie Rice, and J. Zico Kolter. "Fast is better than free: Revisiting adversarial training." International Conference on Learning Representations. 2020.

[3] Li, Lin, and Michael W. Spratling. "Data augmentation alone can improve adversarial training." The Eleventh International Conference on Learning Representations. 2023

[4] Wang, Zekai, et al. "Better diffusion models further improve adversarial training." International Conference on Machine Learning. PMLR, 2023.

questions:
Please refer to the above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730464522585

thread_type:
official_review

<thread_object>
creator: reviewer_agzn

<content>
summary:
This paper proposes a method to address adversarial overfitting issues by optimizing only the parameters in critical layers during adversarial training, instead of all parameters. The authors argue that this approach can effectively alleviate the problem. Specifically, they identify critical layers as those that exhibit a stronger tendency to learn non-robust features or demonstrate higher sensitivity to adversarial input perturbations compared to other layers in the model. The authors empirically demonstrate the effectiveness of their method through experiments conducted on CIFAR10 and CIFAR100 datasets.

strengths:
The numerical experiments in the paper yield promising results, and the method's general applicability is validated across multiple models. The approach presented in this paper is intuitive, and the writing is clear and easy to understand.

weaknesses:
The method proposed by the authors appears to be empirically feasible but lacks theoretical justification. I have certain reservations about the authors' approach, as it appears to be not much different from simply reducing the learning rate. Notably, the authors have fix the maximum number of epochs in their experimental setup, rather than training until convergence, which deviates from conventional practices. If the method merely prevents the model from learning to fit, in order to achieve the so-called mitigation of overfitting, I question the significance of this paper.

questions:
I suggest the authors further investigate the performance of the model after increasing the number of epochs, as well as compare the performance of the model after reducing the learning rate. Additionally, I recommend including the performance of the model without adversarial training in the table, to facilitate comparison for the readers.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730571765020

thread_type:
official_review

<thread_object>
creator: reviewer_jmsg

<content>
summary:
The paper presents an adversarial-finetuning algorithm that selectively updates specific parameters in a deep network. The algorithm relies on iteratively finding `critically vulnerable' parameters and then leveraging variants of adversarial training to robustly train deep networks. The fraction of critical parameters is often less than 4% of the total network, leading to better regularization and reduced overfitting. The paper presents empirical results for a variety of adversarial training approaches, as well as several ablation studies.

strengths:
1. The CLAT formulation expands upon existing intuition of constraining the local lipschitz constant of layers as an adversarial defense, by restricting the number of trainable layers. The idea of using smarter regularization to reduce overfitting is quite interesting, and proves to be useful in the settings that the paper describes.

2. The results and ablations are quite thorough, and support the claims in the paper. I especially appreciate the results in Fig. 2 for clearly showing the advantages and disadvantages of CLAT over just pure adversarial training.

weaknesses:
Some possible improvements are listed below:

1. The criticality metric is essentially an estimate of the ratio of local lipschitz constants across layers. While the use of the metric as a way to rank layers by vulnerability is interesting, it would be better to use standard terminology here rather than coming up with new ones. I also suggest adding a discussion comparing other lipschitz regularization approaches used as an adversarial defense.

2. The paper mostly conducts experiments on CIFAR-10/100 which are generally standard in the field of robustness. However, an experiment on ImageNet or a dataset of similar size and complexity would be a valuable addition. One hypothesis to test is if a dataset of additional complexity also requires a similar number of CLAT parameters. This would also further strengthen the claims from Table 8 which suggests that criticality is an architecture-specific phenomenon.

3. Another ablation that I find would be useful is to track the change in critical layer distributions across training. Since the algorithm re-identifies critical layers every 10 epochs, does the distribution change significantly over training? Specifically, are critical layers more prominent in earlier or later layers as training progresses. This may help further reduce the hyperparameter exploration required for CLAT.

4. Finally, a comparison of epoch wall-time for PGD-AT versus PGD-AT + CLAT would be interesting.

Overall, the results are quite interesting and encouraging, but I would appreciate it if the authors addressed other lipschitz regularization approaches for adversarial defense in the draft.

Minor nit: The plots often have very thin lines close to each other. I suggest increasing the line thickness and fontsize in the figures to make them more clear.

questions:
1. What was the reasoning behind approximating $\frac{\partial F(x')}{\partial x'}$ with a uniform vector in Eq. 5?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730626991184

thread_type:
official_review

<thread_object>
creator: reviewer_s6zb

<content>
summary:
This paper proposes a **Criticality-Leveraged Adversarial Training (CLAT)** method to enhance adversarial robustness by keeping certain parts of the neural network fixed and only fine-tuning the layers that are fragile in terms of robustness. Experiments are conducted across various datasets and neural network structures.

strengths:
1. The explanation of the concept is clear.

2. The presentation in the experimental section is well-organized and easy to follow.

weaknesses:
1. This method is cumbersome as it requires inputting both clean and adversarial images to compute the difference.

2. In Section 3.1, numerous mathematical tools are introduced but ultimately seem unnecessary, as the final objective function minimizes the initial starting point. This makes much of the content in Section 3.1 redundant.

3. Experiments are conducted only on the CIFAR dataset with CNN structures. The experimental design should be expanded to include both small- and large-scale datasets, as well as ViT models.

questions:
1. Some tables could use bold font to highlight the best results.

2. To demonstrate the effectiveness of the proposed Layer Criticality Index, an ablation study on fine-tuning different critical layers with small or large $C_{f_i}$ values would be helpful.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730673595570

thread_type:
official_review

<thread_object>
creator: reviewer_swy4

<content>
summary:
The paper tackles the important issue of adversarial vulnerability of neural networks. Instead of updating all parameters during adversarial training, CLAT focuses on layers that predominantly learn non-robust features, identified by a "criticality index." This selective focus is designed to improve adversarial robustness, can be used with on top of standard adversarial training approaches. The paper offers detailed results on multiple architectures and also evaluated under both white-box and black-box adversarial attacks

strengths:
- CLAT presents an alternative to traditional adversarial training, which requires updating the entire network, and can be used on top of any adversarial training method
- The analysis is detailed and covers many network architectures 
- The method is evaluated on various attacks and also includes the strong attacks such as Auto Attack
- The ablation study and the overhead analysis is good and useful

weaknesses:
- I am unclear about the problem statement and objective. Because, in the Adversarial domain, there are two main adversarial training challenges: (1) the trade-off between clean and adversarial accuracy and (2) adversarial overfitting, and there are methods that cater to either or both, but the evaluation criteria are different. For ex, for trade-off we need to check how clean acc and adv acc are impacted, and for the latter, we need to report both the peak and final accuracy scores, highlighting the impact of overfitting. 
- The concept of robust and non-robust features is important in the paper, and needs more context and explanation. What is the criteria to differentiate, and further an empirical example or analysis will provide more clarity.
Additionally, the theoretical explanation behind the criticality index would benefit from more intuition. Is the motivation based on the paper [Robustness via curvature regularization, and vice versa], that small curvature correlated to higher robustness? How does this relate to criticality or non-robust features? The main focus of the paper about critical layers and its criterion is unclear to me.
- The paper uses multiple terms—such as "weakness," "robust-critical layers," "sensitive layers," and "non-robust features". This variation makes it difficult to understand the precise characteristics, need consistent definitions. Also are we talking about NR layers (weights) or features, there's multiple usages of these terms.
- Since CLAT fine-tunes layers that supposedly learn non-robust features (which aid clean accuracy), one would expect that this selective fine-tuning might reduce clean accuracy, however the clean accuracy improves as well ? I am curious to understand more on this.
- Line 430: Additional intuition or analysis is needed to support this claim. What architectural properties might influence these findings? Additionally, is there a discernible pattern to the critical layers? Identifying such patterns could provide valuable insights into what different layers learn and how this understanding could be leveraged to improve robustness further.

questions:
- Could you elaborate on different hyperparameters and how they were chosen? Ex. the 5% threshold for selecting critical layers. Is there any consideration for adjusting the percentage based on the depth or structure of the network?
- Robust/Non-Robust - is it only layer wise, or also makes sense to see sub-networks or neuron-wise? Introduction could use more context
- Results - (related to point 1 in weaknesses), if the the paper claims to address overfitting in adversarial training, it would be essential to compare CLAT with explicit overfitting reduction techniques like AWP (Adversarial Weight Perturbation) or SWA (Stochastic Weight Averaging). Including metrics such as best vs. final accuracy scores for these methods would help demonstrate whether CLAT effectively mitigates overfitting.
- Computational overhead on bigger and more complex datasets? (other than CIFAR)
-  Line 461 - Did not understand the hypothesis for the behavior of Figure 3. Please elaborate with supporting evidence. Further, how is the behavior on clean accuracy? In general, the ablation study is interesting, but the claims and intuitions presented would benefit from more supporting evidence to strengthen their validity.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732319035853

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of paper revision

comment:
We thank the constructive feedback from all reviewers. The submitted paper has been revised to incorporate your suggestions. In the main paper, we fixed wording issues that lead to the confusion of reviewers sWy4 and oD5s. We also updated the formatting of some tables following the suggestions of the reviewer s6zB. We added a new **Appendix D** to the paper, which consolidated all additional experimental results suggested by the reviewers for easier reference. We will integrate these results and our discussions in the rebuttal text into the finalized version.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733165382034

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank the help and support from all reviewers in improving our paper in this discussion period. We would like to kindly remind the reviewer that the discussion period will conclude soon. We hope our responses and revisions help addressing your concerns. If further clarification is needed, please do not hesitate to ask.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733280860967

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Thanks, final note on revisions and clarifications

comment:
**We sincerely appreciate your thoughtful feedback and the opportunity to refine our work during the review process.** We believe that CLAT marks a significant advancement in adversarial training and has the potential to benefit the research community substantially. The primary motivation of this work is to mitigate robust overfitting during adversarial training, thereby improving both clean and adversarial accuracy. Moreover, CLAT integrates seamlessly with existing adversarial training methodologies to deliver state-of-the-art performance. We are pleased that all reviewers acknowledged the strong empirical performance of CLAT and appreciated the thoroughness and organization of our experiments and ablation studies.

In response to your feedback, we have implemented several key revisions. We ensured consistent terminology throughout the paper and added time comparisons with RiFT and AutoLoRA. Additionally, we conducted an ablation study to further validate the effectiveness of our index metric. To improve clarity, mathematical derivations were reorganized, with some details moved to the appendix. Section 3.1 was updated to incorporate your input, and we clarified the relationship between our "feature weakness" formulation and local Lipschitz constants. Concerns regarding extended training and reduced learning rates were also addressed through comprehensive empirical analysis.

We are confident that these revisions address all concerns and **would be deeply grateful if you find them satisfactory and consider awarding higher scores.** Thank you again for your valuable insights and engagement.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a technique to improve the robust overfitting using a technique called CLAT. Instead of updating all parameters during adversarial training, CLAT focuses on layers that predominantly learn non-robust features, identified by a "criticality index.". The fraction of critical parameters is often less than 4% of the total network, leading to better regularization and reduced overfitting

During the review period, there are several concerns raised about lack of thorough experimentation including datasets and models. Some reviewers also point to the lack of novelty of the approach, and that the paper is unclear to follow.

During the rebuttal, authors have improved the writing taking into account the reviewer's feedback. However, I think the authors still need to improve the paper and make it easy to understand. While the reviewers added experiments on imagenet during the rebuttal, I still feel a solid comprehensive comparisin and ablation studies are still missing, The fact that the method doesn't work with ViTs yet is another issue.

Overall, I feel the authors need to improve the experimentation, positioning and paper writing. So, I feel the paper is not ready for publication yet and I encourage the authors to take into account the invaluable suggestions given by the reviewers to improve the paper.