Normalization Gradients are Least-squares Residuals

Batch Normalization ( BN ) and its variants have seen widespread adoption in the deep learning community because they improve the training of deep neural networks . Discussions of why this normalization works so well remain unsettled . We make explicit the relationship between ordinary least squares and partial derivatives computed when back-propagating through BN . We recast the backpropagation of BN as a least squares fit , which zero-centers and decorrelates partial derivatives from normalized activations . This view , which we term gradientleast-squares , is an extensible and arithmetically accurate description of BN . To further explore this perspective , we motivate , interpret , and evaluate two adjustments to BN . 1 INTRODUCTION . Training deep neural networks has become central to many machine learning tasks in computer vision , speech , and many other application areas . Ioffe & Szegedy ( 2015 ) showed empirically that Batch Normalization ( BN ) enables deep networks to attain faster convergence and lower loss . Reasons for the effectiveness of BN remain an open question ( Lipton & Steinhardt , 2018 ) . Existing work towards explaining this have focused on covariate shift ; Santurkar et al . ( 2018 ) described how BN makes the loss function smoother . This work examines the details of the back-propagation of BN , and recasts it as a least squares fit . This gradient regression zero-centers and decorrelates partial derivatives from the normalized activations ; it passes on a scaled residual during back-propagation . Our view provides novel insight into the effectiveness of BN and several existing alternative normalization approaches in the literature . 1.1 CONTRIBUTIONS . Foremost , we draw an unexpected connection between least squares and the gradient computation of BN . This motivates a novel view that complements earlier investigations into why BN is so effective . Our view is consistent with recent empirical surprises regarding ordering of layers within ResNet residual maps ( He et al. , 2016b ) and within shake-shake regularization branches ( Huang & Narayanan , 2018 ) . Finally , to demonstrate the extensibility of our view , we motivate and evaluate two variants of BN from the perspective of gradient-least-squares . In the first variant , a least squares explanation motivates the serial chaining of BN and Layer Normalization ( LN ) ( Ba et al. , 2016 ) . In the second variant , regularization of the least-squares leads to a version of BN that performs better on batch size two . In both variants , we provide empirical support on CIFAR-10 . In summary , our work presents a view , which we term gradient-least-squares , through which the back-propagation of BN and related work in a neural network can be recast as least squares regression . This regression decomposes gradients into an explained portion and a residual portion ; BN back-propagation will be shown to remove the explained portion . Hopefully , gradient-least-squares will be broadly useful in the future design and understanding of neural network components . Figure 1 reviews normalization with batch statistics , and illustrates our main theorem . { xi } i=1 ... N be activations . Let µ = N∑ i=1 xi N and σ2 = N∑ i=1 ( xi − µ ) 2 N . Let L be a function dependent on the normalized activations zi defined for each j by zj = ( xj − µ ) σ This , along with partial derivatives , are shown in the left figure . Our work establishes a novel identity on the quantities shown in hexagons . The right figure illustrates our main result in a scatter plot , in which each pair ( zi , ∂L ∂zi ) is shown as a data point in the regression . 2 NORMALIZATION GRADIENTS ARE LEAST-SQUARES RESIDUALS . Consider any particular channel within which { xi } are activations to be normalized in BN moment calculations . Ioffe & Szegedy ( 2015 ) defined BN as BN ( x ) = ( x− µ ) σ · c+ b ( 1 ) where σ , µ are batch moments , but b and c are learned per-channel parameters persistent across batches . In BN , the batch dimension and spatial dimensions are marginalized out in the computation of batch moments . For clarity , we consider a simplified version of BN . We ignore the variables b and c in equation 1 responsible for a downstream channel-wise affine transformation . Ignoring b and c is done without loss of generality , since the main observation in this work will focus on the Gaussian normalization and remains agnostic to downstream computations . We also ignore a numerical stability hyperparameter . We examine back-propagation of partial derivatives through this normalization , where µ and σ are viewed as functions of x . Notably , µ and σ are functions of each xi , and thus the division by σ is not affine . We write the normalized output as z = ( x− µ ) σ ( 2 ) We review ordinary least squares of a single variable with intercept ( Friedman et al. , 2001 ) . Let gj = α+βzj + j where α and β are parameters , z and g are observations . zj and gj are entries in z and g respectively . j are i.i.d . Gaussian residuals . We wish to fit α and β α̂ , β̂ = arg min α , β Ej ( ‖g − α− βz‖2 ) ( 3 ) The least-squares problem in equation 3 is satisfied by β̂ = Cov ( z , g ) Var ( z ) and α̂ = E ( g ) − β̂E ( z ) When z are normalized activations and g are partial derivatives , then Ez = 0 and Var ( z ) = 1 . In this special case , the solution simplifies into β̂ = Cov ( z , g ) ( 4 ) α̂ = E ( g ) ( 5 ) Theorem 1 ( Normalization gradients are least-squares residuals ) . Let i ∈ { 1 . . . N } be indices over some set of activations { xi } . Then the moment statistics are defined by µ = N∑ i=1 xi N and σ2 = N∑ i=1 ( xi − µ ) 2 N . Let L be a function dependent on the normalized activations zi defined for each j by zj = ( xj − µ ) σ . Then , the gradients of L satisfy , for all j ∈ { 1 , . . . , N } , the following : σ ∂L ∂xj = ∂L ∂zj − ∂̂L ∂zj ( 6 ) where ∂̂L ∂zj = α̂+ β̂zj ( 7 ) α̂ , β̂ = arg min α , β N∑ i=1 ( ∂L ∂zi − α− βzi ) 2 ( 8 ) Proof : Normalization gradients are least-squares residuals . The proof involves a derivation of partial derivatives by repeated applications of the chain rule and rules of total derivative . Because { zi } normalized over i has mean 0 and variance 1 , the partial derivatives can be rearranged to satisfy the single variable ordinary least squares framework . Fix j . We expand ∂L ∂xj as a linear combination of { ∂L ∂zi } i=1 ... N ∂L ∂xj = ∂L ∂zj ∂zj ∂xj + N∑ i6=j ∂L ∂zi ∂zi ∂xj ( 9 ) We state ∂zi ∂xj directly . Steps are in Appendix A under Lemma 1 . ∂zi ∂xj = −1− zjzi σN if i 6= j N − 1− z2j σN if i = j ( 10 ) Through substitution of equations 10 into 9 , we get ∂L ∂xj = ∂L ∂zj N − 1− z2j σN + N∑ i 6=j [ ∂L ∂zi · −1− zjzi σN ] ( 11 ) σ ∂L ∂xj = ∂L ∂zj + 1 N N∑ i=1 [ ( −1− zizj ) ∂L ∂zi ] ( 12 ) σ ∂L ∂xj = ∂L ∂zj − ( 1 N N∑ i=1 ∂L ∂zi ) − zj N N∑ i=1 zi ∂L ∂zi ( 13 ) Noting that { zi } normalized over i has mean 0 and variance 1 , we recover β̂ and α̂ , in the sense of equations 4 and 5 , from equation 13 . 1 N N∑ i=1 zi ∂L ∂zi = Covi ( zi , ∂L ∂zi ) = β̂ ( 14 ) 1 N N∑ i=1 ∂L ∂zi = Ei [ ∂L ∂zi ] − β̂ · 0 = α̂ ( 15 ) Finally , we rearrange equations 15 and 14 into 13 to conclude , as desired , σ ∂L ∂xj = ∂L ∂zj − α̂− β̂zj = ∂L ∂zj − ∂̂L ∂zj ( 16 ) During back-propagation of a single batch , the normalization function takes in partial derivatives ∂L ∂z ( · ) , and removes that which can be explained by least squares of ∂L ∂z ( · ) against z ( · ) . As illustrated in Figure 1 , during back-propagation , the residual then divides away σ to become ∂L ∂x ( · ) , the gradient for the unnormalized activations . 3 RELATED DEEP LEARNING COMPONENTS VIEWED AS GRADIENT CALCULATIONS . BN aims to control its output to have mean near 0 and variance near 1 , normalized over the dataset ; this is related to the original explanation termed internal covariate shift ( Ioffe & Szegedy , 2015 ) . Most existing work that improve or re-purpose BN have focused on describing the distribution of activations . Definition 1 . In the context of normalization layers inside a neural network , activations are split into partitions , within which means and variances are computed . We refer to these partitions as normalization partitions . Definition 2 . Within the context of a normalization partition , we refer to the moments calculated on the partitions as partition statistics . Theorem 1 shows that BN has least squares fitting built into the gradient computation . Gradients of the activations being normalized in each batch moment calculation are fit with a single-variable with-intercept least squares model , and only a rescaled residual is kept during back-propagation . We emphasize that the data on which the regression is trained and applied is a subset of empirical activations within a batch , corresponding to the normalization partitions of BN . To show extensibility , we recast several popular normalization techniques into the gradient-leastsquares view . We refer to activations arising from a single member of a particular batch as an item . BHWC refers to dimensions corresponding to items , height , width , and channels respectively . In non-image applications or fully connected layers , H and W are 1 . BN marginalizes out the items and spatial dimensions , but statistics for each channel are kept separate . In the subsequent sections , we revisit several normalization methods from the perspective of the gradient . Figure 2 reviews the normalization partitions of these methods , and places our main theorem about gradient-least-squares into context . 3.1 LAYER NORMALIZATION , INSTANCE NORMALIZATION , GROUP NORMALIZATION . Ba et al . ( 2016 ) introduced Layer Normalization ( LN ) in the context of large LSTM models and recurrent networks . Only the ( H , W , C ) dimensions are marginalized in LN , whereas BN marginalizes out the ( B , H , W ) dimensions . In our regression framework , the distinction can be understood as changing the data point partitions in which least squares are fit during back-propagation . LN marginalizes out the channels , but computes separate statistics for each batch item . To summarize , the regression setup in the back-propagation of LN is performed against other channels , rather than against other batch items . Huang & Belongie ( 2017 ) introduced Instance Normalization ( IN ) in the context of transferring styles across images . IN is is closely related to contrast normalization , an older technique used in image processing . IN emphasizes end-to-end training with derivatives passing through the moments . Only the ( H , W ) dimensions are marginalized in IN , whereas BN marginalizes ( B , H , W ) dimensions . In our framework , this can be understood as using fewer data points and a finer binning to fit the least squares during back-propagation , as each batch item now falls into its own normalization partition . Wu & He ( 2018 ) introduced Group Normalization ( GN ) to improve performance on image-related tasks when memory constrains the batch size . Similar to LN , GN also marginalizes out the ( H , W , C ) dimensions in the moment computations . The partitions of GN are finer : the channels are grouped into disjoint sub-partitions , and the moments are computed for each sub-partition . When the number of groups is one , GN reduces to LN . In future normalization methods that involve normalizing with respect to different normalization partitions ; such methods can pattern match with BN , LN , IN , or GN ; the back-propagation can be formulated as a least-squares fit , in which the partial derivatives at normalized activations ∂L ∂z ( · ) are fitted against the normalized z ( · ) , and then the residual of the fit is rescaled to become ∂L ∂x ( · ) . Figure 2 summarize the normalization partitions for BN , LN , IN , and GN ; the figure visualizes , as an example , a one-to-one correspondence between an activation in BN , and a data point in the gradient regression . Theorem 1 is agnostic to the precise nature of how activations are partitioned before being normalized ; thus , equation 9 applies directly to any method that partitions activations and performs Gaussian normalization on each partition . The partitioning of BN , LN , IN , and GN are performed in different respective manners , and each partition is individually subject to Gaussian normalization . Thus , the gradients of BN , LN , IN , and GN are residuals of regressions in the sense of Theorem 1 .

Summary: The primary technical contribution comes from Section 2, where it is demonstrated that the normalized back-propagated gradients obtained from a BN layer can be viewed as the residuals of the gradients obtained without BN regressed via a simple two-parameter model of the activations.  In some sense though this result is to be expected, since centering data (i.e., removing the mean as in BN) can be generically viewed as computing the residuals after a least squares fit of a single constant, and similarly for de-trending with respect to a single independent variable, in this case the activations.  So I'm not sure that Theorem 1 is really that much of an insightful breakthrough, even if it may be nice to work through the precise details in the specific case of a BN layer and the relationship to gradients.