Extreme normalization: approximating full-data batch normalization with single examples
Authors: Sergey Ioffe
Venue: ICLR 2022 Submitted (2022)
Keywords: batch normalization
optimization

Abstract: While batch normalization has been successful in speeding up the training of neural networks, it is not well understood. We cast batch normalization as an approximation of the limiting case where the entire dataset is normalized jointly, and explore other ways to approximate the gradient from this limiting case. We demonstrate an approximation that removes the need to keep more than one example in memory at any given time, at the cost of a small factor increase in the training step computation, as well as a fully per-example training procedure, which removes the extra computation at the cost of a small drop in the final model accuracy. We further use our insights to improve batch renormalization for very small minibatches. Unlike previously proposed methods, our normalization does not change the function class of the inference model, and performs well in the absence of identity shortcuts.

Under review as a conference paper at ICLR 2022
EXTREME NORMALIZATION: APPROXIMATING FULL- DATA BATCH NORMALIZATION WITH SINGLE EXAM- PLES
Anonymous authors Paper under double-blind review
ABSTRACT
While batch normalization has been successful in speeding up the training of neural networks, it is not well understood. We cast batch normalization as an approxima- tion of the limiting case where the entire dataset is normalized jointly, and explore other ways to approximate the gradient from this limiting case. We demonstrate an approximation that removes the need to keep more than one example in memory any given time, at the cost of a small factor increase in the training step computa- tion, as well as a fully per-example training procedure, which removes the extra computation at the cost of small drop in the final model accuracy. We further use our insights to improve batch renormalization for very small minibatches. Unlike previously proposed methods, normalization does not change the function class the inference model, and performs well in the absence of identity shortcuts.
INTRODUCTION
Batch normalization (BN) is an offen-effective tool in training deep models. It achieves such effectiveness by normalizing the elements of a minibatch jointly, and doing pervasively throughout the model. When BN used, ensured that there are representations throughout the model where the per-dimension first and second moments (means and variances) are fixed throughout training.
Batch normalization is not well-understood theoretically, and known to introduce undesired, and often surprising, artifacts into training. Because BN relies on the interaction between examples during training, it causes models to perform poorly when training minibatches are small or contain dependencies between examples. This shown e.g. in (loffe 2017). which also demonstrated that this can be alleviated by carefully using the moving averages of minibatch statistics during training.
The adverse effects of BN, well the engineering complexity needed for large minibatches, have led researchers to tackle the problem normalization that depends only on individual examples, such as (Ba al. 2016 Wu He 2018). These methods, however, change the architecture used for inference, as the normalization needs to be performed during inference as well. While this may lead to difference in performance, such interventions do not change only the optimization, but also the hypothesis class. On the other hand, inference- time BN can be folded into the rest of the model weights yield a regular model (e.g. using only convolutions, bias additions, and ReLU). Hence batch normalization does not change the class of models and can be viewed as solely an optimization device, Approaches that change only the optimization and not the architecture include (Zhang et al.. 2019 Brock et al.) 2221) which remove the normalization entirely, but depend on identity shortcuts in the model.
In this work, we aim develop method for performing per-example normalization in way that does not not modify the inference-timg architecture. The resulting method combines per-example gradient computation, the maintenance of moving first and second moments, and aggregation step that joins the information from multiple examples similar to the way that the gradients with respect to the model parameters are commonly averaged over the minibatch in SGD.

---

Under review as a conference paper at ICLR 2022
2 APPROXIMATING POPULATION NORMALIZATION
Consider model that involves a normalization step, such as Vi = 2; - E or y:= VE Here, 2; is an activation within model computed for a particular (ith) model input, y/ is the normalized activa- tion, and E- is the expectation taken over the training data distribution. The examples above achieve zero mean and unit second moment, respectively. We decouple these two types of normalization in what follows.
Performing normalization that involves the expectation over the data distribution needs take into account the fact that the normalizers such as E x and Ex² depend on the distribution of the activation I and therefore depend on the model parameters that affect æ. For example, if is an output of a layer in a model, then its value depends on the parameters of this layer as well as of the preceding layers. Stochastic optimization relies on computing a stochastic estimate of the gradient aw loss with respect to weights. Without cross-exampel interactions, an unbiased estimate is given by the gradient for a randomly sampled example or minibatch. With cross-example normalization, however, this does not hold since we no longer have per-example gradients: the computation for an example is influenced by all other examples. This work aims at obtaining an estimate, albeit not necessarily unbiased, of the overall gradient from subset of the examples, whose computation requires minimal (ideally none) interaction between examples.
at of the
Let us consider a general normalization layer of the form 2; = fa,, S), where S is defined as the expectation of some statistic of the activation, S = E sx)) Observe that if the model parameters fixed and the model is structured as a directed acyelic graph, then the normalizers S can be computed for various layers by considering them in the topological order. However, just computing the normalizers is not enough; as they depend on the model parameters, dependence in the gradient calculation. Batch Normalization (loffe & Szegedy 2015) proposes solution: approximate each expectation S with its minibatch estimate, through which we can backpropagate during training, using the population averages for inference. However, the fact that no layer other than the input observes the activations computed the same way during training as during inference is likely culprit for the catastrophic failures of batch normalization, when they happen (Ioffe 2017). 7). We would like to use the same normalizers during training and inference.
need to capture this
We could approximate the normalizers by computing exponential moving averages (EMAs) through- out training. The update rate of the EMA allows us to trade off between the variance (lower weight given to the last minibatch means lower variance) and staleness of the parameters (higher weight for the last minibatch means the statistics being averaged are computed for more-similar model parameters). The question is now how to backpropagate through the normalizers, once the examples other than the ones in the current minibatch are no longer available. One ap- proach was proposed in (Yao et al. 2020). In another, Batch Renormalization backpropagates using only the current minibatch, while using the EMAs in the forward pass. Specifically, batch renormalization normalizes its inputs (with vn, denoting the value for the nth training example) as Zn l StopGradient( ()) ) + StopGradient( BB where μ and a are the mean and standard deviation of N n and p. ō are their EMAs.
2.1 COMPUTATION GRAPHS, GRADIENTS, AND JACOBIANS
To better reason about what we would like the minibatch normalization scheme do, let consider the graph representing the computation of the full-batch loss. This is a directed, acyclic graph, whose inputs are the input examples (such as images and labels) and the model weights, and whose output is the batch loss. The other nodes represent intermediate computations. such as hidden layer activations. Each node's ' value is computed as a function (the node's operation) applied to the node's 's inputs (nodes from which there is an edge to the given node). Each fixed sequence of operations can be represented with such a graph. Gradient-based learning relies on the computation of gradients in this graph, typically computed by backpropagation. As alternative formulation, it is easy show that for any DL weight node W and loss L, the gradient aw can be computed by considering all directed paths in the graph from W to L, computing for each such path the product of the Jacobians corresponding to all edges in that path, and adding the products over all the paths. Buckpropagation computes this sum of products efficiently via dynamic programming.
2

---

Under review as a conference paper at ICLR 2022
Batch nodes and shared nodes. For ease of exposition, and without loss of generality, will consider a full-batch computation graph applied to a finite dataset (D1...)) where the nthh example D, could contain, for example, on computation graphs each of whose nodes is either a batch node or a shared node. Batch nodes compute value for each of the data points, while the shared nodes compute quantities that aggregated over the dataset. Specifically:
image and its ground-truth label. We will focus our attention
Batch nodes. The output of a batch node contains one value per training point. One of the batch nodes is the input node that contains the training data set: D = (D)... Dx). Other batch nodes take batch nodes and. possibly, shared nodes as inputs, and compute the outputs such that each output value depends only on the corresponding values from the input batch nodes.
Shared nodes. Shared nodes compute values that are shared by the entire dataset. Shared nodes include the model parameters. The other type of a shared node is an reduction node, which has batch node as the input and computes the average of the corresponding per-example values. Consider a batch node Q = (Q ... ... containing a value for each example. Then reduction node S would compute S N Σ On = En Qn To refer to the batch node Q that is the input the shared node S, we will denote Q Sbatch (Statch .. We are considering the setting in which the model loss L is also shared node, with the total loss computed as the average of per-example losses.
As an example, if the loss computed as = En WW - (4) - yol- where W the parameter to be learned, (In, are the training inputs and outputs, and μr IEn In, then the input batch node would have elements Dn = (In, yn). the shared nodes are L, fr and W, and the other nodes in the computation graph are batch nodes.
Graph over shared nodes. For a computation graph consisting of batch and shared nodes, it is easy to prove (by induction over the graph nodes) that the output of a batch node Q can be written as Q = (f S))==..... where f is some function and S is a subset of the shared nodes in the computation graph. In other words, each element of the batch of values computed by the batch node can be computed from the corresponding input element D, as well as some of the shared nodes. Constructively, such computation is represented by the subgraph of the computation graph consisting of the nodes and edges that can be reached by walking backwards from the given batch node without crossing shared nodes. follows that the value of a reduction shared node S can be computed as S = En, Sbatchh = E,, f(D,..)) We will say that S is S S-dependent ("S" is for "shared") on D and on each of R E S. A shared node S is S-dependent on R if R is either a shared node or the input batch node, and there is directed path from R to S in the computation graph that does not contain intermediate shared nodes.
We can represent the computation of I. (which is reduction node) graph whose nodes are the shared nodes and the input batch node, and which has an edge R -> S if in the original computation graph S is S-dependent on R. We will refer to this graph as the S-graph, and to the original computation graph as the base graph. See Figure [ Each node in the S-graph performs an operation on its inputs that is represented with a subgraph of the original computation graph. Computing the values of all the nodes in the S-graph is no more expensive than a single forward pass in the original computation graph. Similarly, computing the gradient of a linear combination of nodes in the S-graph, with respect to the model parameters, has the same cost as a single backpropagation in the original computation graph.
Jacobians between shared nodes. Consider an edge S + T in the S-graph, where S and T are shared nodes and T = f(D, S) (for some subset of the shared nodes S 3 S and the input batch D). We define the Jacobian TS, af [T From the chain rule, we have the gradient of the loss with respect to a parameter W: AL aw
= -.. J(Sx)1 S&) Sk-1)...JJ J(s. S1) where the summation is over all directed paths (So - T. St = I.) from W to I. in the S-graph.
Because the target T of an edge S + T is a reduction node, we have, T E, f(D,, S), thus J(S,T) = En. (..)). Therefore, the single-example Jacobian (s T) = S) is an unbiased estimate df (S, and can be computed from only the 7ihh training example and the values of some shared nodes. However, a product of such unbiased estimates will not give an unbiased estimate of the product of the Jacobians. For example, consider an S-graph with nodes = J(W, L) + J(S, LW S), but the
W. S, L, edges W + S, W + L, S + L. Then aL aw
3

---

Under review as a conference paper at ICLR 2022
L 1,-MMX m Muneybyy
1% x
Exx,
Wxxz))
E w,00
E
1% 1
Figure 1: An example computation graph (left) and the corresponding S-graph (right). Shared nodes are ovals, batch nodes are rectangles. The S-graph contains the shared nodes and the input batch node, with edge between two nodes if path exists between them in the computation graph that does not pass through other shared nodes. A node in the S-graph computes E. f(Dn, S) where (Dn) are the training examples and S is subset of the shared nodes, An edge S + T in the S-graph corresponds to the Jacobian J(S T). and the desired gradient of the loss with respect to model parameters is a sum of products of such Jacobians. This work aims construct better estimates of such sum of products using a minibatch of examples, or single example.
approximation of this gradient using a single example (or a minibatch) is not in general unbiased: En (W, L) + (S,
= J(W, L) + E., J. (S, (S, L)J,, S) + dL DW.
2.2 DEPENDENCIES BETWEEN JACOBIANS LEAD TO BIASED GRADIENTS
Below, we note some of the ways that dependence between the Jacobian estimates in the computation graph manifests itself, and some ways to reduce these dependencies.
Scaling the second moment gradients. Consider the divisive normalization yn. Sg. where the activation in the nth example, and S is the second moment of the inputs. For shared nodes R and T with S-graph edges R + S - T, the single-example estimates Jn (R, 5) and depend on the corresponding Hn and therefore on each other, resulting in (S, J,,,, S) being a biased estimate of J(S, J S). To ameliorate this, we will modify the estimates of the form (R, S) reduce their dependence on the nith example. Observe that xn is an output of a linear layer 2n = W u, with weights W, then both Qn = xa and &Q,, of un. Based on this, we propose scaling (W, by Q
VS*
()
OW scale quadratically with the scaling $ such that it is independent of the scale of
Un (and the expected inverse scaling amount of scaling is 1). We apply the same scaling to Jn, S) for all R, which amounts to scaling the gradient as it backpropagates through S. If a minibatch of B examples are allowed to interact, and the Jacobians are estimated E J. then we scale the gradient by E Q critical).
In both cases, the amount of scaling is capped by fmax (that value is not
Reducing dependence due to per-node normalization and nonlinearity. Another source of dependence between J(R, S) and J(S. T), where S is the second moment, is the normalization followed by a nonlinearity. For simplicity but without loss of generality, consider a ReLU, As above, the single-example estimate (R, S) depends on Fn. even with the gradient scaling above. So does the single-example estimate of (S, T), which has (RoLU(Tyx+t9) as factor, which is 0 whenever " + 3 < O. We alleviate that dependence by reducing the degree to which a particular nonlinearity input is affected by any normalized activation. To do so, consider an entire vector of node values X,, e.g. the output of a neural network layer for the nith example. Instead of applying both the normalization and the nonlinearity per node, we will compute ReLU(R (7y) + (3) for some fixed matrix R. with denoting matrix product. We can set R a random and fixed rotation matrix, preserving the total energy in both the forward and back pass. Note that if x is computed by linear layer (or convolution), then the rotation by R can be fused into that linear layer and preserves the inference model class. For some matrix sizes, we can also set R to Hadamard matrix, which has the smallest maximum absolute value of the elements. In the convolutional case, we apply a x convolution between the normalization and the nonlinearity, multiplying the Hadamard matrix (or random rotation) by the stack of channel values at each spatial location.
Shrinking the gradients. An additional modification to the gradients that we have found very useful is shrinkage. The intuition is that the variance in the estimation of J(R, S) can be reduced by multiplying the estimate by scalar < 1, and this may in turn reduce the estimation error by finding a more favorable bias- barsinaiancee tradeoff. When scaling the gradients, we apply the same scaling to J (R, S) for all R, and therefore can accomplish this by simply modifying the backpropagation
4

---

Under review as a conference paper at ICLR 2022
through S. When processing the nth example, replace S with ScaleGrad( (S S S Shatch ()) which returns S but scales the gradients in backpropagation.
Averaging Jacobians over the minibatch. Consider minibatch of B examples, from which we want to estimate the gradient al which in turn is the sum of the Jacobian products of the (yy form 1, ... (So. S1) (with So W, S L). Each of Jacobians could be estimated as the minibatch average. Although these estimates are not independent, this can be effective in practice when combined with gradient scaling, and is what Batch Renormalization does. However, computing the minibatch estimates of all of the Jacobians forces interactions among the examples in the minibatch at each layer during backpropagation. We could estimate all of the Jacobians based on individual examples to avoid the interactions, but this increases the bias in the gradient estimation due to the dependence between the Jacobians. As a compromise between these two extremes, we break the product -[ Sk) ... J [So, S1) into subproducts, such that each of these factors is estimated as the product of per-example estimates of the constituent Jacobians for each example in the minibatch, averaged over the minibatch. A decomposition that both gives good empirical performance and lends itself to efficient computation is into two groups, one of which is the single Jacobian J( L)
and the other containing all the others:
1( .)..) (So. S1) biased Σb Jo (s EB Jo 2. Sk-1) (S..
(1)
B
B
The first factor is the minibatch estimate of the partial derivatives of the L node in the S-graph, can be computed via backpropagation, respect to the shared nodes, without backpropagating through the shared nodes. The second factor the biased estimate of the gradient of the linear combination of the shared nodes with respect to the model weights. The product of the two factors is computed by backpropagation from the shared nodes, feeding in the first decomposition factor, and averaging the resulting per-example gradients. The estimate in Eq. (I) is biased even as B - oo because e.g. En Jo(S,, S2) (o S1) / S1). To alleviate the bias, we need to use the rescaling of the gradients through the second moments, and may need apply the gradient shrinkage, as described above.
the minibatch average of per-example gradients with
3 TWO-STAGE BACKPROPAGATION
We present a complete description of the two-stage backpropagation procedure here. We consider a model that contains the subtraction of the mean and division by the square root of the second moment. We estimate the gradients using minibatch of size B, use the estimates (such as exponentially moving averages, EMAs) of the shared nodes in the forward pass, and modify the backpropagation of the gradients.
3.1 MODIFYING THE NORMALIZATION
For generality we consider the convolutional case where multiple values in the model share the normalization statistics. Let x,, (f5.1...a) be the vector being normalized corresponding to the bth example with d being the number of units normalized jointly, such as the number of spatial locations in the convolutional map. Below we show the computation of the subtractive normalization that subtracts the data mean Jr, and divisive normalization that divides by the square root of the data second moment v. Let define ScaleGrad( S, Q,r) for scalar as the operation returning its first argument S, but propagating the gradient to Q after scaling it by r, so that ScaleGrad S, Q, = StopGradien((r)Q + StopGradient(S rQ).
Subtractive: Define per-example mean 14 = EE Tb,i, and current population mean = (MMA(())) Define n ScaleGrad(T, Mb: Tm)) for gradient shrinkage T.. (1 by default). Output (x, u!). E =8,,- D EMA(p)), vf ScaleGrad(z, ry min( vn fm.x)). Output
Divisive: Let v4-
Such normalizations do not involve example interactions in backpropagation, If the examples in the minibatch are allowed to interact (as in Batch Renormalization) the subtractive normalization would be modified to use f = E (4 instead of pa, and the divisive normalization to use v = E. v4 instead of Ub-
5

---

Under review as a conference paper at ICLR 2022
3.2 MODIFYING THE BACKPROPAGATION
We use the estimates from Eq. ① by breaking the backpropagation into two parts, interleaved with gradient aggregation. Let La is the loss computed on 6th example in the minibatch, with the minibatch loss defined as & = Eτ Lg- We will refer to the set of the means in the model as {u}, the set of the second moments as {2), and the weights {W}.
In Sec. 3.1 we modified the normalization such that backpropagation does not require inter-example interactions. Now let us add extra nodes that cause the per-example gradients to interact. For the subtractive normalization, let μr² ScaleGrad( E, (b:Tm) Tm). For the divisive normalization, let v = ScaleGrad(1, p, F. Vb. min( F E. u, ). Note that the gradients of μ and v can be computed by combining per- per-example gradients computed separately from each other.
(a) Backpropagate from the minibatch loss & to all ( [v/]}, and (W). without backpropagating through any of those nodes. This can be done for each example independently. Denote the resulting gradients {yt") ; {9h} and {g"). respectively. Define g" = Σb gh 9" = Σb 96- These are estimates of Jacobians of the form J(S, T.).
(b) Feed in the gradients from the previous step into backpropagation starting from {ui} and {ut}, iee. compute the gradients, with respect to [W), of Σye(n) g" + Eve(v> g" ot where g are considered constant. This can be done for each example separately followed by adding the gradients. Denote the gradients from this step (h").
(c) The total gradients are {o" + 3.
Compared to standard backpropagation in model with batch normalization (or batch renormaliza- tion), the above procedure uses roughly twice the amount computation because the backward pass is performed twice. What is gained is the ability to process one example at a time.
3.3 NORMALIZING CONVOLUTIONAL MODELS
We compute the mean before the linear operation and the second moment after it, which removes the dependence of the mean on the weights of the given layer. With batch normalization, subtracting the mean before or after the linear operation is equivalent, but in our case they are not due to the staleness of the moving averages; computation of mean μ before the linear transform removes the dependence of the EMA estimate of the mean on the layer weights W. If the inference architecture is to add the bias after the linear operation and not before, equivalently replace W(x (1) Wx WW. μ. In convolutional models with circular padding, we can similarly replace W (x - -4) with W + W where W is the sum of the convolutional kernel W along the spatial dimensions. We apply the same modification even with other padding types; in such cases subtracting W μ does not make the normalized activations zero-mean, but we have not found this to be an issue in practice.
3.4
TWO TOWERS
The combination of the normalization modification (to avoid gradient aggregation) and the two-step gradient computation has another interpretation, in which two model towers are created. The first tower uses the modified normalization and also computes the aggregated moments {u²) and {vē}. Those are used to perform the normalization in the second tower, and the minibatch loss is set to the loss computed from the second tower. In backpropagation, gradients for different examples interact only when adding the gradients with respects to the weights, and when backpropagating through {u'} and {v'}. Therefore the gradient computation can be broken up into two stages, the first backpropagating (in the second tower) to {ut) and {ut} and the second continuing from there (in the first tower), which amounts to two sets of per-example computations, with the gradient aggregation after each of the two sets.
We show the construction of a convolutional layer in such two towers in Alg. There the layer receives the minibatches in the first tower and x' in the second. The two sets of values, well as the outputs of the corresponding layers in the two towers, are equal in the forward pass, but the gradients propagate differently between the towers. We also show how the corresponding layer is constructed in the inference model, such that at the inference time acts like a standard convolution followed by bias and the ReL.U. One additional modification the algorithm introduces is scaling the

---

Under review as a conference paper at ICLR 2022
Trainable: W (size H HxHxC;xCo, for convolution spatial support H);); B, 7 (size Co) scalars mm J maxe, Umax (bound on second moment of normalized activations during training); rotation (size C, Coi optional)
Const:
Moving averages: μ (size Ci), D (size Co)
Function Serilayer_Train(e (x, x'): D Input sizes: BxSx5xC; (for spatial size Sx
/ Eij Ni:: /1 ScaleGrra(TT. (m, D size B x C:
p conv(x, W) 1e (Zis W W.gis)))
V Ei Pik..- + E; v = ScaleGrad(D,, M, min( t, fmax))
q P. StopGradient (minieennttmmnn.,. Umax ))
voi F ScaleGrad( pr, Es Fm):
Dt ScaleGrad(V,, EE Vo,,; Tv min( Es Pa, fmax))
q' = conv(x' W) 1² EE., Wasne))
StopGradienn((iin(.. Umax V B)) UpdateEMA(TI, Es :: UpdateEMA((z, Es Vb,.)
return ReLU((Rtatie
7) + (3), ReLU(Rotate(( ① 7) + (3)
Function ConvLayer__nfereneeiis
K
Won R; return ReLU(conv(x, K) + (- 1--matuul pr Eig Kighe) B))
Function Rotate(y):
return convly, reshape(R, R. x 1 x Co x &..))
Algorithm 1: Convolutional layer for training and inference. The training inputs are minibatches x and x' which have the same forward values but differently defined gradients. We build two towers, where the first has the gradients defined to enable per-example gradient computation, and computes the moments used to perform the normalization in the second tower. The loss used for backpropagation is computed for the second tower. The operator ScaleGrad (S, Q, behaves like StopGradien((1 )Q in backpropagation, but returns in the forward pass. The training loss is computed from the second tower, corresponding to x'.
activations to prevent their per-example norm from becoming too large, which happens rarely but can cause instability in training. No such scaling is applied at the inference time.
RELATED WORK
There has been a substantial body of work attempting to reduce or remove the minibatch dependence that is required by batch normalization. One class of approaches changes the architecture such that the normalization is confined to a single example, and the inference architecture contains the same normalization. This includes group normalization (Wu & He, 2018) and layer normalization (Ba et al. 2016), among many others, that normalize sets of activations for an example relative to each other, rather than normalizing the values of the same activation across approaches are applicable in situations that batch normalization is not, such as recurrent models. The normalization affects the model class, which can be a disadvantage, as it conflates two distinct properties of the model architecture - on the one hand, being able to capture the function of interest and containing the appropriate inductive biases, and on the other hand being amenable gradient optimization. This differs from batch normalization, which does not change the function class.
of examples. Such
Another class approaches eliminates the normalization altogether, via combination of architectural building blocks and appropriate initialization. An example of this is the fixup initialization (Zhang et al. 2019), combining the residual connections (He et al.) 2015) with the initialization that causes the residual blocks to act as identity at initialization. This tightly coupled to the choice of the function class. We aim to provide a method that, like batch normalization, does not change the function class, does not rely on skip connections, and enables training that behaves similarly (in terms of accuracy as the function of the training steps) as batch normalization.
Yet another class of approaches, which ours is closest aims recover the properties of batch normalization while reducing the dependence on the minibatch size. One such work is (Yao et al. 2020), which approximates the gradients via the normalization statistics using the information
7

---

Under review as a conference paper at ICLR 2022
collected from multiple examples, while processing only one example at a time and, in addition, enables the use of that information even when the model has been updated since the other examples have been seen. The experiments in that work have been conducted using ResNets (He et al., 2013), while in this work we deliberately avoid the models with skip connections to avoid dependence on the advantages they provide.
The approach taken by (Chiley et al. 2019) also aims at using single examples, combined with the moving average history of previous example statistics, to approximate the behavior of batchnorm. However, they include an extra step of normalizing the activations in an example, relative to the other activations, changing the function class, which we aim avoid.
5 EXPERIMENTS
We train the Inception-v3 (Szegedy et al. 2015) model (without the side head) on the ImageNet ILSVRC-2012 dataset (Russakovsky et al. 2015). The learning rate for the final classification layer is reduced by 10x to account for the potentially larger activations without the minibatch normalization.
The training uses 50 workers, each processing minibatches of size 32 (32 x 50 1600 examples per training step). After each step the moving averages are updated at the rate of 0.2. The moving averages are initialized by using the first 200 steps only update the EMAs but not the weights. For 50 updates after that, each weight update is followed by 50 steps of only updaing the EMAs, to avoid their staleness very early in training when the weights change rapidly (an alternative is ramp up the learning rate (You et al. 2017). During training, we clip the per-example second moment of the normalized activations by umax 5, and cap the scaling of the gradient through second moments by fmax= 2 (these values are not critical). Gradient shrinkage is not used (rm. Tv 1) unless noted. The test accuracy is increased slightly by decaying learning rate half the speed for the biases as compared to convolution weights, injecting multiplicative (uniform from [0.9, 1.1)) and additive (uniform from [-0.1, 0.1) noise after normalization, and setting the dropout keep probability to 0.6.
For the experiments below, we provide the peak validation accuracy along with the number of training steps required to achieve it, as well as the accuracy achieved after the same number of training steps as it takes the baseline model to achieve its peak accuracy.
Baseline. We train the Inception model on ImageNet, using batch normalization. The model reaches the peak validation accuracy of 78.5% % at 120k steps. For the experiments that follow, we report the validation accuracy at 120k steps, as well as the peak validation accuracy.
Two-stage gradient computation. We replaced batch normalization with the method in sec. 3 which allows the gradients to be estimated processing each example separately, followed by the aggregation of gradients, followed by the second stage of gradient computation processing each example separately and using the above aggregation results. The model accuracy is 77.8% at 120k steps, increasing to 78.2% % at 150k steps.
Different minibatches for two stages. While the method in sec. 3 replicates a minibatch, the minibatches used for the two parts of gradient computation do not need to be the same. We have found that the only a small reduction in the model performance results from using a much smaller minibatch for the second stage of the gradient computation (backpropagating from the shared nodes other than the loss). We augmented a minibatch of 32 examples with an extra minibatch of 4 examples, for the resulting minibatch (x1::: xx All 4 + 32 examples are sampled from the training dataset, but the labels for X1:4 are not used. This gives the accuracy of 77% at 120k steps, which increases to 77.6% % at 150k steps.
Improving batch renormalization. Batch renormalization (Toffe] 2017) allows the use of smaller minibatches for normalization, compared to BN. Here, we show how to improve it. Instead of computing the gradient in two passes in sec. at every layer. To enable this, we use the statistics jet and D² instead of je' and v². This removes the need for two minibatches and the doubling of computation, but causes the example gradients to interact at every layer.
perform cross-example normalization
We small microbatches (of size 2 or 4) for normalizaoii,,,,,, with the gradients from multiple microbatches (total 1600 examples) added to perform an update. When normalizing over microbatches of size 4, we get the accuracy of 77.4% % at 170k steps (76.2% at 120k). For this microbatch size,
8

---

Under review as a conference paper at ICLR 2022
(loffe 2017) reports the smaller peak accuracy of 76.5% (at 130k steps, at which point get 76.7%). For a microbatch of 2, we obtain the accuracy of 72% at 180k steps (70.5% at 120k). This can be improved by using the Hadamard rotation between the normalization and ReLU as described in sec. [.2] raising the accuracy to 74.8% % at 180k steps (73.5% % at 120k). The rotation and the preceding normalization can be fused for inference with the convolution before them, thus the inference architecture remains unchanged.
Further improvements the accuracy at small microbatches can be achieved by shrinking the gradients through the means by I'm and through the second moments by rv. For a microbatch of size 2, with Hadamard rotation and I'm = rv 0.5 we get 77.8% at 155k steps (77.2% % 120k).
Single examples. Finally, to push the algorithm to its extreme, we considered the true per-example gradient computation, with no interaction between examples. This is the limiting case of our Batch Renormalization modification, with the microbatch size of I (which the original batch renormalization does not allow). Without gradient shrinkage, we trained the model both with and without the Hadamard rotation, with the rotation increasing the validation accuracy by about 5%: 68.3% vs. 63.8% 5..8% vs, 60% at 120k), with vs. without the rotation, respectively. The rotation is critical in this setting, and the performance can be improved further via gradient shrinkage. With coarse search for Tm: Tv, for 7' = 0, rv = 0.8 we achieve the validation accuracy of 75% at 170k steps (73.7% % at 120k). Although at 120k steps this model performs 5% worse than the baseline model (with BN), our method requires no interaction between the examples, and has the same computation complexity during training as the batch-normalized model.
Residual networks. In addition to the experiments with the Inception architecture, we evaluated the performance of our method on the Resnet-50 architecture (He et 2015). We applied batch normalization to each convolutional layer, and also to the linear skip-connections that change the number of channels. The baseline achieved 76.3% % top-1 accuracy. Two-stage gradient computation gave a similar 7..6%. The improved batch renormalization (described above), normalizing over microbatches of 2 examples, gave 75% when using the gradient scaling factors rm 0, ry = 0.8 and Hadamard rotation.
DISCUSSION
We have presented an interpretation of batch normalization as a method for approximating population normalization similar to the way that stochastic gradient computation aims to approximate the full gradients. We seek other solutions with more favorable properties: making the training model and inference model equivalent in the forward pass, and reducing the cross-example interaction in the gradient computation. We propose a method for estimating the gradients that, by reducing this interaction, enables the gradients to be estimated without holding more than one example in memory at any given time, at the cost of doubling the computation, although this can be reduced (at small cost to the accuracy) by using a smaller minibatch for portions of the gradient estimation. We show how to achieve better performance with batch renormalization for extremely small (2 examples) minibatches, and demonstrate only a modest degradation in performance when avoiding cross evoss-example interactions altogether, wherein the parameter gradient is simple average of the per-example gradients. Our method does not change the inference architecture of convolutional models, in that the model used for inference is still a convnet without any extra operations such inference-time normalization. Our experiments are conducted on Inception-v3, which, unlike the state-of-the-art image analysis architectures, lacks skip connections and enables us to show that our proposed approach does not take advantage of the beneficial properties of such shortcuts. We did not conduct experiments on the state-of-the-as architectures that benefit from skip connections.
We proposed several building blocks motivated by reducing the detrimental effects of the bias introduced when the gradients are estimated from minibatches, and products of Jacobians are estimated as products of their stochastic and non-independent samples. This includes scaling the gradients to compensate for the easily-identifible sources of dependence, introducing the rotation between the normalization and nonlinearity to reduce the gradient dependence, and a scaling applied to gradients as they backpropagate through certain nodes in the computation graph. We suspect that at least some of these building blocks not only enable the reduction in the minibatch size required for normalization, but are applicable to other use cases in which the population gradient is not a mere expectation of per-example gradient over the data, and needs to be estimated from minibatches.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
Andrew Brock, Soham De, Samuel L. Smith, and Karen Simonyan. High pegh-pperformance large-scale image recognition without normalization, 2021.
Vitaliy Chiley, Ilya Sharapov, Atli Kosson, Urs Koster, Ryan Recce, Sofia Samaniego de la Fuente, Vishal Subbiah, and Michael James. Online normalization for training neural networks, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015.
Sergey loffe. Batch renormalization: Towards reducing minibatch dependence in batch-normalized models, 2017.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. CoRR, 1b3/1502.03167, 2015.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bemstein, Alexander C. Berg, and Li Fei-Fei. Imagenet large scale visual recognition challenge, 2015.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jonathon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision, 2015.
Yuxin Wu and Kaiming He. Group normalization, 2018.
Zhuliang Yao, Yue Cao, Shuxin Zheng, Gao Huang, and Stephen Lin. Cross-iteration batch normalization, 2020.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks, 2017.
Hongyi Zhang, Yann N. Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization, 2019.
10


Reviews:
[{"review_id": "UPWa11hB9X", "reviewer": "Reviewer_cEEb", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "While the problem of improving BN's performance on small training batch sizes is important, I find that this paper can be significantly improved in several aspects.\n\n1. The approach is rather ad hoc and not properly motivated. As far as I understand, the forward pass of the proposed method is very similar to batch renorm. The main contribution is a series of techniques to modify the back propagation. This makes it both hard to understand as well difficult to adopt in practice. \n\n2. Empirical evaluation is underwhelming. The majority of the experiments are conducted on Inception V3, with the only two baselines being BN and batch renorm. Also, on large batch sizes, the proposed method underperforms BN despite the additional complexity. On small batch sizes, it marginally outperforms batch renorm, but given the scope of the experiments, it's hard to draw a convincing conclusion. \n\n3. The presentation can be improved. First of all, I find the reliance on the computational graph terminology to be quite unnecessary, and it adds difficulty for the precise understanding of the details. Second, the experiments section can be much better organized, with tabulated results and clear comparisons to baselines. ", "summary_of_the_paper": "This paper proposes an improved version of batch normalization that's intended to work well on small training batch sizes. The technique is similar to batch renormalization, but adopts several techniques to manipulate the backprop procedure. Experiments are are conducted primarily on the Inception V3 network, with comparisons to BN and batch renorm. ", "summary_of_the_review": "Overall, I find that this paper does not meet the bar for ICLR. I suggest the authors work on improving the methodology, adding more evaluations and improving the presentation. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "NJDN5UuSV8D", "reviewer": "Reviewer_mQAF", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength: the idea of altering the computational graph is novel, especially to encompass several samples in a single graph where the weights are 'shared'\n\nWeakness: The computational experiments as I read them show that the new approach is subpar. It seems that the baseline always outperforms. The experiments in general are not presented in a good way. It comes across as 'running out of space' in the experimental section. I suggest to remove or shorten Sections 2 and 3 and improve Section 5. The experiments are also limited to a single dataset. ", "summary_of_the_paper": "In batch normalization (BN) it is common practice to consider statistics over the mini-batch. The authors propose a two step approach for per-sample based normalization that is more accurate (they actually don't prove this) than standard BN. They do this by augmenting the computational graph (in essence the computational graph involves several samples and not just a single sample).  ", "summary_of_the_review": "The idea is sound and it should definitely be further studied. The experimental section begs for improvements. Not just the exposition, but the entire study. More datasets should be included. Inception-v3 is a rather old architecture and thus more modern architectures should be considered. ResNet-50 is included which is positive. \n\nMinor remark: The hyper parameters are listed however there is no word how they were selected. Usually with BN gradient clipping is not needed so it is surprising that it is used. \nI also do not see a discussion on computational times. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "ElQ3GE8hKNv", "reviewer": "Reviewer_bk8j", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper includes some interesting insights on normalization. However, I find it difficult to catch the practical insights. The proposed method doesn’t improve the accuracy, but adds computations. Why would a practitioner use it? I understand the idea is to perform online normalization in a manner similar to SGD, but it doesn’t seem to offer any practical advantages over the existing batch-normalization method.", "summary_of_the_paper": "The paper offers an interpretation of batch normalization as a method for approximating population normalization. The proposed method keeps training model and inference model equivalent in the forward pass, and reduces the cross-example interaction in the gradient computation. Furthermore, the paper introduces a method for estimating the gradients, by reducing interaction between examples, holding a single example in memory at any given time. This comes at the cost of doubling the computation. The computation can be reduced with some cost to the accuracy by using a small mini-batch for parts of gradient estimation. \n", "summary_of_the_review": "The motivation for the proposed method is weak. Although the overall insights are interesting. Improving the motivation, especially with regards to the practical aspects (as detailed above) will help me to improve the ranking of the paper. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "94gLF2aEiJ3", "reviewer": "Reviewer_CZjS", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths**\n\nThe idea of this work is novel and interesting, I really appreciate it.\n\n\n**Weakness**\n\nFrankly speaking, I think this submission is more like a technical report rather than a complete paper. I will demonstrate the weakness of it one by one:\n\n1. Missing important literature. \n\nThe authors missed a previous work MABN(https://openreview.net/forum?id=SkgGjRVKDS). If you read this paper, you will see MABN also deal with small batch issues of BN without introducing additional operators into the model during the inference time. Similar to this paper's idea, MABN also uses moving average statistics during back propagation to solve the drawback of BRN in an elegant way. MABN brings little extra computation overhead while your method double the times of bp during the training time. In a word, if the proposed method only addresses small batch issues of BN, it's trivial and perform worse than a previous method.\n\n2. The demonstration on intuition in section 2.1, 2.2 is really confusing. \n\nIn section 2.1, there are a large amount of text description on some general concepts like computation graph, batch nodes, shared nodes, etc, and the relations among them. I really struggled to get the point why these concepts were needed. As far as I'm concerned, there are only three types of nodes you actually used: weights, feature map(batch nodes), and batch statistics. So why you invent new terms to make readers struggle to understand? Another small complaint is that why you denote jacobian $\\frac{\\partial T}{\\partial S}$ as $J(S, T)$? It should be easier for readers to realize the chain rule if you denote the jacobian $\\frac{\\partial T}{\\partial S}$ as $J(T, S)$.  In a word, I think it will be easier to understand if these concepts can be mathematically formulated. Proper formulas and illustrations is better than verbal descriptions if you aim to propose a general framework;\n\nBesides, I didn't get the main idea of this paper, until the I saw the last sentence in section 2.1. I thought this claim was very interesting, and somehow reasonable, so I looked forward to seeing justification on this claim in the following part. When I saw the title of the section 2.2, the expectation became greater. But what I have read in section 2.2? You began to introduce specific tricks in neural networks consisting of linear, BN, relu layers. Then why you propose a general framework in section 2.1? Why not just directly demonstrate what method you apply on a neural network only consisting of BN, relu, linear(CNN) modules? Besides, I failed to understand why you propose these tricks, the intuition behind these tricks is too vague, I cannot understand the connection between them and the last claim in section 2.1. \n\n3. The experiment result is not convincing.\n\nThe most confusing part is why the experiment results are presented in text descriptions instead of charts or tables. I think it will be much easier for readers to read and compare.\n\nOnly two experiment setting are examined in this paper: inception-v3 and resent50 on imagenet. And the training setting is really wired: the batch size is $32\\times50$. I understand it's perhaps too time-consuming to set a real small batch setting in imagenet experiments, but why you do not use a more standard settings, like $32\\times 8$ batch size to verify the effectiveness of the proposed method in inception-v3? Besides, a real small batch cases are needed, like detection task or segmentation task, in which the batch size are usually lower than $8$. \n\nI also notice you claim that \" in this work we deliberately avoid the models with skip connections to avoid dependence on\nthe advantages they provide.\" Do you imply the previous work (Yao et al. 2020) relies on the skip connections structure? I don't agree. So you need to justify your claim that it is indeed an advantage of your method to work well on non-residual structure, while the counterparts, CBN(Yao et al. 2020) or even MABN, cannot. It means more sophisticated comparison experiments are necessary to make your paper more convincing.\n\n4. Citation format is not correct.\n\nYou need to clearly show the source (journal, conference, etc) of these cited papers . Not just to show respect to the previous work, but also help readers to properly assess the reliability of your paper, if you use these cited papers as supportive evidence. Please follow a formal citation style.\n\n\n", "summary_of_the_paper": "This paper propose a training method to deal with small batch issues of batch normalization. Their proposed method does not need to introduce any extra operators into the model during the inference time. The authors verified the effectiveness of the method in Inception-v3 trained on imagenet.", "summary_of_the_review": "I like the idea of this work, but obviously the authors do not elaborate or justify their idea in either theoretical or empirical aspects. The results is not significant comparing with existing work, either. Besides, the submission seems not a complement paper due to its writing. I lean to reject it.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies the introduction of a variant of batch normalization (BN) to train deep neural network. The underlying idea is a two-step approach for per-sample based normalization, relying on augmenting the computational graph to handle "several samples" nodes.

The reviewers have mentioned that the idea of altering the computational graph is interesting and potentially novel.
Yet, the numerical experiments were not enough precise or solid to back up the claims by the authors, that their proposed BN alternative is of practical interest.
It was also raised that the paper lacks theoretical supports: no formal analysis, most explanations are ad hoc, etc.
Decision: Reject