Towards Stabilizing Batch Statistics in Backward Propagation of Batch Normalization

1 INTRODUCTION . Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) is one of the most popular techniques for training neural networks . It has been widely proven effective in many applications , and become the indispensable part of many state of the art deep models . Despite the success of BN , it ’ s still challenging to utilize BN when batch size is extremely small1 . The batch statistics with small batch size are highly unstable , leading to slow convergence during training and bad performance during inference . For example , in detection or segmentation tasks , the batch size is often limited to 1 or 2 per GPU due to the requirement of high resolution inputs or complex structure of the model . Directly computing batch statistics without any modification on each GPU will make performance of the model severely degrade . To address such issues , many modified normalization methods have been proposed . They can be roughly divided into two categories : some of them try to improve vanilla BN by correcting batch statistics ( Ioffe , 2017 ; Singh & Shrivastava , 2019 ) , but they all fail to completely restore the performance of vanilla BN ; Other methods get over the instability of BN by using instance-level normalization ( Ulyanov et al. , 2016 ; Ba et al. , 2016 ; Wu & He , 2018 ) , therefore models can avoid the affect ∗Equal Contribution . Work was done when Junjie Yan was an intern at Megvii Technology . †Corresponding author . 1In the context of this paper , we use ” batch size/normalization batch size ” to refer the number of samples used to compute statistics unless otherwise stated . We use ” gradient batch size ” to refer the number of samples used to update weights . of batch statistics . This type of methods can restore the performance in small batch cases to some extent . However , instance-level normalization hardly meet industrial or commercial needs so far , for this type of methods have to compute instance-level statistics both in training and inference , which will introduce additional nonlinear operations in inference procedure and dramatically increase consumption Shao et al . ( 2019 ) . While vanilla BN uses the statistics computed over the whole training data instead of batch of samples when training finished . Thus BN is a linear operator and can be merged with convolution layer during inference procedure . Figure 1 ( a ) shows with ResNet-50 ( He et al. , 2016 ) , instance-level normalization almost double the inference time compared with vanilla BN . Therefore , it ’ s a tough but necessary task to restore the performance of BN in small batch training without introducing any nonlinear operations in inference procedure . In this paper , we first analysis the formulation of vanilla BN , revealing there are actually not only 2 but 4 batch statistics involved in normalization during forward propagation ( FP ) as well as backward propagation ( BP ) . The additional 2 batch statistics involved in BP are associated with gradients of the model , and have never been well discussed before . They play an important role in regularizing gradients of the model during BP . In our experiments ( see Figure 2 ) , variance of the batch statistics associated with gradients in BP , due to small batch size , is even larger than that of the widelyknown batch statistics ( mean , variance of feature maps ) . We believe the instability of batch statistics associated with gradients is one of the key reason why BN performs poorly in small batch cases . Based on our analysis , we propose a novel normalization method named Moving Average Batch Normalization ( MABN ) . MABN can completely get over small batch issues without introducing any nonlinear manipulation in inference procedure . The core idea of MABN is to replace batch statistics with moving average statistics . We substitute batch statistics involved in BP and FP with different type of moving average statistics respectively , and theoretical analysis is given to prove the benefits . However , we observed directly using moving average statistics as substitutes for batch statistics can ’ t make training converge in practice . We think the failure takes place due to the occasional large gradients during training , which has been mentioned in Ioffe ( 2017 ) . To avoid training collapse , we modified the vanilla normalization form by reducing the number of batch statistics , centralizing the weights of convolution kernels , and utilizing renormalizing strategy . We also theoretically prove the modified normalization form is more stable than vanilla form . MABN shows its effectiveness in multiple vision public datasets and tasks , including ImageNet ( Russakovsky et al. , 2015 ) , COCO ( Lin et al. , 2014 ) . All results of experiments show MABN with small batch size ( 1 or 2 ) can achieve comparable performance as BN with regular batch size ( see Figure 1 ( b ) ) . Besides , it has same inference consumption as vanilla BN ( see Figure 1 ( a ) ) . We also conducted sufficient ablation experiments to verify the effectiveness of MABN further . 2 RELATED WORK . Batch normalization ( BN ) ( Ioffe & Szegedy , 2015 ) normalizes the internal feature maps of deep neural network using channel-wise statistics ( mean , standard deviation ) along batch dimension . It has been widely proven effectively in most of tasks . But the vanilla BN heavily relies on sufficient batch size in practice . To restore the performance of BN in small batch cases , many normalization techniques have been proposed : Batch Renormalization ( BRN ) ( Ioffe , 2017 ) introduces renormalizing parameters in BN to correct the batch statistics during training , where the renormalizing parameters are computed using moving average statistics ; Unlike BRN , EvalNorm ( Singh & Shrivastava , 2019 ) corrects the batch statistics during inference procedure . Both BRN and EvalNorm can restore the performance of BN to some extent , but they all fail to get over small batch issues completely . Instance Normalization ( IN ) ( Ulyanov et al. , 2016 ) , Layer Normalization ( LN ) ( Ba et al. , 2016 ) , and Group normalization ( GN ) ( Wu & He , 2018 ) all try to avoid the effect of batch size by utilizing instance level statistics . IN uses channel-wise statistics per instance instead of per batch , while LN uses instance-level statistics along channel dimension . But IN and LN shows no superiority to vanilla BN in most of cases . GN divides all channels in predefined groups , and uses group-wise statistics per instance . It can restore the performance of vanilla BN very well in classification and detection tasks . But it have to introduce extra nonlinear manipulations in inference procedure and severely increase inference consumption , as we have pointed out in Section 1 . SyncBN ( Peng et al. , 2018 ) handle the small batch issues by computing the mean and variance across multiple GPUs . This method doesn ’ t essentially solve the problem , and requires a lot of resource . Online Normalization Chiley et al . ( 2019 ) modifies BP by using moving average statistics , so they can set batch size as 1 without degradation of performance , but Online Normalization still have to use instance-level normalization to cooperate with modification in BP , so its inference efficiency is much lower than original BN . Apart from operating on feature maps , some works exploit to normalize the weights of convolution : Weight Standardization ( Qiao et al. , 2019 ) centralizes weight at first before divides weights by its standard deviation . It still has to combine with GN to handle small batch cases . 3 STATISTICS IN BATCH NORMALIZATION . 3.1 REVIEW OF BATCH NORMALIZATION . First of all , let ’ s review the formulation of batch Normalization ( Ioffe & Szegedy , 2015 ) : assume the input of a BN layer is denoted as X ∈ RB×p , where B denotes the batch size , p denotes number of features . In training procedure , the normalized feature maps Y at iteration t is computed as : Y = X − µBt σBt , ( 1 ) where batch statistics µBt and σ 2 Bt are the sample mean and sample variance computed over the batch of samples Bt at iteration t : µBt = 1 B ∑ b Xb , : , σ 2 Bt = 1 B ∑ b ( Xb , : − µBt ) 2 . ( 2 ) Besides , a pair of parameters γ , β are used to scale and shift normalized value Y : Z = Y γ + β . ( 3 ) The scaling and shifting part is added in all normalization form by default , and will be omitted in the following discussion for simplicity . As Ioffe & Szegedy ( 2015 ) demonstrated , the batch statistics µBt , σ 2 Bt are both involved in backward propagation ( BP ) . We can derive the formulation of BP in BN as follows : let L denote the loss , Θt denote the set of the whole learnable parameters of the model at iteration t. Given the partial gradients ∂L∂Y ∣∣∣∣ Θt , Bt , the partial gradients ∂L∂X ∣∣∣∣ Θt , Bt is computed as ∂L ∂X ∣∣∣∣ Θt , Bt = 1 σBt ( ∂L ∂Y ∣∣∣∣ Θt , Bt − gBt − Y ·ΨBt ) ( 4 ) where · denotes element-wise production , gBt and ΨBt are computed as gBt = 1 B ∑ b ∂L ∂Yb , : ∣∣∣∣ Θt , Bt , ΨBt = 1 B ∑ b Yb , : · ∂L ∂Yb , : ∣∣∣∣ Θt , Bt , ( 5 ) It can be seen from ( 5 ) that gBt and ΨBt are also batch statistics involved in BN during BP . But they have never been well discussed before . 3.2 INSTABILITY OF BATCH STATISTICS . According to Ioffe & Szegedy ( 2015 ) , the ideal normalization is to normalize feature maps X using expectation and variance computed over the whole training data set : Y = X − EX√ V ar [ X ] . ( 6 ) But it ’ s impractical when using stochastic optimization . Therefore , Ioffe & Szegedy ( 2015 ) uses mini-batches in stochastic gradient training , each mini-batch produces estimates the mean and variance of each activation . Such simplification makes it possible to involve mean and variance in BP . From the derivation in section 3.1 , we can see batch statistics µBt , σ 2 Bt are the Monte Carlo ( MC ) estimators of population statistics E [ X|Θt ] , V ar [ X|Θt ] respectively at iteration t. Similarly , batch statistics gBt , ΨBt are MC estimators of population statistics E [ ∂L∂Yb , : |Θt ] , E [ Yb , : · ∂L ∂Yb , : |Θt ] at iteration t. E [ ∂L∂Yb , : |Θt ] , E [ Yb , : · ∂L ∂Yb , : |Θt ] are computed over the whole data set . They contain the information how the mean and the variance of population will change as model updates , so they play an important role to make trade off between the change of individual sample and population . Therefore , it ’ s crucial to estimate the population statistics precisely , in order to regularize the gradients of the model properly as weights update . It ’ s well known the variance of MC estimator is inversely proportional to the number of samples , hence the variance of batch statistics dramatically increases when batch size is small . Figure 2 shows the change of batch statistics from a specific normalization layer of ResNet-50 during training on ImageNet . Regular batch statistics ( orange line ) are regarded as a good approximation for population statistics . We can see small batch statistics ( blue line ) are highly unstable , and contains notable error compared with regular batch statistics during training . In fact , the bias of gBt and ΨBt in BP is more serious than that of µBt and σ 2 Bt ( see Figure 2 ( c ) , 2 ( d ) ) . The instability of small batch statistics can worsen the capacity of the models in two aspects : firstly the instability of small batch statistics will make training unstable , resulting in slow convergence ; Secondly the instability of small batch can produce huge difference between batch statistics and population statistics . Since the model is trained using batch statistics while evaluated using population statistics , the difference between batch statistics and population statistics will cause inconsistency between training and inference procedure , leading to bad performance of the model on evaluation data .

Summary: The paper extends recently proposed BatchRenormalization (BRN) technique which uses exponential moving average (EMA) statistics in forward and backward passes of BatchNorm (BN) instead of vanilla batch statistics. Motivation of the work is to stabilize training neural networks on small batch size setup. Authors propose to replace EMA in backward pass by simple moving average (SMA) and show that under some assumptions such replacement reduces variance. Also they consider slightly different way of normalization without centralizing features X, but centralizing convolutional kernels according to Qiao et al. (2019).