Proxy-Normalizing Activations to Match Batch Normalization while Removing Batch Dependence

1 Introduction . Normalization plays a critical role in deep learning as it allows successful scaling to large and deep models . In vision tasks , the most well-established normalization technique is Batch Normalization ( BN ) [ 1 ] . At every layer in the network , BN normalizes the intermediate activations to have zero mean and unit variance in each channel . While indisputably successful when training with large batch size , BN incurs a performance degradation in the regime of small batch size [ 2 , 3 , 4 , 5 , 6 , 7 ] . This performance degradation is commonly attributed to an excessive or simply unwanted regularization stemming from the noise in the approximation of full-batch statistics by mini-batch statistics . Many techniques have been proposed to avoid this issue , while at the same time retaining BN ’ s benefits . Some techniques mimic BN ’ s principle while decoupling the computational batch from the normalization batch [ 2 , 8 , 7 ] . Other techniques are “ batch-independent ” in that they operate independently of the batch in various modalities : through an explicit normalization either in activation space [ 9 , 10 , 11 , 12 , 13 , 3 , 5 , 14 , 15 , 6 ] or in weight space [ 16 , 17 , 18 , 19 , 20 ] ; through the use of an analytic proxy to track activation statistics [ 21 , 22 , 23 ] ; through a change of activation function [ 24 , 5 , 15 ] . In this paper , we push the endeavor to replace BN with a batch-independent normalization a step further . Our main contributions are as follows : ( i ) we introduce a novel framework to finely characterize the various neural network properties affected by the choice of normalization ; ( ii ) using this framework , we show that while BN ’ s beneficial properties are not retained when solely using the prototypical batch-independent norms , they are retained when combining some of these norms with “ Proxy Normalization ” , a novel technique that we hereby introduce ; ( iii ) we demonstrate on an extensive set of experiments that , by reproducing BN ’ s beneficial properties , our batch-independent normalization approach consistently matches or exceeds BN ’ s performance . As a starting point of our analysis , we need to gain a better understanding of those beneficial properties of BN that we aim at reproducing . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 2 Batch Normalization ’ s beneficial properties . We consider throughout this paper a convolutional neural network with d = 2 spatial axes . This neural network receives an input x ∈ RH×W×C0 which , unless otherwise stated , is assumed sampled from a finite dataset D. The neural network maps this input x to intermediate activations xl ∈ RH×W×Cl of height H , width W and number of channels Cl at each layer l. The value of xl at spatial position α ∈ { 1 , . . . , H } × { 1 , . . . , W } and channel c ∈ { 1 , . . . , Cl } is denoted as xlα , c , with the dependency on x kept implicit to avoid overloading notations . The inclusion of BN at layer l leads in the full-batch setting to adding the following operations ∀α , c : ylα , c = xlα , c − µc ( xl ) σc ( xl ) , ỹlα , c = γ l cy l α , c + β l c , ( 1 ) where µc ( xl ) , σc ( xl ) are the mean and standard deviation of xl in channel c , and γlc , β l c are channelwise scale and shift parameters restoring the degrees of freedom lost in the standardization . In the mini-batch setting , the full-batch statistics µc ( xl ) , σc ( xl ) are approximated by mini-batch statistics . Table 1 summarizes the beneficial properties that result from including BN in the neural network . Below , we provide details on each of these properties , and we discuss whether each property is reproduced with batch-independent norms . Scale invariance . When BN is present , the input-output mapping of the neural network is invariant to the scale of weights preceding any BN layer . With such scale invariance plus weight decay , the scale of weights during training reaches an equilibrium with an “ effective ” learning rate depending on both the learning rate and the weight decay strength [ 25 , 26 , 27 , 8 , 28 , 29 , 30 ] . Such mechanism of “ auto rate-tuning ” has been shown to provide optimization benefits [ 31 , 26 , 32 ] . This property is easy to reproduce . It is already obtained with most existing batch-independent norms . Control of activation scale in residual networks . To be trainable , residual networks require the scale of activations to remain well-behaved at initialization [ 33 , 34 , 35 , 20 , 36 , 37 ] . While this property naturally arises when BN is present on the residual path , when BN is not present it can also be enforced by a proper scaling decaying with the depth of the residual path . This “ dilution ” of the residual path with respect to the skip connection path reduces the effective depth of the network and enables to avoid coarse-grained failures modes [ 38 , 39 , 40 , 41 , 42 , 35 ] . This property is easy to reproduce . It is already obtained with most existing batch-independent norms . Regularizing noise . Due to the stochasticity of the approximation of µc ( xl ) , σc ( xl ) by mini-batch statistics , training a neural network with BN in the mini-batch setting can be seen as equivalent to performing Bayesian inference [ 43 , 44 ] or to adding a regularizing term to the training of the same network with full-batch statistics [ 45 ] . As a result , BN induces a specific form of regularization . This regularization is not reproduced with batch-independent norms , but we leave it out of the scope of this paper . To help minimize the bias in our analysis and “ subtract away ” this effect , we will perform all our experiments without and with extra degrees of regularization . This procedure can be seen as a coarse disentanglement of normalization ’ s effects from regularization ’ s effects . Avoidance of collapse . Unnormalized networks with non-saturating nonlinearities are subject to a phenomenon of “ collapse ” whereby the distribution with respect to x , α of the intermediate activation vectors ( xlα,1 , . . . , x l α , Cl ) T becomes close to zero- or one-dimensional in deep layers [ 39 , 46 , 41 , 47 , 48 , 49 , 37 ] . This means that deep in an unnormalized network : ( i ) layers tend to have their channels imbalanced ; ( ii ) nonlinearities tend to become channel-wise linear with respect to x , α and not add any effective capacity [ 39 , 50 , 51 ] . Consequently , unnormalized networks can neither effectively use their whole width ( imbalanced channels ) nor effectively use their whole depth ( channel-wise linearity ) . Conversely , when BN is used , the standardization at each layer prevents this collapse from happening . Even in deep layers , channels remain balanced and nonlinearities remain channel-wise nonlinear with respect to x , α. Consequently , networks with BN can effectively use their whole width and depth . The collapse is , on the other hand , not always avoided with batch-independent norms [ 39 , 18 , 49 ] . Most notably , it is not avoided with Layer Normalization ( LN ) [ 9 ] or Group Normalization ( GN ) [ 3 ] , as we show both theoretically and experimentally on commonly found networks in Section 4 . To the extent possible , we aim at designing a batch-independent norm that avoids this collapse . Preservation of expressivity . We can always express the identity with Eq . ( 1 ) by choosing βlc = µc ( x l ) and γlc = σc ( x l ) . Conversely , for any choice of βlc , γ l c , we can always “ re-absorb ” Eq . ( 1 ) into a preceding convolution with bias . This means that BN in the full-batch setting does not alter the expressivity compared to an unnormalized network , i.e . it amounts to a plain reparameterization of the hypothesis space . The expressivity is , on the other hand , not always preserved with batch-independent norms . In activation space , the dependence of batch-independent statistics on the input x turns the standardization into a channel-wise nonlinear operation that can not be “ re-absorbed ” into a preceding convolution with bias [ 18 ] . This phenomenon is most pronounced when statistics get computed over few components . This means e.g . that Instance Normalization ( IN ) [ 10 ] induces a greater change of expressivity than GN , which itself induces a greater change of expressivity than LN . In weight space , the expressivity can also be altered , namely by the removal of degrees of freedom . This is the case with Weight Standardization ( WS ) [ 18 , 20 ] and Centered Weight Normalization [ 17 ] that remove degrees of freedom ( one per unit ) that can not be restored in a succeeding affine transformation . This reduction of expressivity could explain the ineffectiveness of these techniques in EfficientNets [ 52 ] , as previous works observed [ 20 ] and as we confirm in Section 6 . To the extent possible , we aim at designing a batch-independent norm that preserves expressivity . 3 Theoretical framework of analysis . We specified the different properties that we wish to retain in our design of batch-independent normalization : ( i ) scale invariance , ( ii ) control of activation scale ; ( iii ) avoidance of collapse ; ( iv ) preservation of expressivity . We now introduce a framework to quantify the presence or absence of the specific properties ( iii ) and ( iv ) with various choices of normalization . Propagation . For simplicity , we assume in our theoretical setup that any layer l up to depth L consists of the following three steps : ( i ) convolution step with weights ωl ∈ RKl×Kl×Cl−1×Cl ; ( ii ) normalization step ; ( iii ) activation step sub-decomposed into an affine transformation with scale and shift parameters γl , βl ∈ RCl and an activation function φ which , unless otherwise stated , is assumed positive homogeneous and nonzero ( e.g . φ = ReLU ) . If we denote xl , yl , zl ∈ RH×W×Cl the intermediate activations situated just after ( i ) , ( ii ) , ( iii ) with the convention z0 ≡ x , we may write the propagation through layer l as xl = Conv ( zl−1 ) , ∀α , c : Conv ( zl−1 ) α , c = ( ωl ∗ zl−1 ) α , c , ( 2 ) yl = Norm ( xl ) , ∀α , c : Norm ( xl ) α , c = xlα , c − µIx , c ( xl ) σIx , c ( x l ) , 1 ( 3 ) zl = Act ( yl ) , ∀α , c : Act ( yl ) α , c = φ ( ỹlα , c ) = φ ( γlcy l α , c + β l c ) , ( 4 ) where µIx , c ( x l ) , σIx , c ( x l ) denote the mean and standard deviation of xl conditionally on Ix , c ≡ { c } , { x } , { x , c } , { x , c mod G } for the respective cases Norm = BN , LN , IN , GN with G groups . 1We omit the numerical stability constant and adopt the convention Norm ( xl ) α , c = 0 , ∀α if σIx , c ( x l ) = 0 . Moments . Extending the previous notations , we use µ , σ , P indexed with a ( possibly empty ) subset of variables to denote the operators of conditional mean , standard deviation and power . If we apply these operators to the intermediate activations yl , that implicitly depend on the input x and that explicitly depend on the spatial position α and the channel c , we get e.g . µc ( y l ) = Ex , α [ ylα , c ] , σc ( y l ) = √ Varx , α [ ylα , c ] , Pc ( yl ) = Ex , α [ ( ylα , c ) 2 ] ︸ ︷︷ ︸ µc ( yl ) 2+σc ( yl ) 2 , µx , c ( y l ) = Eα [ ylα , c ] , σx , c ( y l ) = √ Varα [ ylα , c ] , Px , c ( yl ) = Eα [ ( ylα , c ) 2 ] ︸ ︷︷ ︸ µx , c ( yl ) 2+σx , c ( yl ) 2 , where , by convention , x , α , c are considered uniformly sampled among inputs of D , spatial positions and channels , whenever they are considered as random . Power decomposition . Using these notations , we may gain important insights by decomposing the power in channel c of yl , just after the normalization step , as Pc ( yl ) = Ex [ µx , c ( y l ) ] 2 ︸ ︷︷ ︸ P ( 1 ) c ( yl ) + Varx [ µx , c ( y l ) ] ︸ ︷︷ ︸ P ( 2 ) c ( yl ) +Ex [ σx , c ( y l ) ] 2 ︸ ︷︷ ︸ P ( 3 ) c ( yl ) + Varx [ σx , c ( y l ) ] . ︸ ︷︷ ︸ P ( 4 ) c ( yl ) ( 5 ) Since this four-terms power decomposition will be at the core of our analysis , we detail two useful views of it . The first view is that of a hierarchy of scales : P ( 1 ) c ( yl ) measures the power of µc ( yl ) at the dataset scale ; P ( 2 ) c ( yl ) measures the power of µx , c ( yl ) − µc ( yl ) at the instance scale ; the sum of P ( 3 ) c ( yl ) and P ( 4 ) c ( yl ) measures the power of ylα , c−µx , c ( yl ) at the pixel scale . A particular situation where the power would be concentrated at the dataset scale withP ( 1 ) c ( yl ) equal toPc ( yl ) would imply that yl has its distribution fully “ collapsed ” in channel c , i.e . that yl is constant in channel c. The second view is that of a two-level binary tree : on one half of the tree , the sum of P ( 1 ) c ( yl ) and P ( 2 ) c ( yl ) measures the power coming from µx , c ( yl ) , with the relative proportions of P ( 1 ) c ( yl ) and P ( 2 ) c ( yl ) functions of the inter-x similarity and inter-x variability of µx , c ( yl ) ; on the other half of the tree , the sum of P ( 3 ) c ( yl ) and P ( 4 ) c ( yl ) measures the power coming from σx , c ( yl ) , with the relative proportions of P ( 3 ) c ( yl ) and P ( 4 ) c ( yl ) functions of the inter-x similarity and inter-x variability of σx , c ( y l ) . A particular situation where P ( 2 ) c ( yl ) , P ( 4 ) c ( yl ) would be equal to zero would imply that µx , c ( y l ) , σx , c ( yl ) have zero inter-x variability , i.e . that µx , c ( yl ) , σx , c ( yl ) are constant for all x . A version of Eq . ( 5 ) at the layer level instead of channel level will be easier to work with . Defining P ( i ) ( yl ) as the averages of P ( i ) c ( yl ) over c for i ∈ { 1 , 2 , 3 , 4 } , we obtain P ( yl ) = P ( 1 ) ( yl ) + P ( 2 ) ( yl ) + P ( 3 ) ( yl ) + P ( 4 ) ( yl ) . It should be noted that P ( yl ) = 1 for any choice of Norm ∈ { BN , LN , IN , GN } as long as the denominator of Eq . ( 3 ) is nonzero for all x , c [ C.1 ] . Consequently , the terms P ( i ) ( yl ) sum to one , meaning they can be conveniently seen as the proportion of each term i ∈ { 1 , 2 , 3 , 4 } into P ( yl ) . Revisiting BN ’ s avoidance of collapse . When BN is used , yl is normalized not only layer-wise but also channel-wise with P ( 1 ) c ( yl ) = 0 and Pc ( yl ) = 1 . As a first consequence , ỹl ( that is only one affine transformation away from yl ) is unlikely to have its channel-wise distributions collapsed . This means that the nonlinearity φ acting on ỹl is likely to be effectively nonlinear with respect to ỹl ’ s channel-wise distributions.2 As a result , each layer adds capacity and the network effectively uses its whole depth . This is opposite to the situation where ỹl has its channel-wise distributions collapsed with Pc ( ỹl ) − P ( 1 ) c ( ỹl ) Pc ( ỹl ) for all c , which results in φ being close to linear with respect to ỹl ’ s channel-wise distributions . This is illustrated in Figure 1 and formalized in Appendix C.2 . 2Note that : ( i ) the effective nonlinearity of φ with respect to ỹl ’ s channel-wise distributions could be quantified in the context of random nets of Definition 1 with “ reasonable ” choices of β , γ ; ( ii ) BN only guarantees an intra-distribution nonlinearity and not an intra-mode nonlinearity in contexts such as adversarial training [ 54 , 55 ] or conditional GANs [ 56 ] , unless modes are decoupled in BN ’ s computation [ 57 , 58 , 59 , 60 , 61 ] . Pc ( ỹl ) = 1 , 1 2 , 1 4 , 1 8 for ( a ) , ( b ) , ( c ) , ( d ) , respectively . When progressing from ( a ) to ( d ) , the part of the distribution corresponding to active ReLU ( shaded region ) becomes either overly dominant ( blue distribution ) or negligible ( red distribution ) . In either case , the channel-wise distribution ends up concentrated on only one side of piece-wise linearity . As an additional consequence , yl is guaranteed to have its channels well balanced with equal power Pc ( yl ) for all c. As a result , the network effectively uses its whole width . This is opposite to the situation where a single channel c becomes overly dominant over the others with Pc ( yl ) Pc′ ( yl ) for c 6= c′ , which results in downstream layers only “ seeing ” this channel c and the network behaving as if it had a width equal to one at layer l. Revisiting BN ’ s preservation of expressivity . When BN is used , P ( 1 ) c ( yl ) = 0 implies for all c that the terms P ( 2 ) c ( yl ) , P ( 3 ) c ( yl ) , P ( 4 ) c ( yl ) sum to one . Apart from that , BN does not impose any particular constraints on the relative proportions of each term into the sum . This means that the relative proportions of P ( i ) c ( yl ) and P ( i ) ( yl ) for i ∈ { 2 , 3 , 4 } are free to evolve as naturally dictated by the task and the optimizer during learning . This absence of constraints seems sensible . Indeed , imposing constraints on these relative proportions would alter the expressivity , which would not have any obvious justification in general and could even be detrimental in some cases , as we discuss in Section 4 .

Summary: This paper studies normalization strategies for deep residual networks, both theoretically and empirically, and finds the strategies that don't rely on batch-statistics to have particular failure modes, the two most notable ones being channel-wise collapse, where instances collapse to a constant in each channel, and reduction in instance-variability, where the first and second order statistics per channel are identical for each instance, thereby reducing expressivity and the ability to recognize higher order concepts. The paper then proceeds to propose proxy-normalization, a normalization method suitable for small-batch learning (i.e., not relying on batch-statistics) that, when combined with other normalization strategies, such as Layer Norm or Group Norm, alleviates these failure modes. On ImageNet and CIFAR experiments with deep residual networks (RN50 and larger), the work finds that Proxy-Normalization (with sufficient additional regularization such as label-smoothing, dropout, and stochastic depth), can be used to match the performance of BN without relying on batch statistics. The proposed method standardizes the activations of the network according to the mean and variance of a channel-wise gaussian random variable, with learnable mean and variance parameters.