Feature prioritization and regularization improve standard accuracy and adversarial robustness

1 INTRODUCTION . Deep learning models have demonstrated impressive performance in a wide variety of applications ( Goodfellow et al. , 2016 ; Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2014 ; Mnih et al. , 2015 ) . However , recent works have shown that these models are susceptible to adversarial attacks : imperceptible but carefully chosen perturbation added to the input can cause the model to make highly confident but incorrect predictions ( Szegedy et al. , 2013 ; Goodfellow et al. , 2015 ; Kurakin et al. , 2016 ) . Exploring the adversarial robustness of neural networks has recently gained significant attention and there is a rapidly growing body of work related to this topic ( Kurakin et al. , 2016 ; Tramèr et al. , 2017 ; Fawzi et al. , 2018 ; Athalye & Sutskever , 2017 ; Carlini & Wagner , 2017 ; Kolter & Wong , 2017 ; Wong & Kolter , 2018 ; Madry et al. , 2017 ) . A wide variety of methods are proposed to defend a model against adversarial attacks ( Prakash et al. , 2018 ; Liao et al. , 2018 ; Song et al. , 2017 ; Samangouei et al. , 2018 ) . Despite these advances , many techniques are subsequently shown to be ineffective ( Athalye et al. , 2018 ; Athalye & Carlini , 2018 ) , and adversarial training which uses adversarial samples in addition to clean images during the training process has been shown to be able to build relatively robust neural networks ( Madry et al. , 2017 ; Athalye et al. , 2018 ; Dvijotham et al. , 2018 ) . With strong adversaries such as the Projected Gradient Descent ( PGD ) ( Madry et al. , 2017 ) or the Iterative Fast Gradient Sign Method ( I-FGSM ) ( Kurakin et al. , 2016 ) adversarially trained models are able to achieve state-of-the-art performance against a wide range of attacks . Recent advances in the understanding of adversarial training provide insights of its effectiveness . It is shown that standard and robust models depend on very different sets of features ( Tsipras et al. , 2018 ; Tanay et al. , 2018 ) . While standard models utilize features including non-robust ones that are weakly correlated with class labels and easily manipulated by small input perturbations , robust models only use robust features that are highly correlated with class labels and invariant to those perturbations . Although adversarial training learns robust features , there are no explicit design components to encourage a model to depend solely on robust features . Therefore , to further improve the robustness of a model , we propose feature regularization and prioritization schemes . We first propose to use attention mechanism to bias a model towards robust features that are highly correlated with the class labels . We call the learned features at the final layer of a network the global features , and the ones at lower level layers the local features . In our attention module , we use the global features as a way to assign weights to the local features by a non-linear compatibility function . Since global features are directly used to produce class label prediction , we are effectively assigning weights to local features depending on their correlation to the label . Robust features have higher correlation and therefore will be assigned larger weights which in turn contribute to the model ’ s robustness . Next , we propose to use feature regularization to learn robust features that are invariant to input perturbations . We add a L2 regularization term that penalizes the distance between the learned features of a clean sample x and that of its perturbed adversarial counterpart x′ to the training objective . By optimizing this regularizer , we are pushing the model to extract very similar features from the original image and the adversarial image , thus only features that are invariant to the perturbations are learned and we effectively ignore the added noise . From another point of view , a model with small L2 feature distance maps the two nearby points in the image space to nearby points in the learned high dimensional manifold which is a desirable behavior . In this paper , we propose an approach that enhances adversarial training with feature prioritization and regularization to improve the robustness of a model . We use extensive experiments to demonstrate that the attention module focuses on the area of an image which contains the actual object and helps the classifier to only rely on features extracted from those areas . The background clutter and irrelevant features which could be misleading are suppressed . The feature regularization further encourages the model to extract robust features that are not manipulated by the adversarial perturbations . The resulting model has a highly interpretable gradient map that aligns perfectly with salient data characteristics . The main contributions of this paper include : • A method based on feature prioritization and regularization , which significantly outperforms adversarial training . Our model is evaluated on the MNIST , CIFAR-10 , and CIFAR100 datasets , and demonstrates superior performance relative to both standard classification accuracy and adversarial robustness . • We provide justification to show that the attention module helps the model to rely on robust features by assigning larger weights to them . Through qualitative inspection , we show that the attention maps generated by our non-linear attention estimator focus sharply on the regions of interest while suppressing irrelevant background clutter . • In addition to qualitative evaluation of the gradient maps , we propose a novel experimental strategy that quantitatively demonstrates the better alignment of the gradient maps generated by our model with salient data characteristics . 2 RELATED WORK . Due to the extensive amount of literature in this area and the limited length of this paper , we only review some of the most related works in this section . For a more comprehensive survey please refer to Akhtar & Mian ( 2018 ) . Adversarial training . Kurakin et al . ( 2016 ) use adversarial training as a form of data augmentation where it injects adversarial examples during training . In every training mini batch , a mixture of clean images and adversarial images generated by one step Fast Gradient Sign Method ( FGSM ) are used to update the network ’ s parameters . It is then improved by Na et al . ( 2017 ) by adding adversarial examples generated by iterative methods . In Madry et al . ( 2017 ) , it is proposed to replace all clean images with adversarial images which is a direct result of optimizing a saddle point ( min-max ) formulation . By studying the loss landscape of the problem , they suggest that PGD is a universal first-order adversary which is then used in their adversary generating process . L2 regularization . A similar idea with feature regularization is proposed in Kannan et al . ( 2018 ) which they call logit pairing , to prevent a model from being over-confident when making predictions . Compared with logit pairing , feature regularization is more intuitive as it motivates a model to learn very robust features that are invariant to input perturbations , which leads to a robust model . In addition , we propose to also use attention module to further encourage the model to favor robust features which will improve the robustness . Attention Models . Attention in CNN is most commonly deployed for query-based tasks ( Seo et al. , 2016 ; Xu et al. , 2015 ; Jetley et al. , 2018 ) . In Jetley et al . ( 2018 ) a method is presented to use a learned representation of the global image as a query to leverage multiple attention maps at different scales , which allows the expression of a complementary focus on different parts of the image . However , the application of attention to the adversarial robustness aspects has not been seriously explored . To the best of our knowledge , we are the first to employ an attention mechanism in training a robust deep neural network . In our application , we use a ReLU activated neural network instead of the linear-based method as the attention estimator . It allows highly non-linear compatibility between the learned global features and the lower-level local features . 3 APPROACH . We now present our model that combines the attention module and L2 feature regularization , and show how it can be applied to enhance the adversarial training to improve the adversarial robustness of a model and its accuracy . Figure 1 provides an overview of our method . We adopt the terminology in Jetley et al . ( 2018 ) , in particular that local features and global features refer to features extracted at certain layers of a network whose effective receptive fields are subsets of the image and the entire image . We start by forwarding each of the clean and adversarial images and computing the attention weights by a non-linear estimator . Then the individual attention feature is defined to be the weighted combination of the corresponding local features . Next , we define an L2 regularization loss to be the Euclidean distance between the two sets of learned attention features . The attention features of the adversarial image are then used to produce the logits , which is followed by softmax layer to produce the cross-entropy loss . The final loss function of our model is a combination of cross-entropy loss and the regularization loss . 3.1 ADVERSARIAL TRAINING . We adopt the adversarial training described in Madry et al . ( 2017 ) as the basic training approach . It replaces natural training examples by PGD examples , which is suggested to represent a universal first-order adversary . So far PGD has been shown to represent the strongest attack method ( Athalye et al. , 2018 ; Athalye & Carlini , 2018 ) . A model that is trained with PGD adversaries is also robust against a wide range of other attacks and not yet outperformed by any other approach . The adversarial training has a saddle point formulation : min θ E ( x , y ) ∼D [ max δ∈S L ( θ , x+ δ , y ) ] ( 1 ) where D is the distribution of data x and class labels y , L is the cross-entropy loss function for a model with parameters θ , δ is the additive adversarial perturbation with bound S. In this paper we consider l∞ bound as in Madry et al . ( 2017 ) . Our adversarial samples x′ = x + δ are created by PGD : xt+1 = Πx+S ( xt + α sgn ( ∇xL ( θ , x , y ) ) ) ( 2 ) PGD adversaries are computed at each iteration as an approximated optimum of the inner maximization in equation ( 1 ) and an update of the parameters θ is made according to the outer minimization formulation . 3.2 ATTENTION MODEL . As we discussed in Section 1 , our goal of attention model is to favor robust features in making predictions . We propose a non-linear attention model that acts as a feature prioritizing scheme , which is able to put more weight on robust features and less weight on non-robust features to increase the robustness of a classifier . Let lin be the learned feature vector at layer i ∈ { 1 , 2 , ... , I } of a neural network at spatial location n ∈ { 1 , 2 , ... , N } , and let g be the feature vector of the layer just before the final fully connected layer which produces the class label prediction scores ( logits ) . We use a small ReLU activated neural network to generate compatibility scores between the global feature g and local features lin : cin = f ( l i n , g ) ( 3 ) where f is the neural network and the concatenation of g and lin is fed to the network to produce the compatibility scores cin . We then normalize the scores with a softmax operation to get the attention weights : win = exp cin∑ m exp c i m ( 4 ) Next , we compute the weighted sum of local feature vectors which is the attention feature vector at layer i : hi = ∑ n winl i n ( 5 ) We use the outputs of the last residual block as the local feature for computing attention , and replace the global feature g with the corresponding attention descriptor hi for final classification . By using a small neural network instead of the linear alignment models as in Jetley et al . ( 2018 ) , we are able to capture non-linear compatibility between the local and global features when producing the attention weights , which is beneficial considering the multiple non-linear function activated layers between the local and global features .

Summary: This paper studies adversarial training of robust classification models. It is based on PGD training in [madry17]. It proposes two points: 1) add attention schemes, 2) add a feature regularization loss. The results on MNIST and CIFAR10 demonstrate the effectiveness. At last, it did some diagnostic study and visualization on the attention maps and gradient maps.