Improving Adversarial Robustness of Vision Transformers
Anonymous Author(s)
AffiliationAddress email
Abstract
Vision Transformers (ViTs) have recently achieved competitive performance inbroad vision tasks.
Unfortunately, on popular threat models, naturally trained ViTsare proven to provide no more adversarial robustness than convolutional neuralnetworks (CNNs).
Adversarial training is still required by ViTs to defend againstsuch adversarial attacks.
This paper provides a comprehensive evaluation of varioustraining techniques across several datasets, thus bringing the first implementationbenchmark for adversarial training of ViTs.
In addition, regarding ViT as a newtype of model architecture, we further investigate its adversarial robustness fromthe perspective of its unique architectural components.
We find, when we randomlymask the perturbation on some of the patches or the gradient from some attentionblocks during adversarial training, the adversarial robustness of ViTs can be further improved.
Our work may potentially open up a line of work to explore thearchitectural information inside new models like ViTs.
Introduction
Recent years have witnessed the overwhelming advances of Vision Transformers (ViTs) [1, 2, 3].
Different from widely deployed Convolutional Neural Networks (CNNs) [4, 5, 6, 7] that adopt aseries of local convolutional operations to input images, ViTs adopt self-attention mechanisms [8]to sequences of image patches.
Based on large scale pre-training, ViTs have achieved competitiveand even better performance compared to CNNs in several fields such as semantic segmentation [9],object detection [10] and image generation [11].
Despite their great success in a growing number of vision tasks
ViTs fail in providing more adversarialrobustness than CNNs on popular threat models [12].
To defend against such adversarial examplescrafted by adding human-imperceptible perturbations to images [13, 14], ViTs still demand adversarialtraining (AT) [15, 16, 17, 18], an approach that incorporates adversarial examples into training andobtains notable empirical robustness.
Note that implementation details make a big difference in AT[19], which motivates us to comprehensively evaluate various training techniques on benchmarkdatasets, and thus provide the first implementation benchmark for AT of ViTs.
We find that gradientclipping and pre-training on large datasets are necessary for improving AT of ViTs, and advanceddata augmentations indeed improve robustness in contrary to CNNs.
SGD behaves better for ViTsthan AdamW in term of adversarial robustness.
Robust overfitting occurs much easier during AT of
ViTs.
Based on our empirical findings, we provide a practical recipe for adversarial training of ViTsto help researchers in their future work.
Although ViTs can achieve reasonable adversarial robustness via AT just like CNNs, it is unexploredwhether or not we are able to further improve the robustness based on the specific architecturalinformation from ViTs, a model totally different from CNNs.
Interestingly, when we randomly maskthe perturbation on some of the patches or the gradient from some attention blocks during adversarial training, the robustness of ViTs can be further improved.
We term these two simple methods as
Perturbation Random Masking (PRM) and Attention Random Dropping (ARD) respectively, andconduct extensive experiments to verify their effectiveness.
Our main contributions are summarizedas follows:
• We comprehensively evaluate training techniques for AT of ViTs on popular threat models,which indicates the necessity of gradient clipping and pre-training, proper choices ofadvanced data augmentations and optimizers, and the higher propensity to robust overfittingcompared to CNNs.•
Based on the unique architectural information from ViTs, we propose Perturbation RandomMasking (PRM) and Attention Random Dropping (ARD) as a warming-up strategy to
improve the adversarial robustness.
In particular, PRM randomly masks the perturbation on
some of the patches, while ARD randomly masks the gradient from some attention blocks.
• Our work not only provides the first implementation benchmark for AT of ViTs, but also
reminds researchers of the potential of architectural information inside brand new modelslike ViTs.
Related Work
2.1 Training Strategies for ViTs
Due to the lack of inductive bias [1], ViTs require more training techniques to reach their performancepotential.
Dosovitskiy et al [1] first pretrained ViTs on the large-scale datasets (ImageNet-21k or
JFT-300M [20]) and achieved competitive performance on downstream tasks.
Touvron et al
[3]applied strong data augmentations, such as Random Augment [21] and Mixup [22] to train ViTsfrom scratch.
They also find AdamW [23] performs better than SGD [24] for training ViTs.
Besides,several studies [1, 25, 26] adopted gradient clipping to stabilize and accelerate the convergence of
ViTs.
Most related to our work, Steiner et al [27] conducted a comprehensive study on data augmentationand model regularization of ViTs to improve the accuracy.
Unlike them, we focus on investigatingthe strategies to enhance the adversarial robustness.
Adversarial
Robustness
Given a clean example x with its class label y and a model f θ ( · ) with its parameters θ , the goal ofan adversary is to find an adversarial perturbation δ that fools the network into making an incorrectprediction for the perturbed image ( i.e. , f θ ( x + δ ) ̸ = y ), while the perturbation norm does not exceedϵ ( i.e. , ∥ δ ∥ ∞ ≤ ϵ ),
[Equation]
while L ( · , · ) is usually the cross entropy loss.
Previous studies [14, 28, 29, 15, 30] proposed severalmethods to generate adversarial examples, and Project Gradient Descent (PGD)
[15] becomes themost popular one, which is
[Equation]
To defend against adversarial attacks, while several defense methods [31, 32] are proposed, onlyadversarial training (AT) [15] and its variants [16, 17, 18, 33] are demonstrated to show notableempirical robustness [34, 35, 36], which directly incorporates adversarial examples into training asfollows,
[Equation]
Recently, AT of CNNs has become more understood [15, 37, 19].
However, few robustness studiesfocus on ViT, the brand new architecture in computer vision.
For example, several studies [38, 39, 12]have shown ViTs are more robust against common corruptions, which are simpler than adversarialperturbations.
Naturally trained ViTs are only able to provide limited robustness against extremelysmall adversarial perturbations ( e.g. , under threat model of ϵ = 0 .
001 on ImageNet 1 ) [38, 39], while adversarially trained ones fail in providing more adversarial robustness on popular threat modelscompared to CNNs [12].
Further, AT is extremely sensitive to training settings [19], while thereis no training strategy evaluation for AT of ViTs.
This motivates us to comprehensively evaluatevarious training strategies across several datasets, thus bringing the first implementation benchmarkfor adversarial training of ViTs.
Strategies for Adversarial Training of ViTs
In this section, we comprehensively evaluate several training techniques in terms of data (pre-trainingand data augmentation) and training (optimizer, learning rate schedule, and gradient clipping) acrossvarious architectures (vanilla ViT and Swin) and datasets (CIFAR-10 and Imagenette), so as tobenchmark the state-of-the-art (SOTA) adversarial robustness of ViTs.
All experiments are performedon GPUs of Tesla V100.
Datasets.
Since AT is time-consuming (usually 10 × slower than natural training), small datasetsremain the most popular on adversarial robustness [36].
We adopt CIFAR-10 [40] and Imagenette 2 (asubset of 10 classes from ImageNet-1K) datasets.
Architectures.
We use the vanilla ViT (ViT-B) [1] and Swin-B [2] as baseline models.
ViT firstsuccessfully introduced Transformers from natural language processing [8, 41, 42, 43] to computervision and achieved competitive recognition performance compared to CNNs.
With the help ofmulti-stage hierarchical architecture, Swin became a general-purpose backbone and achieved SOTAperformance on several downstream tasks 3 .
By default, all models are pretrained in ImageNet-1K 4 .
Considering that the image size on ImageNet is 224 × 224 and the image size on CIFAR-10 is 32 × 32 ,we downsample the weights for patch embeddings and use the patch size 4 × 4 (the length of thepatch sequence becomes 8 × 8 = 64 on CIFAR-10) to keep the input size on CIFAR-10.
Perturbation Threat Models.
We apply the commonly-used ℓ ∞ attack, that is, the perturbation δsatisfies ∥ δ ∥ ∞ ≤ ϵ .
Following the most popular threat models [36], we set ϵ = 8 / 255 for CIFAR-10and ϵ = 8 / 255 for Imagenette.
When evaluating the adversarial robustness, we apply 20-step PGD(PGD-20)
[15] with the step size 2 / 255 and Auto-Attack (AA) [35] that is the strongest attack toverify the empirical robustness via an ensemble of diverse attacks ( i.e. , two variants of PGD attack,
FAB attack [44], and Square Attack [45]).
Training Settings.
All models are pre-trained on ImageNet-1K and are adversarially trained for 40epochs using SGD with momentum 0.9, weight decay 1e-4, and an initial learning rate 0.1.
Thelearning rate is divided by 10 at the 36-th and 38-th epoch.
Simple data augmentations such asrandom crop with padding and random horizontal flip are applied.
During adversarial training, weuse PGD-10 with step size 2 / 255 to craft adversarial examples.
3.1 The Necessity of Gradient Clipping
[Figure]
First of all, we find that it is difficult to adversarially train ViTs onsmall datasets.
For example, the final robustness of ViT-B is only10% on CIFAR-10 and 9.86% on Imagenette in Table 1.
To explorethe problem, we visualize the norm of gradient during adversarialtraining in Figure 1, and find the gradient explodes at the very beginning of training.
We adopt gradient clipping [46] to AT of ViTs.
Specifically, we ensure the maximum of gradient vector is 1 underthe l 2 -norm.
In Table 1, with the help of gradient clipping, ViT-Bachieves the robustness of 48.33% and 56.00% on CIFAR-10 and
Imagenette.
Although Swin-B can be adversarially trained directly,gradient clipping still improves the performance by a notable margin,i.e. , improving natural accuracy by 3.52% and robustness by 3.87% on Imagenette.
Interestingly, gradient clipping has almost no effect on the results of CNNs (refer to Appendix B), which demonstratesthe difference between ViTs and CNNs in terms of training dynamics.
Since gradient clipping isalmost necessary and beneficial, we always apply it by default in our following experiments.
[Table]
[Figure] [Figure]
3.2 The Necessity of Pre-training
Despite that ViTs perform competitively on ImageNet-1K, which is a difficult recognition task with1000 classes, their performance heavily relies on the large-scale pre-training [1, 47] and strong dataaugmentations [3] due to the lack of CNN-like inductive biases.
Considering that the input size on
CIFAR-10 is much smaller ( 32 × 32 ) and the task itself is much simpler with just 10 classes, we arewondering if pre-training is still necessary for ViTs on such a simple recognition task (although underadversarial attacks).
We evaluate 3 pre-training strategies: 1) pre-training on ImageNet-21K (an extremely large datasetwith 14,197,122 images from 21,841 classes), 2) pre-training on ImageNet-1K (a large dataset withmore than 1.2M images from 1,000 classes), and 3) no pre-training ( i.e. , adversarial training fromscratch).
The final robustness under AA of different AT variants is shown in Figure 2(a).
Comparedto AT from scratch, pre-training-based AT can improve the robustness of ViTs by a notable margin( ∼ 16% on average).
Interestingly, the benefit from pre-training to ViTs is much more significant thanthose to CNNs [48] (+11.6%), which reflects the stronger necessity of ViTs for pre-training comparedto CNNs.
Moreover, pre-training on a larger dataset is unable to bring higher final robustness.
To explore the reason why pre-training can help AT, we plot the learning curve of ViT-B with andwithout pre-training.
In Figure 2(b), we find that ViT without pre-training fails in fitting the trainingdata during adversarial training (the training robustness is only 27.96%).
This suggests that thecurrent random initialization for weight parameters of ViTs is not suitable for AT, even on smalldatasets like CIFAR-10, while pre-training on large datasets could provide a good initialization.
3.3 The Effect of Data Augmentations
In AT of CNNs, a majority of advanced data augmentations like Mixup [22] and RandAugment[21] are proven to be unsuccessful (without considering model weight averaging)
[37, 49], which iscontrary to natural training.
Here, we explore if data augmentations can improve robustness for AT of
ViTs.
Specifically, we incorporate CutMix [50], Mixup, and Randaugment to AT of ViTs respectively.
In Figure 3(a), ViT-B obtains higher robustness on CIFAR-10 with CutMix or Mixup or both, whilebehaving worse with RandAugment.
We also find the similar phenomena of Swin-B on Imagenette.
We conjecture this is because RandAugment is a more advanced technique, which is too difficultfor AT of ViTs.
In conclusion, a suitable combination of data augmentations can directly improvethe adversarial robustness of ViTs, which is different from CNNs.
In the following experiments, wealways adopt CutMix and Mixup to AT of ViTs by default.
[Figure]
[Figure] [Figure]
[Table]
3.4 The Effect of the Optimizer
Recalling that SGD achieves the best adversarial robustness in CNNs [15, 37].
By contrast, the defaultoptimizer of ViTs is AdamW [23].
In this part, we study the choice of optimizers on adversarialrobustness.
We also consider the influence of different combination of optimizer and learning ratescheduler in Appendix C. Part of results on two popular settings are shown in Table 2: 1) AdamW+ cyclic learning rate schedule, and 2) SGD + piecewise learning rate schedule.
We find SGD +piecewise learning rate schedule helps ViT-B to improve the natural accuracy by 4.51% and therobustness by 2.90% over AdamW + cyclic learning rate schedule on CIFAR-10.
In addition, from
Appendix C, we can see that AdamW + piecewise learning rate schedule for AT of ViTs obtainsworse robustness compared to SGD + piecewise learning rate, which indicates the performanceimprovements come from the optimizer itself, rather than the learning rate schedule, i.e., SGDconsistently improves robustness for ViTs.
3.5 Training for More Epochs
Recent research [51, 52] demonstrated that training for longer can also be viewed as increasing modelcomplexity.
Considering that AT requires models with larger complexity [15], we explore if traininglonger is able to improve adversarial robustness.
We adversarially train models for 80 epochs (40epochs in our default setting) and decay the learning rate at 75-th and 78-th epoch.
Figure 3(b)illustrate that robustness has already decreased before the first learning rate decay.
Because the robustoverfitting of ViTs happens much earlier than those of CNNs, training longer fails in improving therobustness and hurts the final robustness instead.
3.6 Recipe for Adversarial Training of ViTs
After conducting a comprehensive evaluation of different training techniques, we provide a practicalrecipe for AT of ViTs.
First, gradient clipping and pre-training are almost necessary.
Second, it is agood choice to use advanced data augmentations ( e.g. , Cutmix and Mixup), because a suitable combination of data augmentations could improve robustness effectively.
Third, in terms of adversarialrobustness of ViTs, it is better to use SGD with piecewise learning rate schedule, since it outperforms
AdamW with cyclic learning rate schedule (the widely used optimizer for natural training of ViTs).
Fourth, since robust overfitting on ViTs happens easier compared to CNNs, it is beneficial to useearly stopping.
Multi-Head Attention
Norm
MLP
Norm
EmbeddedPatches
Norm
MLP
Gradient
F o r w a r d P r op a g a ti on
B ac k w a r d P r op a g a ti on × 𝐿 × 𝐿
[Figure]
+ =
[Figure]
An Architectural Perspective for Adversarially Robust ViTs
In contrast to CNNs, ViTs are a completely different architecture.
For example, it divides theimages into sequences of image patches before feeding them into the model, and it directly integratesglobal information across the entire images even in the lowest layers using self-attention [1].
Thisfundamental difference naturally raises a question whether or not we are able to further improve therobustness based on such specific architectural information.
In this section, we explore the impact ofthese two most obvious architectural changes (multi-head attention and patch-based image splitting)inside ViTs on the adversarial robustness.
4.1 Attention Random Droppping (ARD)
In ViTs, embedded patches 5 are repeatedly processed by encoder blocks, any of which consist of amulti-head attention and a MLP layer, as shown in Figure 4(a).
For the i -th block, its input z i − 1 isfirst processed by a multi-head attention A ( i ) ( · ) , and intermediate result 6 z ′ i
=
A ( i ) ( z i − 1 )
+ z i − 1 isfurther processed by the remaining parts (including a MLP layer and a skip connection) of this blockf ( i ) ( · ) .
For a ViT with L blocks, its output of the last ( i.e. , L -th) block is
[Equation]
According to the chain rule in calculus, the gradient of a loss function L with respect to the input ofthe last block, z L − 1 , can then be decomposed as
[Equation]
Here, we propose the Attention Random Dropping (ARD), which randomly drops the gradient flowingthe multi-head attention
A ( L ) ( · ) ,
[Equation]
where we set u L to 0 with the probability of p and set it to 1 otherwise.
We continue to use the chainrule to obtain the gradient of ARD with respect to the input image x ,
[Equation]
where z 0 is the embedded pathes and fed into the 1 -st Transformer encoder, and all random variablesu i are i.i.d .
Figure 4(a) also illustrates the backward propagation using ARD.
When ARD meets adversarial training, we apply ARD as a warming-up strategy.
After the warmingup periods, the training becomes the vanilla AT.
In particular, the parameter p in ARD decreaseslinearly from 1 to 0 during the first n w epochs, i.e. , p = 1 − n/n w where n is the current epoch( n ≤ n w ).
Here, p is a dynamic parameter along with the training process.
At the beginning, relativelyweak adversarial samples (large p ) are used to warm up while as the training continues, the probabilityp should be gradually smaller to 0 to ensure the strength of the generated adversarial samples duringthe whole adversarial training process.
4.2 Perturbation Random Masking (PRM)
ViTs always split the RGB images into non-overlapping patches, and then process them further.
Considering this patch-based image splitting, we introduce the Perturbation Random Masking (PRM).
Specifically, the perturbation on a patch will be masked with probability k during adversarial examplegeneration:
[Equation]
where δ is the adversarial perturbation, and M is the mask to remove the perturbations on somepatches.
The second equation is exactly the update rule of PGD attack.
Similar to ARD, we incorporate PRM as a warming-up strategy.
Considering an image that issplitted into J patches, the perturbations on ⌊ J · k ⌋ patches will be masked for each iteration duringadversarial example generation within AT, and we linearly decrease k from 1 to 0 at the first n wepochs, i.e., k = 1 − n/n w where n is the current epoch ( n ≤ n w ).
Here, same to p , k is also adynamic parameter along with the training process.
4.3 Overall Algorithm
Obviously, ARD and PRM are two individual methods that have impact on two different aspects,i.e., ARD influences the backward propagation while PRM influences the update of adversarialperturbations.
Thus, we can easily combine them as a mixed warming-up strategy 7 :
[Equation]
The details of ARD or PRM based adversarial training for ViTs are summarized in the Algorithm 1in Appendix D.
4.4 Evaluation on Benchmark Datasets
In this section, we evaluate the effectiveness of ARD and PRM for the adversarial robustness of ViTson benchmark datasets.
Experimental Settings.
For models, we apply 4 variants of ViT (vanilla ViT [1], DeiT [3], ConViT[53], and Swin [2]) and 3 kinds of scales (base, small, and tiny) for each respectively.
For robustnessevaluation, we also adopt 20 steps CW [29] ( ℓ ∞ version of CW loss optimized by PGD-20) and 100steps PGD ∞ [15] in addition to PGD-20 and AA.
We set n w to 10 .
We evaluate the performance on:
1) vanilla setting (the basic setting determined in Section 3), 2) the setting with ARD, 3) the settingwith PRM, and 4) the setting with ARD and PRM
Experimental Results.
In Table 3, after comparing the same kind of ViTs with different scales,we observe that increasing the size of ViTs can improve both the accuracy on clean examples andthe robustness against adversarial attacks.
Moreover, diving into each ViT model with differentarchitectural approaches, we find that our proposed ARD or PRM can boost both robustness andaccuracy.
When only considering the robustness improvement, applying ARD or PRM (one of them)is enough.
Taking Swin-S on CIFAR-10 as an example, we can see that the gain for ARD is 1.16% [Table]
while 1.13% for PRM under AA.
When additionally considering accuracy improvement, it is betterto use both ARD and PRM.
For example, on CIFAR-10 with ViT-B model, applying both ARD and
PRM improves accuracy of 1.60% compared with vanilla, while only applying ARD or PRM, thegains are at most 1.15%.
Similar results can also be observed on Imagenette.
4.5 Combination with Other Defense Methods
Although the above discussions are mainly concentrated on standard adversarial training, we furtherdemonstrate that our methods: ARD and PRM, can be easily combined with other stronger defensemethods.
We adopt TRADES [16] and MART [54] here because of their outstanding performanceevaluated in [36].
The settings are the same as the Section 4.4 except that we do not apply CutMixand Mixup for data augmentations to keep in line with their original papers.
The experiments areconducted on DeiT-Ti, ConViT-Ti and Swin-Ti.
As shown in Table 4, the robustness and accuracy of
TRADES and MART are improved on both datasets.
For example, for DeiT-Ti on CIFAR-10, AArobustness increases by 0.77% on TRADES and 1.63% on MART.
It demonstrates that ARD and
PRM can further improve the performance of existing defense methods.
[Table]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
4.6 Hyper-Parameter Analysis
For simplicity of our proposed warming-up strategy, we use the same hyper-parameter n w as thenumber of warming-up epochs for both ARD and PRM.
We firstly analyze this shared hyper-parametern w ∈ { 0 , 5 , 10 , 15 , 20 , 25 } .
The experiments are based on four ViT variants (ViT, DeiT, ConViT, and
Swin), and the settings are the same as the Section 4.4.
AA is used to evaluate the robustness.
In
Figure 5(a), we significantly improve the robustness as long as adopting our proposed warming-upstrategy into AT, i.e. , n w > 0 .
We obtain the highest robustness at the best n w ( e.g. , n w = 15 for
Swin-S).
Besides, the improvements are not sensitive to the hyper-parameter, which allows us to tunen w easily.
To further exploit the proposed method, we use two individual hyper-parameters n aw and n pw for
ARD and PRM respectively and we peform our experiments on Swin-S. When fixing one ( e.g.,n aw ) as 15 (the best one in the last experiment), we change another hyper-parameter ( e.g. , n pw ) as{ 0 , 5 , 10 , 15 , 20 , 25 } .
In Figure 5 (b), for robustness, these two situations show different trends, i.e. ,a longer warm-up with ARD (the orange line) increase the final robustness, a longer warm-up with
PRM (the blue line) decrease the final robustness.
We can find a good balance when n aw = n pw = 15 .
The former has improved the performance enough, and the latter has not decreased too much.
As aresult, we set n aw = n pw = n w by default.
Conclusion
In this paper, we conducted comprehensive experiments to investigate training techniques acrossdatasets to bring a first implementation benchmark for AT of ViTs.
We also provided some practicalsuggestions, including using gradient clipping, pre-training on large datasets, employing stronger dataaugmentations, and preferring SGD rather than AdamW. Besides, taking the architectural informationinto consideration, we incorporated two simple but effective methods, ARD and PRM, into AT toimprove the adversarial robustness further.
Extensive experiments demonstrated the effectiveness ofour proposed warming-up strategy.
We hope our work not only provides the
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Appendix A.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] SeeAppendix A.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[No] We will release
upon acceptance.
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Section 3.
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[No]
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[Yes] See Section 3.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] See Section 3.
(b) Did you mention the license of the assets?
[Yes] See Section 3.
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[No]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]