Vision Transformer (ViT) models have achieved remarkable performance across various vision tasks, with scalability being a key advantage when applied to large datasets. This scalability enables ViT models to exhibit strong generalization capabilities. However, as the number of parameters increases, the robustness of ViT models to adversarial examples does not scale proportionally. Adversarial training (AT), one of the most effective methods for enhancing robustness, typically requires fine-tuning the entire model, leading to prohibitively high computational costs, especially for large ViT architectures. In this paper, we aim to robustly fine-tune only a small subset of parameters to achieve robustness comparable to standard AT. To accomplish this, we introduce Criticality-Aware Adversarial Training (CAAT), a novel method that adaptively allocates resources to the most robustness-critical parameters, fine-tuning only selected modules. Specifically, CAAT efficiently identifies parameters that contribute most to adversarial robustness. It then leverages parameter-efficient fine-tuning (PEFT) to robustly adjust weight matrices where the number of critical parameters exceeds a predefined threshold. CAAT exhibits favorable generalization when scaled to larger vision transformer architectures, potentially paving the way for adversarial training at scale, e.g, compared with plain adversarial training, CAAT incurs only a 4.3% decrease in adversarial robustness while tuning approximately 6% of its parameters. Extensive experiments on three widely used adversarial learning datasets demonstrate that CAAT outperforms state-of-the-art lightweight AT methods with fewer trainable parameters.
Figures & tables
Figure 1: Illustration of standard AT and our CAAT. The parameters highlighted in blue denote frozen parameters, while those in rhodo represent trainable parameters. We only tune the top- τ critical parameters.
Figure 2: The pipeline of our proposed Criticality-Aware Adversarial Training. Given a vulnerable ViT, we first identify the top- τ critical parameters for adversarial robustness. Instead of directly applying AT to these parameters, we leverage the representational capability of PEFT to enhance adversarial robustness.
Dataset
Method
Tuned params (M)
Tuned / Total (%)
Clean Acc (%)
CW-20 (%)
PGD-10 (%)
AutoAttack (%)
Average (%)
CIFAR-10
Full
85.15
100
82.65
50.69
53.66
47.14
50.50
Adapter-8
1.07
1.25
80.00
48.51
47.82
44.15
46.83
Adapter-32
1.22
1.43
80.06
48.58
47.74
42.47
46.26
LoRA-16
1.24
1.46
83.24
46.71
48.96
44.96
46.88
Aurora
1.15
1.35
81.25
45.71
43.86
42.88
44.15
FullLoRA-AT
2.46
2.87
84.93
47.67
50.30
45.28
47.75
Table 1: Comparisons of different robust fine-tuning methods using a pre-trained ViT-B [ 1 ] backbone across various datasets. ’Total params’ represents the ratio of the total number of parameters involved in adversarial training, while ’Tuned/Total’ denotes the fraction of trainable parameters. The standard adversarial training is highlighted with a gray background. Additionally, the top-1 (%) accuracy is reported in bold , and the top-2 (%) accuracy is reported in underline .
Method
Inference Memory (GB)
Fine-tuning Memory (GB)
Train time (h)
Full
1.3
11.8
13.10
Adapter
2.0
14.0
11.30
LoRA
1.4
10.0
10.48
FullLoRA-AT
1.4
8.5
8.39
HyperAT
1.5
9.2
8.33
CAAT
1.3
7.3
6.00
Table 2: Cost comparisons with ViT-B backbone on the CIFAR10 robustness when faced PGD-10 attack. The best result is in bold .
Arch
Method
Paras
Clean Acc
CW-20
PGD-10
ViT-L
Full
304.33
74.82
52.77
55.39
LoRA
23.58
70.53
48.90
54.86
FullLoRA-AT
81.14
61.76
50.07
44.90
CAAT
11.42
73.72
51.49
54.83
ViT-H
Full
631.51
76.58
50.76
54.37
LoRA
40.36
69.32
44.39
45.97
Table 3: Adversarial robustness result of our CAAT when scaled to larger vision transformer architecture. The best result is in bold .
Figure 3: Visualization of the selected critical parameters of ViT-B on CIFAR10 when subjected to a PGD-10 attack under 0.05M, 0.1M, and 0.2M parameter budgets, respectively.
Dataset
Combined AT
Method
Trainable Params(M)
Clean Acc (%)
PGD-10 (%)
AutoAttack (%)
CIFAR-10
TRADES
Full
85.15
83.57
53.59
48.99
LoRA
4.72
84.73
50.45
45.90
FullLoRA-AT
4.86
86.88
51.66
46.01
CAAT
1.10
86.95
52.15
47.60
MART
Full
85.15
83.66
53.57
49.05
LoRA
4.72
84.80
50.69
46.04
Table 4: Results of combining CAAT with other AT methods. The best result is in bold .
Figure 4: Accuracy vs. parameter efficiency of ViT-B on the CIFAR10 dataset under PGD-10 attack defense.
Method
Tuned / Total (%)
Clean Acc (%)
CW-20 (%)
PGD-10 (%)
AutoAttack (%)
LoRA
1.24
83.24
46.71
50.30
45.28
CAAT(direct tuning)
1.15
79.68
46.82
50.35
45.55
CAAT(indirect tuning)
1.09
87.12
49.33
51.21
46.37
Table 5: Ablation study on direct and indirect tuning. The best result is in bold .
Rank r
Paras
Clean Acc
CW-20
PGD-10
AutoAttack
8
0.57
82.06
46.32
49.64
43.1
16
1.09
87.12
49.33
51.21
46.37
32
2.34
86.77
50.15
50.38
46.94
64
4.93
84.25
47.66
50.5
45.86
Table 6: Results of our CAAT under different rank r .
Sample Num.
200
400
600
800
PGD-10 Acc.
50.10
50.13
50.15
50.15
Table 7: Effect of the number of adversarial samples used to calculate parameter criticality. The best result is in bold .
Dataset
Method
Tuned params (M)
Tuned / Total (%)
Clean Acc (%)
CW-20 (%)
PGD-10 (%)
AutoAttack (%)
CIFAR-10
Full
86.69
100
81.28
48.66
50.69
46.27
Adapter-8
1.72
1.98
75.68
42.10
38.85
39.37
Adapter-32
2.00
2.31
76.53
42.53
39.88
39.47
LoRA-16
3.15
3.63
80.84
45.73
46.59
43.81
FullLoRA-AT
3.17
3.66
82.51
45.75
47.58
44.27
CAAT-Adapter
1.15
1.33
82.89
46.62
47.72
43.95
Table 8: Comparisons of different robust fine-tuning methods using a pre-trained Swin-B [ 45 ] backbone across various datasets. The standard adversarial training is highlighted with a gray background. Additionally, the top-1 (%) accuracy is reported in bold , and the top-2 (%) accuracy is reported in underline .
Division of CEMSE, King Abdullah University of Science and Technology · School of Statistics, Renmin University of China · Department of Computer Science and Engineering, State University of New York at Buffalo +1
School of Computer Science and Engineering, and Key Laboratory of New Generation Artificial Intelligence Technology and Its Interdisciplinary Applications, Southeast University, China. · School of Intelligence Science and Engineering, Southeast University, China.