Adversarial Masking: Towards Understanding Robustness Trade-off for Generalization

1 INTRODUCTION . Deep neural networks have achieved unprecedented success over a variety of tasks and across different domains . However , studies have shown that neural networks are inherently vulnerable to adversarial examples ( Biggio et al. , 2013 ; Szegedy et al. , 2014 ) . To enhance model robustness against adversarial examples , adversarial training ( Goodfellow et al. , 2015 ; Madry et al. , 2018 ) has become one of the most effective and widely applied defense methods , which employs specific attacking algorithms to generate adversarial examples during training in order to learn robust models . Albeit effective in countering adversarial examples , adversarial training often suffers from inferior performance on clean data ( Zhang et al. , 2019 ; Balaji et al. , 2019 ) . This observation has led prior work to extrapolate that a trade-off between robustness and accuracy may be inevitable , particularly for image classification tasks ( Zhang et al. , 2019 ; Tsipras et al. , 2019 ) . However , Yang et al . ( 2020 ) recently suggests that it is possible to learn classifiers both robust and highly accurate on real image data . The current state of adversarial training methods falls short of this prediction , and the discrepancy remains poorly understood . In this paper , we conduct an in-depth study on understanding the trade-off between robustness and clean accuracy in adversarial training , and introduce Adversarial Masking , a new hypothesis stating that a widely used technique , batch normalization ( BN ) , has a significant impact on the trade-off between robustness and natural accuracy . Specifically , we break down BN into normalization and rescaling operations , and find that the rescaling operation has a significant impact on the robustness trade-off while normalization only has marginal influence . Built upon this observation , we hypothesize that adversarial masking ( i.e. , the combination of the rescaling operation and the follow-up ReLU activation fucntion ) acts as a feature masking layer that can magnify or block feature maps to influence the performance of robust or clean generalization . In this hypothesis , different rescaling parameters in BN contribute to different adversarial maskings learned through training . By using a simple linear combination of two adversarial maskings , rather than using robust features learned by adversarial training ( Madry et al. , 2018 ; Ilyas et al. , 2019 ; Zhang et al. , 2019 ) , we show that a well-balanced trade-off can be readily achieved . Based on the Adversarial Masking hypothesis , we further propose RobMask ( Robust Masking ) , a new training scheme that learns an adaptive feature masking for different perturbation strengths . We use the learned adaptive feature masking to incorporate different features so that we could improve model generalization with a better robustness trade-off . Specifically , each perturbation strength is encoded as a low-dimensional vector , and we take this vector as input to a learnable linear projection layer together with ReLU activation , to obtain the adversarial masking for the corresponding perturbation strength . Therefore , for different perturbation strengths , we learn different maskings accordingly . By doing so , rather than hurting the performance on clean test data , we use adversarial examples as powerful regularization to boost model generalization . Experiments on multiple benchmarks demonstrate that RobMask achieves not only significantly better natural accuracy , but also a better trade-off between robustness and generalization . Our contributions are summarized as follows . ( i ) We conduct a detailed analysis to demonstrate that the rescaling operation in batch normalization has a significant impact on the trade-off between robustness and natural accuracy . ( ii ) We introduce Adversarial Masking , a new hypothesis to explain that this trade-off is caused by different feature maskings applied , and different combinations of maskings can lead to different trade-offs . ( iii ) We propose RobMask , a new training scheme to learn an adaptive masking for different perturbation strengths , in order to utilize adversarial examples to boost generalization on clean data . RobMask also achieves a better trade-off between robust and natural accuracy . 2 PRELIMINARY AND RELATED WORK . Adversarial Training Since the discovery of the vulnerability of deep neural networks , diverse approaches have been proposed to enhance model adversarial robustness . A natural idea is to iteratively generate adversarial examples , add them back to the training data , and then retrain the model . For example , Goodfellow et al . ( 2015 ) uses adversarial examples generated by FGSM to augment the data , and Kurakin et al . ( 2017 ) proposes to use a multi-step FGSM to further improve performance . Madry et al . ( 2018 ) shows that adversarial training can be formulated as a min-max optimization problem , and proposes PGD attack ( similar to multi-step FGSM ) to find adversarial examples for each batch . Specifically , for a K-class classification problem , let us denote D = { ( xi , yi ) } ni=1 as the set of training samples with xi ∈ Rd , yi ∈ { 1 , . . . , K } , where K is the number of classes . Considering a classification model fθ ( x ) : Rd → ∆K parameterized by θ , where ∆K represents a K-dimensional simplex , adversarial training can be formulated as : min θ 1 n n∑ i=1 max x′i∈Bp ( xi , ) ` ( fθ ( x ′ i ) , yi ) , ( 1 ) where Bp ( xi , ) denotes the ` p-norm ball centered at xi with radius , and ` ( · , · ) is the cross-entropy loss . The inner maximization problem aims to find an adversarial version of a given data point xi that yields the highest loss . In general , Bp ( xi , ) can be defined based on the threat model , but the ` ∞ ball is the most popular choice among recent work ( Madry et al. , 2018 ; Zhang et al. , 2019 ) , which is also adopted in this paper . For deep neural networks , the inner maximization does not have a closed-form solution . Thus , adversarial training typically uses a gradient-based iterative solver to approximately solve the inner problem . The most commonly used choice is multi-step PGD ( Madry et al. , 2018 ) and C & W attack ( Carlini & Wagner , 2017 ) . Since then , most defense algorithms ( Zhang et al. , 2019 ; Balaji et al. , 2019 ; Wang et al. , 2019 ; Ding et al. , 2018 ; Cheng et al. , 2020 ) are based on a similar min-max framework . Trade-off between Robustness and Accuracy While effective in improving model robustness , adversarial training is known to bear a performance drop on clean test data . Tsipras et al . ( 2019 ) provides a theoretical example of data distribution where any classifier with high test accuracy must also have low adversarial accuracy under ` ∞ perturbations . They claim that high performance on both accuracy and robustness may be unattainable due to their inherently opposing goals . Zhang et al . ( 2019 ) decomposes the robust error as the sum of natural ( classification ) error and boundary error , and provides a differentiable upper-bound using the theory of classification-calibrated loss , based on which they further propose TRADES to achieve different trade-offs by tuning the regularization term . Most recently , Xie et al . ( 2020 ) proposes AdvProp to assign another batch normalization for generating adversarial examples , and shows improved performance on clean data in image classification tasks . There are also parallel studies on applying adversarial training to improve clean data performance in language understanding and vision-and-language tasks ( Zhu et al. , 2019 ; Jiang et al. , 2019 ; Liu et al. , 2020 ; Gan et al. , 2020 ) . Batch Normalization Batch Normalization is a widely adopted technique that enables faster and more stable training of deep neural networks , below , we provide a brief overview of batch normalization , which paves the way to introduce our method . Specifically , batch normalization ( Ioffe & Szegedy , 2015 ) is proposed to reduce the internal co-variate shift to ease neural network training . Considering a convolutional neural network , we can define the input and output as Ib , c , x , y and Ob , c , x , y , respectively . The dimensions correspond to examples with a batch b , channel c , and two spatial dimensions x , y . A neural network applies the same normalization for all activations in a given channel : Ob , c , x , y ← γ Ib , c , x , y − µc√ σ2c + + β ∀b , c , x , y , ( 2 ) where µc = 1|B| ∑ b , x , y Ib , c , x , y denotes the mean for channel c , and σc denotes the corresponding standard deviation . γ and β are two learnable parameters for the channel-wise affine transformation , i.e. , rescaling operations . is a small number to control numerical stability . 3 ADVERSARIAL MASKING . 3.1 BATCH NORMALIZATION ACTS AS ADVERSARIAL MASKING . Ilyas et al . ( 2019 ) disentangles adversarial examples as a natural consequence of non-robust features . Specifically , they construct robust features from an adversarial trained “ robust model ” directly . Therefore , a common belief is that adversarial robustness comes from feature representations learned through adversarial training ( Ilyas et al. , 2019 ; Santurkar et al. , 2019 ) . An interesting question we would like to ask is : can we learn robust features from a vanilla standard-trained model , or , can we obtain non-robust features from an adversarial trained “ robust model ” ? To answer this question , we design the following experiments . We first train a ResNet-18 model with standard and adversarial training , then finetune the networks by allowing only batch normalization ( BN ) to be changed while freezing other parameters . Specifically , we finetune BN in a standard trained model using adversarial training , and finetune BN in an adversarial trained model with standard training , respectively . Results are summarized in Table 1 . Given a standard trained model , by only per- forming adversarial finetuning of the BN layers , the resulting model can already achieve a reasonably good robust accuracy 26.51 % ( the 1st block ) . Similarly , given an adversarial trained model , by only performing standard finetuning of the BN layers , the clean accuracy of the model increases significantly from 78.47 % to 86.96 % ( the 2nd block ) . This experiment demonstrates that we can control the trade-off between clean and robust errors by only tuning the BN layer , so here comes a natural question : which part in BN has contributed to this performance trade-off ? To investigate this , we take a step further to check the difference of every parameter used in BN . In particular , we check the first BN layer after the first convolution layer . As well known , BN uses a running average of the mean and variance during testing . Figure 1a and 1b illustrate the difference on the running mean µ and running variance σ . The batch statistics with and without further standard finetuning of BN ( under the setting in the 2nd block of Table 1 ) are nearly identical across all the dimensions . Figure 1c and 1d plot the learned rescaling parameters γ and β . We can clearly see that the fine-tuned rescaling parameter γ is completely different from the original one , with β unchanged , indicating that γ has a significant impact on the clean and robust trade-off while still performing similar normalization on both sides . On the other hand , Rectified Linear Unit ( ReLU ) ( Agarap , 2018 ) is the most commonly used activation function in deep neural networks . The function returns 0 if it receives any negative input , and for any positive value x it returns that value back ( i.e. , f ( x ) = max ( 0 , x ) ) . During the rescaling operation , γ would magnify or shrink the magnitude of feature maps . Together with β , after ReLU activation , features that become negative will be blocked as 0 . By combining the rescaling operation with the following ReLU activation function , the resulting layer can be viewed as a masking layer to magnify or block the features maps from the convolution layer ( see Figure 2 ( a ) for illustration ) . To further validate this , we plot the feature maps after ReLU activation functions in Figure 4 in the Appendix . Some feature maps are blocked after finetuning BN ( completely black when all pixels are set to 0 ) as well as some are magnified significantly . We term the above observation as Adversarial Masking , and hypothesize that this leads to the trade-off between robustness and natural accuracy .

Summary: This paper follows the direction of previous work AdvProp (Xie et al. (2020)) and aims to use adversarial training as a regularizer to improve the network generalization on clean data. The authors analyze that the different rescaling operation in the batch normalization layer along with ReLU acts as feature masking/selection layer, which can control the trade-off between adversarial robustness and clean data performance. Unlike AdvProp, which uses different batch normalization layer for clean images and adversarial images at a specific perturbation strength, here the authors propose a technique called RobMask that adapts the rescaling parameters of batch normalization based on the perturbation strength during training. The authors show that such adapting technique is more effective than using different batch normalization layers (as in AdvProp) for each perturbation strength and thus improves the clean data performance on CIFAR10/100.