Vanilla Feature Distillation for Improving the Accuracy-Robustness Trade-Off in
Adversarial Training

Since Adversarial attacks were first introduced in [10], many attack algorithms have been proposed to analyze the vulnerability of deep neural networks. This line of research could be roughly divided into two categories according to the knowledge of the adversary, i.e., white-box attacks and black-box attacks. Some studies focus on white-box attacks, where the adversary has full access to the parameters of the target model. With the full knowledge of the target model, attacks can be conducted by Fast Gradient Sign Method (FGSM)[10], Momentum Iterative Method (MIM)[8], Projected Gradient Descent (PGD)[20], and many other methods [28, chakraborty2018adversarial]. Besides, several works focus on black-box attacks, in which the adversary cannot access the target model[19]. In this case, the adversary carries out the attack with the transferability of adversarial examples, i.e., adversarial examples crafted from one model can effectively attack the other model. Without knowing the parameters of the target model, Carlini Wagner Attack (C$\&$W)[3], GAP [24], Feature Importance-aware Attack (FIA)[31], etc. can be implemented[2, 23]. In this paper, we aim to help the target model achieve a better trade-off between accuracy and robustness in both white-box attacks and black-box attacks.

In order to defend against adversarial attacks, adversarial training[20, 30, 34] was proposed to improve the robustness of a model by augmenting the training dataset with adversarial examples when training. Zhang et al.[34] theoretically proved a trade-off between robustness and accuracy in adversarial training and proposed TRADES to trade adversarial robustness against accuracy. Based on the TRADES, Wang et al.[30] investigated the distinctive influence of misclassified and correctly classified examples and proposed Misclassification Aware adveRsarial Training (MART). These methods have successfully enhanced the robustness of deep models, while both of them failed to provide a model with high robustness and high accuracy. Recently, Some researchers found that adversarial training can be improved with knowledge distillation[5, 7, 25]. Helper-based adversarial training (HAT)[25] and Learnable Boundary Guided Adversarial Training (LBGAT)[7] achieved a better trade-off with the knowledge transferred from third-party models. [5] used smoothed labels from Knowledge Distillation (KD) to calibrate the notorious overconfidence of logits generated by pre-trained models. However, they only distill from the logits of the teacher model, while it is difficult for the training model to learn those high accuracy features from the teacher model. By contrast, our method inherits the high accuracy knowledge from the vanilla model by approximating the extracted features to the vanilla features, achieving a better trade-off between accuracy and robustness.

Since adversarial vulnerability is caused by non-robust and sensitive features (non-semantic and imperceptible features, easily tempered by the adversary), adversarial training serves as a typical method to improve robustness against adversarial attacks by incorporating adversarial examples into the process of model training, forcing the model to focus more on robust features thus become immune to tempered adversarial examples. However, those non-robust features are highly predictive that can be exploited by machines, and omitting them leads to large accuracy degradation on clean examples.

To solve this problem, in this paper, we propose Vanilla Feature Distillation Adversarial Training (VFD-Adv) to realize a better trade-off between accuracy and robustness by distilling predictive representations from a high accuracy vanilla model. The basic idea is that, preserving those non-robust but predictive features in adversarial training can mitigate accuracy degradation. To achieve this goal, we constrain features from the adversarial training model that takes adversarial examples and clean examples as input and make them similar to those from the vanilla model fed with the same clean examples. Therefore, those predictive features will be preserved, achieving high accuracy and robustness at the same time.

The pipeline of VFD-Adv is overviewed in Fig.1, where there are two components: 1) the high accuracy vanilla model $F_{van}$ trained only with merely clean examples, and 2) the adversarial training model $F_{adv}$ that needs to be updated in VFD-Adv. A vanilla model can be one of those off-the-shelf models or pre-trained by ourselves. And the adversarial training model is the model we want after training of VFD-Adv. Please note that the vanilla model and the adversarial training model have the same architecture, while parameters of the vanilla model should be freezed once pre-trained. Sharing the spirit of general knowledge distillation and adversarial training during the training process, the vanilla model provides vanilla features for clean examples, while the adversarial training model extracts features for both clean examples and adversarial examples. By matching features to vanilla features, the adversarial training model can learn to preserve those predictive features in adversarial training, thus achieving high accuracy and robustness at the same time.

In this part, we detail the loss functions of the above mentioned VFD-Adv. As illustrated in the overview, we design three loss functions: 1) the clean loss $\mathcal{L}_{clean}$ for supervision of clean examples, and 2) the robustness loss for adversarial training $\mathcal{L}_{adv}$, and 3) the self-distillation loss $\mathcal{L}_{kd}$ for guiding adversarial training towards higher accuracy.

The clean loss $\mathcal{L}_{clean}$ : In the training of VFD-Adv, since our aim is to get a model $F_{adv}$ with high accuracy and high robustness, we take clean examples and their labels for one of supervisions as general adversarial training. This loss can be expressed as:

where $x_{i}$ denotes a clean example, $y_{i}$ is ground truth for $x_{i}$, $\mathcal{J}(\cdot,\cdot)$ is cross-entropy for multi-label classification.

The robustness loss $\mathcal{L}_{adv}$: Meanwhile, we need to improve the robustness of model $F_{adv}$. A typical way to improve robustness is adversarial training, i.e., incorporating adversarial examples into the process of model training. In this way, the adversarial training model learns to correctly classify the adversarial examples. The loss for adversarial training can be expressed as:

where $\hat{y_{i}}$ can be either soft labels (e.g., logits in TRADES[34]) or hard labels (ground truth in ALP[14]), $\mathcal{\phi}(\cdot,\cdot)$ is a distance metric (e.g., Kullback-Leibler Divergence in TRADES[34], cross-entropy in [14]). $\mathcal{L}_{adv}$ itself tends to force the adversarial model to omit non-robust but predictive features of data, reaching a dilemma between gaining more robustness and preserving higher accuracy.

The self-distillation loss $\mathcal{L}_{KD}$ : The design of $\mathcal{L}_{KD}$ begins with knowledge distillation. Knowledge distillation is one of the most popular techniques used to transfer knowledge from one model to another. In particular, when the architectures are identical, this is called self-distillation. The self-distilled model achieves higher accuracy on held-out data [21, 35]. The basic idea is that if we constrain features from the adversarial training model that takes adversarial examples and clean examples as input and make them similar to those from the vanilla model fed with the same clean examples, non-robust but predictive features will be preserved during adversarial training thus mitigate accuracy degradation. Therefore, to get those predictive features, we pre-trained a vanilla model that is optimized towards high accuracy only with clean examples. Then, we distill knowledge from its intermediate layer. More specifically, we call this knowledge vanilla features (i.e., output at layer $k$ of the model for clean examples in our setting). The design of $\mathcal{L}_{KD}$ is defined as:

where $\Gamma(\cdot,\cdot)$ is a distance metric, in this paper we use Euclidean distance. $k$ denotes layer number, $F_{adv}^{k}$ is the output of adversarial training model at layer $k$, $F_{van}^{k}$ is the output of clean examples in vanilla model at layer $k$ (i.e., vanilla features). The first term aims to match features of adversarial examples extracted by the adversarial model to clean counterparts in feature space, while the second term will force features of clean examples extracted by the adversarial model to approximate vanilla features. Different from traditional self-distillation, $\mathcal{L}_{KD}$ forces both features of clean examples $F_{adv}^{k}(x)$ and features of adversarial examples $F_{adv}^{k}(x^{adv})$ to match vanilla features, thus adversarial model will not omit non-robust features and mitigate accuracy degradation.

In this part, we demonstrate the training strategy of VFD-Adv. Instead of proposing a new adversarial training paradigm, our method is a plug-in for existing adversarial training methods. Based on  Eq. 3, Eq. 4 and Eq. 5, we first pre-train a high accuracy vanilla model $F_{van}$ with only clean examples, which architecture is the same as the adversarial training model $F_{adv}$. Then, we get the output at the $k_{th}$ layer of clean examples as vanilla features $F^{k}_{van}(x)$. At the same time, the corresponding output for adversarial training model of adversarial examples $F^{k}_{adv}(x_{adv})$ and clean examples $F^{k}_{adv}(x)$ are extracted. Then, we update the model under the supervision of $\mathcal{L}_{clean}$ to achieve better accuracy, adversarial training loss $\mathcal{L}_{adv}$ to improve robustness and self-distillation loss $\mathcal{L}_{kd}$ to mitigate accuracy degradation. Overall, the objectives of VFD-Adv can be formulated as follows:

where hyperparameters $\beta$ and $\lambda$ are used to adjust the contribution of each term in the total loss. A large value of $\beta$ will force the model to update towards minimizing loss on adversarial examples, and the robustness of the model will increase while the accuracy on clean examples will decrease. $\lambda$ controls feature difference between extracted features and vanilla features. Detailed training algorithm of VFD-Adv can be found in Algorithm 1.

Datasets and Classifiers. For a fair comparison, we follow the previous works[30, 7, 5] to conduct experiments on CIFAR10 and CIFAR100[15] . The CIFAR10 dataset has 10 classes and per class contains 6,000 images, while the CIFAR100 dataset has 100 classes and per class contains 600 images. For each dataset, We split it into one training set and one test set in a ratio of 5:1. In addition, we evaluate our method on ResNet18[11] and VGG16[27] which are widely used in adversarial robustness benchmarks.

Attack Methods. We consider three typical attack methods, i.e., PGD (step=20, $\ell_{\infty}$, $\varepsilon$=8/255), FGSM ($\ell_{\infty}$, $\varepsilon$=8/255) and C$\&$W (kappa=0, step=1000, lr=0.01). Moreover, we evaluate the robustness of the model in two attack settings (i.e., white-box setting and black-box setting), and we use WidResNet[33] as the surrogate model to generate adversarial examples in black-box setting.

Baselines and Parameter Settings. To verify the effectiveness of the proposed method, we compare it with four related state-of-the-art adversarial training methods (i.e., ALP[14], TRADES[25], MART[30], HAT[25]) and three enhanced methods (i.e., KD[5], BGAT and LBGAT[7]) which can further improve the performance. For these methods, we follow their default parameter settings and use the SGD optimizer[6] with a learning rate of 0.1 and a batch size of 128 to train corresponding robust models. We conduct all the experiments on a server with one Nvidia Tesla V100 16GB GPU card.

Fig. 2 shows the comparison of the accuracy vs. robustness trade-off between the original TRADES and its corresponding enhanced versions. We vary the robustness parameter $\beta$ in TRADES to obtain models with different robustness. We can observe that the models trained by different methods exhibit similar robustness when using the same $\beta$, but our method (i.e., TRADES+VFD-Adv) outperforms the others significantly in clean accuracy, achieving a better trade-off between accuracy and robustness.

To quantitatively compare the performance between the proposed VFD-Adv and the baselines, we present the robust and clean accuracy of corresponding trained robust ResNet18 under several typical attack methods. Table 1 shows the results on CIFAR10 and Table 2 shows the results on CIFAR100.

As shown in Table 1, our method significantly improves the robust and clean accuracy of the original adversarial training methods in most cases (white-box attack or black-box attack). For example, compared with the original TRADES, the TRADES+VFD-Adv (i.e., equipped with our proposed method) can further improve the robust accuracy by 4.3% and clean accuracy by 6.4% on average. Meanwhile, the TRADES+VFD-Adv also outperforms those state-of-the-art methods (e.g., equipped with LBGAT, etc.) 3.1% in robust accuracy and 4.4% in clean accuracy on average, achieving the best performance. Moreover, we can draw the same conclusions when combining our method with other adversarial training methods (i.e., ALP, MART, HAT), which indicates the scalability and effectiveness of our method.

In addition, we also present the comparison results on a more complex classification task (i.e., CIFAR100) in Table 2. The results show that our method still can improve the robust accuracy by 2.3% and clean accuracy by 4.5% on average compared with the other state-of-the-art methods, demonstrating the superior performance and the generality of our proposed method on different datasets. In conclusion, the proposed method mitigates the accuracy degradation significantly and improve the robustness. Meanwhile, the proposed method also can be adapted to existing related works flexibly as an extra effective regularization item. We show the experimental results of VGG16 in supplementary file, which show better improvement compared to that on ResNet18.

Furthermore, to qualitatively analyze the effectiveness of our proposed VFD-Adv, we randomly select 5 classes in CIFAR10 and utilize t-SNE[29] to visualize features of 500 clean examples per class and corresponding adversarial examples generated by PGD. As shown in Fig. 3, the first row presents the t-SNE results of clean examples and the second row presents the results of corresponding adversarial examples. For the vanilla model, we can observe that there are 5 clear separable clusters in the case of clean examples but 10 clusters in the case of adversarial examples. The phenomenon indicates that the vanilla model cannot extract right discriminate features from these adversarial examples hence classify them into wrong classes. For all existing defense methods (e.g., TRADE), features of adversarial examples are separated into 5 clusters to some degree, but we can observe that there would appear unexpected obscure boundaries in clean examples, thus leading to degradation in clean accuracy. In contrast, as shown in the most right column, our model can force features of both adversarial examples and clean examples of the adversarial model close to vanilla features, thus realizing clear boundaries in these clusters for both clean examples and adversarial examples. These results are consistent with the rationality of the proposed method.

There are two parameters, i.e., the choice of vanilla feature layer $k$ and the distillation parameter $\lambda$, which affect the performance of the proposed VFD-Adv. Here, we adopt ResNet18 trained on CIFAR10 as the vanilla model to explore the effects of the two parameters. Since ResNet18 has four residual blocks, we choose the output of the last Conv layer of each block (e.g., Layer_1, etc.) and logit layer respectively, and modify $\lambda$ from 0 to 0.04 with a step of 0.005. As shown in Fig. 4, we can observe that the proposed method would achieve the best performance in both robust and clean accuracy when using the layer of the last residual block (i.e., Layer_4). Compared to early layers, the Layer_4 contains more high-level and task-specific information while keeping more spatial information than the logit features, hence the Layer_4 comes as a better choice. In addition, the distillation parameter $\lambda$ also influences the robust and clean accuracy significantly. The distillation parameter $\lambda$ is expected to guide the robust model to learn more predictive features like the vanilla model, but also make the robust model less robust due to relying on such non-robust features. Therefore, as shown in Fig. 4, a larger $\lambda$ tends to yield higher clean accuracy but diminishes the robustness while a smaller $\lambda$ would achieve better robustness with compromising the clean accuracy.