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: Paper summary: This paper proposes a new normalization technique specially designed for settings with small mini-batch sizes (where previous methods like BatchNorm are known to suffer). The approach aggregates mean/variance statistics from previous iterations, weighted based on the Taylor expansion, to get a better estimate of population statistics. The authors evaluate their approach on ImageNet classification, and object detection and instance segmentation on the COCO dataset.