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: The paper introduces a meta adversarial training method to find a good starting point for test-time fine-tuning. It incorporates the test-time fine-tuning procedure into the training phase and strengthens the correlation between the self-supervised and classification tasks. Results on CIFAR-10, STL10 and Tiny-ImageNet show consistent robustness improvement.