Cross-Iteration Batch Normalization

1 INTRODUCTION . Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) has played a significant role in the success of deep neural networks . It was introduced to address the issue of internal covariate shift , where the distribution of network activations changes during training iterations due to the updates of network parameters . This shift is commonly believed to be disruptive to network training , and BN alleviates this problem through normalization of the network activations by their mean and variance , computed over the examples within the mini-batch at each iteration . With this normalization , network training can be performed at much higher learning rates and with less sensitivity to weight initialization . In BN , it is assumed that the distribution statistics for the examples within each mini-batch reflect the statistics over the full training set . While this assumption is generally valid for large batch sizes , it breaks down in the small batch size regime ( Peng et al. , 2018 ; Wu & He , 2018 ; Ioffe , 2017 ) , where noisy statistics computed from small sets of examples can lead to a dramatic drop in performance . This problem hinders the application of BN to memory-consuming tasks such as object detection ( Ren et al. , 2015 ; Dai et al. , 2017 ) , semantic segmentation ( Long et al. , 2015 ; Chen et al. , 2017 ) and action recognition ( Wang et al. , 2018b ) , where batch sizes are limited due to memory constraints . Towards improving estimation of statistics in the small batch size regime , alternative normalizers have been proposed . Several of them , including Layer Normalization ( LN ) ( Ba et al. , 2016 ) , Instance Normalization ( IN ) ( Ulyanov et al. , 2016 ) , and Group Normalization ( GN ) ( Wu & He , 2018 ) , compute the mean and variance over the channel dimension , independent of batch size . Different channel-wise normalization techniques , however , tend to be suitable for different tasks , depending on the set of channels involved . On the other hand , synchronized BN ( SyncBN ) ( Peng et al. , 2018 ) yields consistent improvements by processing larger batch sizes across multiple GPUs . These gains in performance come at the cost of additional overhead needed for synchronization across the devices . A seldom explored direction for estimating better statistics is to compute them over the examples from multiple recent training iterations , instead of from only the current iteration as done in previous techniques . This can substantially enlarge the pool of data from which the mean and variance are obtained . However , there exists an obvious drawback to this approach , in that the activation values from different iterations are not comparable to each other due to the changes in network weights . As shown in Figure 1 , directly calculating the statistics over multiple iterations , which we refer to as Naive CBN , results in lower accuracy . In this paper , we present a method that compensates for the network weight changes among iterations , so that examples from preceding iterations can be effectively used to improve batch normalization . Our method , called Cross-Iteration Batch Normalization ( CBN ) , is motivated by the observation that network weights change gradually , instead of abruptly , between consecutive training iterations , thanks to the iterative nature of Stochastic Gradient Descent ( SGD ) . As a result , the mean and variance of examples from recent iterations can be well approximated for the current network weights via a low-order Taylor polynomial , defined on gradients of the statistics with respect to the network weights . The compensated means and variances from multiple recent iterations are averaged with those of the current iteration to produce better estimates of the statistics . In the small batch size regime , CBN leads to appreciable performance improvements over the original BN , as exhibited in Figure 1 . The superiority of our proposed approach is further demonstrated through more extensive experiments on ImageNet classification and object detection on COCO . These gains are obtained with negligible overhead , as the statistics from previous iterations have already been computed and Taylor polynomials are simple to evaluate . With this work , it is shown that cues for batch normalization can successfully be extracted along the time dimension , opening a new direction for investigation . 2 RELATED WORK . The importance of normalization in training neural networks has been recognized for decades ( LeCun et al. , 1998 ) . In general , normalization can be performed on three components : input data , hidden activations , and network parameters . Among them , input data normalization is used most commonly because of its simplicity and effectiveness ( Sola & Sevilla , 1997 ; LeCun et al. , 1998 ) . After the introduction of Batch Normalization ( Ioffe & Szegedy , 2015 ) , the normalization of activations has become nearly as prevalent . By normalizing hidden activations by their statistics within each mini-batch , BN effectively alleviates the vanishing gradient problem and significantly speeds up the training of deep networks . To mitigate the mini-batch size dependency of BN , a number of variants have been proposed , including Layer Normalization ( LN ) ( Ba et al. , 2016 ) , Instance Normalization ( IN ) ( Ulyanov et al. , 2016 ) , Group Normalization ( GN ) ( Wu & He , 2018 ) , and Batch Instance Normalization ( BIN ) ( Nam & Kim , 2018 ) . The motivation of LN is to explore more suitable statistics for sequential models , while IN performs normalization in a manner similar to BN but with statistics only for each instance . GN achieves a balance between IN and LN , by dividing features into multiple groups along the channel dimension and computing the mean and variance within each group for normalization . BIN introduces a learnable method for automatically switching between normalizing and maintaining style information , enjoying the advantages of both BN and IN on style transfer tasks . Cross-GPU Batch Normalization ( CGBN or SyncBN ) ( Peng et al. , 2018 ) extends BN across multiple GPUs for the purpose of increasing the effective batch size . Though providing higher accuracy , it introduces synchronization overhead to the training process . Kalman Normalization ( KN ) ( Wang et al. , 2018a ) presents a Kalman filtering procedure for estimating the statistics for a network layer from the layer ’ s observed statistics and the computed statistics of previous layers . Batch Renormalization ( BRN ) ( Ioffe , 2017 ) is the first attempt to utilize the statistics of recent iterations for normalization . It does not compensate for the statistics from recent iterations , but rather it down-weights the importance of statistics from distant iterations . This down-weighting heuristic , however , does not make the resulting statistics “ correct '' , as the statistics from recent iterations are not of the current network weights . BRN can be deemed as a special version of our Naive CBN baseline ( without Taylor polynomial approximation ) , where distant iterations are down-weighted . Recent work have also investigated the normalization of network parameters . In Weight Normalization ( WN ) ( Salimans & Kingma , 2016 ) , the optimization of network weights is improved through a reparameterization of weight vectors into their length and direction . Weight Standardization ( WS ) ( Qiao et al. , 2019 ) instead reparameterizes weights based on their first and second moments for the purpose of smoothing the loss landscape of the optimization problem . To combine the advantages of multiple normalization techniques , Switchable Normalization ( SN ) ( Luo et al. , 2018 ) and Sparse Switchable Normalization ( SSN ) ( Shao et al. , 2019 ) make use of differentiable learning to switch among different normalization methods . The proposed CBN takes an activation normalization approach that aims to mitigate the mini-batch dependency of BN . Different from existing techniques , it provides a way to effectively aggregate statistics across multiple training iterations . 3 METHOD . 3.1 REVISITING BATCH NORMALIZATION . The original batch normalization ( BN ) ( Ioffe & Szegedy , 2015 ) whitens the activations of each layer by the statistics computed within a mini-batch . Denote θt and xt , i ( θt ) as the network weights and the feature response of a certain layer for the i-th example in the t-th mini-batch . With these values , BN conducts the following normalization : x̂t , i ( θt ) = xt , i ( θt ) −µt ( θt ) √ σt ( θt ) 2 + ε , ( 1 ) where x̂t , i ( θt ) is the whitened activation with zero mean and unit variance , ε is a small constant added for numerical stability , and µt ( θt ) and σt ( θt ) are the mean and variance computed for all the examples from the current mini-batch , i.e. , µt ( θt ) = 1 m m ∑ i=1 xt , i ( θt ) , ( 2 ) σt ( θt ) = √ 1 m m ∑ i=1 ( xt , i ( θt ) −µt ( θt ) ) 2 = √ νt ( θt ) −µt ( θt ) 2 , ( 3 ) where νt ( θt ) = 1m ∑ m i=1 xt , i ( θt ) 2 , and m denotes the number of examples in the current mini-batch . The whitened activation x̂t , i ( θt ) further undergoes a linear transform with learnable weights , to increase its expressive power : yt , i ( θt ) = γ x̂t , i ( θt ) +β , ( 4 ) where γ and β are the learnable parameters ( initialized to γ = 1 and β = 0 in this work ) . When the batch size m is small , the statistics µt ( θt ) and σt ( θt ) become noisy estimates of the training set statistics , thus degrading the effects of batch normalization . In the ImageNet classification task for which the BN module was originally designed , a batch size of 32 is typical . However , for other tasks requiring larger models and/or higher image resolution , such as object detection , semantic segmentation and video recognition , the typical batch size may be as small as 1 or 2 due to GPU memory limitations . The original BN becomes considerably less effective in such cases . 3.2 LEVERAGING STATISTICS FROM PREVIOUS ITERATIONS . To address the issue of BN with small mini-batches , a naive approach is to compute the mean and variance over the current and previous iterations . However , the statistics µt−τ ( θt−τ ) and νt−τ ( θt−τ ) of the ( t−τ ) -th iteration are computed under the network weights θt−τ , making them obsolete for the current iteration . As a consequence , directly aggregating statistics from multiple iterations produces inaccurate estimates of the mean and variance , leading to significantly worse performance . We observe that the network weights change smoothly between consecutive iterations , due to the nature of gradient-based training . This allows us to approximate µt−τ ( θt ) and νt−τ ( θt ) from the readily available µt−τ ( θt−τ ) and νt−τ ( θt−τ ) via a Taylor polynomial , i.e. , µt−τ ( θt ) = µt−τ ( θt−τ ) + ∂ µt−τ ( θt−τ ) ∂θt−τ ( θt −θt−τ ) +O ( ||θt −θt−τ ||2 ) , ( 5 ) νt−τ ( θt ) = νt−τ ( θt−τ ) + ∂νt−τ ( θt−τ ) ∂θt−τ ( θt −θt−τ ) +O ( ||θt −θt−τ ||2 ) , ( 6 ) where ∂ µt−τ ( θt−τ ) /∂θt−τ and ∂νt−τ ( θt−τ ) /∂θt−τ are gradients of the statistics with respect to the network weights , and O ( ||θt −θt−τ ||2 ) denotes higher-order terms of the Taylor polynomial , which can be omitted since the first-order term dominates when ( θt −θt−τ ) is small . In Eq . ( 5 ) and Eq . ( 6 ) , the gradients ∂ µt−τ ( θt−τ ) /∂θt−τ and ∂νt−τ ( θt−τ ) /∂θt−τ can not be precisely determined at a negligible cost because the statistics µ lt−τ ( θt−τ ) and ν lt−τ ( θt−τ ) for a node at the l-th network layer depend on all the network weights prior to the l-th layer , i.e. , ∂ µ lt−τ ( θt−τ ) /∂θ rt−τ 6= 0 and ∂ν lt−τ ( θt−τ ) /∂θ rt−τ 6= 0 for r ≤ l , where θ rt−τ denotes the network weights at the r-th layer . Only when r = l can these gradients be derived in closed form efficiently . Empirically , we find that as the layer index r decreases ( r≤ l ) , the partial gradients ∂ µ l t ( θt ) θ rt and ∂ν l t ( θt ) θ rt rapidly diminish . These reduced effects of network weight changes at earlier layers on the activation distributions in later layers may perhaps be explained by the reduced internal covariate shift of BN . Motivated by this phenomenon , which is studied in Appendix C , we propose to truncate these partial gradients at layer l. Thus , we further approximate Eq . ( 5 ) and Eq . ( 6 ) by µ lt−τ ( θt ) ≈ µ lt−τ ( θt−τ ) + ∂ µ lt−τ ( θt−τ ) ∂θ lt−τ ( θ lt −θ lt−τ ) , ( 7 ) ν lt−τ ( θt ) ≈ ν lt−τ ( θt−τ ) + ∂ν lt−τ ( θt−τ ) ∂θ lt−τ ( θ lt −θ lt−τ ) . ( 8 ) A naive implementation of ∂ µ lt−τ ( θt−τ ) /∂θ lt−τ and ∂ν lt−τ ( θt−τ ) /∂θ lt−τ involves computational overhead of O ( Cl ×Cl ×Cl−1×K ) , where Cl and Cl−1 denote the channel dimension of the l-th layer and the ( l−1 ) -th layer , respectively , and K denotes the kernel size of θ lt−τ . Here we find that the operation can be implemented efficiently in O ( Cl×Cl−1×K ) , thanks to the averaging over feature responses of µ and ν . See Appendix B for the details .

Summary: This paper proposes a novel Cross-Iteration Batch Normalization (CBN) to address the limitation of BN in the case of small mini-batch sizes. Different from existing methods, CBN exploits the statistics cross different iterations to obtain more accurate estimates of the data statistics. Specifically, the proposed CBN uses Taylor polynomials to approximate the statistics using the information of multiple recent iterations. The experiments on both image classification and object detection tasks demonstrate the effectiveness of the proposed method.