Introduction to Layer Normalization

Batch Normalization (BN), a technique used in deep learning, effectively addresses internal covariate shift by normalizing activations across a mini-batch. However, its reliance on batch statistics can sometimes be a limitation. For instance, BN's performance can degrade with very small mini-batch sizes, as the batch statistics become noisy estimates of the true population statistics. Furthermore, applying BN directly to recurrent neural networks (RNNs) can be tricky because the statistics need to be computed differently for each time step.

Layer Normalization (LN) offers an alternative approach that overcomes these specific limitations. Instead of normalizing across the batch dimension, Layer Normalization computes the mean and variance across all the hidden units in the same layer for a single training example.

How Layer Normalization Works

Imagine the activations $a$ for a specific layer $l$ generated from a single input example $x$. Let $H$ be the number of hidden units in that layer. Layer Normalization calculates the mean $\mu$ and variance $\sigma^2$ using all the summed inputs to the neurons in that layer for that single example:

$$\mu = \frac{1}{H} \sum_{i=1}^{H} a_i$$ $$\sigma^2 = \frac{1}{H} \sum_{i=1}^{H} (a_i - \mu)^2$$

Note that these calculations are performed independently for each training example and do not involve interactions across the batch.

Once the mean and variance are computed, the activations $a_i$ for each hidden unit $i$ in that layer are normalized:

$$\hat{a}_i = \frac{a_i - \mu}{\sqrt{\sigma^2 + \epsilon}}$$

Here, $\epsilon$ is a small constant added for numerical stability, similar to its use in Batch Normalization.

Finally, just like Batch Normalization, Layer Normalization introduces learnable parameters: a per-neuron scale factor $\gamma$ (gamma) and a shift factor $\beta$ (beta). These parameters allow the network to learn the optimal scale and mean of the normalized activations, potentially even recovering the original activations if needed:

$$\text{LN}(a_i) = \gamma_i \hat{a}_i + \beta_i$$

These parameters $\gamma$ and $\beta$ are learned during training along with the network's other weights.

Differences and Advantages

The primary distinction lies in the normalization axis:

• Batch Normalization: Normalizes across the batch dimension, for each feature/channel independently. Statistics depend on the mini-batch.
• Layer Normalization: Normalizes across the feature/hidden unit dimension, for each example independently. Statistics are independent of the mini-batch.

This difference leads to several important properties of Layer Normalization:

1. Batch Size Independence: LN's computations are identical during training and inference and do not depend on the mini-batch size. This makes it suitable for scenarios where small batches are necessary (e.g., due to memory constraints) or where the batch size might vary.
2. Effectiveness in RNNs: LN is often more straightforward to apply and more effective in recurrent networks than BN. Since the normalization is done per time step per example across the hidden units, it naturally handles variable sequence lengths without requiring complex modifications.
3. No Need for Running Statistics: Unlike BN, LN doesn't need to maintain running averages of mean and variance for use during inference, simplifying its implementation.

The following diagram illustrates the different normalization dimensions:

Comparison of normalization dimensions for Batch Normalization (blue, column-wise across batch) and Layer Normalization (green, row-wise across features).

While BN is often the default choice for Convolutional Neural Networks (CNNs), LN has proven particularly useful in NLP tasks involving transformers and RNNs, and in situations where batch statistics might be unreliable. It serves as another valuable tool for stabilizing and potentially accelerating the training of deep neural networks. We won't explore the implementation details here as deeply as we did for Batch Normalization, but most deep learning frameworks provide simple ways to add Layer Normalization layers.