Title: NORMALIZATION GRADIENTS ARE LEAST-SQUARES RESIDUALS

Abstract: 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. 

2. NORMALIZATION GRADIENTS ARE LEAST-SQUARES RESIDUALS:

Consider any particular channel within which {x i } 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 x i , 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 g j = α + βz j + j where α and β are parameters, z and g are observations. z j and g j are entries in z and g respectively. j are i.i.d. Gaussian residuals. We wish to fit α and β

α, β = arg min α,β E j ( 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

EQUATION

Theorem 1 (Normalization gradients are least-squares residuals). Let i ∈ {1 . . . N } be indices over some set of activations {x i }. Then the moment statistics are defined by µ = N i=1

x i N and

σ 2 = N i=1 (x i -µ) 2 N .

Let L be a function dependent on the normalized activations z i defined for each j by z j = (x j -µ) σ . Then, the gradients of L satisfy, for all j ∈ {1, . . . , N }, the following:

σ ∂L ∂x j = ∂L ∂z j - ∂L ∂z j (6)

where

EQUATION

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 {z i } 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 ∂x j as a linear combination of

EQUATION

We state ∂z i ∂x j directly. Steps are in Appendix A under Lemma 1.

∂z i ∂x j =          -1 -z j z i σN if i = j N -1 -z 2 j σN if i = j (10)

Through substitution of equations 10 into 9, we get

EQUATION

Noting that {z i } 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 z i ∂L ∂z i = Cov i z i , ∂L ∂z i = β (14) 1 N N i=1 ∂L ∂z i = E i ∂L ∂z i -β • 0 = α (15)

Finally, we rearrange equations 15 and 14 into 13 to conclude, as desired, , the gradient for the unnormalized activations.

EQUATION

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. BHW C 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. 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.

3.2. WEIGHT NORMALIZATION:

Salimans & Kingma (2016) introduced Weight Normalization (WN) in LSTMs, and noted improvements in the condition number of deep networks; WN divides each weight tensor by their respective vector 2-norms. In the view of gradient-least-squares, WN has a single-variable intercept-0 regression interpretation in back-propagation, analogous to BN. A raw weight vector v, is normalized and scaled before being used as coefficient weights w = c v v, where c is a learned downstream linear scaling parameter.

In this regression setup, the length normalized weights of WN are analogous to the Gaussian normalized activations in BN. We write that z

= v v = w c

, and state directly an analogous relationship between each v ∂L ∂v j and the regression on

z i , ∂L ∂z i i=1...N

. See Appendix B Lemma 2 for steps that derive the following identity: for loss L and for each component j, we have

EQUATION

where

EQUATION

The L2 normalization of weights in WN appears distinct from the Gaussian normalization of activations in BN; nevertheless, WN can also be recast as a least squares regression. The improved ordering has BN coming first, and thus constrains that gradients of the residual map must be decorrelated with respect to some normalized activations inside the residual mapping.

3.3. IDENTITY MAPPINGS IN RESNET, AND SHAKE-SHAKE RESNEXT REGULARIZATION:

An update to the popular ResNet architecture showed that the network's residual mappings can be dramatically improved with a new ordering (He et al., 2016b) . The improvement moved BN operations into early positions and surprised the authors; we support the change from the perspective of gradient-least-squares. Figure 3 reviews the precise ordering in the two versions. Huang & Narayanan (2018) provides independent empirical support for the BN-early order, in shake-shake regularization (Gastaldi, 2017) architectures. We believe that the surprise arises from a perspective that views BN only as a way to control the distribution of activations; one would place BN after a sequence of convolution layers. In the gradient-least-squares perspective, the first layer of each residual mapping is also the final calculation for these gradients before they are added back into the main trunk. The improved residual branch constrains the gradients returning from the residual mappings to be zero-centered and decorrelated with respect to some activations inside the branch.

We illustrate this idea in Figure 3 .

4. NORMALIZATION APPROACHES MOTIVATED BY LEAST SQUARES:

Gradient-least-squares views back-propagation in deep neural networks as a solution to a regression problem. Thus, formulations and ideas from a regression perspective would motivate improvements and alternatives to BN. We pursue and evaluate two of these ideas.

4.1. BN AND LN AS TWO-STEP GRADIENT REGRESSION:

BN and LN are similar to each other, but they normalize over different partitioning of the activations; in back-propagation, the regressions occur respectively with respect to different partitions of the activations. Suppose that a BN and a LN layer are chained serially in either order. This results in a two-step regression during back-propagation; in reversed order, the residual from the first regression is further explained by a second regression on a different partitioning. In principle, whether this helps would depend on the empirical characteristics of the gradients encountered during training.

The second regression could further decorrelate partial gradients from activations. Empirically, we show improvement in a reference ResNet-34-v2 implementation on CIFAR-10 relative to BN with batch size 128. In all cases, only a single per-channel downstream affine transformation is applied, after both normalization layers, for consistency in the number of parameters. See table 1 for CIFAR-10 validation performances. We kept all default hyperparameters from the reference implementation: learning schedules, batch sizes, and optimizer settings.

4.2. ADDRESSING SMALL BATCHES WITH LEAST-SQUARES REGULARIZATION:

BN performs less well on small batches (Ioffe, 2017) . Gradient-least-squares interprets this as gradient regressions failing on correlated data, an issue typically addressed by regularization. We pursue this idea to recover some performance on small batches by use of regularization. Our regularization uses streaming estimates of past gradients to create virtual data in the regression. This performed better than standard BN on the same batch size, but we did not recover the performance of large batches; this is consistent with the idea that regularization could not in general compensate for having much less data. See Appendix C for CIFAR-10 validation performances. to rescale the contributions to the batch mean for each normalization scheme. It uses an analogous set of parameters λ k and activations w k for variances. We sketch the back-propagation of a simplified version of SN in the perspective of gradient-least-squares. We ignore both the division and downstream affine

5. LIMITATIONS AND RELATED WORK:

z → c • z + b.

The normalization calculation inside SwN can be written as:

EQUATION

where Ω = {BN, LN, IN }. There is potentially a unique mean and variance used for each activation. Equation 19 bears similarities to the setup in Theorem 1, but we leave unresolved whether there is a gradient-least-squares regression interpretation for SN.

5.2. DECORRELATED BATCH NORMALIZATION AND SPECTRAL NORMALIZATION:

Decorrelated Batch Normalization (DBN) (Huang et al., 2018) is a generalization of BN that performs Mahalanobis ZCA whitening to decorrelate the channels, using differentiable operations. On some level, the matrix gradient equation resemble the least squares formulation in Theorem 1.

Spectral Normalization (SpN) (Miyato et al., 2018) is an approximate spectral generalization of WN.

For DBN and SpN, the regression interpretations remain unresolved.

5.3. RELATED WORK:

BN has been instrumental in the training of deeper networks (Ioffe & Szegedy, 2015) . Subsequent work resulted in Batch Renormalization (Ioffe, 2017) , and further emphasized the importance of passing gradients through the minibatch moments, instead of a gradient-free exponential running average. In gradient-least-squares, use of running accumulators in the training forward pass would stop the gradients from flowing through them during training, and there would be no least-squares. He et al. (2016b) demonstrate empirically the unexpected advantages of placing BN early in residual mappings of ResNet.

Santurkar et al. ( 2018) showed that BN makes the loss landscape smoother, and gradients more predictable across stochastic gradient descent steps. Balduzzi et al. (2017) found evidence that spatial correlation of gradients explains why ResNet outperforms earlier designs of deep neural networks. Kohler et al. (2018) proved that BN accelerates convergence on least squares loss, but did not consider back-propagation of BN as a least squares residual. Luo et al. (2018b) has recast BN as a stochastic process, resulting in a novel treatment of regularization.

6. DISCUSSION, AND FUTURE WORK:

This work makes explicit how BN back-propagation regresses partial derivatives against the normalized activations and keeps the residual. This view, in conjunction with the empirical success of BN, suggests an interpretation of BN as a gradient regression calculation. BN and its variants decorrelate and zero-center the gradients with respect to the normalized activations. Subjectively, this can be viewed as removing systematic errors from the gradients. Our view also support empirical results in literature preferring early BN placement within neural network branches.

Leveraging gradient-least-squares considerations, we ran two sets of normalization experiments, applicable to large batch and small batch settings. Placing a LN layer either before or after BN can be viewed as two-step regression that better explains the residual. We show empirically on CIFAR-10 that BN and LN together are better than either individually. In a second set of experiments, we address BN's performance degradation with small batch size. We regularize the gradient regression with streaming gradient statistics, which empirically recovers some performance on CIFAR-10 relative to basic BN, on batch size two.

Why do empirical improvements in neural networks with BN keep the gradient-least-squares residuals and drop the explained portion? We propose two open approaches for investigating this in future work. A first approach focuses on how changes to the gradient regression result in different formulations; the two empirical experiments in our work contribute to this. A second approach examines the empirical relationships between gradients of activations evaluated on the same parameter values; we can search for a shared noisy component arising from gradients in the same normalization partition. Suppose that the gradient noise correlates with the activations -this is plausible because the population of internal activations arise from using shared weights -then normalizations could be viewed as a layer that removes systematic noise during back-propagation.

In conclusion, we have presented a novel view that reorganizes the back-propagation of BN as a least squares residual calculation. This view generates novel descriptions of normalization techniques related to BN, and comments on the ordering of layers inside the residual mappings of ResNet. This view is extensible and will motivate novel designs of neural network components in future work. 

{z i }, for i ∈ 1 . . . N . We define µ = 1 N N j=1 x i and σ 2 = 1 N N j=1 (x i -µ) 2 with σ > 0.

The normalized activations z i are defined as z i = x i -µ σ . Then, the partial derivatives satisfy

EQUATION

Proof. In deriving ∂z j ∂x i , we will treat the cases of when j = i and when j = i separately. We start by examining intermediate quantities of interest as a matter of convenience for later use. We define helper quantities u i = x i -µ. Note that each u j depends on all of x i via µ. Next, we write out useful identities

EQUATION

We prepare to differentiate with rule of total derivative:

EQUATION

Making use of equations 21, 22, 23 and 25, We simplify ∂σ ∂x i for any i as follows.

EQUATION

We apply the quotient rule on ∂z j ∂x i when j = i, then substitute equation 33

EQUATION

Similarly, when i = j, 

EQUATION

:



Figures and Tables:

Figure 1: The left figure reviews, for a single channel at a particular layer within a single batch, notable quantities computed during the forward pass and during back-propagation of BN. Let {x i } i=1...N be activations. Let µ =

Figure 2: We review the normalization partitions of BN, LN, GN, and IN. Each normalization partition contains a separate set of data points on which the gradient regression is performed. One partition for each method is illustrated in blue. This figure also shows the correspondence between a single activation and a gradient regression data point for BN.

Figure 3: This figure illustrates the original(He et al., 2016a)  and improved(He et al., 2016b)   residual mappings in ResNets. Arrows point in the direction of the forward pass. Dotted lines indicate that gradients are zero-centered and decorrelated with respect to downstream activations in the residual mapping. The improved ordering has BN coming first, and thus constrains that gradients of the residual map must be decorrelated with respect to some normalized activations inside the residual mapping.

. (2018a) introduced Switch Normalization (SwN), a hybrid strategy for combining moment calculations from LN, BN, and IN. SwN uses learnable scalar logits λ k for k ∈ Ω = {BN, IN, LN } with corresponding softmax weighting activations w k = exp(λ k ) k exp(λ k )

batch size 128)

Table content: <table><tr><td>Normalization</td><td colspan="2">CIFAR-10 Accuracy CIFAR-10 Cross Entropy</td></tr><tr><td>BN, LN</td><td>0.9259</td><td>0.3087</td></tr><tr><td>LN, BN</td><td>0.9245</td><td>0.3389</td></tr><tr><td colspan="2">BN (Ioffe &amp; Szegedy, 2015) 0.9209</td><td>0.3969</td></tr><tr><td>LN (Ba et al., 2016)</td><td>0.9102</td><td>0.3548</td></tr></table>

Shibani Santurkar, Dimitris Tsipras, Andrew Ilyas, and Aleksander Madry. How does batch normalization help optimization?(no, it is not about internal covariate shift). arXiv preprint arXiv:1805.11604, 2018. Consider the Gaussian normalization function that maps all {x i } into corresponding

Table content: <table><tr><td>Yuxin Wu and Kaiming He. Group normalization. In The European Conference on Computer Vision</td></tr><tr><td>(ECCV), September 2018.</td></tr><tr><td>A PARTIAL DERIVATIVES OF THE NORMALIZATION FUNCTION</td></tr><tr><td>Lemma 1.</td></tr></table>

Streaming regularization is less affected by small batch sizes, batch size 2) In our work, we keep track of am exponential running estimates across batches,α * ∼ E b E i ∂L ∂z i X=X (b) and β * ∼ E b E i z i ∂L ∂z i X=X (b)that marginalize the (B, H, W ) dimensions into accumulators of shape C. The b subscript of the outer expectation is slightly abusive notation indicating that α * and β * are running averages across recent batches with momentum as a hyperparameter that determines the weighting. We regularize the gradient regression with virtual activations and virtual gradients, defined as follows. We append two virtual batch items, broadcast to an appropriate shape,x + = µ b + σ b and x -= µ b -σ b .Here, µ b and σ b are batch statistics of the real activations. The concatenated tensor undergoes standard BN, which outputs the usual {z i } for the real activations, but z + = 1 and z -= -1 for the virtual items. The z + and z -do not affect the feed forward calculations, but they receive virtual gradients during back-propagation: eventually modify the gradients received by the real x i activations. The virtual data can be weighted with hyperparameters. In our experiments, we see improvements, robust to a hyperparameter cross-product search over the weightings and the momentum for α * and β * . The momentum for α * and β * were in {.997, .5} and the virtual item weights were in {2 i-1 } i∈{0,1,2,3} . The performance of larger batches are not recovered; regularized regression could not be reasonably expected to recover the performance of regressing with more data. See table2for final validation performances with a reference Tensorflow ResNet-34-v2 implementation on batch size of two. The baseline evaluation with identity (no normalization) experienced noticeable overfitting in terms of cross entropy but not accuracy. The base learning rate was multiplied by 1 64 relative to the baseline rate used in runs with batch size 128.

Table content: <table><tr><td/><td colspan="3">Normalization</td><td colspan="3">CIFAR-10 Accuracy</td><td>CIFAR-10 Cross Entropy</td></tr><tr><td/><td colspan="3">Our Best Hyperparameter</td><td colspan="2">0.9091</td><td>0.3627</td></tr><tr><td/><td colspan="3">Our Worst Hyperparameter</td><td colspan="2">0.9005</td><td>0.4086</td></tr><tr><td/><td colspan="5">BN (Ioffe &amp; Szegedy, 2015) 0.8903</td><td>0.4624</td></tr><tr><td/><td colspan="3">Renorm (Ioffe, 2017)</td><td colspan="2">0.9033</td><td>0.3823</td></tr><tr><td/><td>Identity</td><td/><td/><td colspan="2">0.9086</td><td>0.6934 (0.4229 at best point)</td></tr><tr><td colspan="6">inputs in batch b. ∂L ∂z +</td><td>= α *  + β *</td><td>(53)</td></tr><tr><td/><td/><td/><td/><td/><td>∂L ∂z -</td><td>= α *  -β *</td><td>(54)</td></tr><tr><td colspan="2">Virtual data z + ,</td><td>∂L ∂z +</td><td>and z -,</td><td>∂L ∂z -</td><td colspan="2">regularizes the gradient-least-squares regression.</td><td>∂L ∂z +</td></tr><tr><td>and</td><td>∂L ∂z -</td><td/><td/><td/><td/></tr></table>

Peer review (iclr2019):
 The authors propose a new interpretation of the batch normalization step inside a neural network. The main result shows that the backpropagation of the gradient of some loss function through a batch normalization can be seen as a scaled residual of a least square linear fit. This new interpretation is extended to other normalization technics used in the literature and thus give a "unified" view of such methods. The idea is simple yet very interesting and well introduced. The theoretical results are good and the proofs are well written and easy to follow. However the arguments brought forward by this new vision of batch normalization in applications look light (see sections 3.3, 4.1, 4.2). A more detailed interpretation of this new vision on a single application and its impact would have been preferred than numerous applications as it is done in this paper. Not all the existing normalization methods have been extended with success yet, this makes this unified vision a bit less convincing.