Layer Normalization

Topics: Neural Networks

Motivation

In traditional neural network training, the activations of each layer can vary drastically, which is known as internal covariance shift. This leads to issues like exploding or vanishing gradients. Similarly to Batch Normalization, Layer Normalization addresses these issues, as well as some of the shortcomings of BN.
The working principle is similar to BN (compute metrics, normalize, and scale); however, in LN, the mean and variance are computed across the features for every data vector in the batch instead of computing them across the batch for every feature.

Details of Layer Normalization

Step 1: Compute Mean and Variance for Each Feature

The mean and variance are computed across features (or neurons) for every element of the batch:
$\mu = \frac{1}{H} \sum_{i=1}^H x_i$
$\sigma^2 = \frac{1}{H} \sum_{i=1}^H (x_i - \mu)^2$
where $H$ is the number of features.

For example, if the first layer has $256$ outputs, and a batch size of $32$ is used, then, $32$ means and variances are computed across the $256$ features.

Step 2: Normalize the Input

Each feature is normalized by the calculated layer mean and variance to obtain features with mean $0$ and unit variance:
$\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}$
where $\epsilon$ is a small constant added for numerical stability.

Step 3: Apply Scaling and Shifting

Similarly to BN, scaling and shifting is applied using learnable parameters $\gamma$ and $\beta$ to ensure that the normalized activations can still cover a wide range of values:
$y_i = \gamma \hat{x}_i + \beta$
In practice, $\gamma$ and $\beta$ are not single scalars, but vectors of parameters, where each feature has its own parameters: $\gamma \in \mathbb{R}^H, \; \beta \in \mathbb{R}^H$.

PyTorch Example

References

1. https://www.geeksforgeeks.org/deep-learning/what-is-layer-normalization/