Improving Adversarial Defense with Self-supervised Test-time Fine-tuning

1 INTRODUCTION . Adversarial training and its variants ( Madry et al. , 2018 ; Wang et al. , 2019 ; Zhang et al. , 2019 ) are currently recognized as the most effective defense mechanism against adversarial attacks . However , adversarial training generalizes poorly ; the robust accuracy gap between the training and test set in adversarial training is much larger than the training-test gap in standard training of deep networks ( Neyshabur et al. , 2017 ; Zhang et al. , 2017 ) . Unfortunately , classical techniques to overcome overfitting in standard training , including regularization and data augmentation , only have little effect in adversarial training ( Rice et al. , 2020 ) . The poor generalization of adversarial training is a consequence of adversarial test examples lying far away from the training set ( Huan et al. , 2019 ) . To adapt to these examples , we propose to finetune the network for each test mini-batch . Since the labels of the test images are not available , we exploit self-supervision , which is widely used in the standard training of networks ( Chen et al. , 2020b ; Gidaris et al. , 2018 ; He et al. , 2020 ) . Fine-tuning self-supervised tasks is a substitute of fine-tuning the classification loss at the inference time . Thus , minimizing this self-supervised loss function yields better generalization on the test set . To make our test-time fine-tuning strategy effective , we need to search for a good starting point during training . As will be shown in our experiments , adversarial training does not provide the optimal starting point . We therefore formulate the search for the network parameters that can be effectively fine-tuned as a bilevel optimization problem . Specifically , we introduce a meta adversarial training strategy dedicated to our self-supervised fine-tuning inspired by the model-agnostic metalearning ( MAML ) framework ( Finn et al. , 2017 ) . To this end , we treat the classification of each batch of adversarial images as one task and minimize the corresponding classification error of the self-supervised fine-tuned network . This meta adversarial training strategy strengthens the correlation between the self-supervised and classification tasks so that self-supervised test-time fine-tuning further improves the robust accuracy . In order to reliably evaluate our method , we follow the suggestions of ( Tramer et al. , 2020 ) and design an adaptive attack strategy that is fully aware of the test-time fine-tuning . We empirically demonstrate the effectiveness of our method on the commonly used CIFAR10 ( Krizhevsky et al. , 2009 ) , STL10 ( Coates et al. , 2011 ) and Tiny ImageNet ( Le & Yang , 2015 ) datasets under both standard ( Andriushchenko et al. , 2020 ; Croce & Hein , 2020a ; Madry et al. , 2018 ) and adaptive attacks in the scenarios of both white-box and black-box attacks . The experiments evidence that our method consistently improves the robust accuracy under all attacks . Our contributions can be summarized as follows : 1 . We introduce the framework of self-supervised test-time fine-tuning for adversarially-trained networks and show that it improves the robust accuracy of the test data . 2 . We propose a meta adversarial training strategy based on the MAML framework to find a good starting point and strengthen the correlation between the self-supervised and classification tasks . 3 . Following the suggestions of ( Tramer et al. , 2020 ) , we design an adaptive attack strategy that is fully aware of our test-time fine-tuning . The experiments show that our approach is valid on diverse attack strategies , including standard and adaptive ones in both the white-box and black-box attacks . Our code is available through this link . 2 RELATED WORK . Adversarial Training . In recent years , many approaches have been proposed to defend networks against adversarial attacks ( Guo et al. , 2018 ; Liao et al. , 2018 ; Song et al. , 2018 ) . Among them , adversarial training ( Madry et al. , 2018 ) stands out as one of the most robust and popular methods , even under various strong attacks ( Athalye et al. , 2018 ; Croce & Hein , 2020a ) . Adversarial training optimizes the loss of adversarial examples to find parameters that are robust to adversarial attacks . Several variants of adversarial training have been proposed ( Wang et al. , 2019 ; Zhang et al. , 2019 ) , and their performance is similar to the original version of adversarial training ( Rice et al. , 2020 ) . One important problem that limits the robust accuracy of adversarial training is overfitting . Compared with training on clean images , the gap of robust accuracy between the training and test set is much larger in adversarial training ( Rice et al. , 2020 ) . Moreover , traditional techniques to prevent overfitting , such as data augmentation , have little effect . Recently , some methods have attempted to flatten the weight loss landscape to improve the generalization of adversarial training . In particular , Adversarial Weight Perturbation ( AWP ) ( Wu et al. , 2020 ) achieves this by designing a double-perturbation mechanism that adversarially perturbs both inputs and weights . In addition , learning-based smoothing can flatten the landscape and improve the performance ( Chen et al. , 2021 ) . Self-supervised Learning . In the context of non-adversarial training , many self-supervised strategies have been proposed , such as rotation prediction ( Gidaris et al. , 2018 ) , region/component filling ( Criminisi et al. , 2004 ) , patch-base spatial composition prediction ( Trinh et al. , 2019 ) and contrastive learning ( Chen et al. , 2020b ; He et al. , 2020 ) . While self-supervision has also been employed in adversarial training ( Chen et al. , 2020a ; Kim et al. , 2020 ; Yang & Vondrick , 2020 ; Hendrycks et al. , 2019 ) , their methods only use self-supervised learning at training time to regularize the parameters and improve the robust accuracy . By contrast , we propose to perform self-supervised fine-tuning at test time , which we demonstrate to significantly improve the robust accuracy on test images . As will be shown in the experiments , the self-supervised test-time fine-tuning has larger and complementary improvements over the training time self-supervision . Test-time Fine-tuning . Test-time fine-tuning has been used in various fields , such as image superresolution ( Shocher et al. , 2018 ) and domain adaption ( Sun et al. , 2020 ; Wang et al. , 2021 ) . While our work is thus closely related to Test-Time Training ( TTT ) in ( Sun et al. , 2020 ) , we target a significantly different scenario . TTT assumes that all test samples have been subject to the same distribution shift compared to the training data . As a consequence , it incrementally updates the model parameters when receiving new test images . By contrast , in our scenario , there is no systematic distribution shift , and it is therefore more effective to fine-tune the parameters of the original model for every new test mini-batch . This motivates our meta adversarial training strategy , which searches for the initial model parameters that can be effectively fine-tuned in a self-supervised manner . 3 METHODOLOGY . We follow the traditional multi-task learning formulation ( Caruana , 1997 ) and consider a neural network with a backbone z = E ( x ; θE ) and K + 1 heads . One head f ( z ; θf ) outputs the classification result while the other K heads g1 ( z ; θg1 ) , ... , gK ( z ; θgK ) correspond to K auxiliary self-supervised tasks . θ = ( θE , θf , θg1 , · · · , θgk ) encompasses all trainable parameters , and we further define F = f ◦ E ; Gk = gk ◦ E , k = 1 , 2 , ... , K. ( 1 ) Furthermore , letD = { ( x1 , y1 ) , ( x2 , y2 ) , · · · , ( xn , yn ) } denote the training set containing n images with corresponding labels , and D̃ = { ( x̃1 , ỹ1 ) , ( x̃2 , ỹ2 ) , · · · , ( x̃m , ỹm ) } be the test set containing m images . For further illustration , the labels of the test set are shown . However , they are unknown to the networks at test time . We denote the adversarial examples of x and x̃ as xadv and x̃adv , respectively . They satisfy the constraints ‖xadv − x‖ ≤ ε and ‖x̃adv − x̃‖ ≤ ε , and ε is the size of the adversarial budget . For any set S , we represent the average loss over the elements in S as L ( S ) = 1 |S| ∑ si∈S L ( si ) ( 2 ) where |S| is the number of elements in S. We use F ( · ; θ ) and L ( · ; θ ) to denote a network F and a loss function L parametrized by θ , but we often omit the symbol θ to simplify the notation . The general classification loss , such as the cross-entropy , is denoted by Lcls . We use the superscript “ AT ” to denote the adversarial training loss . For example , we define LATcls ( S ) = 1 |S| ∑ xi , yi∈S max ‖xadvi −xi‖≤ε Lcls ( F ( xadvi ) , yi ) . ( 3 ) Because of the poor generalization of adversarial training , the adversarially-trained model may not give the correct label for a large portion of the adversarial test examples . That is , for a test sample x̃adv with ground-truth label ỹ , we may have arg maxj F ( x̃ adv ; θ ) j 6= ỹ , where F ( x̃adv ; θ ) j denotes the output probability of the j-th class . Since we do not know the label ỹ at test time , we can not directly optimize the classification loss Lcls ( F ( x̃adv ) , ỹ ) . We therefore use self-supervision as a substitute to fine-tune the backboneE . We provide more details on our self-supervised test-time fine-tuning scheme in Section 3.1 and introduce our meta adversarial training method in Section 3.2 . In Section 3.3 , we discuss the specific self-supervised tasks used in the experiments , and further analyze our algorithm in Section 3.4 . 3.1 TEST-TIME FINE-TUNING . Our goal is to perform self-supervised learning on the test examples to overcome the overfitting problem of adversarial training . To this end , let us suppose that an adversarially-trained network with parameters θ0 receives a mini-batch of b adversarial test examples B̃adv = { ( x̃adv1 , ỹ1 ) , ( x̃adv2 , ỹ2 ) , · · · , ( x̃advb , ỹb ) } , As the labels { ỹi } bi=1 are not available , we propose to fine-tune the backbone parameters θE by optimizing the loss function LSS ( B̃adv ) = 1 b K∑ k=1 λk b∑ i=1 LSS , k ( Gk ( x̃advi ) ; θE , θgk ) , ( 4 ) which encompasses K self-supervised tasks . Here , LSS , k represents the loss function of the k-th task and { λk } Kk=1 are the weights balancing the contribution of each task . The number of images b may vary from 1 to m. b = 1 corresponds to the online setting , where only one adversarial image is available at a time , and the backbone parameters θE are adapted to every new image . The online setting is the most practical one , as it does not make any assumptions on the number of adversarial test images the network receives . By contrast , b = m corresponds to the offline setting where all adversarial test examples are available at once . It is similar to the transductive learning ( Gammerman et al. , 1998 ; Vapnik , 2013 ) . Note that our online setting differs from the online test-time training described in TTT ( Sun et al. , 2020 ) ; we do not incrementally update the network parameters as new samples come , but instead initialize fine-tuning from the same starting point θ0 for each new test image . Eqn ( 4 ) encourages θE to update in favor of the self-supervised tasks . However , as the classification head f was only optimized for the old backbone E ( · ; θ0E ) , it will typically be ill-adapted to the new parameters θ∗E , resulting in a degraded robust accuracy . Furthermore , for a small b , the model tends to overfit to the test data , reducing LSS to 0 but extracting features that are only useful for the self-supervised tasks . To overcome these problems , we add an additional loss function acting on the training data that both regularizes the backbone E and optimizes the classification head f so that f remains adapted to the fine-tuned backbone E ( · ; θ∗E ) . Specifically , let B ⊂ D denote a subset of the training set . We then add the regularizer LR ( B ) = LATcls ( B ) = 1 |B| ∑ xi , yi∈B max ‖xadvi −xi‖≤ε Lcls ( F ( xadvi ) , yi ) ( 5 ) to the fine-tuning process . In short , Eqn ( 5 ) evaluates the adversarial training loss on the training set to fine-tune the parameters θf of the classification head . It also forces the backbone E to extract features that can be used to make correct prediction , i.e. , to prevent θE from being misled by LSS when b is small . Combining Eqn ( 4 ) and Eqn ( 5 ) , our final test-time fine-tuning loss is Ltest ( B̃adv , B ) = LSS ( B̃adv ) + λLR ( B ) ( 6 ) where λ sets the influence of LR . Algorithm 1 describes our test-time self-supervised learning algorithm in detail . As SGD is more efficient for more large amount of data , we use SGD to optimize θ when b is large ( e.g . offline setting ) . This version of algorithm is deferred to the Appendix B. Algorithm 1 Self-supervised Test-time Fine-tuning Input : Initial parameters θ0 ; Adversarial test images B̃adv = { x̃advi } bi=1 ; Training data D ; Learning rate η ; Steps T ; Weights λk and λ Output : Prediction of x̃advi : ŷi 1 : for t = 1 to T do 2 : Sample a batch of training images B ⊂ D 3 : Find adversarial xadvi of training image xi ∈ B by PGD attack . 4 : Calculate Ltest in Eqn ( 6 ) 5 : θt = θt−1 − η∇θt−1Ltest ( B̃adv , B ; θt−1 ) 6 : end for 7 : return Prediction ŷi = arg maxj F ( x̃advi ; θT ) j

Summary: Fine-tuning the parameters of an adversarially-trained model at test-time with self-supervised tasks improves robustness to standard adversarial attacks. By incorporating meta-learning into the adversarial training stage, the test-time fine-tuning achieves marginally higher accuracy (though not reduced optimization time, or at least this is not reported). The test-time optimization approach requires the training data be kept available for joint updates on training batches and test data. The self-supervised tasks considered are rotation prediction, a common choice, and vertical flip recognition, which is a close relative of it. Results are reported on the standard adversarial defense benchmark of CIFAR-10, and additionally models are also evaluated on STL and Tiny ImageNet. Design choices are ablated and an adaptive attack is designed and experimented with, but there other checks against obfuscated gradients (for instance those suggested by Athalye et al. 2018) are not pursued, such as attacks in expectation, or decision-based attacks. No comment is made about the computational cost of test-time fine-tuning.