LayerNorm

BatchNorm normalises a feature across the examples in a batch. LayerNorm turns the tensor ninety degrees and normalises an example across its features. Everything else about the two is the same; that one choice is the entire difference, and it is the reason one of them is in every transformer and the other is not.

Because $\mu$ and $\sigma$ depend only on the vector in front of it, the operation is identical at training and inference, identical at batch size 1 and 1024, and identical for the fifth token of a sequence and the five-thousandth. There is no running average to maintain and nothing to synchronise across devices.

What the two terms do

The re-scaling is doing the work. The re-centring is mostly not, and this was established well before RMSNorm made it official — Xu et al. found that LayerNorm’s benefit comes from how it shapes gradients rather than from the forward statistics, and specifically from the derivative of the $1/\sigma$ term, which projects out the component of the gradient along $x$.

Two components are removed: the mean of the gradient, and its projection onto the normalised activation. The second is the important one. It means a gradient step cannot change the magnitude of $x$ — only its direction — which is why normalised networks are insensitive to the scale of their weights and why the learning rate does not need retuning as activations grow with depth.

import torch
from torch import Tensor, nn


class LayerNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5, bias: bool = True):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.bias = nn.Parameter(torch.zeros(dim)) if bias else None
        self.eps = eps

    def forward(self, x: Tensor) -> Tensor:
        dtype = x.dtype
        x = x.float()                                  # statistics in fp32, always
        mu = x.mean(-1, keepdim=True)
        var = x.var(-1, keepdim=True, unbiased=False)  # biased: divide by d, not d-1
        x = (x - mu) * torch.rsqrt(var + self.eps)

        x = x.to(dtype) * self.weight
        return x if self.bias is None else x + self.bias

unbiased=False is not a detail. Bessel’s correction divides by $d-1$, every reference implementation divides by $d$, and using the corrected estimator produces weights that are subtly incompatible with every other stack — a discrepancy of $d/(d-1)$ in the norm, small enough to train through and large enough to shift outputs.

Two simplifications are now standard, and both are subtractive.

The bias $b$ goes first. It is redundant with the bias of the linear layer that follows, and removing it costs nothing measurable while removing a parameter tensor per norm — most models since GPT-J ship without it.

Then the mean subtraction goes, which is RMSNorm: one reduction pass instead of two, one term fewer in the backward pass, and no measured cost. That LayerNorm survived a decade of being progressively hollowed out without anyone finding a task where the removed parts mattered is the most interesting fact about it.