Sandwich Batch Normalization

1 INTRODUCTION . This paper presents a simple , light-weight , and easy-to-implement modification of Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) , yet strongly motivated by various observations ( Zając et al. , 2019 ; Deecke et al. , 2018 ; Xie et al. , 2019 ; Xie & Yuille , 2019 ) drawn from a number of application fields , that BN has troubles standardizing hidden features with very heterogeneous structures , e.g. , from a multi-modal distribution . We call the phenomenon feature distribution heterogeneity . Such heterogeneity of hidden features could arise from multiple causes , often application-dependent : • One straightforward cause is due to input data heterogeneity . For example , when training a deep network on a diverse set of visual domains , that possess significantly different statistics , BN is found to be ineffective at normalizing the activations with only a single mean and variance ( Deecke et al. , 2018 ) , and often needs to be re-set or adapted ( Li et al. , 2016 ) . • Another intrinsic cause could arise from model heterogeneity , i.e. , when the training is , or could be equivalently viewed as , on a set of different models . For instance , in neural architecture search ( NAS ) using weight sharing ( Liu et al. , 2018 ; Dong & Yang , 2019 ) , training the super-network during the search phase could be considered as training a large set of sub-models ( with many overlapped weights ) simultaneously . As another example , for conditional image generation ( Miyato et al. , 2018 ) , the generative model could be treated as a set of category-specific sub-models packed together , one of which would be “ activated ” by the conditional input each time . The vanilla BN ( Figure 1 ( a ) ) fails to perform well when there is data or model heterogeneity . Recent trends split the affine layer into multiple ones and leverage input signals to modulate or select between them ( De Vries et al. , 2017 ; Deecke et al. , 2018 ) ( Figure 1 ( b ) ) ; or even further , utilize several independent BNs to address such disparity ( Zając et al. , 2019 ; Xie et al. , 2019 ; Xie & Yuille , 2019 ; Yu et al. , 2018 ) ( Figure 1 ( c ) ) . While those relaxations alleviate the data or model heterogeneity , we suggest that they might be “ too loose ” in terms of the normalization or regularization effect . Let us take adversarial training ( AT ) ( Madry et al. , 2017 ) as a concrete motivating example to illustrate our rationale . AT is by far the most effective approach to improve a deep model ’ s adversarial robustness . The model is trained by a mixture of the original training set ( “ clean examples ” ) and Norm ! . . . Norm '' Normalization Affine Normalization Af'ine ! . . . Af'ine '' Meta Affine Af'ine ! . . . Af'ine '' Feature map Feature map Conditional information Feature map Conditional information ( a ) Batch Norm ( b ) Categorical Conditional Batch Norm ( d ) Meta Batch Norm Meta Affine Af'ine ! . . . Af'ine '' Norm ! . . . Norm '' Feature map Conditional information ( e ) Meta Batch Norm++ Normalization Af'ine ! . . . Af'ine '' Feature map Conditional information ( c ) Auxiliary Batch Norm Normalization Affine Normalization Af'ine ! . . . Af'ine '' Feature map Feature map Conditional information ( a ) Batch Norm ( b ) Categorical Conditional Batch Norm Sandwich Affine Af'ine ! . . . Af'ine '' Feature map ( c ) Sandwich Batch Norm Normalization Conditional information Figure 1 : Illustration of ( a ) the original batch normalization ( BN ) , composed of one normalization layer and one affine layer ; ( b ) Categorical Conditional BN , composed of one normalization layer following a set of independent affine layers to intake conditional information ; ( c ) our proposed Sandwich BN , sequentially composed of one normalization layer , one shared sandwich affine layer , and a a set of independent affine layers . its attacked counterpart with some small perturbations applied ( “ adversarial examples ” ) . Yet , latest works ( Xie et al. , 2019 ; Xie & Yuille , 2019 ) pointed out that clean and adversarial examples behave like two different domains with distinct statistics on the feature level ( Li & Li , 2017 ; Pang et al. , 2018 ) . Such data heterogeneity puts vanilla BN in jeopardy for adversarial training , where the two domains are treated as one . ( Xie et al. , 2019 ; Xie & Yuille , 2019 ) demonstrated a helpful remedy to improve AT performance by using two separate BNs for clean and adversarial examples respectively , which allows either BN to learn more stable and noiseless statistics over its own focused domain . But what may be missing ? Unfortunately , using two separate BNs ignores one important fact that the two domains , while being different , are not totally independent . Considering that all adversarial images are generated by perturbing clean counterparts only minimally , it is convincing to hypothesize the two domains to be largely overlapped at least ( i.e. , they still share some hidden features despite the different statistics ) . To put it simple : while it is oversimplified to normalize the two domains as “ same one ” , it is also unfair and unnecessary to treat them as “ disparate two ” . More application examples can be found that all share this important structural feature prior , that we ( informally ) call as “ harmony in diversity ” . For instance , weight-sharing NAS algorithms ( Liu et al. , 2018 ; Dong & Yang , 2019 ; Yu et al. , 2018 ) train a large variety of child models , constituting model heterogeneity ; but most child architectures inevitably have many weights in common since they are sampled from the same super net . Similarly , while a conditional GAN ( Miyato et al. , 2018 ) has to produce diverse images classes , those classes often share the same resolution and many other dataset-specific characteristics ( e.g. , the object-centric bias for CIFAR images ) ; that is even more true when the GAN is trained to produce classes of one super-category , e.g. , dogs and cats . Our Contributions : Recognizing the need to address feature normalization with “ harmony in diversity ” , we propose a new SaBN as illustrated in Fig 1 ( c ) . SaBN modifies BN in a “ frustratingly simple ” way : it is equipped with two cascaded affine layers : a shared unconditional sandwich affine layer , followed by a set of independent affine layers that can be conditioned . Compared to Categorical Conditional BN , the new sandwich affine layer is designed to inject an inductive bias , that all re-scaling transformations will have a shared factor , indicating the commodity . Experiments on the applications of NAS and conditional generation demonstrate that SaBN addresses the model heterogeneity issue elegantly , and improves their performance in a plug-and-play fashion . To better address the data heterogeneity altogether , SaBN could further integrate the idea of split/auxiliary BNs ( Zając et al. , 2019 ; Xie et al. , 2019 ; Xie & Yuille , 2019 ; Yu et al. , 2018 ) , to decompose the normalization layer into multiple parallel ones . That yields the new variant called SaAuxBN . We demonstrate it using the application example of adversarial training . Lastly , we extend the idea of SaBN to Adaptive Instance Normalization ( AdaIN ) ( Huang & Belongie , 2017 ) , and show the resulting SaAdaIN to improve arbitrary style transfer . 2 RELATED WORK . 2.1 NORMALIZATION IN DEEP LEARNING . Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) made critical contributions to training deep convolutional networks and nowadays becomes a cornerstone of the latter for numerous tasks . BN normalizes the input mini-batch of samples by the mean and variance , and then re-scale them with learnable affine parameters . The success of BNs was initially attributed to overcoming internal co- variate shift ( Ioffe & Szegedy , 2015 ) , but later on raises many open discussions on its effect of improving landscape smoothness ( Santurkar et al. , 2018 ) ; enabling larger learning rates ( Bjorck et al. , 2018 ) and reducing gradient sensitivity ( Arora et al. , 2018 ) ; preserving the rank of pre-activation weight matrices ( Daneshmand et al. , 2020 ) ; decoupling feature length and direction ( Kohler et al. , 2018 ) ; capturing domain-specific artifacts ( Li et al. , 2016 ) ; reducing BN ’ s dependency on batch sizeIoffe ( 2017 ) ; Singh & Krishnan ( 2020 ) ; Preventing model from elimination singularities Qiao et al . ( 2019 ) ; and even characterizing an important portion of network expressivity ( Frankle et al. , 2020 ) . Inspired by BN , a number of task-specific modifications are proposed by exploiting different normalization axes , such as Instance Normalization ( IN ) ( Ulyanov et al. , 2016 ) for style transfer ; Layer Normalization ( LN ) ( Ba et al. , 2016 ) for recurrent networks ; Group Normalization ( GN ) ( Wu & He , 2018 ) for tackling small batch sizes ; StochNorm Kou et al . ( 2020 ) for fine-tuning task ; Passport-aware NormalizationZhang et al . ( 2020 ) for model IP protection ; and Li et al . ( 2019a ) ; Wang et al . ( 2020 ) ; Zheng et al . ( 2020 ) for image generation . Several normalization variants have been proposed by modulating BN parameters , mostly the affine layer ( mean and variance ) , to improve the controlling flexibility for more sophisticated usage . For example , Harm et al . ( De Vries et al. , 2017 ) presents Conditional BN , whose affine parameters are generated as a function of the input . Similarly , Conditional IN ( Dumoulin et al. , 2016 ) assigns each style with independent IN affine parameters . In ( Miyato et al. , 2018 ) , the authors developed Categorical Conditional BN for conditional GAN image generation , where each generated class has its independent affine parameters . Huang & Belongie ( Huang & Belongie , 2017 ) presented Adaptive IN ( AdaIN ) , which used the mean and variance of style image to replace the original affine parameter , achieving arbitrary style transfer . Spatial adaptivity ( Park et al. , 2019 ) and channel attention ( Li et al. , 2019b ) managed to modulate BN with higher complexities . A few latest works investigate to use multiple normalization layer instead of one in BN . ( Deecke et al. , 2018 ) developed mode normalization by employing a mixture-of-experts to separate incoming data into several modes and separately normalizing each mode . ( Zając et al. , 2019 ) used two separate BNs to address the domain shift between labeled and unlabeled data in semi-supervised learning . Very recently , ( Xie & Yuille , 2019 ; Xie et al. , 2019 ) reveal the two-domain issue in adversarial training and find improvements by using two separate BNs ( AuxBN ) . 2.2 BRIEF BACKGROUNDS FOR RELATED APPLICATIONS . We leverage four important applications as test beds . All of them appear to be oversimplified by using the vanilla BN , where the feature homogeneity and heterogeneity are not properly handled . We briefly introduce them below , and will concretely illustrate where the heterogeneity comes from and how our methods resolve the bottlenecks in Sec . 3 . Generative Adversarial Network Generative adversarial networks ( GANs ) have been prevailing since its origin ( Goodfellow et al. , 2014a ) for image generation . Many efforts have been made to improve GANs , such as modifying loss function ( Arjovsky et al. , 2017 ; Gulrajani et al. , 2017 ; Jolicoeur-Martineau , 2018 ) , improving network architecture ( Zhang et al. , 2018 ; Karras et al. , 2019 ; Gong et al. , 2019 ) and adjusting training procedure ( Karras et al. , 2017 ) . Recent works also tried to improve the generated image quality by proposing new normalization modules , such as Categorical Conditional BN and spectral normalization ( Miyato et al. , 2018 ) . Neural Architecture Search ( NAS ) The goal of NAS is to automatically search for an optimal model architecture for the given task and dataset . It was first proposed in ( Zoph & Le , 2016 ) where a reinforcement learning algorithm iteratively samples , trains and evaluates candidate models from the search space . Due to its prohibitive time cost , weight-sharing mechanism was introduced ( Pham et al. , 2018 ) and becomes a popular strategy to accelerate the training of sampled models ( Liu et al. , 2018 ) . However , weight-sharing causes performance deterioration due to unfair training ( Chu et al. , 2019 ) . In addition , a few NAS benchmarks ( Ying et al. , 2019 ; Dong & Yang , 2020 ; Zela et al. , 2020 ) were recently released , with ground-truth accuracy for candidate models pre-recorded , enabling researchers to evaluate the performance of search method more easily . Adversarial Robustness Deep networks are notorious for the vulnerability to adversarial attacks ( Goodfellow et al. , 2014b ) . In order to enhance adversarial robustness , countless training approaches have been proposed . ( Dhillon et al. , 2018 ; Papernot & McDaniel , 2017 ; Xu et al. , 2017 ; Meng & Chen , 2017 ; Liao et al. , 2018 ; Madry et al. , 2017 ) . Among them , adversarial training ( AT ) ( Madry et al. , 2017 ) are arguably the strongest , which train the model over a mixture of clean and perturbed data . Overall , the normalization in AT has , to our best knowledge , not been studied in depth . A pioneer work ( Xie et al. , 2019 ) introduce an auxiliary batch norm ( AuxBN ) to improve the clean image recognition accuracy . Neural Style Transfer Style transfer is a technique generating a stylized image , by combining the content of one image with the style of another . Various improvements are made on the normalization methods in this area . Ulyanov et al . ( 2016 ) proposed Instance Normalization , improving the stylized quality of generated images . Conditional Instance Normalization ( Dumoulin et al. , 2016 ) and Adaptive Instance Normalization ( Huang & Belongie , 2017 ) are proposed , enabling the network to perform arbitrary style transfer .

Summary: In this paper, the authors propose Sandwich Affine strategy to separate the affine layer in BN into one shared sandwich affine layer, cascaded by several parallel independent affine layers. Such method should well address the inherent feature distribution heterogeneity in many tasks. Following this idea, the SaAuxBN and SaIN have also been introduced in this paper. The extensive experiments demonstrate the effectiveness of such methods in neural architecture search (NAS), image generation, adversarial training, and style transfer.