Encoding Robustness to Image Style via Adversarial Feature Perturbations

1 Introduction . Robust optimization for neural networks has been a major focus of recent research . A mainstream approach to reducing the brittleness of classifiers is adversarial training , which solves a min-max 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . optimization problem in which an adversary makes perturbations to images to degrade network performance , while the network adapts its parameters to resist degradation [ 9 , 18 , 25 ] . The result is a hardened network that is no longer brittle to small perturbations to input pixels . While adversarial training makes networks robust to adversarial perturbations , it does not address other forms of brittleness that plague vision systems . For example , shifts in image style , lighting , color mapping , and domain shifts can still severely degrade the performance of neural networks [ 12 ] . We propose adapting adversarial training to make neural networks robust to changes in image style and appearance , rather than small perturbations at the pixel level . We formulate a min-max game in which an adversary chooses adversarial feature statistics , and network parameters are then updated to resist these changes in feature space that correspond to appearance differences of input images . This game is played until the network is robust to a variety of changes in image space including texture , color , brightness , etc . The idea of adversarial feature statistics is inspired by the observation that the mean and variance of feature maps encode style information , and thus enable the transfer of style information from a source image to a target image through normalization [ 15 , 40 ] . Unlike standard approaches that rely on feature statistics from auxiliary images to define an image style , we use adversarial optimization of feature statistics to prepare classifiers for the worst-case style that they might encounter . We propose training with Adversarial Batch Normalization ( AdvBN ) layers . Before each gradient update , the AdvBN layer performs an adversarial feature shift by re-normalizing with the most damaging mean and variance . By using this layer in a robust optimization framework , we create networks that are resistant to various domain shifts representable by shifts in feature statistics . An advantage of this method is that it does not require additional auxiliary data from new domains . We show that robust training with AdvBN layer hardens classifiers against changes in image appearance and style [ 6 , 45 ] , as well as common image corruptions [ 12 ] . Besides classification , the effectiveness of AdvBN is also shown in the task of semantic segmentation , where it improves cross-domain generalization . 2 Related work . Adversarial training . Adversarial training and its variants [ 7 , 25 , 32 , 49 ] have been widely studied for producing models that are robust to adversarial examples [ 26 , 36 ] through solving min-max optimization problems . Besides defending against adversarial attacks , recent work has shown that adversarial training can be effectively applied to many other tasks [ 8 , 17 , 30 , 53 , 55 ] . Adversarial data augmentation [ 34 , 42 , 54 ] proposes to generate worst-case unseen domains using data augmentation and an adversarial loss , thus improving the generalization of neural networks . Another work [ 46 ] interprets the original formula of adversarial training as a type of data augmentation and reveals the distributional discrepancy between the feature representations of adversarial examples and clean data . The combination of adversarial training and feature statistics has been studied in previous work [ 24 , 27 ] , where it has been used to defend against adversarial attacks [ 27 ] , or to generate feature distributions [ 24 ] . Our method differs from previous work in both the target of perturbation and the objective function . We craft adversarial feature distributions by directly perturbing the mean and variance of feature maps instead of through a variational auto-encoder [ 24 ] , and our objective does not include regularization terms . Feature Perturbation . Feature perturbation has been an effective approach to generate novel data distributions from a given source domain [ 14 , 19 , 22 , 39 ] . Feature perturbations can be applied in different ways , such as adding spatial noise to the feature maps [ 37 ] or transforming feature maps using classical signal processing [ 22 ] . We will specifically discuss feature perturbation through re-normalization . While feature normalization [ 16 , 40 ] is first proposed to accelerate neural network training [ 1 ] , the mean and variance of deep feature representations have been shown to effectively capture image style information , and style transfer can be realized through feature re-normalization [ 15 ] . The idea of feature re-normalization has also been adopted to help neural networks adapt from the source domain to a target domain , using feature statistics obtained from the latter [ 21 , 31 ] . Recent work [ 20 ] also proposes to use re-normalization for feature interpolation to improve the generalization capability of neural networks . Instead of re-normalizing features with statistics of other samples or from other domains , we simulate the worst-case scenario , encouraging models to be less sensitive to style information and thus to generalize better to images of varying appearances . Robustness to distribution shifts . The definition of “ distribution shifts ” varies from one topic to another . For example , distribution shifts in the meta-learning literature [ 19 , 37 ] often refer to the discrepancy in discriminatory feature distributions of novel tasks from different domains . Another definition is the subtle difference between training and testing data that are sampled from the same underlying distribution [ 38 , 50 ] . This work mainly focuses on distribution shifts across image data that amount to major “ style ” discrepancy , including variations in illumination [ 6 , 45 ] , weather condition [ 29 ] , and image quality degradation [ 12 ] . Various methods have been proposed to produce neural networks that generalize well to this type of distribution shift , including test time training [ 35 ] , test time adaptation [ 43 ] , training with noise [ 30 ] , and novel network architectures [ 28 ] , etc . Data augmentation is a popular method which is designed to increase the diversity of training data and prevent neural networks from over-fitting . Recent studies in data augmentation have proposed to use novel augmentation operations [ 6 , 13 , 48 , 51 ] , policy searching [ 3 , 10 , 23 ] , and adversarial approaches [ 17 , 42 , 55 ] , etc . However , the potency of data augmentation can be limited by the choice of applicable augmentations [ 2 , 5 ] . Feature space augmentation through feature interpolation [ 20 , 41 ] , on the other hand , is not restricted to the family of image transformations . Our method also works in feature space , but instead of interpolating , we perturb feature statistics adversarially . 3 Adversarial Batch Normalization . We propose Adversarial Batch Normalization ( AdvBN ) , a module that adversarially perturbs deep feature distributions such that the features confuse CNN classifiers . We iteratively compute adversarial directions in feature space by taking PGD steps on batch statistics . In the next section , we will train on these perturbed feature distributions in order to produce models robust to domain shifts . Consider a pre-trained classification network , g , with L layers . We divide g into two parts , g1 , l and gl+1 , L , where gm , n denotes layersm through n. Now , consider a batch of data , x , with corresponding labels , y . Formally , the AdvBN module is defined by BNδadv ( x ; g , l , y ) = δ ′ σ · σ ( f ) · ( f − µ ( f ) σ ( f ) ) + δ′µ · µ ( f ) , where f = g1 , l ( x ) , ( 1 ) ( δ′µ , δ ′ σ ) = argmax ( δµ , δσ ) L [ gl+1 , L ( δσ · ( f − µ ( f ) ) + δµ · µ ( f ) ) , y ] , subject to ‖δµ − 1‖∞ ≤ , ‖δσ − 1‖∞ ≤ , ( 2 ) where L is the cross-entropy loss , and the maximization problem is solved with projected gradient descent . Note that the feature input of gl+1 , L is a simplified form of the AdvBN formulation in Eq ( 1 ) , where the two σ ( f ) ’ s cancel out . Simply put , the AdvBN module is a PGD attack on batch norm statistics which can be inserted inside a network . δµ , δσ are vectors with length equal to the number of channels in the output of layer l , and we multiply by them entry-wise , one scalar entry per channel , similarly to Batch Normalization . Additionally , note that this module acts on a per-batch basis so that features corresponding to the same image are perturbed differently across training epochs as training samples are randomly shuffled between epochs during training . In Eq ( 1 ) , we re-normalize the feature by adding adversarial statistics δµ · µ ( f ) , rather than simply δµ , so that ` ∞ bounds and steps size do not need to depend on µ ( f ) . Intuitively , we permit the mean of adversarial features to vary more when the clean features have a mean of higher magnitude . Visualizing feature shifts . We adopt the VGG [ 33 ] based autoencoder from AdaIN [ 15 ] , in which the encoder is the first few layers of a pre-trained VGG-19 image classification network . The decoder is trained to restore the input image from the output of the encoder . After we obtain an autoencoder that can perform the identity mapping on input images , we plug in an AdvBN module after the encoder . To compute cross-entropy loss for the AdvBN module , the remaining downstream layers of the aforementioned VGG-19 classifier are used , which takes the encoded feature as input and outputs the class prediction . The features perturbed by AdvBN are then fed into the decoder to create our visualizations . In Figure 2 , images with adversarial feature distributions exhibit differences in color , texture , and edges . We draw two major conclusions from these visualizations which highlight the adversarial properties of these domains . The first one concerns textures : CNNs have been shown to rely heavily on image textures for classification [ 6 ] . Images from the adversarial domain , on the other hand , have inconsistent textures across samples . For example , the furry texture of a dog is smoothed in Figure 2 ( b ) , and the stripes disappear from a zebra in Figure 2 ( d ) , whereas visible textures appear in ( f ) and ( g ) of Figure 2 . The second conclusion pertains to color . Previous study [ 52 ] suggests that colors serve as important information for CNNs . In the adversarial domain , we find suppressed colors ( Figure 2 ( a ) , ( c ) ) and unnatural hue ( Figure 2 ( e ) , ( g ) ) . Figure 3 illustrates how the appearance of reconstructed images shifts as adversarial perturbations to feature statistics become larger . See Appendix B for additional examples . We use this visualization technique to process the entire ImageNet validation set and denote it as ImageNet-AdvBN in Figure 1 . By evaluating different methods on this dataset , we observe that performance on ImageNet-AdvBN is consistently degraded , which validates the adversarial property of features generated by AdvBN . Experiments concerning performance on ImageNet-AdvBN are listed in Appendix B .

Summary: This paper proposes to improve robustness on ImageNet classification by adversarially optimizing the batch-norm statistics of a pretrained classifier. The initial layers of the classifier are frozen, and only the upper subnetwork of the classifier is finetuned after the addition of the adversarial batch-norm layer. The benefits of the method often appear in conjunction with image augmentation methods, where applying AdvBN on top of an image augmentation method tends to result in better performance.