CAT: Customized Adversarial Training for Improved Robustness

1 Introduction

Deep neural networks (DNNs) have proved their effectiveness on a variety of domains and tasks. However, it has been found that DNNs are highly vulnerable to adversarial examples [Szegedy et al., 2014]. To enhance the robustness of DNNs against adversarial examples, adversarial training [Goodfellow et al., 2015; Madry et al., 2018] has become one of the most effective and widely used methods. Given a pre-defined perturbation tolerance, denoted as ϵ, adversarial training aims to minimize the robust loss, defined as the worst-case loss within ϵ-ball around each example, leading to a min-max optimization problem. [Madry et al., 2018] shows that applying a multi-step projected gradient descent (PGD) attack to approximately solve the inner maximization leads to a robust model, and several recent research has proposed various ways to improve adversarial training [Zhang et al., 2019;

Wang, 2019; Wang et al., 2019a; Balaji et al., 2019; Ding et al., 2018].

However, standard adversarial training methods still have a hypothetical and possibly problematic assumption: the perturbation tolerance ϵ is a large and fixed constant throughout the training process, which ignores the fact that every data point may have different intrinsic robustness. Intuitively, some examples are naturally closer to the decision boundary, and enforcing large margin on those examples will force the classifier to give up on those examples, leading to a distorted decision surface. This intuition may explain the known issue of the undesirable robustnessaccuracy tradeoff in adversarial robustness [Su et al., 2018; Tsipras et al., 2019]. Furthermore, with a different perturbation tolerance, it is questionable whether we should still force the model to learn to fit the one-hot label as in the original adversarial training formulation. In the extreme case, if an example is perturbed to the decision boundary, a good classifier yielding the binary class prediction probabilities should output [0.5, 0.5] instead of [1, 0]. This aspect becomes crucial when each example is associated with a different level of perturbation. Although some recent papers have started to address the uniform ϵ issue by treating correctly and incorrectly classified examples differently [Ding et al., 2018] or assigning non-uniform perturbation level [Balaji et al., 2019], none of them have tried to incorporate customized training labels in this process.

Motivated by these ideas, we propose a novel Customized Adversarial Training (CAT) framework that can substantially improve the performance of adversarial training. Throughout the adversarial training process, our algorithm dynamically finds a non-uniform and effective perturbation level and the corresponding customized target label for each example. This leads to better generalization performance and furthermore, with a careful design on adaptive ϵ tuning, our algorithm has only negligible computational overhead and runs as fast as the original adversarial training algorithm. Furthermore, we theoretically explain why the proposed method could lead to improved generalization performance.

Our method significantly outperforms existing adversarial training methods on the standard CIFAR-10 defense task.

With Wide-ResNet structure on CIFAR-10, under 8/255 ℓ∞ perturbation, our method achieves 73% robust accuracy under PGD attack and 71% robust accuracy under Carlini and Wagner (C&W) attack [Carlini and Wagner, 2017], while the current best model only achieves 58.6% under PGD attack and 56.8% under C&W attack. Furthermore, our method only degrades the clean accuracy from 95.93% (standard test accuracy) to 93.48%, while other adversarial training methods have clean accuracy below 91.34%.

2 Background andMotivation

2.1 Preliminaries

Adversarial training can be formulated as a min-max optimization problem. For a K-class classification problem, let D = {(xi, yi)}i=1,...,n denote the set of training samples in the dataset with xi ∈ Rd , yi ∈ {1, . . . ,K} := [K], we consider a classfication model fθ(x) : Rd → [K] parameterized by θ. We denote by hθ(x) : Rd → [0, 1]K as the prediction output for each class, i.e., fθ(x) = argmaxi[hθ(x)]i.

Adversarial training can be formulated as:

where Bp(xi, ϵ) denotes the ℓp-norm ball centered at xi with radius ϵ. The inner maximization problem aims to find an adversarial version of a given data point xi that achieves the highest loss. In general, one can define Bp(xi, ϵ) based on the threat model, but the ℓ∞ ball is the most popular choice adopted by recent works [Madry et al., 2018; Zhang et al., 2019; Wang, 2019; Ding et al., 2018; Wang et al., 2019b], which we also use in this paper. For a deep neural network model, the inner maximization does not have a closed form solution, so adversarial training methods typically use a gradient-based iterative solver to approximately solve the inner problem. The most commonly used choice is the multi-step PGD [Madry et al., 2018] and C&W attack [Carlini and Wagner, 2017].

2.2 Motivation

Intuitively, if adversarial training can always find a model with close-to-zero robust error, one should always use a large ϵ for training because it will automatically imply robustness to any smaller ϵ. Unfortunately, in practice a uniformly large ϵ is often harmful. In the following we empirically explain this problem and use it to motivate our proposed algorithm.

We use a simple linear classification case to demonstrate why a uniformly large ϵ is harmful. In Figure 1a, we generate a synthetic linearly separable dataset with the margin set to be 1.75 for both classes, where the correct linear boundary can be easily obtained by standard training. In Figure 1b, we run adversarial training with ϵ = 1, and since this ϵ is smaller than the margin, the algorithm can still obtain near-optimal results. However, when we use a large ϵ = 4 for adversarial training in Figure 1c, the resulting decision boundary becomes significantly worse. It is because adversarial training cannot correctly fit all the samples with a margin up to 4, so it will sacrifice some data samples, leading to distorted and undesirable decision boundary. This motivates the following two problems:

• We shouldn't set the same large ϵ uniformly for all samples. Some samples are intrinsically closer to the decision boundary and they should use a smaller ϵ. Without doing this, adversarial training will give up on those samples, which leads to worse training and generalization error (see more discussions in Section 3.3 on the generalization bounds).
• The adversarial training loss is trying to force the prediction to match the one-hot label (e.g., [1, 0] in the binary classification case) even after large perturbations. However, if a sample is perturbed, the prediction shouldn't remain one-hot. For instance, if a sample is perturbed to the decision boundary of a binary classification problem, the prediction of a perfect model should be [0.5, 0.5] instead of [1, 0], which also makes adversarial training fail to recover a good decision hyperplane.

Furthermore, we observe that even if adversarial training can obtain close-to-zero training error with large ϵ (e.g., [Gao et al., 2019] proves that this will happen for overparameterized network with large-enough margin), a uniformly large ϵ will lead to larger generalization gap. This could be partially explained by the theoretical results provided by [Yin et al., 2018], which shows that the adversarial Rademacher complexity has a lower bound with an explicit dependence on the perturbation tolerance. The empirical results in Table 1 also illustrate this problem. When conducting adversarial training with ϵ = 0.3 on CIFAR10 VGG-16, we found that the model achieves close-to-zero robust training error on all ϵ ≤ 0.3, but it suffers larger generalization gap compared to training with smaller ϵ. This also demonstrates that a uniformly large ϵ is harmful even when it achieves perfect training error.

3 CAT (Customized Adversarial Training)

In this section, we propose the Customized Adversarial Training (CAT) framework that improves adversarial training by addressing the above-mentioned problems. First, our algorithm has an auto-tuning method to customize the ϵ used for each training example. Second, instead of forcing the model to fit the original label, we customize the target label for each example based on its own ϵ. In the following we will describe these two components in more detail.

3.1 Auto-tuning Perturbation Strength for Adversarial Training

The first component of our algorithm is an ϵ auto-tuning method which adaptively assigns a suitable ϵ for each example during the adversarial training procedure. Let ϵi be the

(a) Standard training

perturbation level assigned to example i. Based on the intuition mentioned in Section 2.2, we do not want to further increase ϵ if we find the classifier does not have capacity to robustly classify the example, which means we should set

and the adversarial training objective becomes

Note that ϵi in (2) depends on θ, while θ in (3) also depends on ϵi. We thus propose an alternative update scheme - conducting one SGD update on θ, and then updating the ϵi in the current batch. However, finding ϵi exactly requires bruteforce search for every possible value, which adds significant computational overhead to adversarial training.

Therefore, we only conduct a simplified update rule on ϵi as follows. Starting from an initial perturbation level of zero, at each iteration we conduct adversarial attack (e.g., PGD attack) with perturbation tolerance ϵi + η where η is a constant. If the attack is successful, then we reset current ϵi to 0 to encourage model learning a more robust classifier towards those examples. While if the attack is unsuccessful, which means an attacker still cannot find an adversarial example that satisfies maxx′i∈Bp(xi,ϵi+η) fθ(x′i ) ̸= yi, then we increase ϵi = ϵi + η. The attack results will also be used to update the model parameter θ, so this adaptive scheme does not require any additional cost. In practice, we also have an upper bound on the final perturbation to ensure that ϵi remains bounded for each i.

3.2 Adaptive Label Uncertainty for Adversarial Training

As mentioned in Section 2.2, the standard adversarial training loss is trying to enforce a sample being classified as the original one-hot label after ϵ perturbation. However, this may not be ideal. In the extreme case, if a sample is perturbed to the decision boundary, the prediction must be far away from onehot. This problem is more severe when using non-uniform ϵi, since each different ϵi will introduce a different bias to the loss, and that may be one of the reasons that purely adaptive ϵ-scheduling does not work well (see our ablation study in Section 5 and also the results reported in [Balaji et al., 2019]).

In the following, we propose an adaptive label smoothing approach to reflect different perturbation tolerance on each example. Szegedy [2016] introduced label smoothing that converts one-hot label vectors into one-warm vectors representing low-confidence classification, in order to prevent the model from making over-confident predictions. Specifically, with a one-hot encoded label y, the smoothed version is

where α ∈ [0, 1] is the hyperparameter to control the smoothing level. In the adaptive setting, we set α = cϵi so that a larger perturbation tolerance would receive a higher label uncertainty and c is a hyperparameter. A common choice of u is u = 1K . However, this strict requirement tries to enforce every other labels having the same probability, which may not make sense in practice. On the other hand, as shown Section 2.2, adversarial training is easy to overfit and generate a large generalization gap. To better address these issues, we sample from a distribution instead. Specifically, we use u = Dirichlet(β) where Dirichlet(·) refers to the Dirichlet distribution and β ∈ RK is concentration hyperparamter. With different perturbation tolerance, the adaptive version of label smoothing is

The Final Objective Function. Combining the two aforementioned techniques, our Customized Adversarial Training (CAT) method attempts to minimize the following objective:

where ˜ yi is defined in (4). As described in Section 3.1, we approximately minimize this objective with an alternative update scheme, which incurs almost no additional cost compared to the original adversarial training algorithm. The detailed algorithm is presented in Algorithm 1.

Algorithm 1 CAT algorithm

Input: Training dataset (X, Y ), cross entropy loss or mix loss ℓ, scheduling parameter η, weighting factor c, perturbation upperbound ϵmax Initial every sample's ϵi with 0 for epoch=1, . . . , N do for i=1, . . . , B do ˜ yi ← (1 - cϵi)yi + cϵiDirichlet(β) ϵi ← ϵi + η δi ← 0 for j = 1 . . .m do δi ← δi + α · sign(∇δℓ(fθ(xi + δi), ˜ yi) δi ← max(min(δi, ϵi), -ϵi) end for if fθ(xi + δi) ̸= yi then ϵi ← 0 end if ϵi ← min(ϵmax, ϵi) ˜ yi ← (1 - cϵi)yi + cϵiDirichlet(β) θ ← θ - γθ∇θℓ(fθ(xi + δi), ˜ yi) end for end for return θ

Choice of Loss Function. In general, our framework can be used with any loss function ℓ(·). In the previous works, cross entropy loss is commonly used for ℓ. However, the model trained by smoothing techniques tends to have a smaller logit gap between true label and other labels. Therefore, in order to encourage model to generate a larger logit gap, we propose a mixed loss to enhance the defense performance towards C&W∞ attack. That is,

where Z(x) ∈ RK is the final (logit) layer output, and [Z(x)]i is the prediction score for the i-th class and y0 is the original label. The parameter κ encourages the adversary to find higher confident adversarial examples in training.

3.3 Theoretical Analysis

To better understand how our scheme improves generalization, we now provide some theoretical analysis. Recall we denote by hθ(x) : Rd → [0, 1]K as the prediction probability for the K classes. We define the bilateral margin that our paper is essentially maximizing over as follows.

Definition 1 (Bilateral margin). We define the bilateral perturbed network output by Hθ(x, δi , δo ):

The bilateral margin is now defined as the minimum norm of (δi , δo ) required to cause the classifier to make false predictions:

This margin captures both the relative perturbation on the input layer δi and the soft-max output δo .

Theorem 2. Suppose the parameter space Θ we optimize over has covering number that scales as logN‖·‖op (η,Θ) ≤ ⌊C 2 /η2 ⌋ for some complexity C. Then with probability 1 - δ over the draw of the training data, any classifer fθ, θ ∈ Θ which achieves training error zero satisfies:

where ζ is of small order O ( 1n log(1/δ)).

We defer the proof to the Appendix, which is adapted from Theorem 2.1 of [Wei and Ma, 2019]. We observe the population risk is bounded by two key factors, the average of 1 mF (xi,yi) and C, the covering number of the parameter space. On one side, the average of 1 mF (xi,yi) is dominated by the samples with the smallest margin. Therefore when we do adversarial training, it is important that we not only achieve higher overall accuracy, but also make sure the samples closer to the decision boundary have large enough margin. This can not be achieved by simply using constant and large ϵ that will maintain a large margin for most samples but sacrifice the accuracy of a small portion of data. On the other hand, the covering number of the network's parameter space can be roughly captured by a bound of product of all layers' weight norms. We hypothesize that with more flexibility in choosing ϵ, our algorithm will converge faster than using larger constant ϵ and will have more implicit regularization effect. To testify this hypothesis, we roughly measure the model complexity C by the product of the weight norms of different models. In comparison to our model, when training with constant ϵ = 0.01, 0.02 and 0.03, it respectively yields C as large as 2.54, 3.53 and 1.39 times of that of our model, which means our model indeed has more implicit regularization effect among others.

4 RelatedWork

Adversarial Training. To enhance the adversarial robustness of a neural network model, a natural idea is to iteratively generate adversarial examples, add them back to the training data, and retrain the model. For example, [Goodfellow et al., 2015] use adversarial examples generated by FGSM to augment the data, and [Kurakin et al., 2017] propose to use a multiple-step FGSM to further improve the performance. [Madry et al., 2018] show that adversarial training can be formulated as a min-max optimization problem, and propose to use PGD attack (similar to multi-step FGSM) to find adversarial examples for each batch. After that, many defense algorithms are based on a similar min-max framework. However, each of them uses slightly different loss functions.

We see that except for natural training which directly minimizes the cross entropy loss (denoted as CE), all training techniques involve the use of the min-max framework. TRADES and MMA use the unperturbed data's cross entropy loss as an additional regularization term to achieve a better trade-off between clean and robust error.

Similar to our method, both MMA and IAAT have samplewise adaptive ϵ during training. They also utilize the adaptive ϵ to find the largest possible ϵi for every sample xi. However, they do not consider the adaptive label technique mentioned in Section 3.2. As a result, they can only achieve better clean accuracy while the improvements in robust accuracy are limited. Our CAT algorithm (CAT with CE loss) is more general than IAAT and MMA. CAT reduce to IAAT when we set c = 0 in adaptive label smoothing. Moreover, MMA could be treated as a special case of CAT when we use a line search scheme to find the ϵi and c = 0. Also, in Section 5, we will show the importance of the adaptive label uncertainty step in CAT. Recently, a concurrent work [Stutz et al., 2019] combines label smoothing with adversarial training. However, in the adversarial training process, they still use the same ϵ for all the examples, which is quite different from our instancewise auto-tuning ϵ. Since their model jointly performs detection (dropping low confident examples) and prediction, the results and formulation are not directly comparable to other adversarial training methods.

5 Performance Evaluation

In this section, we conduct extensive experiments to show that CAT achieves a strong result on both clean and robust accuracy. We include the following methods into our comparison:

• Customized Adversarial Training (CAT): Our proposed method.
• Adversarial training: The adversarial training method proposed in [Madry et al., 2018] where they use a Kstep PGD attack as adversary.
• TRADES: TRADES [Zhang et al., 2019] improves adversarial training by an additional loss on the clean examples and achieves the state-of-art performance on robust accuracy.
• Natural: the natural training which only minimizes the cross entropy loss.

Furthermore, since many recently proposed adversarial training methods have considered CIFAR-10 with Wide-ResNet structure as the standard setting for reporting their numbers, we also compare our performance with 7 previous methods on this specific setting.

5.1 Experimental Setup

Dataset and Model Structure. We use CIFAR-10 dataset for performance evaluation. We use both standard VGG16 [Simonyan and Zisserman, 2015] and Wide ResNet that is used in both vanilla adversarial training [Madry et al., 2018] and TRADES [Zhang et al., 2019]. For VGG-16, we implement adversarial training with the standard hyperparameters and train TRADES with the official implementation. For Wide ResNet, since the model has become standard for testing adversarial training methods, we use exactly the same model structure provided by [Madry et al., 2018; Zhang et al., 2019]. We use the models' checkpoint released by TRADES official repository and implement the Madry's adversarial training using the standard hyper-parameters. All our experiments were implemented in Pytorch-1.4 and conducted using a GTX 2080 TI GPU.

Implementation Details. We set the number of iterations in adversarial attack to be 10 for all methods during training. Adversarial training and TRADES are trained on PGD attacks setting ϵ = 8/255 with cross entropy loss (CE). All the models are trained using SGD with momentum 0.9, weight decay 5 × 10-4 . For VGG-16/Wide ResNet models, we use the initial learning rate of 0.01/0.1, and we decay the learning rate by 90% at the 80th, 140th, and 180th epoch. For CAT, we set epsilon scheduling parameter η = 0.005, ϵmax = 8/255 and weighting parameter c = 10. We set β = 1 for the distribution Dirichlet(β), which is equal to a uniform distribution. Also, we set κ = 10. Our code is publicly available at https://github.com/cmhcbb/ CAT-Customized-Adversarial-Training-for-Improved-Robustness.

5.2 Robustness Evaluation and Analysis

White-box Attacks Results. For CIFAR10, we evaluate all the models under white-box ϵ = 8/255 ℓ∞-norm bounded non-targeted PGD and C&W attack. Specifically, we use both PGDX (X-step PGD with step size ϵ/5) and C&W∞. As suggested, we test our model under different steps PGD and multiple random restarts.

The experimental results are shown in Table 2, where we can easily see that CAT clearly outperforms other methods. CAT achieves a significant better robust accuracy at the standard 8/255 perturbation threshold considered in the literature, and also have better clean accuracy. We also test the performance of CAT under attacks with 20 restarts and 1,000 iterations to confirm the robustness of the model. Futhermore, we visualize the loss landscape in the ablation study.

Wide-ResNet has become a standard structure for comparing adversarial training methods, and it's standard to train and evaluate with 8/255 ℓ∞ norm perturbation. For this setting, we collect the reported accuracy from 7 other adversarial training methods, with several of them published very recently, to have a detailed full comparison. As shown in Table 3, our method achieves state-of-art robust accuracy while maintaining a high clean accuracy.

Black-box Transfer Attacks Results. We follow the criterion of evaluating transfer attacks as suggested by [Athalye et al., 2018] to inspect whether the models trained by CAT will cause the issue of obfuscated gradients and give a false sense of model robustness. We generate 10,000 adversarial examples of CIFAR-10 from natural models with ϵ = 8/255 and evaluate their attack performance on the target model. Table 4 shows that CAT achieves the best accuracy compared with adversarial training and TRADES, suggesting the effectiveness of CAT in defending both white-box and transfer attacks.

5.3 Ablation Study

The Importance of Adaptive Label Uncertainty. Here we discuss and perform an ablation study using VGG-16 and CIFAR-10 on the importance of adaptive label uncertainty and adaptive instance-wise ϵ. In Table 5, Adv train denotes the original adversarial training, Adv+LS denotes adversarial training with label smoothing (setting y by Eq (4)), AdpAdv denotes adversarial training with adaptive instance-wise

ϵ, and CAT is the proposed method which is a combination of these two tricks. We found that only applying adaptive instance-wise ϵ or label smoothing cannot significantly boost the robust accuracy over standard adversarial training, but the proposed method, by nicely combining these two ideas, can significantly improve the performance.

Loss Landscape Exploration. To further verify the superior robustness using CAT, we visualize the loss landscape of different training methods in Figure 2. Following the implementation in [Engstrom et al., 2018], we divide the data input along a linear space grid defined by the sign of the input gradient and a random Rademacher vector, where the x- and y- axes represent the magnitude of the perturbation added in each direction and the z-axis represents the loss. As shown in Figure 2, CAT generates a model with a lower and smoother loss landscape. Also, it could be taken as another strong evidence that we have found a robust model through CAT.

6 Conclusions

In this paper, we propose CAT, a customized adversarial training method that is designed to have better generalization for both clean and robust performance. We also provide a theoretical analysis that explains the performance of our algorithm. Experimental results show that CAT has achieved state-of-art robust accuracy and a high clean accuracy while keeping similar running time as standard adversarial training. The success of CAT indicates that it is crucial to customize the perturbation level on both data sample side and its label in adversarial training.