Boosting Adversarial Robustness with CLAT: Criticality Leveraged Adversarial Training

W1: Robust overfitting

Sorry for the confusion. We are discussing and tackling the same phenomena as mentioned in the two papers you provided. We have unified our wording to “robust overfitting”, or simply “overfitting”, in the revised paper.

W2: Motivation

We hypothesize that only specific critical layers in a network should be further robustified via adversarial training, and focusing on these layers, rather than the entire model or a larger subset, improves results and mitigates overfitting. With the proposed criticality metric in Equ. (6), we identify layers that are still learning non-robust features in the adversarial training process, and propose to sepcifically reduce their criticality with the rest of the model frozen. This layer-selective tuning enables the critical layers to learn robust features, while preventing the other already-robust layers to overfit on the adversarial training objective.

Our hypothesis is supported by our results and reinforced by the ablation study, which shows that fine-tuning either too many or too few layers negatively impacts performance. Figure 3 demonstrates the importance of selecting the optimal number of critical layers for adversarial accuracy. Clean accuracy trends similarly, as the same number of critical layers is optimal for both adversarial and clean accuracy, and fine-tuning too many layers degrades performance in both cases. These results are shown in Appendix D.4, Figure 7.

W3: TRADES

CLAT differs significantly from the TRADES objective in both methodology and ultimate goal.

CLAT focuses on fine-tuning specific layers identified as critical, while freezing the rest of the model. This selective optimization mitigates overfitting and enhances robustness without making global changes.

In contrast, TRADES optimizes the entire network to reduce the difference between clean and adversarial outputs. This is a broader approach aimed at achieving robustness across the entire network, not just specific critical layers.

Mathematical Differences

The two objectives are fundamentally different:

• In CLAT, the additional term $\lambda \mathcal{L}_C(f_S)$ targets the identified critical layers to reduce their sensitivity to perturbations. The objective $\mathcal{L}_C(f_S)$ is formulated in Equ. (8). This targets feature robustness specifically at the critical layers most prone to learning non-robust features, reducing their criticality without modifying the rest of the model.

• In TRADES, the second term

  $\lambda \cdot \max_{x' \in \mathcal{B}(x)} \mathcal{L}(f_\theta(x'), f_\theta(x))$

  focuses on minimizing the output discrepancy between clean and adversarial examples. This involves finding adversarial examples through an inner maximization and ensuring that the network's response to these examples is similar to its response to clean data, thereby enhancing adversarial robustness.

W4: Stronger attacks

To verify the reliability of our evaluation, we provide additional results with increased steps of PGD, as well as additional attack methods, in Appendix D.7 (Table 17). CLAT shows consistent performance improvement over baseline under all attacking scenarios. Furthermore, we already show the performance of CLAT under varying epsilon values on different networks in Appendix B and include AutoAttack evaluations in the main paper, showing our performance improvement is universal.

W5: Data augmentation

While augmentation techniques have proven effective in mitigating adversarial training, we emphasize that this represents an orthogonal line of research, focusing on data, whereas our work is centered on model and objective optimization. Nonetheless, we provide results demonstrating the performance of CLAT after applying augmentation techniques to illustrate its versatility, as shown in Appendix D.8.2. Applying CLAT with the Augmented data further improves model robustness compared to baseline methods.

Thank you for raising these insightful points. We address your concerns in two parts:

1. Extended Training Epochs and Convergence: In response to your suggestion regarding convergence, we extended the training duration to 150 adversarial epochs, using a cosine learning rate scheduler that decays the learning rate to zero by epoch 150. As shown in Figure 5 (Appendix D.1), models trained with PGD-AT show a consistent decline in performance with more adversarial training epochs, highlighting the overfitting effect. In contrast, models trained with CLAT improves both performance and robustness across all epochs, demonstrating the mitigation of overfitting and allowing the model to fit more with more epochs being trained.
2. Learning Rate Reduction Comparison: We also conducted experiments with a reduced learning rate, reducing it by a factor of 10 at the 70th epoch, as presented in Appendix D.2 (Figure 6). This setup included comparisons between CLAT with the original learning rate, CLAT with the reduced learning rate, and PGD-AT with the reduced learning rate, all trained to 150 epochs with eventual decay of the learning rate to zero. Across both clean and adversarial accuracies, CLAT consistently outperformed the other configurations, reinforcing its effectiveness beyond merely reducing the learning rate.

To add more discussion on the learning rate reduction, reducing learning rate may boost the performance in the specific epoch, but does not prevent the model from further overfitting. CLAT, on the other hand, enables continued performance improvement throughout training epochs regardless of the learning rate used.

Lastly, the orange line in Fig. 5 indicates the performance of CLAT from scratch without any prior adversarial training. We also show the performance of CLAT on clean pre-trained models in Appendix Section C.1 (Table 11). These results show CLAT mitigates overfitting and improves both accuracy and robustness throughout the finetuning epochs, no matter of the starting point.

Thank you so much for your review, we really appreciate your attention to detail and have carefully addressed all of your concerns/ questions.

W1: Cumbersome

We respectfully disagree. CLAT training does not increase the computational complexity compared to the standard adversarial training methods, so the total training time does not differ. The only additional cost of CLAT involves identifying critical layers with the criticality indices. The computational overhead for our critical indices is negligible.

• In our investigation in Sec 4.4, we present the average time required to compute critical layers across various batch sizes using randomized data on DenseNet121. It's worth noting that smaller networks exhibit even shorter computation times. Overall, we find that criticality indices can be computed with batch size as small as 10 (See Table 10). Conservatively, if we consider a batch size of 50, our critical layers are derived from a mere 0.0008% of the training data. Adopting criticality estimation every 10 epochs, as proposed in our method, introduces neglectable additional complexity to the training process.
• Lastly, estimating the criticality for all layers involves two forward passes: one with a batch of clean input and one with corresponding PGD adversarial examples against the model output. Importantly, both clean inputs and adversarial examples are readily available during the adversarial training process. As we empirically demonstrate that the criticality estimation process is insensitive to the selection of training data, we can simply record the hidden features of the final batch of adversarial training data for the criticality computation in the new layer selection cycle.

W2: Math in Sec. 3.1

Intuitively, a layer learning non-robust features will be influenced more by adversarial examples. This impact can be measured by the difference between the output features when facing a clean input vs. facing an adversarial input. In Sec. 3.1,we further make a theoretical connection between the output feature difference and the local curvature around a layer’s output, showing that the output MSE under attack, as in Equ. (5), is an approximation to the local loss curvature, which represents the robustness as derived in [Robustness via curvature regularization, and vice versa]. This derivation ensures Equ. (5) is a valid estimation for local feature weakness/robustness, which builds the foundation for deriving the criticality metrics in Equ. (6).

It should also be noted that Equ. (3), which comes with the first-order gradient term, is a more accurate measurement to feature robustness and would lead to a more accurate criticality metrics. However, as we use the criticality metrics as both layer selection criteria and layer finetuning objective (see Equ. (8)), having a first-order gradient term in the metric will make the optimization significantly more costly. Here we empirically find that the approximation in Equ. (5) is good enough for improving clean accuracy and robustness of CNN models. Yet the full derivation may be useful for future work to apply CLAT-like methods on more complicated architectures that require a more precise criticality measurement.

W3: Other datasets and models

As noted by reviewer jMsg, our experiments primarily focus on CIFAR-10/100, which is a standard benchmark in robustness research. However, we have extended our results to include performance on Imagenette and ImageNet, as presented in revised Appendix D.6

ViT models, due to its larger size and diverse layer types, bring unique challenges to the feature-based layer selection process. This is out of the scope of this paper and we plan to address it in future work. Initial experiments on applying competing techniques such as RiFT and AutoLoRA on DeiT, ViT, and SwinB leads to more than a 3% reduction in adversarial accuracy (on PGD-10) compared to full adversarial training, showcasing the different characteristic of ViTs and CNNs.

Q1: Tables

Thank you for the feedback. We have revised accordingly

Q2: Layer selection ablation

We provide comprehensive ablations in the manuscript to demonstrate the effectiveness of our Layer Criticality Index. Notably, we compare the performance of selecting layers based on our critical indices versus random selection in Tab 12 (Appendix C.2), as well as dynamic versus static critical indices during CLAT in Tab 14 (Appendix D.5). These results clearly indicate the effectiveness of our layer selection approach

In our experiments, we fine-tune layers with the largest Cfi values, as these are identified as the most critical. We do not perform ablations on fine-tuning layers with the smallest Cfi values, as this aligns with the premise of RiFT, which targets the most robust layers and finetune them with clean training objectives. Notably, our method consistently outperforms RiFT across all benchmarks, validating the superiority of our approach

Thank you for your insightful review!

W1: Local lipschitz constants

We thank the reviewer for pointing us to the line of Lipschitz regularization. We agree that the formulation we reached in Equation (5) is similar to the local Lipschitz constants, and would love to make the connection in our revised paper.

Meanwhile, from the adversarial robustness perspective, there seems to be less exploration on Lipschitz regularization approaches. We take inspiration from the curvature formulation as in Equ. (3), which is a more accurate measurement to feature robustness and would lead to a more accurate criticality metrics. However, as we use the criticality metrics as both layer selection criteria and layer finetuning objective (see Equ. (8)), having a first-order gradient term in the metric will make the optimization significantly more costly. Here we empirically find that the approximation in Equ. (5) is good enough for improving clean accuracy and robustness of CNN models. Yet the full derivation may be useful for future work to apply CLAT-like methods on more complicated architectures that require a more precise criticality measurement.

W2: Larger dataset

We have provided results for Imagenette (Appendix D.6.1, Table 15) and Imagenet (Appendix D.6.2, Table 16) in our rebuttal. We used the same setting as in the CIFAR training, such as selecting 5% of the layers and using a random batch of 50 data points for criticality measurements. We note that our improvements are consistent as the CIFAR experiments, strengthening the observation that criticality seems to be an architecture specific phenomenon.

W3: Critical layer change*

Table 13 (Appendix D.3) presents the computed critical indices for DN121, RN50, and RN18. We recompute the critical indices at epochs 70, 80, and 90 during the fine-tuning phase (30 epochs). Prior to this, we conduct adversarial training without freezing any layers. We observe that the distribution of critical layers changes over the course of training, although there is no clear trend indicating whether they are more prominent in earlier or later layers as training progresses. The results from this experiment, combined with the results from the ablation study where we do not change critical layers at all indicate that dynamic layer selection is crucial for CLAT.

Additionally, Table 14 (Appendix D.5) bolsters our reasoning for utilizing dynamic selection. Across the board, CLAT conducted with dynamic layers outperforms CLAT with fixed layer selection throughout finetuning.

W4: Wall time

Expanding on results from Table 10, one PGD-AT epoch on DN121 takes 67 seconds on CIFAR and one epoch for CLAT finetuning takes the same amount of time. This is not including critical index recomputation time, which is minimal compared to the adversarial training time. Even if we take a batch size of 100 here (Critical indices can be computed accurately for batches as small as 10) for computing critical indices (and we do this thrice), it adds only 3.15s to the training time every 10 training epochs. Consider a typical PGD-AT+CLAT training process with 70 epochs of PGD-AT pre training followed by 30 epochs of CLAT finetuning, the total training time would take 10.5 extra seconds over the 1.86 hours standard adversarial training time. This timing is fairly consistent with most other networks as well.

Line and font in figures

Thanks for pointing this out, we will be sure to fix this in the final version.

Q1: Approxmation in Equ. (5)

As we use the criticality metrics as both layer selection criteria and layer finetuning objective (see Equ. (8)), having a first-order gradient term in the metric will make the optimization significantly more costly. Here we empirically find that the approximation in Equ. (5) is good enough for improving clean accuracy and robustness of CNN models.

Problem statement and objective

Overfitting is the primary challenge we address, but we caveat that this encapsulates the degradation of both clean and adversarial accuracy. Overfitting occurs when the model becomes excessively tailored to adversarial examples, reducing its ability to generalize and impacting both clean accuracy and robustness over time.To accentuate this, we extend the learning curve experiments in Figure 2 to 150 adversarial training epochs, using a cosine learning rate scheduler that decays the learning rate to 0 by epoch 150. As shown in Figure 5 in the updated Appendix D.1, models trained with PGD-AT exhibit a consistent decline in performance with additional adversarial training epochs, clearly illustrating the effects of overfitting. In contrast, models trained with CLAT maintain consistent performance and robustness improvements across all epochs. We note that CLAT models also achieve higher peak and final accuracies as a result of being significantly less prone to overfitting.

Criticality and non-robust features

The concept of robust and non-robust features is indeed crucial, and we appreciate the opportunity to provide more context. Non-robust features are those that significantly impact adversarial robustness, as they are highly sensitive to adversarial perturbations. Intuitively, a layer learning non-robust features will be more affected by adversarial examples. This sensitivity can be measured by the difference between the layer's output features when faced with a clean input versus an adversarial input.

In Section 3.1, we further establish a theoretical connection between this output feature difference and the local curvature around a layer's output. Specifically, we show that the output MSE under attack, as defined in Equation (5), serves as an approximation to the local loss curvature, which correlates with robustness as derived in [Robustness via curvature regularization, and vice versa].

To isolate the contribution of each layer to the learning of non-robust features, we define a layer criticality measure in Equation (6), which represents the ratio of the feature MSE computed before the layer to the feature MSE computed after the layer. Layers with higher criticality contribute more significantly to the learning of non-robust features, aligning with our definition of "critical layers" in Definition 3.1. Empirical results in Table 7 demonstrate that selecting these critical layers is crucial for the robustness and accuracy improvements achieved by CLAT.

Weakness, critical, and sensitive

We appreciate the reviewer's feedback and apologize for the confusion caused by inconsistent terminology. To clarify:

• "Weakness" refers to a property of the feature, as defined in Equation (5), whereas "criticality" refers to a property of the layer, as defined in Equation (6). We use the term "weakness" as a concise noun to describe how non-robust a feature is. Since "non-robust" is an adjective and "non-robustness" feels too cumbersome, "weakness" serves as a fitting alternative as it is the opposite of "robustness."
• We acknowledge that the term "sensitive" was unnecessary and led to ambiguity. As a result, we revised Section 3.1 to eliminate the use of "sensitive" for greater clarity.
• We now use consistent terminology throughout the paper: in adjective form, we refer to a "non-robust feature" and a "critical layer," while in noun form, we use "feature weakness" and "layer criticality." These changes are clearly reflected in the revised Section 3.1 and aim to eliminate any ambiguity regarding the terms.

Additionally, we have ensured consistent usage of the terms “layer” and “feature” to avoid confusion between non-robust features and critical layers. The revised manuscript now clearly differentiates between features and layers and uses these terms consistently.

Clean accuracy improvement

To achieve an adversarially robust model, we want all layers to learn robust features. Note that well-learnt robust features should also be useful, which leads to both robust and accurate classification. We believe an ideal conversion from learning non-robust features to robust features should not lead to clean accuracy drop. A model can be both robust and accurate, much like the human brain.

We hypothesize that overfitting is the true cause of loss of clean accuracy in standard adversarial training. In other words, we believe that overfitting to adversarial examples can degrade the model’s performance on clean data. CLAT is designed to counter this issue by selectively fine-tuning only the critical layers that need robustness improvements, while keeping already-robust layers fixed. We propose that this approach prevents overfitting to adversarial data and ensures that the model retains its ability to perform well on clean inputs.

Our empirical results suggest that this selective fine-tuning approach effectively mitigates overfitting during adversarial training, allowing the critical layers to transition from learning non-robust features to learning robust ones without any loss of clean accuracy—and even leading to improvements. This outcome supports our hypothesis that improving robustness can indeed go hand-in-hand with improving clean accuracy, distinguishing CLAT from standard adversarial training methods that typically face a tradeoff between accuracy and robustness.

Architectural property

We appreciate this insightful question. Line 430 is what we observe empirically, and we acknowledge that the exact architectural properties influencing these findings remain unclear. Identifying patterns in critical layers and understanding the architectural properties that contribute to these results is indeed an important area for future research. We explicitly highlight this as an open question (in Conclusions) and a promising direction for future work.

Hyperparameters

The three primary hyperparameters in our approach are:

1. Percentage of Critical Layers: For each network, we conducted experiments where the algorithm was applied to 1, 2, ..., up to the total number of layers that could be selected as critical (entire fine-tuning). As stated in Lines 460–464, “Allowing more layers to be fine-tuned enhances the model flexibility, which initially improves CLAT performance. However, fine-tuning more layers diminishes CLAT's effectiveness. This is likely due to the diversion of attention towards fine-tuning less-critical layers, detracting from more critical ones.” Based on this, 5% was selected as a general threshold that balances performance across different architectures. We show the robustness-layer number curve in Fig 3 and the accuracy-layer number curve in Fig. 7 (revised Appendix D.4), which reflects our observations.
2. Number of Epochs for Fine-Tuning: Based on Figure 2, fine-tuning for 30 epochs appears to be optimal. This observation is consistent with results obtained when fine-tuning is extended to 150 epochs, where the learning rate decays to 0, across all tested networks.
3. Frequency of Computing Critical Indices: The frequency of recomputing critical indices (every 𝑋 epochs) was determined by testing the variance in critical layers after every epoch of training. We found that critical layers typically take approximately 10 epochs to become less critical through CLAT, making 10 epochs an effective interval for recomputation.

Subnetworks

The observation regarding robustness/non-robustness at sub-network or neuron-wise levels is an interesting direction. However, we focus on layerwise adjustments because they represent the outermost lever in the optimization space, offering significant performance improvements with manageable complexity.

We initially explored alternative granularities, such as blockwise and channelwise adjustments. Blockwise adjustments did not provide any gains, while channelwise adjustments introduced substantial computational overhead due to the significantly higher number of structures involved. Given these findings and the strong performance achieved with layerwise adjustments, we chose not to pursue finer-grained methods further.

Overfitting reduction techniques

CLAT proposes a way to mitigate overfitting by selecting and updating only the critical layers. We believe CLAT is orthogonal to other explicit overfitting reduction techniques that are applied in the full model optimization process. Results for combining CLAT with other explicit overfitting reduction techniques are shown in Appendix D.8.1, Table 18, where final performance is reported. CLAT works together with the mentioned methods and improves performance universally.

Computational overhead

We want to clarify that our method does not actually have any computational overhead compared to standard PGD-AT except for the times required to compute critical indices (which is trivial as shown in Table 10). No matter the dataset or model used, criticality can be stably computed with a single training batch. In our ImageNet experiments in Appendix D.6.2 we find a batch size of 50 is adequate for criticality measurements. This cost is negligible compared to the adversarial training cost of the full training set.

Behavior of Figure 3

This is addressed above in the discussion on hyperparameter selection. To reiterate, Figure 3 demonstrates the importance of selecting the optimal number of critical layers for adversarial accuracy. Clean accuracy trends similarly, as the same number of critical layers is optimal for both adversarial and clean accuracy, and fine-tuning too many layers degrades performance in both cases. These results are shown in Appendix D.4, Figure 7.

We hypothesize that only specific critical layers in a network should be further robustified via adversarial training, and focusing on these layers, rather than the entire model or a larger subset, improves results and mitigates overfitting. This hypothesis is supported by our results and reinforced by the ablation study, which shows that fine-tuning either too many or too few layers negatively impacts performance.