Inside LayerNorm

Each pair of batch and token indices selects its own row of d features. With LayerNorm(d), the mean and variance reduce only that last axis. Variance divides by d, which is correction=0 in PyTorch. Keeping the reduced axis gives one statistic per row with shape (B, S, 1).

μ  = mean(X, dim=-1, keepdim=True)
σ² = mean((X − μ)², dim=-1, keepdim=True)
r  = 1 / √(σ² + ε)
X̂  = (X − μ) × r
Y  = X̂ × γ + β

The learned scale γ and shift β each have shape (d,) and are shared across batch items and token positions. Statistics broadcast across features; parameters broadcast across batch items and tokens. Both input and output keep the same (B, S, d) addresses.

This follows the PyTorch LayerNorm definition and the original GPT-2 normalization code.

The default symbolic view uses a 2 × 4 × 8 tensor so every entry can be seen. Model presets show a representative window of a larger tensor; shape labels and the cost ledger refer to the full selected dimensions. Values are synthetic examples, not activations or weights taken from a trained model; rows deliberately sit at different levels with different spreads so that normalization visibly brings them onto one scale.

The memory ledger models one forward pass. Tensor storage follows the selected precision, with FP32 statistics. The numerical examples use JavaScript numbers; changing storage precision does not simulate rounding. The teaching peak retains input, centered values, squares, normalized values, and output for inspection, plus two statistics per row (the mean, and the variance, which is converted into r in place) and both parameters when affine is enabled. With N = B·S·d, R = B·S, and w bytes per value, affine enabled uses 5Nw + 8R + 2dw bytes. Disabling affine removes both parameters and the separate output buffer, giving 4Nw + 8R bytes.

Broadcasting reuses a value along a compatible axis; the repeated cells in the animation do not imply an allocated copy. Each addition, subtraction, multiplication, and division counts as one basic FLOP. A complete forward pass with affine costs 7N + R basic FLOPs plus R reciprocal square roots, counted separately. The R extra FLOPs are the ε additions; they and the reciprocal square roots are charged to the variance broadcast step, where r is computed once per row and then multiplied in. Disabling affine gives 5N + R basic FLOPs, with the same R reciprocal square roots.

A fused implementation can keep intermediates inside a kernel instead of allocating separate tensors. These estimates exclude gradients, optimizer state, allocator overhead, and kernel workspace. Actual kernel instructions, memory traffic, and runtime depend on the implementation and hardware.