Training Deep Networks
When working with data, we often preprocess before training.
Choices regarding data preprocessing often make an enormous difference in the final results.
Recall our application of MLPs to predicting house prices (:numref:sec_kaggle_house).
Our first step when working with real data
was to standardize our input features to have
zero mean \boldsymbol{\mu} = 0 and unit variance \boldsymbol{\Sigma} = \boldsymbol{1} across multiple observations :cite:friedman1987exploratory, frequently rescaling the latter so that the diagonal is unity, i.e., \Sigma_{ii} = 1.
Yet another strategy is to rescale vectors to unit length, possibly zero mean per observation.
This can work well, e.g., for spatial sensor data. These preprocessing techniques and many others, are
beneficial for keeping the estimation problem well controlled.
For a review of feature selection and extraction see the article of :citet:guyon2008feature, for example.
Standardizing vectors also has the nice side-effect of constraining the function complexity of functions that act upon it. For instance, the celebrated radius-margin bound :cite:Vapnik95 in support vector machines and the Perceptron Convergence Theorem :cite:Novikoff62 rely on inputs of bounded norm.
Intuitively, this standardization plays nicely with our optimizers
since it puts the parameters a priori on a similar scale.
As such, it is only natural to ask whether a corresponding normalization step inside a deep network
might not be beneficial. While this is not quite the reasoning that led to the invention of batch normalization :cite:Ioffe.Szegedy.2015, it is a useful way of understanding it and its cousin, layer normalization :cite:Ba.Kiros.Hinton.2016, within a unified framework.
Second, for a typical MLP or CNN, as we train,
the variables
in intermediate layers (e.g., affine transformation outputs in MLP)
may take values with widely varying magnitudes:
whether along the layers from input to output, across units in the same layer,
and over time due to our updates to the model parameters.
The inventors of batch normalization postulated informally
that this drift in the distribution of such variables could hamper the convergence of the network.
Intuitively, we might conjecture that if one
layer has variable activations that are 100 times that of another layer,
this might necessitate compensatory adjustments in the learning rates. Adaptive solvers
such as AdaGrad :cite:Duchi.Hazan.Singer.2011, Adam :cite:Kingma.Ba.2014, Yogi :cite:Zaheer.Reddi.Sachan.ea.2018, or Distributed Shampoo :cite:anil2020scalable aim to address this from the viewpoint of optimization, e.g., by adding aspects of second-order methods.
The alternative is to prevent the problem from occurring, simply by adaptive normalization.
Third, deeper networks are complex and tend to be more liable to overfitting.
This means that regularization becomes more critical. A common technique for regularization is noise
injection. This has been known for a long time, e.g., with regard to noise injection for the
inputs :cite:Bishop.1995. It also forms the basis of dropout in :numref:sec_dropout. As it turns out, quite serendipitously, batch normalization conveys all three benefits: preprocessing, numerical stability, and regularization.
Batch normalization is applied to individual layers, or optionally, to all of them:
In each training iteration,
we first normalize the inputs (of batch normalization)
by subtracting their mean and
dividing by their standard deviation,
where both are estimated based on the statistics of the current minibatch.
Next, we apply a scale coefficient and an offset to recover the lost degrees
of freedom. It is precisely due to this normalization based on batch statistics
that batch normalization derives its name.
Note that if we tried to apply batch normalization with minibatches of size 1,
we would not be able to learn anything.
That is because after subtracting the means,
each hidden unit would take value 0.
As you might guess, since we are devoting a whole section to batch normalization,
with large enough minibatches the approach proves effective and stable.
One takeaway here is that when applying batch normalization,
the choice of batch size is
even more significant than without batch normalization, or at least,
suitable calibration is needed as we might adjust batch size.
Denote by \mathcal{B} a minibatch and let \mathbf{x} \in \mathcal{B} be an input to
batch normalization (\textrm{BN}). In this case the batch normalization is defined as follows:
\textrm{BN}(\mathbf{x}) = \boldsymbol{\gamma} \odot \frac{\mathbf{x} - \hat{\boldsymbol{\mu}}_\mathcal{B}}{\hat{\boldsymbol{\sigma}}_\mathcal{B}} + \boldsymbol{\beta}.
:eqlabel:eq_batchnorm
In :eqref:eq_batchnorm,
\hat{\boldsymbol{\mu}}_\mathcal{B} is the sample mean
and \hat{\boldsymbol{\sigma}}_\mathcal{B} is the sample standard deviation of the minibatch \mathcal{B}.
After applying standardization,
the resulting minibatch
has zero mean and unit variance.
The choice of unit variance
(rather than some other magic number) is arbitrary. We recover this degree of freedom
by including an elementwise
scale parameter \boldsymbol{\gamma} and shift parameter \boldsymbol{\beta}
that have the same shape as \mathbf{x}. Both are parameters that
need to be learned as part of model training.
The variable magnitudes
for intermediate layers cannot diverge during training
since batch normalization actively centers and rescales them back
to a given mean and size (via \hat{\boldsymbol{\mu}}_\mathcal{B} and {\hat{\boldsymbol{\sigma}}_\mathcal{B}}).
Practical experience confirms that, as alluded to when discussing feature rescaling, batch normalization seems to allow for more aggressive learning rates.
We calculate \hat{\boldsymbol{\mu}}_\mathcal{B} and {\hat{\boldsymbol{\sigma}}_\mathcal{B}} in :eqref:eq_batchnorm as follows:
$$\hat{\boldsymbol{\mu}}\mathcal{B} = \frac{1}{|\mathcal{B}|} \sum{\mathbf{x} \in \mathcal{B}} \mathbf{x}
\textrm{ and }
\hat{\boldsymbol{\sigma}}\mathcal{B}^2 = \frac{1}{|\mathcal{B}|} \sum{\mathbf{x} \in \mathcal{B}} (\mathbf{x} - \hat{\boldsymbol{\mu}}_{\mathcal{B}})^2 + \epsilon.$$
Note that we add a small constant \epsilon > 0
to the variance estimate
to ensure that we never attempt division by zero,
even in cases where the empirical variance estimate might be very small or vanish.
The estimates \hat{\boldsymbol{\mu}}_\mathcal{B} and {\hat{\boldsymbol{\sigma}}_\mathcal{B}} counteract the scaling issue
by using noisy estimates of mean and variance.
You might think that this noisiness should be a problem.
On the contrary, it is actually beneficial.
This turns out to be a recurring theme in deep learning.
For reasons that are not yet well-characterized theoretically,
various sources of noise in optimization
often lead to faster training and less overfitting:
this variation appears to act as a form of regularization.
:citet:Teye.Azizpour.Smith.2018 and :citet:Luo.Wang.Shao.ea.2018
related the properties of batch normalization to Bayesian priors and penalties, respectively.
In particular, this sheds some light on the puzzle
of why batch normalization works best for moderate minibatch sizes in the 50--100 range.
This particular size of minibatch seems to inject just the "right amount" of noise per layer, both in terms of scale via \hat{\boldsymbol{\sigma}}, and in terms of offset via \hat{\boldsymbol{\mu}}: a
larger minibatch regularizes less due to the more stable estimates, whereas tiny minibatches
destroy useful signal due to high variance. Exploring this direction further, considering alternative types
of preprocessing and filtering may yet lead to other effective types of regularization.
Fixing a trained model, you might think
that we would prefer using the entire dataset
to estimate the mean and variance.
Once training is complete, why would we want
the same image to be classified differently,
depending on the batch in which it happens to reside?
During training, such exact calculation is infeasible
because the intermediate variables
for all data examples
change every time we update our model.
However, once the model is trained,
we can calculate the means and variances
of each layer's variables based on the entire dataset.
Indeed this is standard practice for
models employing batch normalization;
thus batch normalization layers function differently
in training mode (normalizing by minibatch statistics)
than in prediction mode (normalizing by dataset statistics).
In this form they closely resemble the behavior of dropout regularization of :numref:sec_dropout,
where noise is only injected during training.
Batch Normalization Layers
Batch normalization implementations for fully connected layers
and convolutional layers are slightly different.
One key difference between batch normalization and other layers
is that because the former operates on a full minibatch at a time,
we cannot just ignore the batch dimension
as we did before when introducing other layers.
Fully Connected Layers
When applying batch normalization to fully connected layers,
:citet:Ioffe.Szegedy.2015, in their original paper inserted batch normalization after the affine transformation
and before the nonlinear activation function. Later applications experimented with
inserting batch normalization right after activation functions.
Denoting the input to the fully connected layer by \mathbf{x},
the affine transformation
by \mathbf{W}\mathbf{x} + \mathbf{b} (with the weight parameter \mathbf{W} and the bias parameter \mathbf{b}),
and the activation function by \phi,
we can express the computation of a batch-normalization-enabled,
fully connected layer output \mathbf{h} as follows:
\mathbf{h} = \phi(\textrm{BN}(\mathbf{W}\mathbf{x} + \mathbf{b}) ).
Recall that mean and variance are computed
on the same minibatch
on which the transformation is applied.
Convolutional Layers
Similarly, with convolutional layers,
we can apply batch normalization after the convolution
but before the nonlinear activation function. The key difference from batch normalization
in fully connected layers is that we apply the operation on a per-channel basis
across all locations. This is compatible with our assumption of translation
invariance that led to convolutions: we assumed that the specific location of a pattern
within an image was not critical for the purpose of understanding.
Assume that our minibatches contain m examples
and that for each channel,
the output of the convolution has height p and width q.
For convolutional layers, we carry out each batch normalization
over the m \cdot p \cdot q elements per output channel simultaneously.
Thus, we collect the values over all spatial locations
when computing the mean and variance
and consequently
apply the same mean and variance
within a given channel
to normalize the value at each spatial location.
Each channel has its own scale and shift parameters,
both of which are scalars.
Layer Normalization
🏷️subsec_layer-normalization-in-bn
Note that in the context of convolutions the batch normalization is well defined even for
minibatches of size 1: after all, we have all the locations across an image to average. Consequently,
mean and variance are well defined, even if it is just within a single observation. This consideration
led :citet:Ba.Kiros.Hinton.2016 to introduce the notion of layer normalization. It works just like
a batch norm, only that it is applied to one observation at a time. Consequently both the offset and the scaling factor are scalars. For an $n$-dimensional vector \mathbf{x}, layer norms are given by
\mathbf{x} \rightarrow \textrm{LN}(\mathbf{x}) = \frac{\mathbf{x} - \hat{\mu}}{\hat\sigma},
where scaling and offset are applied coefficient-wise
and given by
$$\hat{\mu} \stackrel{\textrm{def}}{=} \frac{1}{n} \sum_{i=1}^n x_i \textrm{ and }
\hat{\sigma}^2 \stackrel{\textrm{def}}{=} \frac{1}{n} \sum_{i=1}^n (x_i - \hat{\mu})^2 + \epsilon.$$
As before we add a small offset \epsilon > 0 to prevent division by zero. One of the major benefits of using layer normalization is that it prevents divergence. After all, ignoring \epsilon, the output of the layer normalization is scale independent. That is, we have \textrm{LN}(\mathbf{x}) \approx \textrm{LN}(\alpha \mathbf{x}) for any choice of \alpha \neq 0. This becomes an equality for |\alpha| \to \infty (the approximate equality is due to the offset \epsilon for the variance).
Another advantage of the layer normalization is that it does not depend on the minibatch size. It is also independent of whether we are in training or test regime. In other words, it is simply a deterministic transformation that standardizes the activations to a given scale. This can be very beneficial in preventing divergence in optimization. We skip further details and recommend that interested readers consult the original paper.
Batch Normalization During Prediction
As we mentioned earlier, batch normalization typically behaves differently
in training mode than in prediction mode.
First, the noise in the sample mean and the sample variance
arising from estimating each on minibatches
is no longer desirable once we have trained the model.
Second, we might not have the luxury
of computing per-batch normalization statistics.
For example,
we might need to apply our model to make one prediction at a time.
Typically, after training, we use the entire dataset
to compute stable estimates of the variable statistics
and then fix them at prediction time.
Hence, batch normalization behaves differently during training than at test time.
Recall that dropout also exhibits this characteristic.
(Implementation from Scratch)
To see how batch normalization works in practice, we implement one from scratch below.
