Covariate Shift is a phenomenon in machine learning where the distribution of the input features (covariates) changes between the training and testing phases, while the relationship between the inputs and the outputs remains the same.
Covariate shift to the intermediate layers of a neural network. Instead of the input data distribution changing between training and testing, the distribution of the activations (outputs) produced by one layer changes during training, causing the next layer to receive inputs from a constantly changing distribution. This was the primary motivations behind the development of Batch Normalization. The idea was that if the activations could be normalized to maintain a consistent distribution throughout training, each layer would learn under more stable conditions, leading to faster and more reliable convergence. Reducing ICS is not the primary reason batch normalization improves performance. The current understanding is that Batch Normalization mainly helps by smoothing the optimization landscape, making gradient based optimization more stable and allowing larger learning rates.
Used to stabilize and accelerate the training of deep neural networks Normalize the activations of a layer using the statistics (mean and variance) computed from the current mini batch. BatchNorm is usually applied before the activation function because its goal is to stabilize the distribution of the values entering the nonlinearity. If normalization were applied after ReLu, it would operate on already-clipped activations, potentially reintroducing negative values and reducing the intended effect of the activation. By normalizing the pre-activations, the activation function consistently receives inputs with a controlled mean and variance, leading to more stable and efficient training.
Consider an input matrix :
Here, rows represents the each sample in the batch, columns represents the features. We can take the mean and variance of each feature ;
[1.63,;1.63,;1.63,;3.27] $$
Now each value is normalized using the statistics of its own feature.
While this makes training more stable, it also restricts the representation learned by the network. Some neurons may actually benefit from producing activations with a different mean or variance. If we always force the outputs to have mean 0 and variance 1, the network loses this flexibility.
To overcome this, Batch Normalization introduces two learnable parameters for every feature (or neuron): γ (Gamma): scales the normalized activations. β (Beta): shifts the normalized activations.
$$ Y=\gamma X_{\text{norm}}+\beta \
X_{\text{norm}}= \begin{bmatrix} -1.22 & -1.22 & 0.00 & -1.22\ 0.00 & 0.00 & -1.22 & 0.00\ 1.22 & 1.22 & 1.22 & 1.22 \end{bmatrix}
\ \gamma=[2,;1,;0.5,;3] \quad \beta=[1,;0,;-2,;0] \ Y= \begin{bmatrix} -1.44 & -1.22 & -2.00 & -3.66\ 1.00 & 0.00 & -2.61 & 0.00\ 3.44 & 1.22 & -1.39 & 3.66 \end{bmatrix} $$
Notice that the outputs are no longer constrained to have mean 0 and variance 1. Instead, the network learns the most suitable distribution for each neuron while still benefiting from the stability provided by normalization.
<u><b>Running Mean and Running Variance Computation in Batch Normalization (BN)</u></b>
During training, Batchnorm normalizes using the current mini-batch statistics. During inference, there is no mini-batch ( or we don’t want the output to depend on other test samples), so it uses statistics accumulated during training.
First linear layer computation;
Computing the mean vector and variance vector we get;
Normalized output becomes;
Now we update the running statistics for the BN(1);
Second Linear Layer Computation & BN(2) Running Statistics Update
After Batch Normalization (and assuming an identity activation function for simplicity),
$$ Z_2 = \hat{Z}_1 W_2 = \begin{bmatrix} -1 & -1 & -1 \ 1 & 1 & 1 \end{bmatrix} \begin{bmatrix} 1 & 0 & 0 & 1 \ 0 & 1 & 1 & 0 \ 1 & 1 & 0 & 0 \end{bmatrix}
\begin{bmatrix} -2 & -2 & -1 & -1 \ 2 & 2 & 1 & 1 \end{bmatrix} $$
Computing the column-wise mean (
$$ \hat{Z}_2 = \frac{Z_2 - \mu_2}{\sqrt{\sigma_2^2 + \epsilon}} = \begin{bmatrix} \frac{-2-0}{\sqrt{4}} & \frac{-2-0}{\sqrt{4}} & \frac{-1-0}{\sqrt{1}} & \frac{-1-0}{\sqrt{1}} \ \frac{2-0}{\sqrt{4}} & \frac{2-0}{\sqrt{4}} & \frac{1-0}{\sqrt{1}} & \frac{1-0}{\sqrt{1}} \end{bmatrix}
\begin{bmatrix} -1 & -1 & -1 & -1 \ 1 & 1 & 1 & 1 \end{bmatrix} $$
Updating Running Statistics for BN(2)
Assuming initial running statistics for
Layer Normalization (LayerNorm) normalizes the activations across the features of each individual sample, rather than across the samples in a mini-batch. Independent of batch size Same behaviour during training and inference There is no running statistics
Unlike Batch Normalization (which normalizes across the batch dimension for each feature channel), Layer Normalization normalizes the activations across the features for each individual sample independently. This makes LayerNorm independent of batch size, which is why it is standard practice in Transformers and RNNs.
Given the input matrix
Assume learnable affine parameters: scale
$$ \hat{x}^{(1)} = \frac{x^{(1)} - \mu_1}{\sqrt{\sigma_1^2 + \epsilon}} = \begin{bmatrix} \frac{1 - 2}{\sqrt{2/3}} & \frac{2 - 2}{\sqrt{2/3}} & \frac{3 - 2}{\sqrt{2/3}} \end{bmatrix}
\begin{bmatrix} -\sqrt{1.5} & 0 & \sqrt{1.5} \end{bmatrix} \approx \begin{bmatrix} -1.2247 & 0 & 1.2247 \end{bmatrix} $$
$$ \hat{x}^{(2)} = \frac{x^{(2)} - \mu_2}{\sqrt{\sigma_2^2 + \epsilon}} = \begin{bmatrix} \frac{5 - 16/3}{\sqrt{14/9}} & \frac{4 - 16/3}{\sqrt{14/9}} & \frac{7 - 16/3}{\sqrt{14/9}} \end{bmatrix}
\begin{bmatrix} -\frac{1}{\sqrt{14}} & -\frac{4}{\sqrt{14}} & \frac{5}{\sqrt{14}} \end{bmatrix} \approx \begin{bmatrix} -0.2673 & -1.0690 & 1.3363 \end{bmatrix} $$
Since
RMS Norm is almost Layer Norm, except it removes one operation.
Taking the Mean sqauared
\frac{14}{3} =\sqrt{4.67} = 2.16 $$
\begin{bmatrix} 0.46&0.93&1.39 \end{bmatrix} $$
Researchers found that removing the mean isn’t actually necessary. Simply controlling the magnitude of the activations is often enough for stable training.
Instance Normalization (InstanceNorm) normalizes each channel of each sample independently by computing the mean and variance over the spatial dimensions (height and width). It does not depend on other samples in the batch. Group Normalization (GroupNorm) divides the feature channels into several groups and normalizes each group independently. It computes statistics within each group for every sample.
BatchNorm → Compare this feature across different samples. LayerNorm → Compare different features within the same sample. RMSNorm → Scale the sample by its overall magnitude. InstanceNorm → Normalize each image channel independently. GroupNorm → Normalize a few channels together as a group.