Layer Normalization (2016) › The Solution

Layer normalization addresses these challenges by computing normalization statistics across the features of each individual example, rather than across the batch dimension. This fundamental shift eliminates all batch dependencies while preserving the core benefits of normalizing activations. The technique is applied to the inputs of each layer, transforming them to have zero mean and unit variance across the feature dimensions.

The mathematical formulation is straightforward. For a layer with input vector $h$ of dimension $d$, layer normalization first computes the mean and variance across all features:

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

These statistics are computed independently for each example in the batch, using only that example's features. The input is then normalized and scaled:

$$\hat{h} = \frac{h - \mu}{\sqrt{\sigma^2 + \epsilon}}$$

where $\epsilon$ is a small constant added for numerical stability. Finally, learnable parameters $\gamma$ and $\beta$ are applied to allow the network to learn the optimal scale and shift:

$$y = \gamma \odot \hat{h} + \beta$$

where $\odot$ represents element-wise multiplication. The parameters $\gamma$ and $\beta$ have the same dimension as $h$ and are learned during training, allowing the network to recover the original activations if normalization proves unhelpful for a particular layer or to adjust the scale to optimal values.

The key insight is that layer normalization computes statistics across what the paper calls the "layer," meaning all the features or dimensions of the input for a single example. For a fully connected layer, this means normalizing across all $d$ hidden units. For a convolutional layer, it means normalizing across all feature maps and spatial locations for a single example. For a recurrent layer, it means normalizing across all hidden units at each time step, independently for each example.

This approach provides several immediate advantages over batch normalization. First, it requires no batch dimension, making it applicable in settings where batch size is one or varies during training. Second, it's naturally suited to variable-length sequences because each sequence is normalized independently. Third, the statistics are deterministic for a given input, eliminating the randomness that comes from batch composition in batch normalization. This determinism makes layer normalization more predictable and easier to reason about.

The normalization operates identically during training and inference, without the need for separate modes or running averages. This eliminates a source of potential bugs and distributional shifts between training and deployment. The technique also integrates smoothly with gradient-based optimization. This provided gradient flow benefits similar to batch normalization while avoiding the complications that batch statistics introduce in recurrent settings.