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: The paper studies the problem of making predictive models robust to distributional shifts. This is achieved using an adversarial training strategy that adversarially perturbs batch norm statistics of a pre-trained model, which effectively captures variations in image style such as color, texture, etc. This worst-case optimization strategy exposes the model to larger variations in image statistics and semantics rather than just just pixel-based perturbations, and as a result is expected to produce better robustness under large distribution shifts.