Improving Generalization of Adversarial Training via
Robust Critical Fine-Tuning
Kaijie Zhu1,2 , Jindong Wang3 , Xixu Hu4 , Xing Xie3 , Ge Yang1,2 *School of Artifical Intelligence, University of Chinese Academy of Sciences
2
Institute of Automation, CAS 3 Microsoft Research 4 City University of Hong Kong
1

Abstract

1. Introduction
The pursuit of accurate and trustworthy artificial intelligence systems is a fundamental objective in the deep learning community. Adversarial examples [43, 14], which perturbs input by a small, human imperceptible noise thatcan cause deep neural networks to make incorrect predictions, pose a significant threat to the security of AI sys* Corresponding author: Ge Yang <ge.yang@ia.ac.cn>.

54.00
53.75
Adversarial Robust Accuracy

Deep neural networks are susceptible to adversarial examples, posing a significant security risk in critical applications. Adversarial Training (AT) is a well-establishedtechnique to enhance adversarial robustness, but it often
comes at the cost of decreased generalization ability. This
paper proposes Robustness Critical Fine-Tuning (RiFT), a
novel approach to enhance generalization without compromising adversarial robustness. The core idea of RiFT is
to exploit the redundant capacity for robustness by finetuning the adversarially trained model on its non-robustcritical module. To do so, we introduce module robust
criticality (MRC), a measure that evaluates the significance of a given module to model robustness under worstcase weight perturbations. Using this measure, we identify
the module with the lowest MRC value as the non-robustcritical module and fine-tune its weights to obtain fine-tuned
weights. Subsequently, we linearly interpolate between theadversarially trained weights and fine-tuned weights to derive the optimal fine-tuned model weights. We demonstrate the efficacy of RiFT on ResNet18, ResNet34, and
WideResNet34-10 models trained on CIFAR10, CIFAR100,
and Tiny-ImageNet datasets. Our experiments show that
RiFT can significantly improve both generalization and outof-distribution robustness by around 1.5% while maintaining or even slightly enhancing adversarial robustness. Codeis available at https://github.com/microsoft/robustlearn.

53.50
53.25
53.00
52.75

Non-robust-critical module
Linear layer
All layers
Robust-critical module
Vertical Line

52.50
52.25

81.5

82.0

82.5
83.0
Standard Accuracy

83.5

Figure 1. Interpolation results of fine-tuning on different modules
of ResNet18 on CIFAR10 dataset. Dots denote different interpolation points between the final fine-tuned weights of RiFT and
the initial adversarially trained weights. All fine-tuning methods improve the generalization ability, but only fine-tuning on the
non-robust-critical module (layer2.1.conv2) can preserverobustness. Additionally, fine-tuning on robust-critical module(layer4.1.conv1) causes the worst trade-off between generalization and robustness. In the initial interpolation stage, finetuning on non-robust-critical modules enhances adversarial robustness by around 0.3%.

tems. Notable experimental and theoretical progress has
been made in defending against such adversarial examples[6, 4, 10, 18, 11, 15, 36]. Among various defense methods[49, 32, 54, 30, 8], adversarial training (AT) [28] has beenshown to be one of the most promising approaches [4, 11]
to enhance the adversarial robustness. However, compared
to standard training, AT severely sacrifices generalization
on in-distribution data [40, 44, 55, 35, 31] and is exceptionally vulnerable to certain out-of-distribution (OOD) examples [13, 50, 21] such as Contrast, Bright and Fog, resulting
in unsatisfactory performance.
Prior arts tend to mitigate the trade-off between gener-

alization and adversarial robustness within the adversarial
training procedure. For example, some approaches have
explored reweighting instances [56], using unlabeled data[35], or redefining the robust loss function [55, 46, 48, 31].
In this paper, we take a different perspective to address sucha trade-off by leveraging the redundant capacity for robustness of neural networks after adversarial training. Recent
research has demonstrated that deep neural networks can
exhibit redundant capacity for generalization due to theircomplex and opaque nature, where specific network modules can be deleted, permuted [45], or reset to their initialvalues [52, 9] with only minor degradation in generalization
performance. Hence, it is intuitive to ask: Do adversariallytrained models have such redundant capacity? If so, how toleverage it to improve the generalization and OOD robustness 1 while maintaining adversarial robustness?
Based on such motivation, we introduce a new concept
called Module Robust Criticality (MRC) 2 to investigate theredundant capacity of adversarially trained models for robustness. MRC aims to quantify the maximum incrementof robustness loss of a module’s parameters under the constrained weight perturbation. As illustrated in Figure 3,
we empirically find that certain modules do exhibit redundant characteristics under such perturbations, resulting in
negligible drops in adversarial robustness. We refer to themodules with the lowest MRC value as the non-robustcritical modules. These findings further inspire us to propose a novel fine-tuning technique called Robust CriticalFine-Tuning (RiFT), which aims to leverage the redundantcapacity of the non-robust-critical module to improve generalization while maintaining adversarial robustness. RiFT
consists of three steps: (1) Module robust criticality characterization, which calculates the MRC value for each module and identifies the non-robust-critical module. (2) Nonrobust-critical module fine-tuning, which exploits the redundant capacity of the non-robust-critical module via finetuning its weights with standard examples. (3) Mitigating
robustness-generalization trade-off via interpolation, whichinterpolates between adversarially trained parameters andfine-tuned parameters to find the best weights that maximize the improvement in generalization while preserving
adversarial robustness.
Experimental results demonstrate that RiFT significantly
improves both the generalization performance and OOD
robustness by around 2% while maintaining or even improving the adversarial robustness of the original models.
Furthermore, we also incorporate RiFT to other adversarial training regimes such as TRADES [55], MART [46],AT-AWP [48], and SCORE [31], and show that such incorporation leads to further enhancements. More importantly,
1 Here, generalization refers to generalization to in-distribution (ID)samples, and OOD robustness refers to generalization to OOD samples.
2 In our paper, a module refers to a layer of the neural network.

our experiments reveal several noteworthy insights. First,
we found that fine-tuning on non-robust-critical modules
can effectively mitigate the trade-off between adversarialrobustness and generalization, showing that these two canboth be improved (Section 5.3). As illustrated in Figure 1,
adversarial robustness increases alongside the generalization in the initial interpolation procedure, indicating that thefeatures learned by fine-tuning can benefit both generalization and adversarial robustness. This contradicts the previous claim [44] that the features learned by optimal standard and robust classifiers are fundamentally different. Second, the existence of non-robust-critical modules suggeststhat current adversarial training regimes do not fully utilize the capacity of DNNs (Section 5.2). This motivates
future work to design more efficient adversarial training approaches using such capacity. Third, while previous study[24] reported that fine-tuning on pre-train models could distort the learned robust features and result in poor performance on OOD samples, we find that fine-tuning adversarially trained models do NOT lead to worse OOD performance (Section 5.3).
The contribution of this work is summarized as follows:
1. Novel approach. We propose the concept of module
robust criticality and verify the existence of redundantcapacity for robustness in adversarially trained models.
We then propose RiFT to exploit such redundancy toimprove the generalization of AT models.
2. Superior experimental results. Our approach improves both generalization and OOD robustness of AT
models by around 2%. It can also be incorporated withprevious AT methods to mitigate the trade-off between
generalization and adversarial robustness.
3. Interesting insights. The findings of our experiments
shed light on the intricate interplay between generalization, adversarial robustness, and OOD robustness.
Our work emphasizes the potential of leveraging theredundant capacity in adversarially trained models toimprove generalization and robustness further, which
may inspire more efficient and effective training methods to fully utilize this redundancy.

2. Related Work
Trade-off between adversarial robustness and generalization The existence of such trade-off has been extensively debated in the adversarial learning community[40, 44, 55, 20, 35, 31]. Despite lingering controversies, theprevalent viewpoint is that this trade-off is inherent. Theoretical analyses [44, 35, 20] demonstrated that the trade-off
provably exists even in simple cases, e.g., binary classification and linear regression. To address this trade-off, various methods have been proposed during adversarial train-ing, such as instance reweighting [56], robust self-training[35], incorporating unlabeled data [7, 18], and redefining
the robust loss function [55, 46, 48, 31]. This paper presents
a novel post-processing approach that exploits the excess
capacity of the model after adversarial training to address
such trade-off. Our RiFT can be used in conjunction withexisting adversarial training techniques, providing a practical and effective way to mitigate the trade-off further.

Redundant Fitting Capacity The over-parameterized
deep neural networks (DNNs) exhibit striking fitting power
even for random labels [52, 3]. Recent studies have shownthat not all modules contribute equally to the generalization
ability of DNNs [45, 38, 53, 9], indicating the redundantfitting capacity for generalization. Veit et al. [45] foundthat some blocks can be deleted or permuted without degrading the test performance too much. Rosenfeld andTsotsos [38] demonstrated that one could achieve comparable performance by training only a small fraction of network parameters. Further, recent studies have identified certain neural network modules, referred to as robust modules[53, 9], rewinding their parameters to initial values results
in a negligible decline in generalization. Previous studies
have proposed methods to reduce the computational andstorage costs of deep neural networks by removing the redundant capacity for generalization while preserving comparable performance, such as compression [16] and distillation [19]. In contrast, our work focuses on the redundantcapacity for robustness of adversarially trained models andtries to exlpoit such redundancy.

Fine-tuning Methods Pre-training on large scale datasets
has been shown to be a powerful approach for developing high-performing deep learning models [5, 12, 34, 22].
Fine-tuning is a widely adopted approach to enhance thetransferability of pre-trained models to downstream tasks
and domain shifts. Typically, fine-tuning methods involve
fine-tuning the last layer (linear probing) [1, 24] or all layers (fully fine-tuning) [1, 18, 29, 24]. Salman et al. [39]demonstrated that both fully fine-tuning and linear probing of adversarially trained models can improve the transfer performance on downstream tasks. Nevertheless, recent
studies [2, 47, 24] have suggested that fine-tuning can degrade pre-trained features and underperformance on out-ofdistribution (OOD) samples. To address this issue, different
fine-tuning techniques are proposed such as WiSE-FT [47]and surgical fine-tuning [27] that either leveraged ensemble learning or selective fine-tuning for better OOD performance. Kumar et al. [24] suggested the two-step strategy oflinear probing then full fine-tuning (LP-FT) combines thebenefits of both fully fine-tuning and linear probing.

3. Module Robust Criticality
Improving the generalization of adversarially trainedmodels requires a thorough understanding of DNNs, which,however, proves to be difficult due to the lack of explainability. Luckily, recent studies show that specific modules
in neural networks, referred to as critical modules [53, 9],significantly impact model generalization if their parameters are rewound to initial values. In this work, we propose
a metric called Module Robust Criticality (MRC) to evaluate the robustness contribution of each module explicitly.

3.1. Preliminaries
We denote a l-layered DNN as f (θ) = ϕ(x(l) ; θ (l) ) ◦. . . ◦ ϕ(x(1) ; θ (1) ), where θ (i) is the parameter of i-th
layer and ϕ(·) denotes the activation function. We use
θAT and θF T to denote the weights of the adversariallytrained and fine-tuned model, respectively. We use D ={(x1 , y1 ), ..., (xn , yn )} to denote a dataset and Dstd meansa standard dataset such as CIFAR10. The cross-entropy loss
is denoted by L and ∥·∥p is denoted as the ℓp norm.
Let ∆x ∈ S denote the adversarial perturbation applied
to a clean input x, where S represents the allowed range
of input perturbations. Given a neural network f (θ) and adataset D, adversarial training aims to minimize the robust
loss [28] as:arg min R(f (θ), D), whereθX
R(f (θ), D) =max L(f (θ, x + ∆x), y).
(x,y)∈D(1)∆x∈S

Here, R(f (θ), D) is the robust loss to find the worstcase input perturbation that maximizes the cross-entropy
classification error.

3.2. Module Robust Criticality
Definition 3.1 (Module Robust Criticality). Given a weight
perturbation scaling factor ϵ > 0 and a neural network f (θ),
the robust criticality of a module i is defined as
MRC (f, θ (i) , D, ϵ) = max R(f (θ + ∆θ), D)∆θ∈Cθ− R(f (θ), D),(2)where ∆θ = {0, . . . , 0, ∆θ (i) , 0, . . . , 0} denotes theweight perturbation with respect to the module weights θ (i) ,Cθ = {∆θ ∥∆θ∥p ≤ ϵ∥θ (i) ∥p }, R(·) is the robust loss defined in Eq. (1).
The MRC value for each module represents how they
are critically contributing to model adversarial robustness.
The module with the lowest MRC value is considered redundant, as changing its weights has a negligible effect onrobustness degradation. We refer to this module as thenon-robust-critical module. Intuitively, MRC serves as an
upper bound for weight changing of a particular module,
as demonstrated in Theorem 3.1. Since we do not know
the optimization directions and how they might affect themodel robustness to adversarial examples, we measure theextent to which worst-case weight perturbations affect therobustness, providing an upper bound loss for optimizing
the weight. Further, the MRC for a module depicts the
sharpness of robust loss landscape [48, 41] around the minima θ (i) . If the MRC score is high, it means that the robust
loss landscape with respect to θ (i) is sharp, and fine-tuning
this module is likely to hurt the adversarial robustness.
Theorem 3.1. The MRC value for a module i serves as an
upper bound for the robust loss increase when we optimize
the module under constraint Cθ :R(f (θ ∗ ), D) − R(f (θ), D) ≤ MRC (f, θ (i) , D, ϵ),X
where θ ∗ = arg min
L(f (θ ′ , x), y).
(3)
θ ′ ,(θ ′ −θ)∈Cθ

(x,y)∈D

Proof. By the definition of MRC, for any weights (θ ′ −θ) ∈Cθ , we have:R(f (θ ′ ), D) − R(f (θ), D) ≤ MRC (f, θ (i) , D, ϵ). (4)Thus, for the optimized weights:Xθ ∗ = arg min
L(f (θ ′ , x), y),θ ′ ,(θ ′ −θ)∈Cθ

(5)(x,y)∈D

Algorithm 1 Module Robust Criticality Characterization
Input: neural network f , adversarially trained model
weights θAT , desired module i’s weights θ (i) , standard
dataset Dstd , weight perturbation scaling factor ϵ, optimization iteration steps T , learning rate γ.
Output: The module robust criticality of module i.
1: Initialize adversarial dataset: Dadv = {}2: for Batch Bk ∈ Dstd do ▷ Generate adversarial dataset
3:
Bkadv = PGD-10(θS AT , Bk )4:
Dadv = Dadv Bkadv
5: end for
6: Freeze all parameters of θAT except for θ (i)7: θ1 = θAT
8: for t = 1, . . . , T do
▷ Iterate T epochs
9:θt+1 = θt
10:
for Batch Bkadv ∈ Dadv do
11:
Calculate Loss: L(f, θt , Bkadv ))12:
θt+1 = θt+1 + γ∇θt (L)▷ Gradient Ascent13:
end for(i)(i)
14:∆θ (i) = θt+1 − θAT ▷ Check perturb constraint(i)15:
if ∥∆θ (i) ∥2 ≥ ϵ∥θAT ∥2 then
16:

∥θ

(i)

∥

∆θ (i) = ϵ ∥∆θAT(i) ∥22 ∆θ (i)θt+1 = θt + ∆θ (i)
break
19:
end if
20: end for
21: MRC (θ (i) ) = L(f, θT , Dadv ) − L(f, θAT , Dadv )22: Return MRC (θ (i) )17:18:

it satisfies
R(f (θ ∗ ), D) − R(f (θ), D) ≤ MRC (f, θ (i) , D, ϵ). (6)Such that the proof ends.
Remark: The definition of MRC is similar in spirit to thework of Zhang et al. [53] and Chatterji et al. [9]. However, MRC differs fundamentally from them in two aspects.
First, MRC aims to capture the influence of a module on
adversarial robustness, while Zhang et al. [53] and Chatterji et al. [9] focus on studying the impact of a module on
generalization. Second, MRC investigates the robustness
characteristics of module weights under worst-case weight
perturbations, whereas Zhang et al. [53] and Chatterji etal. [9] analyzed the properties of a module by rewinding its
weights to their initial values. Similar to [25, 41], we define the weight perturbation constraint Cθ as a multiple ofthe ℓp norm of original parameters, which ensures the scaleinvariant property and allows us to compare the robust criticality of modules across different layers, see Appendix A
for a detailed proof.
Theorem 3.1 establishes a clear upper bound for finetuning particular modules. This theorem assures us that

fine-tuning on non-robust-critical modules shouldn’t harm
the model robustness. However, it does not ascertain if finetuning the robust-critical module will lead to a significant
decline in robust accuracy.

3.3. Relaxation of MRC
Optimizing in Eq. (2) requires simultaneously finding
worst-case weight perturbation ∆θ and worst-case input
perturbation ∆x, which is time-consuming. Thus, we propose a relaxation version by fixing ∆x at the initial optimizing phase. Concretely, we first calculate the adversarial
examples ∆x with respect to θAT . By fixing the adversarial
examples unchanged during the optimization, we iteratively
optimize the ∆θ by gradient ascent method to maximize
the robust loss to find the optimal ∆θ. We set a weight
perturbation constraint and check it after each optimization
step. If the constraint is violated, we project the perturbation onto the constraint set. The pseudo-code is describedin Algorithm 1. In our experiments, if not specified, we
set ∥·∥p = ∥·∥2 and ϵ = 0.1 for Cθ , the iterative step for
optimizing ∆θ is 10.

AT Models: 𝜽𝑨𝑻
Non-robust-critical module
Robust-critical module

𝟏−𝜶

AT Models: 𝜽𝑨𝑻+

Best FT Models: 𝜽∗𝑭𝑻FT Models: 𝜽𝑭𝑻𝜶∗:𝜶

Step 1 :Module robust criticality
characterization.

Step 2 :Fine-tuning on
non-robust critical module.

Finding the best interpolation
factor 𝜶∗ which maximize theincrease of generalization
while preserving robustness.

Step 3 :Mitigating robustnessgeneralization trade-off
via interpolation.

Figure 2. The pipeline of our proposed Robust Critical Fine-Tuning (RiFT).

4. RiFT: Robust Critical Fine-tuning
In this paper, we propose RiFT, a robust critical finetuning approach that leverages MRC to guide the finetuning of a deep neural network to improve both generalization and robustness. Let Padv (x, y) and Pstd (x, y) denotethe distributions of adversarial and standard inputs, respectively. Then, applying an adversarially trained model on
Padv (x, y) to Pstd (x, y) can be viewed as a distributional
shift problem. Thus, it is natural for RiFT to exploit the redundant capacity to fine-tune adversarially trained models
on the standard dataset.
Specifically, RiFT consists of three steps as shown in
Figure 2. First, we calculate the MRC of each module andchoose the module with the lowest MRC score as our nonrobust-critical module. Second, we freeze the parameters of
the adversarially trained model except for our chosen nonrobust-critical module. Then we fine-tune the adversarially trained models on corresponding standard dataset Dstd .
Third, we linearly interpolate the weights of the original adversarially trained model and fine-tuned model to identify
the optimal interpolation point that maximizes generalization improvement while maintaining robustness.
Step 1: Module robust criticality characterization According to the Algorithm 1, we iteratively calculate the
MRC value for each module θ (i) ∈ θAT , then we choose
the module with the lowest MRC value, denoted as θ̃:θ̃ = θ (i) where i = arg min MRC (f, θ (i) , D, ϵ).

(7)iStep 2: Fine-tuning on non-robust-critical modules
Next, we freeze the rest of the parameters and fine-tune on
desired parameters θ̃. We solve the following optimization

problem by SGD with momentum [42]X
L(f (x, (θ̃; θ \ θ̃)), y) + λ∥θ̃∥2 ,arg minθ̃(8)(x,y)∈D

where λ is the ℓ2 weight decay factor.
Step 3: Mitigating robustness-generalization trade-off
via interpolation For a interpolation coefficient α, the interpolated weights is calculated as:θα = (1 − α)θAT + αθF T ,(9)where θAT is the initial adversarially trained weights andθF T is the fine-tuned weights obtained by Eq. (8). Since
our goal is to improve the generalization while preserving
adversarial robustness, thus the best interpolation point ischosen to be the point that most significantly improves thegeneralization while the corresponding adversarial robustness is no less than the original robustness by 0.1.
Remark: Theorem 3.1 establishes an upper bound on
the possible drop in robustness loss that can be achieved
through fine-tuning. It is expected that the second step of
optimization would enforce the parameters to lie within theboundary Cθ in order to satisfy the theorem. However, here
we do not employ constrained optimization but find the optimal point by first optimizing without constraints and theninterpolating. This is because (1) the constraints are empirically given and may not always provide the optimal range
for preserving robustness, and it is possible to fine-tune outside the constraint range and still ensure that there is notmuch loss of robustness. (2) the interpolation procedure
serves as a weight-ensemble, which may benefit both robustness and generalization, as noted in WiSE-FT [47]. The
complete algorithm of RiFT is shown in Appendix B.

5. Experiments

5.2. Empirical Analysis of MRC

5.1. Experimental Setup

Before delving into the main results of RiFT, we first
empirically analyze our proposed MRC metric in Definition 3.1, which serves as the foundation of our RiFT approach. We present the MRC analysis on ResNet18 [17] on
CIFAR-10 in Figure 3, where each column corresponds to
the MRC value and its corresponding robust accuracy dropof a specific module.
Our analysis shows that the impact of worst-case weight
perturbations on model robustness varies across different
modules. Some modules exhibit minimal impact on robustness under perturbation, indicating the presence of redundant capacity for robustness. Conversely, for other modules, the worst-case weight perturbation shows a significant impact, resulting in a substantial decline in robustness.
For example, in module layer2.1.conv2, worst-caseweight perturbations only result in a meager addition of
0.09 robust loss. However, for layer4.1.conv1, theworst-case weight perturbations affect the model’s robust
loss by an additional 12.94, resulting in a substantial decline (53.03%) in robustness accuracy. Such robust-criticaland non-robust-critical modules are verified to exist in various network architectures and datasets, as detailed in Appendix C.4. We also observe that as the network capacity
decreases (e.g., from WRN34-10 to ResNet18) and the task
becomes more challenging (e.g., from CIFAR10 to TinyImageNet), the proportion of non-robust-critical modules
increases, as less complex tasks require less capacity, leading to more non-robust-critical modules.
It is worthy noting that the decrease in robust accuracy
doesn’t directly correlate with MRC. For instance, bothlayer4.0.conv2 and layer4.1.conv1 have a robust accuracy drop of 53.05%, yet their MRC values differ.
This discrepancy can be attributed to the different probability distributions of misclassified samples across modules,
resulting in same accuracy declines but different losses.

Datasets We adopt three popular image classification
datasets: CIFAR10 [23], CIFAR100 [23], and TinyImageNet [26]. CIFAR10 and CIFAR100 comprise 60,000
32 × 32 color images in 10 and 100 classes, respectively.
Tiny-ImageNet is a subset of ImageNet and contains 200
classes, where each class contains 500 colorful images withsize 64 × 64. We use three OOD datasets accordingly toevaluate the OOD robustness: CIFAR10-C, CIFAR100-C,
and Tiny-ImageNet-C [18]. These datasets simulate 15
types of common visual corruptions and are grouped into
four classes: Noise, Blur, Weather, and Digital.
Evaluation metrics We use the test set accuracy of eachstandard dataset to represent the generalization ability. For
evaluating adversarial robustness, we adopt a common setting of PGD-10 [28] with constraint ℓ∞ = 8/255. We run
PGD-10 with three times and select the worst robust accuracy as the final metric. The OOD robustness is evaluatedby the accuracy of the test set of the corrupted dataset corresponding to the standard dataset.
Training details We use ResNet18 [17], ResNet34[17], WideResNet34-10 (WRN34-10) [51] as backbones.
ResNet18 and ResNet34 are 18-layer and 34-layer ResNetmodels, respectively. WideResNet34-10 is a 34-layer
WideResNet model with a widening factor of 10. Similarly,
we adopt PGD-10 [28] with constraint ℓ∞ = 8/255 for adversarial training. Following standard settings [37, 33], wetrain models with adversarial examples for 110 epochs. The
learning rate starts from 0.1 and decays by a factor of 0.1 at
epochs 100 and 105. We select the weights with the highest
test robust accuracy as our adversarially trained models.
We fine-tune the adversarially trained models θAT using SGD with momentum [42] for 10 epochs. The initial
learning rate is set to 0.001.3 We decay the learning rate by
1/10 after fine-tuning for 5 epochs We choose the weights
with the highest test accuracy as fine-tuned model weights,
denoted as θF T . We then interpolate between initial adversarially trained model weights θAT and θF T , the best interpolation point selected by Step 3 in Section 4 is denoted
as θF∗ T . We then compare the generalization, adversarialrobustness, and OOD robustness of θF∗ T and θAT .
We report the average of three different seeds and omit
the standard deviations of 3 runs as they are tiny (< 0.20%),
which hardly effect the results. Refer to Appendix C for
more training details.
3 The best learning rate for fine-tuning vary across architectures anddatasets and is required to be carefully modified.

5.3. Main Results
Table 1 summarizes the main results of our study, fromwhich we have the following findings.
RiFT improves generalization First, RiFT effectively
mitigates the trade-off between generalization and robustness raised by adversarial training. Across different datasets
and network architectures, RiFT improves the generalization of adversarially trained models by approximately 2%.
This result prompts us to rethink the trade-off, as it may becaused by inefficient adversarial training algorithm rather
than the inherent limitation of DNNs. Furthermore, as
demonstrated in Figure 1, both adversarial robustness andgeneralization increase simultaneously in the initial interpolation process, indicating that these two characteristics canRobust Acc Drop -3.16 -18.89 -23.04 -52.01 -9.47 -35.79 -34.49 -51.83 -2.86 -50.32 -49.75 -53.03 -7.22 -52.58 -53.03 -53.03 -7.09 -6.47

lay

lay

er1

.1.

c

co

nv
1
.0.
lay coner1 v1.0
lay .coner1 v2.1
lay .coner1 v1.1
lay .coner2 v2.0lay .coner2 v1.0
lay .coner2 v2lay .1.coer2 nv1.1.
lay conv
er3
2.0
lay .coner3 v1.0
lay .coner3 v2.1
lay .coner3 v1.1
lay .coner4 v2.0lay .coner4 v1lay .0.co
er4 nv2o
er4 nv1
.1.
co
nv
2
lin
ea
r

Module Robust Criticality 0.12 0.57 0.93 7.76 0.27 1.82 1.69 3.94 0.09 5.67 5.78 9.80 0.25 6.83 11.19 12.94 0.23 0.19

Figure 3. Example of module robust criticality (MRC) and its corresponding robust accuracy drop of ResNet18 trained on CIFAR10. Eachcolumn represents an individual module. The first row represents the corresponding robust accuracy drop and the second row represents
the MRC value of each module. The higher the MRC value is, the more robust-critical the module is. Some modules are not criticalto robustness, exhibiting redundant characteristics for contributing to robustness. However, some modules are critical to robustness. For
example, the robust acc drop is only 2.86% for layer2.1.conv2 while for layer4.1.conv1 the robust acc drop is up to 53.03%.
Table 1. Results of RiFT on different datasets and backbones. Std means the standard test accuracy for in distribution generalization, OOD
denotes the OOD robust accuracy of corresponding corruption dataset (e.g., CIFAR10-C). Adv denotes the adversarial robust accuracy. In
each column, we bold the entry with the higher accuracy. RiFT improves both generalization and OOD robustness across architectures anddatasets while maintaining adversarial robustness.

Architecture

CIFAR10

MethodCIFAR100Tiny-ImageNet

Std

OOD

Adv

Std

OOD

Adv

Std

OOD

Adv

ResNet18

AT
AT+RiFT
∆

81.46
83.44
+1.98

73.56
75.69
+2.13

53.63
53.65
+0.02

57.10
58.74
+1.64

46.43
48.06
+1.63

30.15
30.17
+0.02

49.10
50.61
+1.51

27.68
28.73
+1.05

23.28
23.34
+0.06

ResNet34

AT
AT+RiFT
∆

84.23
85.41
+1.18

75.37
77.15
+1.78

55.31
55.34
+0.03

58.67
60.88
+2.21

48.24
49.97
+1.73

30.50
30.58
+0.08

50.96
52.54
+1.58

27.91
30.07
+2.16

24.27
24.37
+0.10

WRN34-10

AT
AT+RiFT
∆

87.41
87.89
+0.48

78.75
79.31
+0.56

55.40
55.41
+0.01

62.35
64.56
+2.21

50.61
52.69
+2.08

31.66
31.64
-0.02

52.78
55.31
+2.53

31.81
33.86
+2.05

26.07
26.17
+0.10

Avg

∆

+1.21

+1.49

+0.02

+2.02

+1.81

+0.02

+1.87

+1.75

+0.08

be improved together. This trend is observed across different datasets and network architectures; see Appendix C.5
for more illustrations. This finding challenges the notion that the features of optimal standard and optimal robust classifiers are fundamentally different, as previously
claimed by Tsipras et al. [44], as fine-tuning procedures canincrease both robustness and generalization.
Fine-tuning improves OOD robustness Second, ourstudy also investigated the out-of-distribution (OOD) robustness of the fine-tuned models and observed an improvement of approximately 2%. This observation is noteworthy
because recent work [2, 24, 47] showed that fine-tuning pretrained models can distort learned features and result in underperformance in OOD samples. Furthermore, Yi et al.
[50] demonstrated that adversarial training enhances OOD
robustness, but it is unclear whether fine-tuning on adversarially trained models distorts robust features. Our results
indicate that fine-tuning adversarially trained models does
not distort the robust features learned by adversarial train-

ing and instead helps improve OOD robustness. We suggest
fine-tuning adversarially trained models may be a promisingavenue for further improving OOD robustness.

5.4. Incorporate RiFT to Other AT Methods
To further validate the effectiveness of RiFT, we conduct experiments on ResNet18 [17] trained on CIFAR10
and CIFAR100 [23] using four different adversarial training techniques: TRADES [55], MART [46], AWP [48], andSCORE [31], and then apply our RiFT to the resulting models. As shown in Table 2, our approach is compatible withvarious adversarial training methods and improves generalization and OOD robustness.

5.5. Ablation Study
Fine-tuning on different modules To evaluate the efficacy of fine-tuning the non-robust-critical module, we conducted further experiments by fine-tuning the adversariallytrained model on different modules. Specifically, we usedfour fine-tuning methods: fully fine-tuning, linear probing

Table 3. Results of fine-tuning on different modules.

Table 2. Results of RiFT + other AT methods.

CIFAR10

MethodCIFAR100

Std

OOD

Adv

Std

OOD

Adv

TRADES
TRADES+RiFT
∆

81.54
81.87
+0.33

73.42
74.09
+0.67

53.31
53.30
-0.01

57.44
57.78
+0.34

47.23
47.52
+0.29

30.20
30.22
+0.02

MART
MART+RiFT
∆

76.77
77.14
+0.37

68.62
69.41
+0.79

56.90
56.92
+0.02

51.46
52.42
+0.96

42.07
43.35
+1.28

31.47
31.48
+0.01

AWP
AWP+RiFT
∆

78.40
78.79
+ 0.39

70.48
71.12
+0.64

53.83
53.84
+0.01

52.85
54.89
+2.04

43.10
45.08
+1.98

31.00
31.05
+0.05

SCORE
SCORE+RiFT
∆

84.20
85.65
+1.45

75.82
77.37
+1.55

54.59
54.62
+0.03

54.83
57.63
+2.80

45.39
47.77
+2.38

29.49
29.50
+0.01

(fine-tuning on the last layer), fine-tuning on the non-robustcritical module, and fine-tuning on the robust-critical module. The experiment was conducted using ResNet18 on
CIFAR-10, and the results are presented in Figure 1 andTable 3. As described in Section 3.2, MRC is an upperbound for weight perturbation, indicating the criticality ofa module in terms of model robustness. Fine-tuning on
a non-robust-critical module can help preserve adversarialrobustness but does not guarantee improvement in generalization. Similarly, fine-tuning on the robust-critical module does not necessarily hurt robustness. However, our experiments observed that all fine-tuning methods improved
generalization ability, but only fine-tuning on non-robustcritical module preserved adversarial robustness. Moreover,fine-tuning on the robust-critical module exhibited the worst
trade-off between generalization and robustness comparedto fine-tuning on all layers.

More non-robust-critical modules, more useful? To investigate whether fine-tuning on more non-critical modules
could further improve generalization, we additionally finetune on the top two, top three, and top five non-robustcritical modules. However, Table 3 reveals that generalization and OOD robustness did not surpass the results
achieved by fine-tuning a singular non-robust-critical module. Notably, performance deteriorated when fine-tuning
multiple non-critical modules compared to fine-tuning alllayers. It’s pivotal to note that this doesn’t negate MRC’s
applicability to several modules. The MRC for module iis evaluated with other module parameters held constant,making it challenging to discern the impact of worst-caseperturbations across multiple modules using the MRC of asingle one. We posit that broadening MRC’s definition toencompass multiple modules might address this problem.

Method

Std

OOD

Adv

All layers
Last layer
Robust-critical
Non-robust-critical

83.56
83.35
83.36
83.44

75.48
75.16
75.42
75.69

52.66
52.75
52.48
53.65

Table 4. Results of fine-tuning on multiple non-robustcritical modules.

Method

Std

OOD

Adv

Top 1
Top 2
Top 3
Top 5

83.44
83.41
83.59
83.70

75.69
75.61
75.77
75.82

53.65
52.47
52.22
52.35

Ablation on interpolation factor α∗ The value of α∗ isclosely related to the fine-tuning learning rate. Specifically,a large learning rate can result in substantial weight updates
that may push the fine-tuned weights θF T away from theiradversarially trained counterparts θAT . Our empirical results indicate that a fine-tuning learning rate of 0.001 is
suitable for most cases and that the corresponding α∗ value
generally ranges between 0.6 to 0.9.

Factors related to the generalization gain of RiFT ”Our
results unveiled patterns and behaviors that offer insights
into the determinants of the generalization gains observedwith RiFT. First, the generalization gain of RiFT is a function of both the neural network’s inherent capacity and theinherent difficulty posed by the classification task. Specifically, as the classification task becomes more challenging,
the robust criticality of each module increases, which in turndecreases the generalization gain of RiFT. This effect canbe mitigated by using a model with a larger capacity. For
instance, we observe that the generalization gain of RiFT
increases as we switch from ResNet18 to ResNet34 andto WRN34-10 when evaluating on CIFAR100 and TinyImageNet. Further, We observed that the generalization
gain of RiFT with WRN34-10 on CIFAR10 is notably
lower, at approximately 0.5%, compared to 2% on otherdatasets. This might be attributed to the minimal generalization disparity between adversarially trained models
and their standard-trained counterparts; specifically, while
WRN34-10’s standard test accuracy stands at around 95%,
its adversarial counterpart registers at 87%. It is evident that
fine-tuning on a single module may not yield significant improvements. Investigating these patterns further could offer
strategies for enhancing the robustness and generalization
capabilities of deep neural networks.

6. Conclusion
In this paper, we aim to exploit the redundant capacityof the adversarially trained models. Our proposed RiFT
leverages the concept of module robust criticality (MRC)
to guide the fine-tuning process, which leads to improved
generalization and OOD robustness. The extensive experiments demonstrate the effectiveness of RiFT across various
network architectures and datasets. Our findings shed lighton the intricate relationship between generalization, adversarial robustness, and OOD robustness. RiFT is a primary
exploration of fine-tuning the adversarially trained models.
We believe that fine-tuning holds great promise, and we callfor more theoretical and empirical analyses to advance ourunderstanding of this important technique.

References[1] Armen Aghajanyan, Akshat Shrivastava, Anchit Gupta, Naman Goyal, Luke Zettlemoyer, and Sonal Gupta. Betterfine-tuning by reducing representational collapse. In International Conference on Learning Representations, 2021. 3[2] Anders Andreassen, Yasaman Bahri, Behnam Neyshabur,and Rebecca Roelofs. The evolution of out-of-distribution
robustness throughout fine-tuning.
arXiv preprint
arXiv:2106.15831, 2021. 3, 7
[3] Devansh Arpit, Stanislaw Jastrzkebski, Nicolas Ballas,
David Krueger, Emmanuel Bengio, Maxinder S Kanwal,
Tegan Maharaj, Asja Fischer, Aaron Courville, Yoshua Bengio, et al. A closer look at memorization in deep networks. In
International Conference on Machine Learning, pages 233–
242. PMLR, 2017. 3[4] Anish Athalye, Nicholas Carlini, and David Wagner. Obfuscated gradients give a false sense of security: Circumventing
defenses to adversarial examples. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International
Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 274–283. PMLR,
10–15 Jul 2018. 1[5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33:1877–1901, 2020. 3[6] N. Carlini and D. Wagner. Towards evaluating the robustnessof neural networks. In 2017 IEEE Symposium on Security
and Privacy, pages 39–57, Los Alamitos, CA, USA, may
2017. IEEE Computer Society. 1[7] Yair Carmon, Aditi Raghunathan, Ludwig Schmidt, John C
Duchi, and Percy S Liang. Unlabeled data improves adversarial robustness. In H. Wallach, H. Larochelle, A.
Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems,
volume 32. Curran Associates, Inc., 2019. 3[8] Alvin Chan, Yi Tay, Yew Soon Ong, and Jie Fu. Jacobian
adversarially regularized networks for robustness. In International Conference on Learning Representations, 2020. 1[9] Niladri Chatterji, Behnam Neyshabur, and Hanie Sedghi.
The intriguing role of module criticality in the generalizationof deep networks. In International Conference on Learning
Representations, 2020. 2, 3, 4[10] Jeremy Cohen, Elan Rosenfeld, and Zico Kolter. Certified
adversarial robustness via randomized smoothing. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine
Learning, volume 97 of Proceedings of Machine Learning
Research, pages 1310–1320. PMLR, 09–15 Jun 2019. 1[11] Francesco Croce and Matthias Hein. Reliable evaluation of adversarial robustness with an ensemble of diverse
parameter-free attacks. In International Conference on Machine Learning, pages 2206–2216. PMLR, 2020. 1[12] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov,
Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner,
Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image isworth 16x16 words: Transformers for image recognition atscale. In International Conference on Learning Representations, 2021. 3[13] Justin Gilmer, Nicolas Ford, Nicholas Carlini, and Ekin
Cubuk. Adversarial examples are a natural consequence oftest error in noise. In International Conference on Machine
Learning, pages 2280–2289. PMLR, 2019. 1[14] Ian Goodfellow, Jonathon Shlens, and Christian Szegedy.
Explaining and harnessing adversarial examples. In International Conference on Learning Representations, 2015. 1[15] Sven Gowal, Sylvestre-Alvise Rebuffi, Olivia Wiles, FlorianStimberg, Dan Andrei Calian, and Timothy A Mann. Improving robustness using generated data. In M. Ranzato,
A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman
Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 4218–4233. Curran Associates, Inc., 2021. 1[16] Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning,trained quantization and huffman coding. arXiv preprint
arXiv:1510.00149, 2015. 3[17] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.
Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern
Recognition, pages 770–778, 2016. 6, 7
[18] Dan Hendrycks, Kimin Lee, and Mantas Mazeika. Using
pre-training can improve model robustness and uncertainty.
In International Conference on Machine Learning, pages
2712–2721. PMLR, 2019. 1, 3, 6
[19] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint
arXiv:1503.02531, 2015. 3[20] Adel Javanmard, Mahdi Soltanolkotabi, and Hamed Hassani.
Precise tradeoffs in adversarial training for linear regression.
In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume
125 of Proceedings of Machine Learning Research, pages
2034–2078. PMLR, 09–12 Jul 2020. 2[21] Klim Kireev, Maksym Andriushchenko, and Nicolas Flammarion. On the effectiveness of adversarial training against

common corruptions. In Uncertainty in Artificial Intelligence, pages 1012–1021. PMLR, 2022. 1[22] Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan
Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby.
Big transfer (bit): General visual representation learning. In
Computer Vision–ECCV 2020: 16th European Conference,
Glasgow, UK, August 23–28, 2020, Proceedings, Part V 16,
pages 491–507. Springer, 2020. 3[23] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple
layers of features from tiny images. 2009. 6, 7
[24] Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones,
Tengyu Ma, and Percy Liang. Fine-tuning can distort pretrained features and underperform out-of-distribution. In International Conference on Learning Representations, 2022.
2, 3, 7[25] Jungmin Kwon, Jeongseop Kim, Hyunseo Park, andIn Kwon Choi. Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In
Marina Meila and Tong Zhang, editors, Proceedings of the
38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research,
pages 5905–5914. PMLR, 18–24 Jul 2021. 4[26] Ya Le and Xuan Yang. Tiny imagenet visual recognition
challenge. CS 231N, 7(7):3, 2015. 6[27] Yoonho Lee, Annie S Chen, Fahim Tajwar, Ananya Kumar,
Huaxiu Yao, Percy Liang, and Chelsea Finn. Surgical finetuning improves adaptation to distribution shifts. In The
Eleventh International Conference on Learning Representations, 2023. 3[28] Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt,
Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. In International
Conference on Learning Representations, 2018. 1, 3, 6
[29] John P Miller, Rohan Taori, Aditi Raghunathan, Shiori
Sagawa, Pang Wei Koh, Vaishaal Shankar, Percy Liang,
Yair Carmon, and Ludwig Schmidt. Accuracy on the line:on the strong correlation between out-of-distribution andin-distribution generalization. In Marina Meila and Tong
Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedingsof Machine Learning Research, pages 7721–7735. PMLR,
18–24 Jul 2021. 3[30] Aamir Mustafa, Salman Khan, Munawar Hayat, Roland
Goecke, Jianbing Shen, and Ling Shao. Adversarial defense
by restricting the hidden space of deep neural networks. In
The IEEE International Conference on Computer Vision, October 2019. 1[31] Tianyu Pang, Min Lin, Xiao Yang, Jun Zhu, and Shuicheng
Yan.
Robustness and accuracy could be reconcilable
by (Proper) definition. In Kamalika Chaudhuri, Stefanie
Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan
Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedingsof Machine Learning Research, pages 17258–17277. PMLR,
17–23 Jul 2022. 1, 2, 3, 7[32] Tianyu Pang, Kun Xu, Yinpeng Dong, Chao Du, Ning Chen,
and Jun Zhu. Rethinking softmax cross-entropy loss for ad-

versarial robustness. In International Conference on Learning Representations, 2020. 1[33] Tianyu Pang, Xiao Yang, Yinpeng Dong, Hang Su, and Jun
Zhu. Bag of tricks for adversarial training. In International
Conference on Learning Representations, 2021. 6[34] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario
Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. 3[35] Aditi Raghunathan, Sang Michael Xie, Fanny Yang, John
Duchi, and Percy Liang. Understanding and mitigating thetradeoff between robustness and accuracy. In Hal Daumé
III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 7909–7919.
PMLR, 13–18 Jul 2020. 1, 2, 3[36] Sylvestre-Alvise Rebuffi, Sven Gowal, Dan A Calian, Florian Stimberg, Olivia Wiles, and Timothy Mann. Fixing
data augmentation to improve adversarial robustness. arXiv
preprint arXiv:2103.01946, 2021. 1[37] Leslie Rice, Eric Wong, and Zico Kolter. Overfitting in adversarially robust deep learning. In Hal Daumé III and Aarti
Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedingsof Machine Learning Research, pages 8093–8104. PMLR,
13–18 Jul 2020. 6[38] Amir Rosenfeld and John K Tsotsos. Intriguing properties ofrandomly weighted networks: Generalizing while learning
next to nothing. In 2019 16th Conference on Computer andRobot Vision (CRV), pages 9–16. IEEE, 2019. 3[39] Hadi Salman, Andrew Ilyas, Logan Engstrom, Ashish
Kapoor, and Aleksander Madry. Do adversarially robust imagenet models transfer better? Advances in Neural Information Processing Systems, 33:3533–3545, 2020. 3[40] Ludwig Schmidt, Shibani Santurkar, Dimitris Tsipras, Kunal
Talwar, and Aleksander Madry. Adversarially robust generalization requires more data. In S. Bengio, H. Wallach, H.
Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett,
editors, Advances in Neural Information Processing Systems,
volume 31. Curran Associates, Inc., 2018. 1, 2[41] David Stutz, Matthias Hein, and Bernt Schiele. Relating adversarially robust generalization to flat minima. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 7807–7817, 2021. 4[42] Ilya Sutskever, James Martens, George Dahl, and Geoffrey
Hinton. On the importance of initialization and momentum
in deep learning. In International Conference on Machine
Learning, pages 1139–1147. PMLR, 2013. 5, 6[43] Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan
Bruna, Dumitru Erhan, Ian Goodfellow, and Rob Fergus.
Intriguing properties of neural networks. arXiv preprint
arXiv:1312.6199, 2013. 1[44] Dimitris Tsipras, Shibani Santurkar, Logan Engstrom,
Alexander Turner, and Aleksander Madry. Robustness may
be at odds with accuracy. In International Conference on
Learning Representations, 2019. 1, 2, 7[45] Andreas Veit, Michael J Wilber, and Serge Belongie. Residual networks behave like ensembles of relatively shallow net-

works. Advances in Neural Information Processing Systems,
29, 2016. 2, 3
[46] Yisen Wang, Difan Zou, Jinfeng Yi, James Bailey, Xingjun
Ma, and Quanquan Gu. Improving adversarial robustness
requires revisiting misclassified examples. In International
Conference on Learning Representations, 2020. 2, 3, 7[47] Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim,
Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok
Namkoong, and Ludwig Schmidt. Robust fine-tuning ofzero-shot models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages
7959–7971, June 2022. 3, 5, 7
[48] Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust generalization. In H.
Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin,
editors, Advances in Neural Information Processing Systems,
volume 33, pages 2958–2969. Curran Associates, Inc., 2020.
2, 3, 4, 7
[49] Cihang Xie, Jianyu Wang, Zhishuai Zhang, Zhou Ren, andAlan Yuille. Mitigating adversarial effects through randomization. In International Conference on Learning Representations, 2018. 1[50] Mingyang Yi, Lu Hou, Jiacheng Sun, Lifeng Shang, Xin
Jiang, Qun Liu, and Zhiming Ma. Improved ood generalization via adversarial training and pretraing. In Marina
Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 ofProceedings of Machine Learning Research, pages 11987–
11997. PMLR, 18–24 Jul 2021. 1, 7
[51] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Edwin R. Hancock Richard C. Wilson and William
A. P. Smith, editors, Proceedings of the British Machine Vision Conference, pages 87.1–87.12. BMVA Press, September 2016. 6[52] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin
Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017. 2, 3
[53] Chiyuan Zhang, Samy Bengio, and Yoram Singer. Are alllayers created equal? Journal of Machine Learning Research, 23(67):1–28, 2022. 3, 4[54] Haichao Zhang and Jianyu Wang. Defense against adversarial attacks using feature scattering-based adversarial training.
In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'AlchéBuc, E. Fox, and R. Garnett, editors, Advances in Neural
Information Processing Systems, volume 32. Curran Associates, Inc., 2019. 1[55] Hongyang Zhang, Yaodong Yu, Jiantao Jiao, Eric Xing, Laurent El Ghaoui, and Michael Jordan. Theoretically principled trade-off between robustness and accuracy. In Kamalika
Chaudhuri and Ruslan Salakhutdinov, editors, Proceedingsof the 36th International Conference on Machine Learning,volume 97 of Proceedings of Machine Learning Research,
pages 7472–7482. PMLR, 09–15 Jun 2019. 1, 2, 3, 7[56] Jingfeng Zhang, Jianing Zhu, Gang Niu, Bo Han, Masashi
Sugiyama, and Mohan Kankanhalli.
Geometry-aware

instance-reweighted adversarial training. In International
Conference on Learning Representations, 2021. 2, 3

A. Proof of the scale-invariant property
Without loss of generality, assume a two layers neuralnetwork f and ϕ is a ReLU-based activation function.
f (θf , x) = θ (2) ϕ(θ (1) x).

(10)

The corresponding scaled neural network g is:
1 (2)θ ϕ(βθ (1) x),βg(θg , x) =Theorem A.4. The Module Robustness Criticality (MRC)
proposed in Definition 3.1 is invariant to the scaling of theparameters.
(1)∀(z, β) ∈ R × R , ϕ(βz) = βϕ(z).

≤ max R(g(θg + ∆θg ), D).

(26)∆θf∗ = arg max R(f (θf + ∆θf ), D),(27)∆θg∗ = arg max R(g(θg + ∆θg ), D).

(28)∆θg ∈CθgLet∆θf ∈Cθf∆θg ∈Cθg(12)

Proof.
ϕ(βz) = max(βz, 0) = β max(z, 0) = βϕ(z).

(25)∆θf ∈Cθf

Theorem A.1. The rectified function ϕ(x) = max(x, 0) is
a homogeneous function where+

max R(f (θf + ∆θf ), D)(11)

where the non-negative β is the scaling factor.
Suppose we calculate the MRC value of the first moduleθ (1) and β1 θ (1) .

(13)

(1)

Proof. Let ∆θf = {∆θf , 0}, ∆θg = {∆θg , 0} be theperturbation of the first layer for network f and g respectively. First, we prove

Consider the perturbation ∆θ˜g = β∆θf∗ for g, it is easy
to show that ∆θ˜g ∈ Cθg ,(1)Cθf = {∆θf ∥∆θf ∥p ≤ ϵ∥θf ∥p },(29)∥∆θg ∥p ≤ ϵ∥θg(1) ∥p }(30)

Cθg = {∆θg

Theorem A.2. ∀x, f (θf , x) ≡ g(θg , x).

(1)= {∆θg ∆θg = β∥∆θf ∥p ≤ ϵβ∥θf ∥p }.

Proof.
1 (2)θ ϕ(βθ (1) x)β
1≡ βθ (2) ϕ(θ (1) x)βg(θg , x) =

≡θ

(2)

ϕ(θ

(1)

x)

≡ f (θf , x)

(14)

Therefore,R(g(θg + ∆θ˜g ), D) ≤

(15)

(17)

g(θg + ∆θ˜g )1= θ (2) ϕ((βθ (1) + β∆θf∗ )x)
β

(18)

(20)

≡ max ℓ(g(θg , x + ∆x), y).

(21)

R(f (θf ), D) =(x,y)∈D≡X(x,y)∈D=θ (2) ϕ((θ (1) + ∆θf∗ )x)(36)≡f (θf + ∆θf∗ ).

(37)

(38)

≤ R(g(θg + ∆θg∗ ), D).

(39)

Similarly, we can prove
max R(g(θg + ∆θg ), D)(40)≤ max R(f (θf + ∆θf ), D).

(41)

max R(f (θf + ∆θf ), D)(42)= max R(g(θg + ∆θg ), D).

(43)∆θg ∈Cθg

max ℓ(f (θf , x + ∆x), y) (22)∆x∈S

max ℓ(g(θg , x + ∆x), y) (23)∆x∈S

= R(g(θg ), D)(35)

R(f (θf + ∆θf∗ ), D) ≡ R(g(θg + ∆θ˜g ), D)

Thus,
X

(34)

According to Theorem A.3,
(19)

max ℓ(f (θf , x + ∆x), y)

∆x∈S

(33)

Repeat the same analysis as presented in Theorem A.2,

Thus,
∆x∈S

(32)

(16)

Proof. According to Theorem A.2,
∀x + ∆x, f (θf , x + ∆x) ≡ g(θg , x + ∆x).

max R(g(θg + ∆θg ), D)∆θg ∈Cθg= R(g(θg + ∆θg∗ ), D).

Theorem A.3. The robust losses of f and g are equal:R(f (θf ), D) ≡ R(g(θg ), D).

(31)

(24)

∆θf ∈CθfThus,∆θf ∈Cθf∆θg ∈Cθg

Such that the proof ends.

B. Algorithm of RiFT
The complete algorithm of RiFT is presented in Algorithm 2.
Algorithm 2 Robust Critical Fine-Tuning
Input: adversarially trained model weights θAT , standard
dataset Dstd , weight perturbation scaling factor α, finetuning optimization iteration steps T and learning rateγ, weight decay facotr λ.
∗Output: The fine-tuned model weights θAT.
1: Step 1: Calculate MRC for each module
2: for Module weight θ (j) do
3:
Calculate MRC value of θ (j) using Algorithm 1.
4: end for
5: Select the module with lowest MRC value, denote as
non-robust critical module θ (i)6: Step 2: Fine-tuning on Non-robust critical module
7: θ1 = θAT
8: for t = 1, . . . , T do
▷ Fine-tuning T epochs
9:
for Batch Bk ∈ Dstd do
10:
Calculate loss: L(f (θt ), Bk ))
(i)
(i)
11:
θt+1 = θt+1 − γ∇θt (L) ▷ Gradient Descent12:
end for13:
θF T = θt if θt obtain highest std test acc.
14: end for
15: Step 3: Interpolation
16: for α ∈ (0, 1, 0.05) do
17:
θα = (1 − α)θAT + αθF T
18:
θF∗ T = θα if it reaches best standard test acc while
preserve the robustness as θAT .
19: end for∗
20: Return Fine-tuned model weights θF
T

C. Training Details

the one with highest robust accuracy. The hyper-parametersettings are either based on their original paper or same as
the vanilla AT, depends on which method achieves betterrobust accuracy.

C.3. Fine-tuning Details
The hyper-parameter that most affects fine-tuning is theinitial learning rate. According to our experience, we find asmall learning rate usually performs better. If the adversarial robustness of the final fine-tuned weights is still higher
than the robustness of the initial adversarial training, wethen increase the learning rate.

C.4. The MRC value of ResNet34 and WRN34-10
Figure C.1 and Figure C.2 shows the Module Robust Criticality (MRC) value of each module in ResNet34
trained on CIFAR100 and WideResNet34 trained on TinyImageNet, respectively. It can be observed that both models exhibit redundant capacity. Additionally, Figure C.3
and Figure C.4 shows the MRC value of each module in
ResNet18 trained on CIFAR100 and Tiny-ImageNet, respectively. As we discussed in Section 5.3 and Section 5.5,
ResNet18 has a lower redundant capacity compared to
ResNet34 and WideResNet34, and the redundant capacity
decreases as the classification task becomes more complex.

C.5. More interpolation results
Figure C.5 shows the interpolation results of different
modules of ResNet18 trained on CIFAR100 dataset. It canbe observed that fine-tuning on robust-critical module canalso help improve generalization and robustness. This does
not mean that our MRC is wrong, as we claimed in Section 5.2, fine-tuning on robust-critical module does not necessarily hurt robustness. The MRC provides guidance on
which module to fine-tune for optimal results, and still, finetuning on non-robust-critical module achieves the highest
test accuracy while preserving robustness.

C.1. Experiment Environment
All experiments are conducted on a workstationequipped with an NVIDIA GeForce RTX 3090 GPU with
24GB memory and NVIDIA A100 with 80GB memory.
The PyTorch version is 1.11.0.

C.2. Adversarial Training Details
For vanilla adversarial training, We set the initial learning rate as 0.1, which decays at 100 and 105 epochs withfactor 10. When generating adversarial examples, we set
BN as train mode since it usually achieves higher robustness.
When incorporating RiFT with other adversarial training methods, the SCORE method is incorporated with
TRADES. For the CIFAR100 training, we ran with threedifferent learning rate and select the best model weights as

D. Analysis of the complexity of MRC algorithm
When identifying the most non-robust-critical module,
it is required to iterate all modules of the model. Supposea model with n modules, for each module, the calculation
complexity depends on the iteration steps in Algorithm 1.
Considering the different overheads for each iterative computation of the modules at different locations, for example,
when calculating the last module’s MRC value, it only requires forward-backward iteration of the last layer of parameters. Thus, the average total forward-backward iteration of each module is n/2. In our experiments, we set thelearning rate as 1 and the iteration step as 10. Thus, in our
experiments, the complexity of MRC algorithm cost 5n total forward-backward propagation.

Robust Acc Drop -0.05 -0.83 -0.71 -13.37 -0.64 -22.84 -0.57 -4.21 -0.61 -1.31 -0.24 -1.92 -0.21 -2.89 -0.20 -1.47 -0.76 -0.17 -0.01 -0.48 -0.18 -5.52 -0.16 -5.52 -0.40 -0.52 -0.24 -18.75 -4.72 -24.44 -4.33 -6.33 -2.41 -10.23

r

2

ea
lin

1

co
nv

.2.

2

co
nv

.2.

er4
lay

1

co
nv

.1.

er4

er4

lay

2

co
nv

.1.

lay

1

co
nv

.0.

er4
lay

2

co
nv

.0.

er4

er4

lay

1

co
nv

.5.

lay

2

co
nv

.5.

er3
lay

1

co
nv

.4.

er3
lay

2

co
nv

.4.

er3

er3

lay

1

co
nv

.3.

lay

2

co
nv

.3.

er3
lay

1

co
nv

.2.

er3
lay

2

co
nv

.2.

er3

er3

lay

1

co
nv

.1.

lay

2

co
nv

.1.

er3
lay

1

co
nv

.0.

er3

er3

lay

2

co
nv

.0.
er3

lay

1

co
nv

.3.

lay

2

co
nv

.3.

er2
lay

1

co
nv

.2.

er2

er2

lay

2

co
nv

.2.

lay

1

co
nv

.1.

er2
lay

2

co
nv

.1.

er2
lay

1

co
nv

.0.

er2

er2

lay

2

co
nv

.0.

lay

1

co
nv

.2.

er2
lay

2

co
nv

.2.

er1
lay

1

co
nv

.1.

er1

er1
lay

lay

2

co
nv

co
nv

.1.

.0.

lay

er1

er1

1
co
nv
lay

lay

er1

.0.

co

nv

1

Module Robust Criticality 0.04 0.06 0.04 1.16 0.03 2.82 0.04 0.34 0.06 0.10 0.02 0.15 0.02 0.17 0.02 0.21 0.13 0.01 0.02 0.07 0.02 0.53 0.03 0.52 0.03 0.02 0.03 3.20 0.63 5.31 0.41 0.58 0.17 0.83

Figure C.1. Example of module robust criticality (MRC) and its corresponding robust accuracy drop of ResNet34 trained on CIFAR100.

2.37

fc

0.14

0.25

0.24

0.01

0.43

0.01

1.64

0.03

4.11

0.07

1.17

1.42

5.27

0.05

0.42

0.02

0.66

0.03

3.70

0.03

5.77

1.25

3.90

1.02

1.35

0.33

4.18

1.77

0.14

blo

blo

ck1

.la
y

0.47

co
nv
er. 1
0
.co
ck1
n
.la
ye v1
blo
r.
ck1 0.co
n
.la
ye v2
blo
r.
ck1 1.co
n
.la
ye v1
blo
r.
ck1 1.co
n
.la
ye v2
blo
r.
ck1 2.co
n
.la
ye v1
blo
r.
ck1 2.co
n
.la
ye v2
blo
r.
ck1 3.co
n
.la
ye v1
blo
r.
ck1 3.co
n
.la
ye v2
blo
r.
ck1 4.co
n
.la
ye v1
blo
r.
ck2 4.co
n
.la
ye v2
blo
r.
ck2 0.co
n
.la
ye v1
blo
r.
ck2 0.co
n
.la
ye v2
blo
r.
ck2 1.co
n
.la
ye v1
blo
r.
ck2 1.co
n
.la
ye v2
blo
r.
ck2 2.co
n
.la
ye v1
blo
r.
ck2 2.co
n
.la
ye v2
blo
r.
ck2 3.co
n
.la
ye v1
blo
r.
ck2 3.co
n
.la
ye v2
blo
r.
ck2 4.co
n
.la
ye v1
blo
r.
ck3 4.co
n
.la
ye v2
blo
r.
ck3 0.co
n
.la
ye v1
blo
r.
ck3 0.co
n
.la
ye v2
blo
r.
ck3 1.co
n
.la
ye v1
blo
r.
ck3 1.co
n
.la
ye v2
blo
r.
ck3 2.co
n
.la
ye v1
blo
r.
ck3 2.co
n
.la
ye v2
blo
r.
ck3 3.co
n
.la
ye v1
blo
r.
ck3 3.co
n
.la
ye v2
blo
r.
ck3 4.co
n
.la
ye v1
r.4
.co
nv
2

Robust Acc Drop -0.29 -2.29 -3.72 -4.03 -0.05 -6.69 -0.05 -19.54 -0.60 -34.77 -1.25 -15.85 -16.27 -34.75 -1.07 -7.64 -0.40 -10.56 -0.69 -30.97 -0.64 -32.87 -14.20 -41.21 -14.90 -40.39 -5.44 -27.12 -17.09 -38.11 -9.28 -30.64
Module Robust Criticality 0.02

Figure C.2. Example of module robust criticality (MRC) and its corresponding robust accuracy drop of WideResNet34 trained on TinyImageNet.

Robust Acc Drop -1.13 -23.00 -14.87 -29.10 -23.87 -27.67 -24.41 -29.09 -17.45 -28.35 -27.03 -29.09 -12.01 -29.02 -28.94 -29.09 -25.51 -26.45

r
ea

v2
co
n

lin

v1
co
n

er4
lay

er4

3.26

.1.

v2
co
n
lay

er4
lay

.1.

v1
co
n

er4
lay

.0.

v2
co
n

er3

.0.

v1
co
n
lay

er3

.1.

v2
co
n
lay

er3
lay

.1.

v1
co
n

er3
lay

.0.

v2
co
n

er2

.0.

v1
co
n
lay

er2

.1.

v2
co
n
lay

er2
lay

.1.

v1

er2
lay

.0.

co
n

er1
lay

.0.

v2

er1
lay

.1.

er1
lay

.1.

er1

.0.

co
n

6.53 13.06 21.28 1.06 21.33 24.24 41.72 4.99

v1

3.66 10.54 1.61

co
n

6.06

v2

2.92

co
n

6.53

v1

1.30

lay

.0.

co

2.31

co
n

0.09

nv
1

Module Robust Criticality

Figure C.3. Example of module robust criticality (MRC) and its corresponding robust accuracy drop of ResNet18 trained on CIFAR100.

Robust Acc Drop -6.45 -34.66 -34.00 -34.93 -34.72 -34.92 -34.67 -34.92 -33.63 -34.93 -34.77 -34.93 -26.73 -34.93 -34.93 -34.93 -34.61 -34.93

r
ea
lin

v2
on

er4

.1.
c

v1
on
lay

er4

.1.
c

v2
on
lay

.0.
c

v1
lay

er4

on
.0.
c

v2
lay

er4

on

er3

.1.
c

v1
on
lay

lay

er3

.1.
c

v2
on

v1

.0.
c

lay

er3

on
.0.
c

v2
on

er3
lay

er2

.1.
c

v1
on
lay

lay

er2

.1.
c

v2
on

v1

.0.
c

lay

er2

on
.0.
c

v2
on

9.85 15.14 15.15 7.27 18.89 14.33 57.21 4.82 80.39 100.67 103.50 11.15 11.29

er2
lay

.1.
c

v1
on

er1
lay

.1.
c
er1

on

v1

.0.
c
er1

lay

er1

.0.
c

on

nv
co

lay

8.99 14.26 7.01

v2

6.33

lay

0.38

1

Module Robust Criticality

Figure C.4. Example of module robust criticality (MRC) and its corresponding robust accuracy drop of ResNet18 trained on Tiny-ImageNet.

Adversarial Robust Accuracy

30.4
30.2
30.0
29.8
29.6
29.4
57.0

Non-robust-critical module
Linear layer
All layers
Robust-critical module
Vertical Line
57.5

58.0
58.5
Standard Accuracy

59.0

Figure C.5. Interpolation results of fine-tuning on different modules of ResNet18 on CIFAR100 dataset. Dots denote different interpolation points between the final fine-tuned weights of RiFT
and the initial adversarially trained weights.