PyTorch LayerNorm: shapes, epsilon, and affine parameters

PyTorch nn.LayerNorm normalizes the trailing dimensions specified by normalized_shape, then optionally applies learned scale and shift parameters. For an input shaped (batch, sequence, features), LayerNorm(features) normalizes each feature vector independently. It preserves the input shape.

The important choice is which values belong to one normalization group. Two LayerNorm configurations can return the same shape while computing different results. This guide uses fixed CPU float64 tensors to make that difference visible, then checks variance, epsilon, affine parameters, evaluation mode, and image-layout errors.

All six examples run independently with Python 3.12.14 and PyTorch 2.15.0+cpu. Our 44 checks include independent scalar arithmetic and a finite-difference gradient comparison. They do not measure GPU speed or model quality. The LayerNorm API reference documents the arguments used below.

Choose normalized_shape from the trailing axes

The first tensor has shape (2, 2, 3). Treat its axes as batch, sequence position, and feature. An integer normalized_shape of 3 means the final dimension must have length three. A tuple (2, 3) selects the final two dimensions together.

import torch

x = torch.tensor([
    [[1., 2., 4.], [10., 20., 40.]],
    [[3., 5., 9.], [2., 3., 8.]]
], dtype=torch.float64)
last = torch.nn.LayerNorm(3, elementwise_affine=False)(x)
last_two = torch.nn.LayerNorm((2, 3), elementwise_affine=False)(x)
print('Input and output shapes:', tuple(x.shape), tuple(last.shape), tuple(last_two.shape))
print('Last-axis means:', last.mean(dim=-1).round(decimals=6).tolist())
print('Two-axis means:', last_two.mean(dim=(-2, -1)).round(decimals=6).tolist())
print('Same values:', torch.allclose(last, last_two))
print('First row, last axis:', [round(v, 4) for v in last[0, 0].tolist()])
print('First row, two axes:', [round(v, 4) for v in last_two[0, 0].tolist()])

Input and output shapes: (2, 2, 3) (2, 2, 3) (2, 2, 3)
Last-axis means: [[-0.0, -0.0], [-0.0, 0.0]]
Two-axis means: [-0.0, 0.0]
Same values: False
First row, last axis: [-1.069, -0.2673, 1.3363]
First row, two axes: [-0.8612, -0.7884, -0.6429]

LayerNorm(3) computes four separate three-value normalizations, one for each batch-position pair. LayerNorm((2, 3)) computes two six-value normalizations, one for each batch member. Neither combines the two batch members in this example.

The first row becomes approximately [-1.069, -0.2673, 1.3363] with the last-axis choice, but [-0.8612, -0.7884, -0.6429] when it shares statistics with the larger values in the second row. The shapes agree, so a shape assertion alone would miss the change in meaning.

The printed -0.0 values are rounded floating-point results close to zero. Our checks compare them with a tolerance. A separate test changes the second batch member and confirms that the first member's two-axis result stays unchanged.

For a (B, L, D) sequence tensor, LayerNorm(D) leaves B and L free to vary. LayerNorm((L, D)) requires those two trailing sizes to match its configuration. Decide whether positions should share statistics before choosing the tuple. Our transformer architecture guide provides the wider model context; the examples here focus on the normalization operation itself.

Reproduce the arithmetic with the correct variance

Within each group, subtract the mean and divide by the square root of variance plus eps. Before the optional affine step, the expression is (x - mean) / sqrt(variance + eps). The following calculation uses the mean squared deviation, dividing by the number of values rather than one fewer:

import torch

x = torch.tensor([[1., 2., 4.], [8., 8., 8.]], dtype=torch.float64)
eps = 1e-5
mean = x.mean(dim=-1, keepdim=True)
variance = ((x - mean) ** 2).mean(dim=-1, keepdim=True)
manual = (x - mean) / torch.sqrt(variance + eps)
layer = torch.nn.LayerNorm(3, eps=eps, elementwise_affine=False)
y = layer(x)
print('Means:', mean.flatten().tolist())
print('Variances:', variance.flatten().tolist())
print('Matches manual:', torch.allclose(y, manual, rtol=1e-12, atol=1e-12))
print('Constant row:', y[1].tolist())
print('Output variances:', y.var(dim=-1, correction=0).tolist())

Means: [2.3333333333333335, 8.0]
Variances: [1.5555555555555554, 0.0]
Matches manual: True
Constant row: [0.0, 0.0, 0.0]
Output variances: [0.9999935714698979, 0.0]

For [1, 2, 4], the mean is 7/3 and the variance is 14/9. Our independent Python calculation matches the module output. Using sample variance with correction=1 does not match; when checking with torch.var, specify correction=0 for this comparison.

The constant [8, 8, 8] group has zero variance and produces zeros with positive eps and no affine transformation. It does not produce unit variance. Even the nonconstant group's output variance is slightly below one: before scale and shift, it is v/(v + eps), where v is the input group's variance.

That distinction matters when writing assertions. Check the actual expression and numerical tolerance rather than requiring every output group to have variance exactly one. Learned scale and shift can change the output statistics further.

Understand the learned scale and shift

With elementwise_affine=True, LayerNorm stores a weight and bias with the configured normalized_shape. The default weight is one and the default bias is zero. This example replaces them with known values so their effect can be checked:

import torch

x = torch.tensor([[1., 2., 4.], [2., 4., 8.]], dtype=torch.float64)
layer = torch.nn.LayerNorm(3, dtype=torch.float64)
with torch.no_grad():
    layer.weight.copy_(torch.tensor([2., 0.5, -1.], dtype=torch.float64))
    layer.bias.copy_(torch.tensor([1., -2., 3.], dtype=torch.float64))
z = torch.nn.LayerNorm(3, elementwise_affine=False)(x)
y = layer(x)
expected = z * layer.weight + layer.bias
scale_only = torch.nn.LayerNorm(3, bias=False, dtype=torch.float64)
no_affine = torch.nn.LayerNorm(3, elementwise_affine=False)
print('Matches scale and shift:', torch.allclose(y, expected, rtol=1e-12, atol=1e-12))
print('Output row means:', [round(v, 4) for v in y.mean(dim=-1).tolist()])
print('Parameter counts:', *(sum(p.numel() for p in m.parameters())
                              for m in (layer, scale_only, no_affine)))
print('Scale-only bias:', scale_only.bias)

Matches scale and shift: True
Output row means: [-0.536, -0.536]
Parameter counts: 6 3 0
Scale-only bias: None

The output equals the normalized values multiplied by [2, 0.5, -1], then shifted by [1, -2, 3]. These three feature parameters are shared across the leading row dimension. The row means are now approximately -0.536, demonstrating that learned affine output need not have zero mean.

The parameter counts separate three configurations: the ordinary length-three layer has six scalars, bias=False retains three scale parameters, and elementwise_affine=False removes both scale and shift. For normalized_shape=(2, 3), our additional check found a (2, 3) weight and bias, totaling twelve scalars.

The functional layer_norm API accepts weight, bias, and eps explicitly. Our functional call matched this configured module when all those values agreed. Use the module when you want it to own and register the learned parameters as part of your model.

These examples specify float64 parameters where affine parameters are present, matching the input tensors. Keep a real model's inputs and parameters on compatible devices and dtypes; changing normalized_shape will not fix a device or dtype mismatch.

See why eps changes near-constant groups

Epsilon is added inside the square root. Its effect is easy to see when the input variance is small:

import torch

x = torch.tensor([[1., 1.001, 0.999], [7., 7., 7.]], dtype=torch.float64)
small = torch.nn.LayerNorm(3, eps=1e-5, elementwise_affine=False)(x)
large = torch.nn.LayerNorm(3, eps=0.1, elementwise_affine=False)(x)
print('Small eps:', [[round(v, 6) for v in row] for row in small.tolist()])
print('Large eps:', [[round(v, 6) for v in row] for row in large.tolist()])
print('Larger eps reduces magnitudes:', bool(large[0].abs().max() < small[0].abs().max()))
print('Both finite:', bool(torch.isfinite(small).all() and torch.isfinite(large).all()))

Small eps: [[0.0, 0.306186, -0.306186], [0.0, 0.0, 0.0]]
Large eps: [[0.0, 0.003162, -0.003162], [0.0, 0.0, 0.0]]
Larger eps reduces magnitudes: True
Both finite: True

The first group differs from its mean by only about 0.001. With eps=1e-5, its nonzero outputs have magnitude about 0.306186. Increasing eps to 0.1 reduces them to about 0.003162. Both results match our independent arithmetic. The constant group stays at zero and both outputs remain finite.

The large value is deliberately chosen to expose the effect, not recommended as a tuning default. Changing eps changes the computation. Record it with the model configuration, especially when comparing implementations or loading a checkpoint whose architecture is defined elsewhere.

Separate evaluation mode from gradient tracking

This layer uses the current input's group statistics in both training and evaluation. The following example compares the modes, then backpropagates through the evaluation output:

import torch

x = torch.tensor([[1., 2., 4.], [3., 7., 8.]], dtype=torch.float64, requires_grad=True)
layer = torch.nn.LayerNorm(3, dtype=torch.float64)
layer.train()
training = layer(x)
layer.eval()
evaluation = layer(x)
weights = torch.tensor([1., 2., 4.], dtype=torch.float64)
(evaluation * weights).sum().backward()
print('Train/eval equal:', torch.equal(training, evaluation))
print('Running-stat buffers:', list(layer.named_buffers()))
print('Eval output requires gradients:', evaluation.requires_grad)
print('Input gradient finite:', bool(torch.isfinite(x.grad).all()))
print('Input gradient nonzero:', bool(x.grad.abs().max() > 0))

Train/eval equal: True
Running-stat buffers: []
Eval output requires gradients: True
Input gradient finite: True
Input gradient nonzero: True

The isolated LayerNorm returns exactly equal train and eval outputs for this fixed input and parameter state. It has no running-statistic buffers in the tested configuration. Calling eval() does not detach the result: the input receives finite, nonzero gradients.

We weighted the output by [1, 2, 4] to avoid using an uninformative sum of centered values as the gradient example. The verification suite compares each input derivative with an independently calculated finite difference, and checks scale and bias gradients separately.

Other layers in a larger model may still behave differently between train and eval. Choose gradient recording separately according to your task; our PyTorch eval, no_grad, and inference_mode guide covers those controls. Equal LayerNorm outputs do not mean eval() is unnecessary for the whole model.

Avoid confusing image channels with the final axis

For a tensor in (N, C, H, W) order, LayerNorm(C) still checks the final axis, which is W. It does not interpret C as a channel name. Our input has two channels and a width of three, so that call fails:

import torch

x = torch.arange(12, dtype=torch.float64).reshape(1, 2, 2, 3)
try:
    torch.nn.LayerNorm(2, elementwise_affine=False)(x)
except RuntimeError as error:
    print('NCHW channel-count attempt:', type(error).__name__)
channels_last = x.permute(0, 2, 3, 1)
normalized = torch.nn.LayerNorm(2, elementwise_affine=False)(channels_last)
back = normalized.permute(0, 3, 1, 2)
print('Restored shape:', tuple(back.shape))
print('Channel means:', back.mean(dim=1).round(decimals=6).tolist())
print('First pixel:', [round(v, 6) for v in back[0, :, 0, 0].tolist()])

NCHW channel-count attempt: RuntimeError
Restored shape: (1, 2, 2, 3)
Channel means: [[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]
First pixel: [-0.999999, 0.999999]

The permutation changes the logical order to (N, H, W, C). LayerNorm(2) then normalizes the pair of channel values at each pixel. Permuting back restores (1, 2, 2, 3). At the first pixel, values 0 and 6 have mean 3 and variance 9, producing approximately [-0.999999, 0.999999].

Our additional LayerNorm(3) call on the original NCHW tensor succeeds but normalizes width, producing different values. If channel count and width happen to be equal, the incorrect call may also pass its shape check. Label the axes rather than relying on an exception to expose every mistake.

The permute reference describes dimension reordering. A reshape is not a substitute for deliberately moving the channel axis. Also distinguish channel-only normalization at each pixel from normalizing all channel and spatial values together with a shape such as (C, H, W); they define different groups.

Check a small input before using the layer in a model

• Write down the input axis meanings and the exact trailing shape that should share statistics.
• Compare one group against a manual calculation with correction=0 and eps inside the square root.
• Inspect parameter shapes and distinguish bias=False from elementwise_affine=False.
• Test the full model's evaluation and gradient behavior separately from this layer's arithmetic.

The examples and verification code are available in the review package. Our checks use dense CPU float64 tensors; CUDA, mixed precision, training improvements, and runtime performance require their own evaluation. LayerNorm is not a universal replacement for every other normalization method.