CONTROLLING COVARIATE SHIFT USING EQUILIBRIUM NORMALIZATION OF WEIGHTS

1 INTRODUCTION . The introduction of normalizing layers to neural networks has in no small part contributed to the deep learning revolution in machine learning . The most successful of these techniques in the image classification domain is the batch normalization ( BatchNorm ) layer ( Ioffe & Szegedy , 2015 ) , which works by normalizing the univariate first and second order statistics between layers . Batchnorm has seen near universal adoption in image classification tasks due to its surprisingly multifaceted benefits . Compared to an unnormalized network , its has been widely observed that using batch norm empirically results in : • Stability over a wide range of step sizes • Faster convergence ( particularly with larger step sizes ) • Improved generalization The multiple effects of BatchNorm make it both hard to replace and hard to analyze . In this paper we introduce Equilibrium Normalization ( EquiNorm ) , a normalization that works in weight space and still uses a form of batch statistics unlike previous weight space approaches . EquiNorm results in very rapid convergence , even more so than BatchNorm , however as we will show in our experiments , this also results in a tendency to overfit . When combined with additional regularisation , EquiNorm can significantly outperform BatchNorm , which benefits less from this additional regularisation . 2 RELATED WORK . A number of normalization layers have been proposed that can be considered alternatives to batch normalization . Batch normalization has also been extended as batch renormalization ( Ioffe , 2017 ) to handle smaller batch sizes . Layer/Instance Normalization A simple modification of BatchNorm involves computing the statistics independently for each instance , so that no averaging is done across each mini-batch , instead averaging either across channels ( layer norm ) or separately for each channel ( instance norm ) . Unfortunately these techniques are known to not generalize as well as batch norm for convolutional neural networks ( Sec 6.7 ; Sec 4.1 . Jimmy Lei Ba , 2016 ; Yuxin Wu , 2018 ) . Group Normalization A middle ground between layer and instance normalization can be found by averaging statistics over small groups of channels . This has been shown empirically to be superior to either approach , although there is still a gap in generalization performance ( Yuxin Wu , 2018 ) . Like the approaches above , it avoids a dependence on batch statistics , allowing for potentially much smaller batches to be used without a degradation in generalization performance . Weight Normalization Additional stability can be introduced into NN training by constraining the norm of the weights corresponding to each output channel/neuron to be one . When this is done by an explicit division operation in the forward pass , rather than via an optimization constraint , this is known as weight normalization ( Salimans & Kingma , 2016 ) . An additional learnable scaling factor is also introduced . Unfortunately , to match the generalization performance of BatchNorm on image classification tasks such as CIFAR-10 , this technique needs to be used together with partial ( additive only ) BatchNorm ( Section 5.1 , Salimans & Kingma , 2016 ) . Local Response Normalization A precursor to batch norm , local normalization methods ( Jarrett et al. , 2009 ; Lyu & Simoncelli , 2008 ) played an important part in the seminal AlexNet architecture ( Krizhevsky et al. , 2012 ) , and were widely used before batch norm was introduced . LR normalization has similarities to group norm in that it uses a group of neighboring channels ( with ordering set arbitrary at initialization ) for normalization . Although it aids generalization in a similar manner to BatchNorm , it does not accelerate convergence or allow for larger step sizes to be used ( Sec 4.2.1 , Ioffe & Szegedy , 2015 ) . 3 ASSUMPTIONS . The EquiNorm method functions by modifying the weights of a convolution before it is applied . For justifying the form of our method , we make the following assumptions about this convolution , which we will discuss relaxing after detailing the method : 1 . All inputs to the convolutional layer are positive , such as when the layer is preceded by a ReLU . 2 . The convolution has stride one . 3 . Cyclic padding is used . 4 . All weights are non-zero , and there exists at least one positive and one negative weight per output channel . 4 METHOD . Consider initially for simplicity a convolutional layer with a single input and output channel . Let w : kernelheight× kernelwidth , be the weight kernel for this neuron , and let x : batchsize× height× width , be the input tensor for a single mini-batch . We will compute scalar quantities s and b that modify the weights as follows : w′′ = sw′ = s ( w + b ) . This transformation will be fully differentiated through during the backwards pass ( using automatic differentiation ) so that the gradient of w is correct . As with BatchNorm , we also include an additional affine transformation after the convolution to ensure no expressivity is lost due to the normalization operation . The core idea of equilibrium normalization is to balance the contribution of positive and negative weights to the output of the convolution . To this end , we introduce additional notation to address the positive and negative weights separately . Let superscripts +/− ( i.e . w+/w− ) indicate sums of the positive/negative elements respectively . Also , let v be the sum of the input data to the layer , v = ∑ i , j , k xijk . As we have two constants to determine , we need two constraints that we wish to be satisfied . The first constraint we introduce is common with batch normalization , a constraint on the mean of the output . Since under the cyclic padding assumption , each weight is multiplied by each input element , we can constrain the mean of the output to be zero by requiring that : vs ∑ j , j ( wjk + b ) = 0 , ∴ b = −mean ( w ) . The second constraint controls the magnitude within the total output of the layer , of the positive weight elements : vsw′+ = r , ∴ s = r vw′+ . The constant r is set so that the contribution is of average 1 per output element , which is achieved by setting r to the product of batch-size , output width and output height . Note that due to the mean constraint , this automatically results in the negative weight contribution also being of magnitude r . 4.1 FULL CASE . When multiple input channels are used , each weight no longer multiples each input , rather they each multiply only inputs from a single channel . To compensate for this we need to compute per-channel sums vc ( where c is the channel index ) and change the second constraint as follows : channelsin∑ c vcsw ′+ c = r. The first constraint changes in the same fashion . When multiple output channels are used , we just duplicate this procedure applying it to each channel ’ s weights separately . We thus maintain a s and b value per output channel , and compute as intermediate values a w′+ of matrix shape . For completeness we give the full equations below . All summations are over the full range of the summed indexes . Tensor shapes w : channelsout× channelsin× kernelheight× kernelwidth x : batchsize× channelsin× heightin× widthin w′+ : channelsout× channelsin v : channelsin s , b : channelsout Updates : vc = ∑ i , j , k xicjk , wdc = ∑ j , k wdcjk , r = batchsize× heightout× widthout× stride2 bd = − ∑ c vcwdc ( kernelheight× kernelwidth ) ∑ c vc , w′+dc = ∑ j , k ( wdcjk + bd ) I [ wdcjk + bd > 0 ] , sd = r∑ c vcw ′+ dc , w′′dcjk = sd ( wdcjk + bd ) . At test time , we follow the technique used in BatchNorm of using a running estimate of the data statistics ( vc in our method ) that is computed during training time . SINGLE PASS FORMULATION . The above calculation requires two passes over the weights , first to compute b , then to compute the sum of positive elements after the addition of b . We can do an approximate computation using only one pass by assuming the sign of each element does not change after the addition of b . The s calculation changes as follows : n+dc = ∑ j , k I [ wdcjk > 0 ] , sd = r∑ c vc ( w+dc + bdn + dc ) . We use this variant in all experiments that follow . 5 DISCUSSION . 5.1 CONTROLLING COVARIATE SHIFT . The original justification for batch normalization is its ability to minimize covariate shift , although there is some debate on whether or not this is the main contributing factor to its effectiveness , see Santurkar et al . ( 2018 ) . In this context , covariate shift refers to the change between steps of the statistics of the outputs of a layer . Like batch normalization , the approach we propose also controls the shift in the outputs of a layer between steps , just in a different way . Our approach is motivated by a hypothesis that it is not necessary to control the mean and variance precisely ; other notions of scale and shift may work as well or better . Santurkar et al . ( 2018 ) show that normalizing by other norms , such as L1 , can work well , supporting this hypothesis . The sum of the output from positive and negative weights is a form of L1 control which can be contrasted with the L2 control that Batchnorm uses . This control can be motivated by Young ’ s convolution inequality , which bounds the output of a convolution operation in terms of the norm of the input : ‖x ∗ w‖r ≤ ‖w‖p ‖x‖q , where 1 p + 1 q = 1 r + 1 . Note that p , q , r ≥ 1 is also required , and that this only applies directly when there is a single input and output channel , which we assume in the remainder of this section for simplicity . For EquiNorm , we have assumed that the input is positive , so that our input sum is equivalent to the L1 norm of the input . Additionally , after subtracting off the mean , the weight vector w′ has L1 norm equal to w′+ − w′− = 2w′+ , so we are also normalizing the weights by the L1 norm . In effect , we are applying Young ’ s convolution inequality with p = q = r = 1 . It is also possible to apply the above inequality with p = 2 , q = 1 and r = 2 . I.e . normalize the weights using the L2 norm , giving a bound on the L2 norm of the output in terms of the L1 norm of the input . This is less satisfying as one convolution ’ s output is the input of another convolution ( after passing through scaling & a nonlinearity ) so we would like to use the same norm for both inputs and outputs . The related weight normalization ( WN , Salimans & Kingma , 2016 ) method normalizes weights by their L2 norm , and differs from our method by centering outputs using an additional mean-only output batchnorm . Additionally , since it doesn ’ t normalize by the input norm , the output norm can be correspondingly large . These differences have a significant effect in practice .

Summary: This manuscript introduces a new layer-wise transform, EquiNorm, to improve upon batch normalization. As with batch normalization and related techniques, the idea is to introduce a simple linear transform at each layer to reduce the dependency of the features to the data. Unlike batch normalization, the procedure does not modify the inputs to the layers but rather the layer weights. For this purpose, a scaling factor and a shift is computed on a mini batch, separating positive and negative weights to compute easily both running estimates of shift and of spread (here in the l1 sense). The method is compared to BatchNorm and GroupNorm on several classic computer vision datasets. Empirically, the method converges faster in the beginning of the optimization, in the sense that in the first few epochs the test accuracy is higher than for BatchNorm. However, this benefit decreases with more epochs and when the results are close to peak performance the difference between methods is a small. In addition, test accuracy can decrease at the end of the optimization, which the authors interpret as a sign of increased overfit, and tackle with clever data augmentation. The paper reads well.