Layer Normalization › The Layer Normalization Formula

To understand layer normalization, let's start with a fundamental question: what does it mean for a neural network layer to have "unstable" activations, and how can we fix it?

Imagine a token's hidden representation as a vector of 768 numbers (a typical transformer dimension). During training, these numbers can drift: some become very large, others very small, and their collective distribution shifts unpredictably. This creates two problems. First, downstream layers must constantly adapt to changing input statistics, making learning inefficient. Second, when values grow too large or too small, gradients either explode or vanish, destabilizing training entirely.

The solution is elegant: before each token's representation moves to the next layer, we transform it to have a predictable, standardized distribution. Specifically, we want the 768 features to have zero mean and unit variance. This "resets" the scale at every layer, preventing drift from accumulating.

The key insight is that we don't need to look outside the current token to achieve this. The token's own 768 values carry all the information we need to compute a mean and a standard deviation. Subtract the mean, divide by the standard deviation, and you have a normalized representation that is stable regardless of how the upstream weights evolved during training.

In practice, this normalization is applied at multiple points within each transformer block. This ensures that neither the attention mechanism nor the feed-forward network ever sees wildly out-of-distribution inputs. The cumulative effect over many layers is dramatic: networks with layer normalization can train stably with learning rates that would cause catastrophic divergence in networks without it.

The first step is computing where the current distribution is centered. Given a hidden state vector $\mathbf{x} = [x_1, x_2, \ldots, x_d]$ representing one token, we calculate its mean:

$$\mu = \frac{1}{d} \sum_{i=1}^{d} x_i$$

where:

• $\mu$: the arithmetic mean of all $d$ features in this token's representation
• $d$: the hidden dimension (e.g., 768 for BERT-base, 4096 for LLaMA-7B)
• $x_i$: the value of the $i$-th feature

This tells us the "center of mass" of the representation. If $\mu = 2.5$, the features are shifted toward positive values; if $\mu = -1.3$, they lean negative. The goal is to shift this center to zero.

Next, we need to know how spread out the values are. A representation where all values cluster tightly around the mean is very different from one where values are scattered widely. We capture this with variance:

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

where:

• $\sigma^2$: the variance, measuring how much the features deviate from their mean
• $(x_i - \mu)^2$: the squared deviation of each feature from the mean

Squaring ensures that positive and negative deviations don't cancel out. If $\sigma^2$ is large, the features are spread out; if small, they're tightly clustered. We'll use the standard deviation $\sigma = \sqrt{\sigma^2}$ to rescale the values to unit variance.

Notice that we use the population variance (dividing by $d$) rather than the sample variance (dividing by $d - 1$). This is a deliberate choice: we're not trying to estimate the variance of an underlying population; we're computing the exact variance of the $d$ values in front of us. Using $d$ rather than $d - 1$ gives us an unbiased estimator for the squared deviation within this specific vector, which is what we need for a deterministic normalization operation.

With mean and variance in hand, we can now standardize each feature:

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

where:

• $\hat{x}_i$: the normalized value of the $i$-th feature
• $\epsilon$: a tiny constant (typically $10^{-5}$) added to prevent division by zero if variance is extremely small

This two-part transformation is exactly what you'd do to standardize any dataset: subtract the mean (centering at zero), then divide by the standard deviation (scaling to unit variance). The result $\hat{x}_i$ has zero mean and approximately unit variance across the $d$ features.

The subtraction by $\mu$ ensures the output is centered at zero, removing any bias in the representation. The division by $\sqrt{\sigma^2 + \epsilon}$ ensures that features with large variance are scaled down to unit range, and features with small variance are scaled up. This compression of the dynamic range is what prevents activations from exploding across layers.

Layer normalization need not be restrictive. Forcing every representation to have exactly zero mean and unit variance might seem limiting: what if the optimal representation for some layer needs a different distribution?

The solution is to add learnable parameters that can undo the normalization if needed:

$$y_i = \gamma_i \cdot \hat{x}_i + \beta_i$$

where:

• $y_i$: the final output for the $i$-th feature
• $\gamma_i$: a learned scale parameter for feature $i$ (initialized to 1)
• $\beta_i$: a learned shift parameter for feature $i$ (initialized to 0)

These parameters are learned during training, just like weights and biases. If the network discovers that feature $i$ should have mean 3.5 and standard deviation 2.0, it can learn $\gamma_i = 2.0$ and $\beta_i = 3.5$ to recover that distribution. This means layer normalization never reduces the network's representational power: it starts from a stable baseline but can learn any distribution it needs.

The initialization of $\gamma = 1$ and $\beta = 0$ is deliberate. At the start of training, layer normalization acts as a pure standardization with no modification: the output is exactly the normalized input. As training progresses, $\gamma$ and $\beta$ adjust to whatever distribution the subsequent layers find most useful. This means the network starts with the stability benefits of normalization and gradually adapts to the specific distribution requirements of the task, without ever losing the ability to represent the original data.

Putting all the pieces together, layer normalization maps a hidden state vector $\mathbf{x}$ of dimension $d$ into:

$$\text{LayerNorm}(\mathbf{x}) = \gamma \odot \frac{\mathbf{x} - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$$

where:

• $\mathbf{x} = [x_1, x_2, \ldots, x_d]$: the input vector representing one token's hidden state
• $\mu = \frac{1}{d} \sum_{i=1}^{d} x_i$: the mean across all features
• $\sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i - \mu)^2$: the variance across all features
• $\epsilon$: a small stability constant (typically $10^{-5}$ or $10^{-6}$)
• $\gamma = [\gamma_1, \gamma_2, \ldots, \gamma_d]$: learned scale parameters, initialized to ones
• $\beta = [\beta_1, \beta_2, \ldots, \beta_d]$: learned shift parameters, initialized to zeros
• $\odot$: element-wise multiplication

The formula reads naturally: subtract the mean, divide by the standard deviation (with a safety epsilon), then apply a learned scale and shift.

The key insight that distinguishes layer normalization from batch normalization is the dimension over which we compute statistics. Batch normalization asks: "What's the typical value of feature $i$ across all samples in this batch?" Layer normalization asks: "What's the typical value across all features for this particular token?"

For transformers, the layer normalization approach is far more natural. Each token is processed independently, and we want stable statistics regardless of what other tokens or samples happen to be in the batch. This independence also means layer normalization works identically during training and inference, with no need for running statistics or batch size considerations.

To make this concrete, consider a transformer with hidden dimension 512. Batch normalization would compute 512 separate statistics (one per feature dimension), each estimated from however many tokens appear at the same position across the batch. Layer normalization computes 2 statistics (one mean and one variance) from the 512 features of a single token. The batch normalization statistics are estimates that improve with larger batches; the layer normalization statistics are exact computations that require no batch at all. During inference when you process one sequence at a time, layer normalization requires no special handling; batch normalization in inference mode must use stored running statistics from training, which may not accurately represent the test distribution.