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: The paper investigates the role of Batch Normalization (BN) in the generalization of deep networks, and its impact in the trade-off between clean and robust accuracy in adversarially trained networks. The authors demonstrate that the rescaling operations in BN when considered in conjunction with the ReLU activation, serve as a feature masking operation. Based on these observations, the authors propose RobMask, which uses a linear combination of such rescaling operations to achieve improved generalization in deep networks.