Year: 2021

Abstract:
Although adversarial training and its variants currently constitute the most effective way to achieve robustness against adversarial attacks, their poor generalization limits their performance on the test samples. In this work, we propose to improve the generalization and robust accuracy of adversarially-trained networks via self-supervised test-time fine-tuning. To this end, we introduce 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. The extensive experiments on CIFAR10, STL10 and Tiny ImageNet using different self-supervised tasks show that our method consistently improves the robust accuracy under different attack strategies for both the white-box and black-box attacks.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 METHODOLOGY
text: 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 cannot 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.

heading: 3.1 TEST-TIME FINE-TUNING
text: 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

heading: 3.2 META ADVERSARIAL TRAINING
text: To make the best out of optimizing Ltest at test time, we should find suitable starting point θ0, i.e., a starting point such that test-time self-supervised learning yields better robust accuracy. We translate this into a meta learning scheme, which entails a bilevel optimization problem.
Specifically, we divide the training data into s small exclusive subsets D = ∪sj=1Bj and let Badvj to be adversaries of Bj . We then formulate meta adversarial learning as the bilevel minimization of
Lmeta(D; θ) = 1
s ∑ Bj⊂D LATcls (Bj ; θ∗j (θ)), θ∗j = arg min θ LSS(Badvj ; θ) , (7)
where LSS is the self-supervised loss function defined in Eqn (4) and LATcls is the loss function of adversarial training defined in Eqn (3). As bilevel optimization is time-consuming, following MAML (Finn et al., 2017), we use a single gradient step of the current model parameters θ to approximate θ∗j . That is, we compute
θ∗j ≈ θ − α∇θLSS(Badvj ; θ) . (8) In essence, this meta adversarial training scheme searches for a starting point such that fine-tuning with LSS will lead to a good robust accuracy. If this holds for all training subsets, then we can expect the robust accuracy after fine-tuning at test time also to increase. Note that, because the meta learning objective of Eqn (7) already accounts for classification accuracy, the regularization by LR is not needed during meta adversarial learning.
Accelerating Training. To compute the gradient ∇θLmeta(D; θ), we need to calculate the timeconsuming second order derivatives −α∇2θLSS(Badvj ; θ)∇θ∗jL AT cls (Bj ; θ ∗ j ) . Considering that adversarial training is already much slower than standard training (Shafahi et al., 2019), we cannot
afford another significant training overhead. Fortunately, as shown in (Finn et al., 2017), second order derivatives have little influence on the performance of MAML. We therefore ignore them and take the gradient to be
∇θLmeta(D; θ) ≈ 1
s ∑ Bj⊂D ∇θ∗jL AT cls (Bj ; θ ∗ j ) . (9)
However, by ignoring the second order gradient, only the parameters on the forward path of the classifier F , i.e., θE and θf , will be updated. In other words, optimizing Eqn (7) in this fashion will not update {θgk}Kk=1. To nonetheless encourage each self-supervised head Gk to output the correct prediction, we incorporate an additional loss function encoding the self-supervised tasks,
LATSS (D) = ∑ k λkLATSS,k(D) = ∑ k λk |D| ∑ xi∈D max ‖xadvi −xi‖≤ε LSS,k(Gk(xadvi )) . (10)
Note that we use the adversarial version of LSS to provide robustness to the self-supervised tasks, which, as shown in (Chen et al., 2020a; Hendrycks et al., 2019; Yang & Vondrick, 2020), is beneficial for the classifier. The final meta adversarial learning objective therefore is
Ltrain(D) = Lmeta(D) + λ′LATSS (D) (11) where λ′ balances the two losses. Algorithm 2 shows complete meta adversarial training algorithm.
Algorithm 2 Meta Adversarial Training Input: Training set D; Learning rate α, β; Iterations T ; Weights λk and λ′ Output: Starting parameters θ0 for the test-time fine-tuning 1: for t = 1 to T do 2: Sample q exclusive batches of training images B1, B2, · · · , Bq ⊂ D 3: Using PGD to find the adversaries Badvj : x adv j,i = arg max‖xadvj,i −xj,i‖≤ε Lcls(F (x adv j,i ), yj,i)
4: for batches B1, B2, · · · , Bq do 5: θ∗j = θ − α∇θLSS(Badvj ; θ) 6: lmeta,j = LATcls (Bj ; θ∗j ) 7: end for 8: θ = θ − βq ∑ Bj [ ∇θ∗j lmeta,j + λ ′∇θLATSS (Bj ; θ) ]
9: end for 10: return Trained parameters θ0 = θ

heading: 3.3 SELF-SUPERVISED TASKS
text: In principle, any self-supervised tasks can be used for test-time fine-tuning, as long as they are positively correlated with the robust accuracy. However, for the test-time fine-tuning to remain efficient, we should not use too many self-supervised tasks. Furthermore, as we aim to support the fully online setting, where only one image is available at a time, we cannot incorporate a contrastive loss (Chen et al., 2020b; He et al., 2020; Kim et al., 2020) to LSS . In our experiments, we therefore use two self-supervised tasks that have been shown to be useful to improve the classification accuracy: Rotation Prediction and Vertical Flip Prediction.
Rotation Prediction is a widely used self-supervision task proposed in (Gidaris et al., 2018) and has been employed in adversarial training as an auxiliary task to improve the robust accuracy (Chen et al., 2020a; Hendrycks et al., 2019). Following (Gidaris et al., 2018), we create 4 copies of the input image by rotating it with Ω = {0◦, 90◦, 180◦, 270◦}. The task then consists of a 4-way classification problem, where the head grotate aims to predict the correct rotation angle. The loss for an image x is the average cross-entropy over the 4 copies, given by
Lrotate(x) = − 1
4 ∑ ω∈Ω log(Grotate(xω)ω) , (12)
where xω is the rotated image with angle ω ∈ Ω, Grotate = grotate ◦ E denotes the classifier for rotation prediction, and Grotate(·)ω is the predicted probability for the ω angle. The head grotate is a fully-connected layer followed by a softmax layer.
Table 1: The statistics ρ(x̃adv) and two self-supervised tasks. The dataset is CIFAR10 and the network is WideResNet34-10. Adversarial budget ε = 0.031
Tasks E(ρ(x̃adv)) P (ρ(x̃adv) > 0) Rotation 0.15 68.51%
VFlip 0.22 72.16%
0.5 0.0 0.5 1.0 (xadv)
0.0
0.2
0.4
0.6
0.8
1.0 Rotation
1.0 0.5 0.0 0.5 1.0 (xadv)
0.0
0.2
0.4
0.6
0.8
1.0 Vertical Flip
Figure 1: Empirical cdf of ρ(x̃advi ) on CIFAR10 and WideResNet-34-10. Adversarial budget ε = 0.031
Vertical Flip (VFlip) Prediction is a self-supervised task similar to rotation prediction and has also been used for self-supervised learning (Saito et al., 2020). In essence, we make two copies of the input image and flip one copy vertically. The head gvflip then contains a 2-way fully-connected layer followed by a softmax layer and predicts whether the image is vertically flipped or not. The corresponding loss for an image x is
Lvflip(x) = − 1
2 ∑ v∈V log(Gvflip(xv)v) , (13)
where V = {flipped, not flipped} is the operation set and Gvflip = gvflip ◦ E. xv denotes and transformed input and Gvflip(·)v is the probability of operation v. Note that we do not flip the image horizontally as it is a common data augmentation technique and classifiers typically seek to be invariant to horizontal flip.

heading: 3.4 METHOD ANALYSIS
text: We now analyze why our test-time self-supervised fine-tuning improves the robust accuracy. To this end, we observe the correlation between the gradient of self-supervised loss LSS and the classification loss Lcls. For an input x̃advi with label ỹi, the gradient correlation with respect to the shared parameters θE can be expressed as
ρ(x̃advi ) = ∇θELcls(x̃advi , ỹi)T∇θELSS(x̃advi )
‖∇θELcls(x̃advi , ỹi)‖2‖∇θELSS(x̃advi )‖2 . (14)
If ρ(x̃advi ) is significantly larger than 0, then gradient descent w.r.t. LSS should act as a good substitute for optimizing Lcls. Let us consider a small step η in the direction of ∇θELSS(x̃advi ). Then, the classification loss Lcls can be approximated by Taylor expansion as
Lcls(x̃advi , ỹi; θE − η∇θELSS(x̃advi ))− Lcls(x̃advi , ỹi; θE) ≈− ηρ(x̃advi )‖∇θELcls(x̃advi , ỹi)‖2‖∇θELSS(x̃advi )‖2 .
(15)
As θE contains millions of parameters, the gradient norm is typically large, even a small η will decrease the classification loss when ρ(x̃advi ) is significantly larger than 0.
Below, we further confirm this empirically. We choose cross-entropy as the classification loss. For all adversarial test inputs x̃adv ∼ D̃adv , where D̃adv is the adversaries of the test data, we regard ρ(x̃adv) as a random variable and calculate its empirical statistics on the test dataset. Table 1 shows the empirical statistics of an adversarially-trained model on CIFAR10, and Figure 1 shows the cumulative distribution of ρ(x̃adv). The mean of ρ(x̃adv) is indeed significantly larger than 0 and P (ρ(x̃adv) > 0) is larger than the robust accuracy of the adversarially-trained network (50%-60%), which implies that self-supervised test-time fine-tuning helps to correctly classify the adversarial test images.

heading: 4 ADAPTIVE ATTACKS
text: In the white-box attacks, the attacker knows every detail of the defense method. Therefore, we need to assume that the attacker is aware of our test-time fine-tuning method and will adjust its strategy for generating adversarial examples accordingly. Here, we discuss one such strong adaptation strategy targeted to our method.
Suppose that the attacker is fully aware of the hyperparameters for test-time fine-tuning. Then, finding adversaries B̃adv of the clean subset B̃ can be achieved by maximizing the adaptive loss
x̃advi = arg max ‖x̃advi −xi‖≤ε Lattack(F (x̃advi ), y; θT (B̃adv)) , (16)
where Lattack refers to the general attack loss, such as the cross-entropy or the difference of logit ratio (DLR) (Croce & Hein, 2020a). Let θT be the fine-tuned test-time parameters using Algorithm 1. At the k-th step of the attack, it depends on the input B̃(k) = {(x̃(k)j , ỹj)}bj=1 via the update
θt+1 = θt − η∇θtLtest(B̃(k), B) , (17) where Ltest and B are the loss function and subset of training images mentioned in Eqn (6). As θT is a function of the input B̃(k), we can calculate the end-to-end gradient of x̃(k)i ∈ B̃(k) as ∇ x̃ (k) i Lattack(F (x̃(k)i ); θT (B̃(k))). However, θT goes through T gradient descent steps, and thus calculating the gradient∇ x̃ (k) i θT (B̃(k)) requires T -th order derivatives of the backbone E, which is virtually impossible if T or the dimension of θE is large. We therefore approximate the gradient as
Grad(x̃ (k) i ) ≈ ∇x̃(k)i Lattack(F (x̃ (k) i ); θ T ) , (18)
which treats θT as a fixed variable so that high-order derivatives from θT (B̃(k)(x̃(k)i )) can be avoided. Although this approximation makes Grad(x̃(k)i ) inaccurate, common white-box attacks use projected gradients, which are robust to such inaccuracies. For example, PGD only uses the sign of the gradient under an `∞ adversarial budget. Note that solving the maximization in Eqn (16) does not necessarily require calculating the gradient Grad(x̃(k)i ). For instance, we will also use Square Attack (Andriushchenko et al., 2020), a strong score-based black-box attack, to maximize Eqn (16) and generate adversaries for B̃.
As another approximation to save time, one can also fixing θT for several iterations. This leverages the intuition that attack strategies often make small changes to the input x̃, and thus, for the intermediate images in the k-th and (k+1)-th steps, θT (B̃(k)) and θT (B̃(k+1)) should be close. Therefore, a general version of our adaptive attacks only updates θT every u iterations, with u a hyperparameter. We provide the algorithms for PGD, AutoPGD, FAB and Square Attack in the Appendix C.

heading: 5 EXPERIMENTS
text: Experimental Settings. Following previous works (Cui et al., 2020; Huang et al., 2020), we consider `∞-norm attacks with an adversarial budget ε = 0.031(≈ 8/255). We evaluate our method on three datasets: CIFAR10 (Krizhevsky et al., 2009), STL10 (Coates et al., 2011) and Tiny ImageNet (Le & Yang, 2015). We also use two different network architectures: WideResNet-3410 (Zagoruyko & Komodakis, 2016) for CIFAR10, and ResNet18 (He et al., 2016) for STL10 and Tiny ImageNet. Following the common settings in adversarial training, we train the network for 100 epochs using SGD with a momentum factor of 0.9 and a weight decay factor of 5× 10−4. The learning rate β starts at 0.1 and is divided by a factor of 10 at the 50-th and the 75-th epochs. The step size α in Eqn (8) the same as β; the factor λ′ in Eqn (11) is 1.0. We set the weight of each self-supervised task λk = 1K where K is the number of tasks. We use 10-iteration PGD with step size 0.007(≈ 0.031/4) to find the adversarial Badvj during training. To prevent the computational burden of meta adversarial training, we use divide the classification of 32 images as one task and sample 8 mini-batches B1, ..., B8 in each iteration.(i.e. |Bj | = 32 and q = 8 in Algorithm 2). More detailed experimental settings are provided in the Appendix B.
Attack Methods. We consider four common white-box and black-box attack methods: PGD20 (Madry et al., 2018), AutoPGD (Croce & Hein, 2020a), FAB (Croce & Hein, 2020b) and Square Attack (Andriushchenko et al., 2020). We apply both the standard and adaptive versions of these methods. Furthermore, for AutoPGD, we use both the cross-entropy and DLR (Croce & Hein, 2020a) loss. More details are provided in the Appendix C.
Baselines. We compare our method with the following methods: 1) Regular adversarial training (Regular AT), which uses LATcls in Eqn (3). 2) Regular adversarial training with an additional selfsupervised loss, i.e., using LATcls + λ′LATSS for adversarial training, where LATSS is given in Eqn (10).
This corresponds to the formulation of (Hendrycks et al., 2019). 3) Meta adversarial learning (Algorithm 2) without test-time fine-tuning.
We include additional figures, experiments with different adversarial budget and experiments about label leaking in Appendix A.

heading: 5.1 ONLINE TEST-TIME FINE-TUNING
text: Let us first evaluate our method in the online setting, where only one adversarial image is available when fine-tuning the network. We fine-tune the network for 10 steps with a momentum factor of 0.9 and a learning rate η = 5× 10−4. We set λk = 1K and λ = 15.0. CIFAR10. Table 2 shows the robust accuracy for different attacks and using two different tasks for fine-tuning. The adaptive attack is not applicable to models without fine-tuning. As we inject different self-supervised tasks into the adversarial training stage, and as different self-supervised tasks may impact the robust accuracy differently (Chen et al., 2020a), the robust accuracy without fine-tuning still varies. The vertical flipping task yields better robust accuracy before fine-tuning but its improvements after fine-tuning is small. By contrast, rotation prediction achieves low robust accuracy before fine-tuning, but its improvement after fine-tuning is the largest. Using both tasks together combines their effect and yields the highest overall accuracy after test-time fine-tuning.
Note that our self-supervised test-time fine-tuning, together with meta adversarial learning, consistently improves the robust accuracy under different attack methods. Under the strongest adaptive AutoPGD, test-time fine-tuning using both tasks achieves a robust accuracy of 57.70%, significantly outperforming regular adversarial training.
STL10 and Tiny ImageNet. As using both the rotation and vertical flip prediction led to the highest overall accuracy on CIFAR10, we focus on this strategy for STL10 and Tiny ImageNet. Table 3 and 4 shows the robust accuracy on STL10 and Tiny ImageNet using a ResNet18. Our approach also significantly outperforms regular adversarial training on these datasets.

heading: 5.2 OFFLINE TEST-TIME FINE-TUNING
text: Let us now study the offline setting, where all test adversarial examples are available at test time. We use SGD to fine-tune the model for 10 epochs. We set λ = 10.0 and keep the other settings the same as for online fine-tuning. As shown in Table 2, 3, 4, the offline fine-tuning further improves the robust accuracy over the online version.

heading: 5.3 ABLATION STUDY
text: Removing LSS or LR. In our previous experiments, test-time fine-tuning was achieved using a combination of two loss functions: LSS and LR. To study the effect of each of these terms separately, we remove either one of them from Ltest. In Table 5, we report the robust accuracy after online fine-tuning using only LR and only LSS . While, as expected, removing LSS tends to reduce more accuracy than removing LR. It shows the benefits of our self-supervised test-time fine-tuning strategy. Nevertheless, the best results are obtained by exploiting both loss terms.
Removing Meta Adversarial Training. Our meta training strategy in Algorithm 2 aims to strengthen the correlation between the self-supervised tasks and classification. To show its effectiveness, we perform an ablation study where we fine-tune the model with regular adversarial training (i.e., setting α = 0 in the line 5 of Algorithm 2). We then perform the same test-time fine-tuning on the model without meta adversarial training, using the same hyperparameters as in the meta adversarial training case. As shown in Table 6, the robust accuracy and the improvements of fine-tuning are consistently worse without meta adversarial training.

heading: 6 CONCLUSION
text: We propose self-supervised test-time fine-tuning on adversarially-trained models to improve their generalization ability. Furthermore, we introduce a meta adversarial training strategy to find a good starting point for our self-supervised fine-tuning process. Our extensive experiments on CIFAR10, STL10 and Tiny ImageNet demonstrate that our method consistently improves the robust accuracy under different attack strategies, including strong adaptive attacks where the attacker is aware of our test-time fine-tuning technique. In these experiments, we utilize three different sources of selfsupervision: rotation prediction, vertical flip prediction and the ensemble of them. While all settings lead to improved robust accuracy, some self-supervised tasks seem to have more impact than others. Finding better self-supervised tasks and their combinations will be the focus of our future work.
REPRODUCIBILITY STATEMENT
We provide pesudocode of our algorithms and the detailed hyperparameters in the main text and Appendix. The datasets we use (CIFAR10, STL10 and Tiny ImageNet) are all publicly available. In addition, we provide the code to reproduce the experiment results here.
ETHICS STATEMENT
Adversarial examples bring security threats to the current machine learning models. Adversarial examples may be maliciously used to attack the safety-critical systems like autonomous driving and face verification. While adversarial training is effective in defending against adversarial attacks, it suffers from poor generalization. We propose to use the self-supervised test-time fine-tuning to increase the robust accuracy of these models. It helps build safer machine learning systems.

heading: A ADDITIONAL EXPERIMENTS
text: Robust Accuracy v.s Fine-tuning Steps. Figure 2 shows the robust accuracy at each step of the test-time fine-tuning for different self-supervised tasks and attack methods. When using the standard version of attacks, the robust accuracy gradually increases as fine-tuning proceeds. When using our adaptive attacks, the adversarial examples are generated to attack the network with θT (T = 10) instead of θ0. Thus, when the parameters gradually change from θ0 to θT , the accuracy drops.
Accuracy Improvement on Clean Images. Our model is able to improve not only the robust accuracy but also the natural accuracy of clean images on adversarially-trained models. To evidence this, we maintain all the components of our model and simply replace the adversarial input images with clean images (i.e. replacing B̃adv with clean inputs B̃ in Algorithm 1) and perform the same self-supervised test-time fine-tuning.
As shown in Table 7, our approach increases the clean image accuracy. This phenomenon further strengthens our conjecture that the improvement of robust accuracy is due to the improvement of generalization instead of perturbing the model parameters, because the randomly perturbing the parameters usually lower the natural accuracy of the model. This also suggests that our method is applicable to the broader problem of improving the accuracy of deep learning models with poor generalization ability.
Accuracy Improvement on Inputs with Different Adversarial Budget. Our method is also able to improve the robust accuracy of inputs with different adversarial budgets. As shown in Table 8, we set `∞ budget of the adversarial inputs to be 0.015 to perform the online test-time fine-tuning. The robust accuracy is ed improved.
Label Leaking. One potential concern about our method is whether the effect of label leaking leads to the improvement of the robust accuracy. The attacker uses the ground truth label to generate the adversarial examples. Someone may worry it leaks the information of the ground truth label to the adversarial images and the test-time fine-tuning can utilize the information to guess the correct label.
Previous experiments on clean images already show that test-time fine-tuning is effective even if there is no information of the ground truth label. We further resolve the concern of the label leaking by the following experiments. The attacker randomly lowers the score of the false label to perform the adversarial attack. And we conduct the self-supervised test-time fine-tuning on these “adversarial” images. If our method uses the information of leaked label to improve the robust accuracy, it will predict the false label and reduce the accuracy. However, as shown in Table 9, the robust accuracy is also improved after the test-time fine-tuning. It demonstrates that the improvement of robust accuracy is not a result of label leaking. It is worth noting that such attack is not targeted at the ground truth label. Therefore, the accuracy of the input images, along with the accuracy improvement of the test-time fine-tuning, are similar to the clean images in Table 7.
Comparison with SOAP (Shi et al., 2021). Our method is different from SOAP as we are finetuning the model to adapt to new examples instead of purifying the input. We apply SOAP-RP to the adversarially-trained model and find that its improvement is marginal. Under AutoPGD, the accuracy is improved from 53.09% to 53.57%. This improvement is much smaller than our method, whose improvement is from 53.09% to 57.93%. SOAP only has little effect when combining with the commonly-used adversarial training.
Transfer Attack. In Table 10, we perform a transfer attack from the static adversarial defense. We use the robust networks with the same architecture as the substitute model, and the test-time fine-tuning also improves the robust accuracy.
Expectation Attack. In Table 11, we show the results of the expectation attack. We modify the adaptive attack and average the gradient from 10 fine-tuned models, whose training batches are different. We evaluate the model using the ensemble of rotation and vertical flip as the selfsupervised task on CIFAR10. We evaluate the model with Adaptive-AutoPGD-EOT and AdaptiveSquareAttack-EOT. One is the strongest attack in our method and the other is a black-box attack that is less likely to be affected by gradient masking. The experiment shows that the expectation attack has little influence on the improvement of our test-time fine-tuning.
Boundary Attack. We use one of the SOTA decision-based attacks: RayS (Chen & Gu, 2020). We test it on CIFAR10 with the ensemble of rotation and vertical flip. Table 12 shows that our method also improves the robust accuracy for the decision-based attack.
Combination with (Gowal et al., 2020). We combine our test-time fine-tuning with the adversarial training using additional data (Gowal et al., 2020). We apply our Meta AT to it with the ensemble of rotation and vertical flip. Using a WideResNet-28-10, it achieves robust accuracy 62.07% under AutoPGD. With our test-time fine-tuning, the robust accuracy is improved to 64.34%. The improvement of robust accuracy is 2.27%.
Inference Time. Table 13 shows the inference time for different methods. While the inference time for our method is larger than SOAP and normal method when the batch size is 1, the inference time gets closer when using larger batch size. And the batch size of 20 or more is a common scenario of the inference.

heading: B DETAILS OF OUR EXPERIMENTAL SETTING
text: Meta Adversarial Training. We consider an `∞ norm with an adversarial budget ε = 0.031. We also use two different network architectures: WideResNet-34-10 for CIFAR10 and ResNet18 for STL10 and Tiny ImageNet. Following the common settings for adversarial training, we train the network for 100 epochs using SGD with a momentum factor of 0.9 and a weight decay factor of 5× 10−4. The learning rate β starts at 0.1 and is divided by a factor of 10 after the 50-th and again after the 75-th epochs. The step size α in Eqn (8) is the same as β. The factor λ′ in Eqn (11) is set to 1.0. We use 10-iteration PGD (PGD-10) with a step size of 0.007 to find the adversarial image Badvj at training time. The weight of each self-supervised task is set to λk = 1 K . We set |Bj | = 32 and sample 8 batches B1, ..., B8 in each iteration. Furthermore, we save the model after the 51-st epoch for further evaluation, as the model obtained right after the first learning rate decay usually yields the best performance (Rice et al., 2020).
We use PGD with the standard cross-entropy loss to generate adversarial examples at training time in the line 3, line 6 and line 8 of Algorithm 2. The hyperparameters of the attacks are as follows:
• Line 3: PGD-10 with step size 0.007.
• Line 6: As θ∗j is similar to θ, the adversarial examples at this step are similar to those at Line 4. To save training time, we therefore choose the starting point of the attack as the adversarial examples in Line 4 and use PGD-2 with a step size of 0.005.
• Line 8: PGD-3 with step size 0.02.
Online Test-time Fine-tuning. We fine-tune the network for T = 10 steps with a momentum of 0.9 and a learning rate of η = 5 × 10−4. We set λk = 1K and λ = 15.0. In line 2 of Algorithm 1, we sample a batch B ⊂ D containing 20 training images. In line 3, we use PGD-10 with a step size of 0.007.
Offline Test-time Fine-tuning. The algorithm for offline fine-tuning is shown in Algorithm 3. As we stochastic gradient descent is more efficient for large amount of data, we use stochastic gradient
descent in the offline fine-tuning. This is the main different between Algorithm 1 (online fine-tuning) and Algorithm 3 (offline fine-tuning). We also fine-tune the network for 10 epochs. The batch size of each B̃advj is 128. The other hyperparameters are the same as in the online version.
Algorithm 3 Self-supervised Test-time Fine-tuning with SGD 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: Divide B̃adv into r subsets B̃adv1 , ..., B̃ adv r
3: for B̃advj in B̃adv1 , ..., B̃advr do 4: Sample a batch of training images B ⊂ D 5: Find adversarial xadvi of training image xi ∈ B by PGD attack. 6: θt = θt−1 − η∇θt−1Ltest(B̃advj , B; θt−1) 7: end for 8: end for 9: return Prediction ŷi = arg maxj F (x̃advi ; θT )j
Attacks. The detailed settings of each attack are provided below:
• PGD-20. We use 20 iterations of PGD with step size γ = 0.003. The attack loss is the cross-entropy.
• AutoPGD. We use both the cross-entropy and the difference of logits ratio (DLR) as the attack loss. The hyperparameters are the same as in (Croce & Hein, 2020a).
• FAB. We use the code from (Croce & Hein, 2020a) and keep the hyperparameters the same. • Square Attack. We set T = 2000 and the initial fraction of the elements p = 0.3. The other
hyperparameters are the same as in (Andriushchenko et al., 2020).
For the adaptive versions, we set the interval u = dT/5e.

heading: C ALGORITHM FOR ADAPTIVE ATTACKS
text: In Algorithm 4, 6, 5 and 7, we show the algorithms for `∞ norm-based adaptive PGD, AutoPGD, Square Attack and FAB, respectively. The main difference between the original and adaptive version is the target loss function for maximization. The reader may refer to (Andriushchenko et al., 2020; Croce & Hein, 2020a;b) for more detailed description of the steps in these algorithms (e.g., the condition for decreasing the learning rate in AutoPGD).
Algorithm 4 `∞ Norm Adaptive PGD Attack
Input: Test images B̃ = {(x̃i, ỹi)}; Attack loss Lattack; Step size γ; Iterations T ; Intervals u; Adversarial budget ε; Trained parameters of the network θ0.
Output: Adversarial images B̃adv = {x̃advi } 1: Add random noise to x̃i in B̃ and get B̃′ 2: for t = 1 to T do 3: if t mod u = 0 then 4: Get final parameters θT by taking B̃′ as input image for Algorithm 1: θ = θT 5: end if 6: for x̃′i in B̃′ do 7: Grad(x̃′i) = ∇x̃′iLattack(F (x̃ ′ i), ỹi; θ)
8: x̃′i = Clip[x̃i−ε,x̃i+ε](x̃ ′ i + γSign(Grad(x̃ ′ i)))
9: end for 10: end for 11: return Adversarial image x̃advi = x̃′i
Algorithm 5 `∞ Norm Adaptive AutoPGD
Input: Test images B̃ = {(x̃i, ỹi)}; Attack loss Lattack; Step size γ; Iterations T ; Intervals u; Adversarial budget ε; Parameter of the adversarially-trained network θ0; Decay iterations W = {w0, ..., wn}; Momentum ξ Output: Adversarial image B̃adv = {x̃advi } 1: Get final parameter θT by taking B̃ as input image for Algorithm 1. 2: θ = θT
3: for x̃i in B̃ do 4: x̃0i = x̃i 5: Grad(x̃i) = ∇x̃iLattack(F (x̃i), ỹi; θ) 6: x̃1i = Clip[x̃i−ε,x̃i+ε](x̃ 0 i + γSign(Grad(x̃i))) 7: l0i = Lattack(F (x̃0i ), ỹi; θ) 8: l1i = Lattack(F (x̃1i ), ỹi; θ) 9: l∗i = max{l0i , l1i }
10: x̃∗i = x̃ 0 i if l∗i = l0i else x̃∗i = x̃1i 11: end for 12: for t = 1 to T − 1 do 13: if t mod u = 0 then 14: Get final parameter θT by taking B̃∗ = {x̃∗i } as input image for Algorithm 1. 15: θ = θT 16: end if 17: for i = 1, ..., |B̃| do 18: Grad(x̃ti) = ∇x̃tiLattack(F (x̃ t i), ỹi; θ) 19: zt+1i = Clip[x̃i−ε,x̃i+ε](x̃ t i + γSign(Grad(x̃ t i))) 20: x̃t+1i = Clip[x̃i−ε,x̃i+ε](x̃ t i + ξ(z t+1 i − zti) + (1− ξ)(x̃ti − x̃ t−1 i )) 21: lt+1i = Lattack(F (x̃ t+1 i ), ỹi; θ) 22: x̃∗i = x̃ t+1 i and l ∗ i = l t+1 i if l t+1 i > l ∗ i 23: if k ∈W and satisfy the condition of dropping learning rate then 24: γ = γ/2 and x̃t+1i = x̃ ∗ i 25: end if 26: end for 27: end for 28: return Adversarial image x̃advi = x̃∗i
Algorithm 6 `∞ Norm Adaptive Square Attack
Input: Test images B̃ = {(x̃i, ỹi)}; Attack loss Lattack; Step size γ; Iterations T ; Intervals u; Image size w; Color channels c; Adversarial budget ε; Parameter of the adversarially-trained network θ0. Output: Adversarial image B̃adv = {x̃advi } 1: Add noise to x̃i in B̃ and get B̃′ 2: for t = 1 to T do 3: if t mod u = 0 then 4: Get final parameter θT by taking B̃′ as input image for Algorithm 1. 5: θ = θT
6: end if 7: for x̃′i in B̃′ do 8: ht ←− side length of the square to modify (according to some schedule) 9: δ ←− array of zeros of size w × w × c
10: Sample uniformly r, s ∈ {0, ..., w − ht} ⊂ N 11: for j = 1, ..., c do 12: ρ←− Uniform(−2ε, 2ε) 13: δr+1:r+ht,s+1:s+ht = ρ · 1ht×ht 14: end for 15: x̃newi = Clip[x̃i−ε,x̃i+ε](x̃ ′ i + δ) 16: lnewi = Lattack(F (x̃newi ), ỹi; θ) 17: if lnew < l∗ then 18: x̃′i = x̃ new i 19: l∗i = l new i 20: end if 21: end for 22: end for 23: return Adversarial image x̃advi = x̃′i
Algorithm 7 `∞ Norm Adaptive FAB
Input: Test images B̃ = {(x̃i, ỹi)}; Step size γ; Iterations T ; Intervals u; Adversarial budget ε; Trained parameters of the network θ0; αmax, η, β. Output: Adversarial images B̃adv = {x̃advi } 1: Add random noise to x̃i in B̃ and get B̃′ 2: v = +∞ 3: for t = 1 to T do 4: if t mod u = 0 then 5: Get final parameters θT by taking B̃′ as input image for Algorithm 1: θ = θT 6: end if 7: for x̃′i in B̃′ do 8: Grad(x̃′i)l = ∇x̃′iF (x̃ ′ i; θ)l
9: s = arg minl 6=yi |F (x̃′i;θ)l−F (x̃ ′ i;θ)yi |
‖Grad(x̃′i)l−Grad(x̃′i)yi‖1 10: δt = proj∞(x̃ ′ i, πs, C) 11: δtorig = proj∞(x̃i, πs, C)
12: α = min
{ ‖δt‖∞
‖δt‖∞+‖δtorig ‖∞ , αmax
} ∈ [0, 1]
13: x̃′i = projC ( (1− α)(x̃′i + ηδt) + α(x̃i + ηδtorig) ) 14: if x̃′i is not classified as yi then 15: if ‖x̃′i − x̃i‖∞ < v then 16: x̃advi = x̃ ′ i 17: v = ‖x̃′i − x̃i‖∞ 18: end if 19: x̃′i = (1− β)x̃i + βx̃′i 20: end if 21: end for 22: end for 23: return Adversarial image x̃advi

Review summaries:
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.

This paper presents a nice framework of test-time fine tuning through self-supervision for adversarially trained networks with the purpose of improving the robust accuracy on test data. A meta adversarial training strategy is also proposed that strengthens the correlation between the self-supervised and classification tasks to provide a good starting point for test-time fine tuning. The proposed method improved robustness performance against different attack methods including adaptive attacks where the attacker has knowledge of the fine tuning technique.

Propose to improve the generalization and robust accuracy of adversarially-trained networks via self-supervised test-time fine-tuning. Introduce a meta adversarial training method to find a good starting point for test-time fine-tuning, which  incorporates the test-time fine-tuning procedure into the training phase. Observe consistent improvement under different attack strategies for both white-box and black-box attacks.

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.