Beyond BatchNorm: Towards a Unified Understanding of Normalization in Deep Learning

1 Introduction . Normalization techniques are often necessary to effectively train deep neural networks ( DNNs ) [ 1 , 2 , 3 ] . Arguably , the most popular of these is BatchNorm [ 1 ] , whose success can be attributed to several beneficial properties that allow it to stabilize a DNN ’ s training dynamics : for example , ability to propagate informative activation patterns in deeper layers [ 4 , 5 ] ; reduced dependence on initialization [ 6 , 7 , 8 ] ; faster convergence via removal of outlier eigenvalues [ 9 , 10 ] ; auto-tuning of learning rates [ 11 ] , equivalent to modern adaptive optimizers [ 12 ] ; and smoothing of loss landscape [ 13 , 14 ] . However , depending on the application scenario , BatchNorm ’ s use can be of limited benefit or even a hindrance : for example , BatchNorm struggles when training with small batch-sizes [ 3 , 15 ] ; in settings with train-test distribution shifts , BatchNorm can undermine a model ’ s accuracy [ 16 , 17 ] ; in meta-learning , it can lead to transductive inference [ 18 ] ; and in adversarial training , it can hamper accuracy on both clean and adversarial examples by estimating incorrect statistics [ 19 , 20 ] . To either address specific shortcomings or to replace BatchNorm in general , several recent works propose alternative normalization layers ( interchangeably called normalizers in this paper ) . For example , Brock et al . [ 23 ] propose to match BatchNorm ’ s forward propagation behavior in Residual Email : { eslubana , dickrp } @ umich.edu , and hidenori.tanaka @ ntt-research.com * Work partially performed during an internship at Physics & Informatics Laboratories , NTT Research . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . networks [ 24 ] by replacing it with Scaled Weight Standardization [ 25 , 26 ] . Wu and He [ 3 ] design GroupNorm , a batch-independent method that groups multiple channels in a layer to perform normalization . Liu et al . [ 27 ] use an evolutionary algorithm to search for both normalizers and activation layers . Given the right training configuration , these works show their proposed normalizers often achieve similar test accuracy to BatchNorm and even outperform it on some benchmarks . This begs the question , are we ready to replace BatchNorm ? To probe this question , we plot training curves for models defined using different combinations of normalizer , network architecture , batch size , and learning rate on CIFAR-100 . As shown in Figure 1 , clear trends begin to emerge . For example , we see LayerNorm [ 2 ] often converges at a relatively slower speed ; Weight Normalization [ 28 ] can not be trained at all for ResNets ( with and without SkipInit [ 6 ] ) ; Instance Normalization [ 22 ] results in unstable training in deeper non-residual networks , especially with small batch-sizes . Overall , evaluating hundreds of models in different settings , we see evident success/failure modes exist for all normalization techniques , including BatchNorm . As we noted before , prior works have established several properties to help explain such success/failure modes for the specific case of BatchNorm . However , given the pursuit of alternative normalizers in recent works , these properties need to be generalized so that one can accurately determine how normalization techniques beyond BatchNorm affect DNN training . In this work , we take a first step towards this goal by extending known properties of BatchNorm at initialization to several alternative normalization techniques . As we show , these properties are highly predictive of a normalizer ’ s influence on DNN training and can help ascertain exactly when an alternative technique is capable of serving as a replacement for BatchNorm . Our contributions follow . • Stable Forward Propagation : In Section 3 , we show activations-based normalizers are provably able to prevent exploding variance of activations in ResNets , similar to BatchNorm [ 5 , 6 ] . Parametric normalizers like Weight Normalization [ 28 ] do not share this property ; however , we explain why architectural modifications proposed in recent works [ 6 , 7 ] can resolve this limitation . • Informative Forward Propagation : In Section 4 , we first show the ability of a normalizer to generate dissimilar activations for different inputs is a strong predictor of optimization speed . We then extend a known result for BatchNorm to demonstrate the rank of representations in the deepest layer of a Group-normalized [ 3 ] model is at least Ω ( √ width/Group Size ) . This helps us illustrate how use of GroupNorm can prevent high similarity of activations for different inputs if the group size is small , i.e. , the number of groups is large . This suggests Instance Normalization [ 22 ] ( viz. , GroupNorm with group size equal to 1 ) is most likely and LayerNorm [ 2 ] ( viz. , GroupNorm with group size equal to layer width ) is least likely to produce informative activations . • Stable Backward Propagation : In Section 5 , we show normalization techniques that rely on individual sample and/or channel statistics ( e.g. , Instance Normalization [ 22 ] ) suffer from an exacerbated case of gradient explosion [ 29 ] , often witnessing unstable backward propagation . We show this behavior is mitigated by grouping of channels in GroupNorm , thus demonstrating a speed–stability trade-off characterized by group size . Related Work : Due to its ubiquity , past work has generally focused on understanding BatchNorm [ 5 , 4 , 6 , 9 , 10 , 7 , 13 , 29 , 30 , 31 ] . A few works have studied LayerNorm [ 32 , 33 ] , due to its relevance in natural language processing . In contrast , we try to analyze normalization methods in deep learning in a general manner . As we show , we can identify properties in BatchNorm that readily generalize to other normalizers and are often predictive of the normalizer ’ s impact on training . Our analysis is inspired by a rich body of work focused on understanding randomly initialized DNNs [ 34 , 35 , 36 , 37 , 38 ] . Most related to us is the contemporary work by Labatie et al . [ 39 ] , who analyze the impact of different normalization layers on expressivity of activations and conclude LayerNorm leads to high similarity of activations in deeper layers . As we discuss , this result is in fact a special case of our Claim 3 . 2 Preliminaries : Normalization Layers for DNNs . We first clarify the notations and operations used by the normalizers discussed in this work . Specifically , we define operators µ { d } ( T ) and σ { d } ( T ) , which calculate the mean and standard deviation of a tensor T along the dimensions specified by set { d } . ‖T ‖ denotes the ` 2 norm of T . RMS { d } ( T ) denotes the root mean square of T along dimensions specified by set { d } . For example , for a vector v ∈ Rn , we have RMS { 1 } ( v ) = √∑ i v 2 i/n . We assume the outputs of these operators broadcast as per requirements . ρ ( . ) denotes the sigmoid function . We define symbols b , c , x to denote the batch , channel , and spatial dimensions . For feature maps in a CNN , x will include both the height and the width dimensions . The notation c/g denotes division of c neurons ( or channels ) into groups of size g. When grouping is performed , each group is normalized independently . Normalization Layers : We analyze ten normalization layers in this work . These layers were chosen to cover a broad range of ideas : e.g. , activations-based layers [ 1 , 40 ] , parametric layers [ 23 , 28 ] , hand-engineered layers [ 3 ] , AutoML designed layers [ 27 ] , and layers [ 22 , 2 , 4 ] that form building blocks of recent techniques [ 41 ] . 1 . Activations-Based Layers : BatchNorm ( BN ) [ 1 ] , Layer- Norm ( LN ) [ 2 ] , Instance Normalization ( IN ) [ 22 ] , GroupNorm ( GN ) [ 3 ] , Filter Response Normalization ( FRN ) [ 40 ] , Variance Normalization ( VN ) [ 4 ] , EvoNormBO [ 27 ] , and EvoNoRMSO [ 27 ] fall in this category . These layers function in the activation space . Note that Variance Normalization is an ablation of BatchNorm that does not use the mean-centering operation . Typically , given activations AL at layer L , these layers use an operation of the form Anorm = φ ( γ σ { d } ( AL ) ( AL − µ { d } ( AL ) ) + β ) ) . Here , γ and β are learned affine parameters used for controlling quantities affected by the normalization operations ( such as mean , standard deviation , and RMS ) and φ is a non-linearity , such as ReLU . The exact operations for these layers , minus the affine parameters , are shown in Table 1 . 2 . Parametric Layers : Weight Normalization ( WN ) [ 28 ] and Scaled Weight Standardization ( SWS ) [ 23 ] fall in this category . Table 1 shows the exact operations . These layers function in the parameter space and act on a filter ’ s weights ( W ) to generate normalized weights ( Wnorm ) . The normalized weightsWnorm are used for processing the input : AL+1 = φ ( Wnorm ∗ AL ) . 3 Stable Forward Propagation . Stable forward propagation is a necessary condition for successful DNN training [ 36 ] . In this section , we identify and demystify the role of normalization layers in preventing the problem of exploding or vanishing activations during forward propagation . These problems can result in training instability due to exploding or vanishing gradients during backward propagation [ 36 , 38 ] . Building on a previous study on BatchNorm , we first show that activations-based normalizers provably avoid exponential growth of variance in ResNets1 , ensuring training stability . Thereafter , we show parametric normalizers do not share this property and ensuring stable training requires explicit remedies . 3.1 Activations-Based Normalizers and Exponential Variance in Residual Networks . Hanin and Rolnick [ 38 ] show that for stable forward propagation in ResNets , the average variance of activations should not grow exponentially ( i.e. , should not explode ) . Interestingly , Figure 1 shows that all activations-based normalizers are able to train the standard ResNet [ 24 ] architecture stably . For BatchNorm , this behavior is provably expected . Specifically , De and Smith [ 6 ] find that to ensure variance along the batch-dimension is 1 , BatchNorm rescales the Lth layer ’ s residual path output by a factor of O ( 1/√L ) . This causes the growth of variance in a Batch-Normalized ResNet to be linear in depth , hence avoiding exponential growth of variance in and ensuring effective training . We now show this result can be extended to other normalization techniques too . Claim 1 . Similar to BatchNorm [ 6 ] , GroupNorm [ 3 ] avoids exponential growth of variance in ResNets , ensuring stable training . Proof . We follow the same setup as De and Smith [ 6 ] . Assume the Lth residual path ( fL ) is processed by a normalization layer N , after which it combines with the skip connection to generate the next output : yL = yL−1 + N ( fL ( yL−1 ) ) . The covariance of layer input and Residual path ’ s output is assumed to be zero . Hence , the output ’ s variance is : Var ( yL ) = Var ( yL−1 ) + Var ( N ( fL ( yL−1 ) ) ) . Now , assume GroupNorm with group size G is used for normalizing the D-dimensional activation signal , i.e. , N = GN ( . ) . This implies for the gth group , σg , x ( GN ( fL ( yL−1 ) ) ) = 1 . Then , for a batch of size N , denoting the ith sample activations as y ( i ) L , and using ( y ( i ) L ) j to index the activations , we note the residual output ’ s variance averaged along the spatial dimension is : 〈Var ( N ( fL ( yL−1 ) ) 〉 = 1D ∑D j=1 ( 1 N ∑N i=1 ( GN ( fL ( y ( i ) L−1 ) ) j ) 2 ) = 1 N ∑N i=1 ( 1 D ∑D j=1 ( GN ( fL ( yL−1 ) ( i ) ) j ) 2 ) = 1N ∑N i=1 G D ( ∑D/G g=1 ( σg , x ( GN ( fL ( yL−1 ) ( i ) ) ) ) 2 ) = 1 . Overall , this implies 〈Var ( yL ) 〉 = 〈Var ( yL−1 ) 〉 + 1 . Recursively applying this relationship for a bounded variance input , we see average variance at the Lth layer is in O ( L ) . Thus , similar to BatchNorm , use of GroupNorm will enable stable forward propagation in ResNets by ensuring signal variance grows linearly with depth . We show empirical demonstrations of Claim 1 in Figure 2 , where the average activation variance is plotted for a ResNet-56 . As can be seen , for all activations-based normalizers , the growth of variance is linear in the number of layers . At the end of a Residual module , which spatially downsamples the signal , the variance plummets . However , the remaining layers follow a pattern of linear growth , as expected by our result . We note our 1The case of non-residual networks is discussed in appendix . In brief , most normalizers help avoid exploding/vanishing activations by enforcing unit activation variance in the batch , channel , or spatial dimensions . theory does not apply to EvoNorms , which are designed via AutoML . However , empirically , we see EvoNorms also avoid exponential growth of variance in ResNets . Thus , our analysis shows , all activations-based normalizers in Table 1 share the beneficial property of stabilizing forward propagation in ResNets , similar to BatchNorm .

Summary: This paper summarizes multiple aspects of a variety of normalization schemes and what makes them work or not work. Specifically, the authors analyze properties that prior works have found to be critical in BatchNormalso in the context of other normalizers (like GroupNorm), thus arguing for a more unified understanding and insight towards these normalization components. The discussion and thus contribution can be broken into three parts: stable activation, discriminative representation and gradient explosion analysis.