Differentiable Learning-to-Normalize via Switchable Normalization

1 INTRODUCTION . Normalization techniques are effective components in deep learning , advancing many research fields such as natural language processing , computer vision , and machine learning . In recent years , many normalization methods such as Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) , Instance Normalization ( IN ) ( Ulyanov et al. , 2016 ) , and Layer Normalization ( LN ) ( Ba et al. , 2016 ) have been developed . Despite their great successes , existing practices often employed the same normalizer in all normalization layers of an entire network , rendering suboptimal performance . Also , different normalizers are used to solve different tasks , making model design cumbersome . To address the above issues , we propose Switchable Normalization ( SN ) , which combines three types of statistics estimated channel-wise , layer-wise , and minibatch-wise by using IN , LN , and BN respectively . SN switches among them by learning their importance weights . By design , SN is adaptable to various deep networks and tasks . For example , the ratios of IN , LN , and BN in SN are compared in multiple tasks as shown in Fig.1 ( a ) . We see that using one normalization method uniformly is not optimal for these tasks . For instance , image classification and object detection prefer the combination of three normalizers . In particular , SN chooses BN more than IN and LN in image classification and the backbone network of object detection , while LN has larger weights in the box and mask heads . For artistic image style transfer ( Johnson et al. , 2016 ) , SN selects IN . For neural architecture search , SN is applied to LSTM where LN is preferable than group normalization ( GN ) ( Wu & He , 2018 ) , which is a variant of IN by dividing channels into groups . The selectivity of normalizers makes SN robust to minibatch size . As shown in Fig.1 ( b ) , when training ResNet50 ( He et al. , 2016 ) on ImageNet ( Deng et al. , 2009 ) with different batch sizes , SN is close to the “ ideal case ” more than BN and GN . For ( 8 , 32 ) as an example1 , ResNet50 trained with SN is able to achieve 76.9 % top-1 accuracy , surpassing BN and GN by 0.5 % and 1.0 % respectively . In general , SN obtains better or comparable results than both BN and GN in all batch settings . ∗The first two authors contribute equally . Corresponding to : pluo.lhi @ gmail.com , { renjiamin , pengzhanglin , zhangruimao , lijingyu } @ sensetime.com . 1In this work , minibatch size refers to the number of samples per GPU , and batch size is ‘ # GPUs ’ times ‘ # samples per GPU ’ . A batch setting is denoted as a 2-tuple , ( # GPUs , # samples per GPU ) . Overall , this work has three key contributions . ( 1 ) We introduce Switchable Normalization ( SN ) , which is applicable in both CNNs and RNNs/LSTMs , and improves the other normalization techniques on many challenging benchmarks and tasks including image recognition in ImageNet ( Russakovsky et al. , 2015 ) , object detection in COCO ( Lin et al. , 2014 ) , scene parsing in Cityscapes ( Cordts et al. , 2016 ) and ADE20K ( Zhou et al. , 2017 ) , artistic image stylization ( Johnson et al. , 2016 ) , neural architecture search ( Pham et al. , 2018 ) , and video recognition in Kinetics ( Kay et al. , 2017 ) . ( 2 ) The analyses of SN are presented where multiple normalizers can be compared and understood with geometric interpretation . ( 3 ) By enabling each normalization layer in a deep network to have its own operation , SN helps ease the usage of normalizers , pushes the frontier of normalization in deep learning , as well as opens up new research direction . We believe that all existing models could be reexamined with this new perspective . We ’ ll make the code of SN available and recommend it as an alternative of existing handcrafted approaches . In the following sections , we first present SN in Sec.2 and then discuss its relationships with previous work in Sec.3 . SN is evaluated extensively in Sec.4 . 2 SWITCHABLE NORMALIZATION ( SN ) . We describe a general formulation of a normalization layer and then present SN . A General Form . We take CNN as an illustrative example . Let h be the input data of an arbitrary normalization layer represented by a 4D tensor ( N , C , H , W ) , indicating number of samples , number of channels , height and width of a channel respectively as shown in Fig.2 . Let hncij and ĥncij be a pixel before and after normalization , where n ∈ [ 1 , N ] , c ∈ [ 1 , C ] , i ∈ [ 1 , H ] , and j ∈ [ 1 , W ] . Let µ and σ be a mean and a standard deviation . We have ĥncij = γ hncij − µ√ σ2 + + β , ( 1 ) where γ and β are a scale and a shift parameter respectively . is a small constant to preserve numerical stability . Eqn . ( 1 ) shows that each pixel is normalized by using µ and σ , and then re-scale and re-shift by γ and β . IN , LN , and BN share the formulation of Eqn . ( 1 ) , but they use different sets of pixels to estimate µ and σ . In other words , the num- bers of their estimated statistics are different . In general , we have µk = 1 |Ik| ∑ ( n , c , i , j ) ∈Ik hncij , σ 2 k = 1 |Ik| ∑ ( n , c , i , j ) ∈Ik ( hncij − µk ) 2 , ( 2 ) where k ∈ { in , ln , bn } is used to distinguish different methods . Ik is a set pixels and |Ik| denotes the number of pixels . Specifically , Iin , Iln , and Ibn are the sets of pixels used to compute statistics in different approaches . IN was established in the task of artistic image style transfer ( Johnson et al. , 2016 ; Huang & Belongie , 2017 ) . In IN , we have µin , σ2in ∈ RN×C and Iin = { ( i , j ) |i ∈ [ 1 , H ] , j ∈ [ 1 , W ] } , meaning that IN has 2NC elements of statistics , where each mean and variance value is computed along ( H , W ) for each channel of each sample . LN ( Ba et al. , 2016 ) was proposed to ease optimization of recurrent neural networks ( RNNs ) . In LN , we have µln , σ2ln ∈ RN×1 and Iln = { ( c , i , j ) |c ∈ [ 1 , C ] , i ∈ [ 1 , H ] , j ∈ [ 1 , W ] } , implying that LN has 2N statistical values , where a mean value and a variance value are computed in ( C , H , W ) for each one of the N samples . BN ( Ioffe & Szegedy , 2015 ) was first demonstrated in the task of image classification ( He et al. , 2016 ; Krizhevsky et al. , 2012 ) by normalizing the hidden feature maps of CNNs . In BN , we have µbn , σ 2 bn ∈ RC×1 and Ibn = { ( n , i , j ) |n ∈ [ 1 , N ] , i ∈ [ 1 , H ] , j ∈ [ 1 , W ] } , in the sense that BN treats each channel independently like IN , but not only normalizes across ( H , W ) , but also the N samples in a minibatch , leading to 2C elements of statistics . 2.1 FORMULATION OF SN . SN has an intuitive expression ĥncij = γ hncij − Σk∈Ωwkµk√ Σk∈Ωw′kσ 2 k + + β , ( 3 ) where Ω is a set of statistics estimated in different ways . In this work , we define Ω = { in , ln , bn } the same as above where µk and σ2k can be calculated by following Eqn. ( 2 ) . However , this strategy leads to large redundant computations . In fact , the three kinds of statistics of SN depend on each other . Therefore we could reduce redundancy by reusing computations , µin = 1 HW H , W∑ i , j hncij , σ 2 in = 1 HW H , W∑ i , j ( hncij − µin ) 2 , µln = 1 C C∑ c=1 µin , σ 2 ln = 1 C C∑ c=1 ( σ2in + µ 2 in ) − µ2ln , µbn = 1 N N∑ n=1 µin , σ 2 bn = 1 N N∑ n=1 ( σ2in + µ 2 in ) − µ2bn , ( 4 ) showing that the means and variances of LN and BN can be computed based on IN . Using Eqn . ( 4 ) , the computational complexity of SN is O ( NCHW ) , which is comparable to previous work . Furthermore , wk and w′k in Eqn . ( 3 ) are importance ratios used to weighted average the means and variances respectively . Eachwk orw′k is a scalar variable , which is shared across all channels . There are 3×2 = 6 importance weights in SN . We have Σk∈Ωwk = 1 , Σk∈Ωw′k = 1 , and ∀wk , w′k ∈ [ 0 , 1 ] , and define wk = eλk Σz∈ { in , ln , bn } eλz and k ∈ { in , ln , bn } . ( 5 ) Here each wk is computed by using a softmax function with λin , λln , and λbn as the control parameters , which can be learned by back-propagation ( BP ) . w′k are defined similarly by using another three control parameters λ′in , λ ′ ln , and λ ′ bn . Training . Let Θ be a set of network parameters ( e.g . filters ) and Φ be a set of control parameters that control the network architecture . In SN , we have Φ = { λin , λln , λbn , λ′in , λ′ln , λ′bn } . Training a deep network with SN is to minimize a loss function L ( Θ , Φ ) , where Θ and Φ can be optimized jointly by back-propagation ( BP ) . This training procedure is different from previous meta-learning algorithms such as network architecture search ( Colson et al. , 2007 ; Liu et al. , 2018 ; Pham et al. , 2018 ) . In previous work , Φ represents as a set of network modules with different learning capacities , where Θ and Φ were optimized in two BP stages iteratively by using two training sets that are nonoverlapped . For example , previous work divided an entire training set into a training and a validation set . However , if Θ and Φ in previous work are optimized in the same set of training data , Φ would choose the module with large complexity to overfit these data . In contrast , SN essentially prevents overfitting by choosing normalizers to improve both learning and generalization ability as discussed below . Analyses of SN . To understand SN , we theoretically compare SN with BN , IN , and LN by representing them using weight normalization ( WN ) ( Salimans & Kingma , 2016 ) that is independent of mean and variance . WN is computed as v w Tx ‖w‖2 , where w and x represent a filter and an image patch . WN normalizes the norm of each filter to 1 and rescales to v. Remark 1 . Let x be an image patch with zero mean and unit variance , wi be a filter of the ith channel , i ∈ { 1 , 2 , ... , C } , and v be the filter norm of WN . Eqn . ( 1 ) can be rewritten by ĥin = γ wi Tx ‖wi‖2 +β ; ĥbn = γ wi Tx ‖wi‖2 +β , s.t . γ ≤ v ; and ĥln = γ wi Tx ‖wi‖2+ ∑C j 6=i ‖wj‖2 +β . By combining them , SN in Eqn . ( 3 ) can be reformulated by ĥsn = winĥin +wbnĥbn +wlnĥln = γ wi Tx ‖wi‖2+wln ∑C j 6=i ‖wj‖2 + β , s.t . wbnγ ≤ v , where win , wln , wbn are the weights and γ ≤ v is a constraint . Remark 1 simplifies SN in Eqn . ( 3 ) , enabling us to compare different normalizers geometrically by formulating them with respect to WN . In Fig.3 , ĥin of IN can be computed similarly to WN with an additional bias β , where the norms of all filters are normalized to 1 and then rescaled to γ . As γ and v have the same learning dynamic , the length of γ would be identically to v ( see w1 , w2 of IN ) . Moreover , ĥbn in BN can be rewritten as WN with regularization over γ , making it shorter than v. Compared to IN and LN , Luo et al . ( 2019 ) shows that the regularization of BN improves generalization and increases angle between filters , preventing them from coadaptation ( see w1 , w2 of BN ) . Furthermore , ĥln in LN normalizes each filter among channels where the filter norm is less constrained than IN and BN . That is , LN allows γ > v to increase learning ability . Finally , ĥsn in SN inherits the benefits from all of them and enables balance between learning and generalization ability . For example , when the batch size is small , the random noise from the batch statistics of BN would be too strong . SN is able to maintain performance by decreasing wbn and increasing wln , such that the regularization from BN is reduced and the learning ability is enhanced by LN . This phenomenon is supported by our experiment . More results are provided in Appendix B. Variants of SN . SN has many extensions . For instance , a pretrained network with SN can be finetuned by applying the argmax function on its control parameters where each normalization layer selects only one normalizer , leading to sparse SN . For ( 8 , 32 ) as an example , SN with sparsity achieves top-1 accuracy of 77.0 % in ImageNet with ResNet50 , which is comparable to 76.9 % of SN without sparsity . Moreover , when the channels are divided into groups , each group could select its own normalizer to increase representation power of SN . Our preliminary results suggest that group SN performs better than SN in some senses . For instance , group SN with only two groups boosts the top-1 accuracy of ResNet50 to 77.2 % in ImageNet . The above two variants will be presented as future work due to the length of paper . This work focuses on SN where the importance weights are tied between channels . Inference . When applying SN in test , the statistics of IN and LN are computed independently for each sample , while BN uses batch average after training without computing moving average in each iteration . Here batch average is performed in two steps . First , we freeze the parameters of the network and all the SN layers , and feed the network with a certain number of mini-batches randomly chosen from the training set . Second , we average the means and variances produced by all these mini-batches in each SN layer . The averaged statistics are used by BN in SN . We find that batch average makes training converged faster than moving average . It can be computed by using a small amount of samples . For example , top-1 accuracies of ResNet50 on ImageNet by using batch average with 50k and all training samples are 76.90 % and 76.92 % respectively . They are trained much faster and slightly better than 76.89 % of moving average . Appendix A shows more results . Implementation . SN can be easily implemented in existing softwares such as PyTorch and TensorFlow . The backward computation of SN can be obtained by automatic differentiation ( AD ) in these softwares . Without AD , we need to implement back-propagation ( BP ) of SN , where the errors are propagated through µk and σ2k . We provide the derivations of BP in Appendix H .

Summary: In this work, the authors propose Switchable Normalization (SN), which *learns* to switch / select different normalization algorithms (including batch normalization (BN), Instance Normalization (IN), Layer Normalization (LN)) in layers of the networks and in different applications. The idea is motivated by observations (shown in Fig 1) that, 1) different tasks tend to have applied different normalization methods; 2) some normalization methods (e.g. BN) are fragile to very small batch size.