Batch and Layer Normalization

Why Neural Networks Need Normalization

Training a deep network is a negotiation between layers that keep changing their own inputs. Layer 10 is trying to learn a function of the activations produced by layer 9 — but layer 9 is simultaneously updating its weights, which shifts the distribution layer 10 sees at every training step. In a 24-layer Transformer or a 100-layer ResNet the effect compounds: by the time gradients flow back from the loss, each layer is chasing a moving target.

Without intervention this leads to three very concrete problems that any practitioner has encountered:

1. Saturating activations. If pre-activations drift to large positive or negative values, sigmoid/tanh flatten out and gradients become numerically zero ($\frac{\partial \sigma}{\partial z} = \sigma(z)(1-\sigma(z)) \approx 0$). The network stops learning.
2. Exploding or vanishing gradients. The Jacobian of a deep network is a product of per-layer Jacobians. If their magnitudes drift above 1, the product explodes; below 1, it vanishes exponentially with depth.
3. Step-size fragility. The optimal learning rate depends on the scale of activations. If that scale drifts during training, a learning rate that worked at epoch 1 is catastrophically wrong at epoch 10.

Normalization layers fix all three by imposing a contract on the activations that flow between layers: no matter what weights come before, the activation entering the next layer will have a predictable mean and variance. That contract is the single idea behind BatchNorm, LayerNorm, InstanceNorm, GroupNorm, and RMSNorm — they differ only in which axis they normalize over.

Mental model. Think of normalization as inserting a variance-control valve between layers. It standardizes the signal, then offers the next layer a pair of knobs ($\gamma, \beta$) to pick whatever distribution it actually wants. The optimizer becomes a thermostat — gradients correct the knobs, not the (already well-conditioned) pre-activations.

Internal Covariate Shift: The Original Motivation

Ioffe and Szegedy coined the term internal covariate shift (ICS) in 2015 to name the phenomenon: as weights update, the distribution of inputs to each internal layer changes. Their paper showed that simply forcing each layer's inputs to be zero-mean, unit-variance dramatically stabilized training and let them use 10× larger learning rates for ImageNet-scale models.

The mathematical reading is subtler. Let $h^{(\ell)}$ be the activations at layer $\ell$ and $W^{(\ell)}$ the weights above it. The gradient update to $W^{(\ell)}$ depends on both $h^{(\ell-1)}$ (its input) and the local Jacobian. If $h^{(\ell-1)}$ has a time-varying distribution, SGD is solving a non-stationary optimization problem — each step is optimizing a slightly different loss surface than the previous step.

Santurkar et al. (2018) later showed that BN's benefits extend beyond reducing ICS: it smooths the loss landscape (smaller Lipschitz constants of both loss and gradients), which is the real reason larger step sizes stay stable. Either framing leads to the same practical recipe — standardize before the next linear layer.

The demo below simulates the statistics of four stacked layers over 50 training epochs, with and without BatchNorm. Toggle BN on and watch the layer means and variances stay tightly around $(0, 1)$ ; toggle it off and watch the distributions drift apart as depth increases.

The Core Idea: Standardize, Then Re-Learn the Distribution

Every normalization layer, without exception, follows the same four-step recipe. The differences between BN, LN, IN, GN and RMSNorm are purely about the set of elements over which the statistics are computed.

1. Choose a reduction axis set $\mathcal{A}$ — the elements that share a (μ, σ) pair.
2. Compute the mean: $\mu = \frac{1}{|\mathcal{A}|} \sum_{i \in \mathcal{A}} x_i$
3. Compute the variance: $\sigma^{2} = \frac{1}{|\mathcal{A}|} \sum_{i \in \mathcal{A}} (x_i - \mu)^{2}$
4. Standardize and re-scale: $\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^{2} + \varepsilon}}$, then $y_i = \gamma_i \hat{x}_i + \beta_i$.

The $\varepsilon$ guards against division by zero for dead features. The learnable affine pair $(\gamma, \beta)$ is the critical escape hatch: it guarantees the network can recover the original distribution if needed ($\gamma = \sigma, \beta = \mu$ makes the layer the identity), so normalization never reduces expressivity — it only conditions the optimization problem.

The only remaining question is which elements share statistics. Answering it differently gives each normalization variant.

Compactly, every variant in the family computes $\text{Norm}(x) = \gamma \cdot \frac{x - \mu_{\mathcal{A}}}{\sqrt{\sigma^{2}_{\mathcal{A}} + \varepsilon}} + \beta$ — only the choice of $\mathcal{A}$ changes.

Batch Normalization: Statistics Along the Batch Axis

BatchNorm was the breakthrough. For a convolutional activation tensor $X \in \mathbb{R}^{N \times C \times H \times W}$, it computes one mean and one variance per channel, reducing across the batch, height and width axes: $\mu_c = \frac{1}{N H W}\sum_{n,h,w} x_{n,c,h,w}$ and $\sigma_c^{2} = \frac{1}{N H W}\sum_{n,h,w} (x_{n,c,h,w} - \mu_c)^{2}$.

Intuitively, every channel is a feature detector; BN normalizes each detector's response across all the locations and samples it sees in the batch, so a single pixel in a single image never dominates the channel's scale.

Why BN needs running statistics

During training, $\mu_c, \sigma_c$ are computed from the current batch. At inference time you often process a single example — there's no batch to average over. BN solves this by maintaining an exponential moving average (EMA) of the mini-batch statistics during training, then using the frozen EMA means/variances at inference:

The running mean is updated as $\mu^{\text{run}}_c \leftarrow (1 - m)\,\mu^{\text{run}}_c + m\,\mu^{\text{batch}}_c$ with momentum $m \approx 0.1$. This train/eval discrepancy is the classic source of BatchNorm footguns: fine-tuning with batch size 1 breaks the EMA updates, and shuffling data differently between train and eval produces subtly different stats.

BatchNorm From Scratch: Code Trace

The math above compresses into one page of NumPy. Below is a class-based implementation that mirrors$\texttt{nn.BatchNorm2d}$ exactly: it holds per-channel learnable parameters$\gamma, \beta$, maintains EMA buffers for inference, and exposes a training-mode forward pass. Every line is annotated with actual numerical values for a small (4, 3, 2, 2) input.

1import numpy as np
2
3class BatchNorm2dScratch:
4    """Training-mode forward pass that matches nn.BatchNorm2d exactly."""
5
6    def __init__(self, C, eps=1e-5, momentum=0.1):
7        self.gamma = np.ones((C,), dtype=np.float32)
8        self.beta  = np.zeros((C,), dtype=np.float32)
9        self.running_mean = np.zeros((C,), dtype=np.float32)
10        self.running_var  = np.ones((C,),  dtype=np.float32)
11        self.eps, self.m = eps, momentum
12
13    def forward_train(self, x):
14        # x shape: (N, C, H, W) — image-style conv activations
15        mu  = x.mean(axis=(0, 2, 3))                    # (C,)
16        var = x.var(axis=(0, 2, 3))                     # (C,)
17        # EMA update — PyTorch uses m on the batch side, (1-m) on running
18        self.running_mean = (1 - self.m) * self.running_mean + self.m * mu
19        self.running_var  = (1 - self.m) * self.running_var  + self.m * var
20        # Normalize with the CURRENT batch stats (not running)
21        x_hat = (x - mu[None, :, None, None]) / np.sqrt(var[None, :, None, None] + self.eps)
22        return self.gamma[None, :, None, None] * x_hat + self.beta[None, :, None, None]
23
24bn = BatchNorm2dScratch(C=3)
25x  = np.random.randn(4, 3, 2, 2).astype(np.float32)
26y  = bn.forward_train(x)
27print("y.shape      =", y.shape)
28print("batch μ[c=0] =", round(float(x[:, 0].mean()), 4))
29print("running μ    =", bn.running_mean)

The same computation in PyTorch is a single line, but the built-in module hides the train/eval switch that powers BN. The next trace makes that switch explicit: we run the SAME input through the SAME layer twice, once in $\texttt{train}$ mode and once in $\texttt{eval}$ mode, and show the outputs disagree.

1import torch
2import torch.nn as nn
3
4bn = nn.BatchNorm2d(num_features=3, eps=1e-5, momentum=0.1)
5x  = torch.randn(4, 3, 2, 2)
6
7bn.train()                       # enable running-stats update
8y_train = bn(x)
9
10bn.eval()                        # freeze running stats, use them for forward
11y_eval  = bn(x)
12
13print("y_train shape   :", y_train.shape)
14print("running_mean    :", bn.running_mean)
15print("running_var     :", bn.running_var)
16print("train ≠ eval?   :", not torch.allclose(y_train, y_eval))

The divergence is the whole point. In training mode BN normalizes with the CURRENT batch's statistics and updates its running buffers; in eval mode it freezes those buffers and uses them instead, so the forward pass becomes a pure function of the input. Early in training the two modes disagree because the running stats lag behind the true dataset distribution — forgetting $\texttt{model.eval()}$before inference is the classic BatchNorm footgun, and this visualizer makes it unmissable.

Layer Normalization: Statistics Along the Feature Axis

Layer Normalization, introduced by Ba, Kiros and Hinton (2016), reduces in the opposite direction from BatchNorm. For a Transformer activation of shape $(N, T, D)$ — batch of $N$ sequences, each of length $T$, each token a $D$-dimensional vector — LayerNorm computes one mean and variance per token, reducing across the feature dimension: $\mu_{n,t} = \frac{1}{D}\sum_{d=1}^{D} x_{n,t,d}$ and $\sigma_{n,t}^{2} = \frac{1}{D}\sum_{d=1}^{D}(x_{n,t,d} - \mu_{n,t})^{2}$.

Critically, the statistics depend only on the token itself — not on the batch, not on other tokens. This has three consequences that make LN the correct choice for Transformers:

1. Batch-size independent. Train with batch 256 or infer with batch 1 — LayerNorm does the same math either way. There is no train/eval mismatch, no running statistics to maintain.
2. Length-independent. Each token is normalized by itself, so sequences of length 4 and 4096 are both handled correctly. BatchNorm, which would average across the sequence dimension, would mix information across positions — a leak attention is designed to avoid.
3. Causal-safe. In an autoregressive decoder, token $t$ must not see information from tokens {>}t. LN's stats come only from position $t$'s own features, so no future leakage.

Geometric reading. LayerNorm projects each sample vector onto the unit sphere (after centering) in $\mathbb{R}^{D}$, then scales it to radius $\sqrt{D}$. Direction is preserved; only magnitude is rewritten. For attention this matters: dot-product scores $Q K^{\top}$ are angle-like quantities, and LayerNorm makes sure each token's query/key contributes the same magnitude regardless of what happened upstream.

BatchNorm vs LayerNorm: The Dimensional Story

The fastest way to internalize the difference is to see which cells of a (N, C, H, W) tensor share their mean and variance. BN colors a whole $(N, H, W)$ slab of one channel; LN colors a whole$(C, H, W)$ slab of one sample; InstanceNorm colors one $(H, W)$ feature map; GroupNorm interpolates between LN and IN.

InstanceNorm and GroupNorm: The Interpolation

Between BatchNorm (statistics pooled over the entire batch per channel) and LayerNorm (statistics over the full feature vector per sample) lie two useful interpolations. InstanceNorm, introduced by Ulyanov, Vedaldi and Lempitsky (2016) in Instance Normalization: The Missing Ingredient for Fast Stylization, computes one $(\mu, \sigma^{2})$ per $(n, c)$ — each feature map is normalized in isolation, independently of the batch and of other channels. It became the default for style-transfer networks and GAN generators because it strips content-specific statistics out of each feature map, letting learned $\gamma, \beta$ reinject the desired style.

GroupNorm, proposed by Wu and He (ECCV 2018, Group Normalization), sits in between. It partitions the $C$ channels into $G$ groups and computes statistics over $(C/G, H, W)$ per sample. The per-group mean is$\mu_{n,g} = \frac{1}{(C/G) H W}\sum_{c \in g, h, w} x_{n,c,h,w}$, with an analogous variance. Two limits recover the neighbors: $G = 1$ reduces to LayerNorm (all channels share stats) and $G = C$ reduces to InstanceNorm (each channel alone). For vision tasks with small batch sizes — detection, segmentation, 3D medical imaging — $G = 32$ is the empirical sweet spot and GroupNorm routinely outperforms BatchNorm.

[[1.7641, 0.4002], [0.9787, 2.2409]]  — channel 0 of sample 0

1import numpy as np
2
3def group_norm(x, gamma, beta, G, eps=1e-5):
4    """x shape (N, C, H, W). Normalize within each (sample, group)."""
5    N, C, H, W = x.shape
6    assert C % G == 0
7    x_g = x.reshape(N, G, C // G, H, W)
8    mu  = x_g.mean(axis=(2, 3, 4), keepdims=True)
9    var = x_g.var(axis=(2, 3, 4),  keepdims=True)
10    x_hat = (x_g - mu) / np.sqrt(var + eps)
11    x_hat = x_hat.reshape(N, C, H, W)
12    return gamma[None, :, None, None] * x_hat + beta[None, :, None, None]
13
14np.random.seed(0)
15x = np.random.randn(2, 6, 2, 2).astype(np.float32)
16gamma = np.ones((6,), dtype=np.float32)
17beta  = np.zeros((6,), dtype=np.float32)
18y_ln = group_norm(x, gamma, beta, G=1)   # equivalent to LayerNorm
19y_gn = group_norm(x, gamma, beta, G=2)   # 2 groups of 3 channels
20y_in = group_norm(x, gamma, beta, G=6)   # equivalent to InstanceNorm
21print("Per-row mean ≈ 0?", y_gn.mean(axis=(1,2,3)).round(4))

LayerNorm From Scratch in Python

The entire LayerNorm forward pass fits in five lines of NumPy. We'll step through it with concrete numbers so there's no black box. Our input is a $3 \times 4$ matrix (three samples, four features). Click any line on the right to see what that line computes and why.

        f0    f1    f2    f3
samp0  2.0   4.0   6.0   8.0
samp1  1.0   3.0   2.0   6.0
samp2  5.0   7.0   3.0   9.0

[[5.0000],
 [3.0000],
 [6.0000]]

[[5.0000],
 [3.5000],
 [5.0000]]

        f0     f1     f2     f3
samp0 -3.0   -1.0   +1.0   +3.0
samp1 -2.0    0.0   -1.0   +3.0
samp2 -1.0   +1.0   -3.0   +3.0

[[5.00001],
 [3.50001],
 [5.00001]]   (eps=1e-5 is added for numerical safety)

[[2.2361],
 [1.8708],
 [2.2361]]

         f0       f1       f2       f3
samp0  -1.3416  -0.4472   0.4472   1.3416
samp1  -1.0690   0.0000  -0.5345   1.6036
samp2  -0.4472   0.4472  -1.3416   1.3416

         f0       f1       f2       f3
samp0  -2.6833  -0.4472   0.2236   1.3416
samp1  -2.1381   0.0000  -0.2673   1.6036
samp2  -0.8944   0.4472  -0.6708   1.3416
(column 0 doubled, column 2 halved, rest unchanged)

         f0       f1       f2       f3
samp0  -2.6833  -0.4472   0.2236   1.8416
samp1  -2.1381   0.0000  -0.2673   2.1036
samp2  -0.8944   0.4472  -0.6708   1.8416
(beta=0.5 added only to the last column)

        f0    f1    f2    f3
samp0  2.0   4.0   6.0   8.0
samp1  1.0   3.0   2.0   6.0
samp2  5.0   7.0   3.0   9.0

         f0       f1       f2       f3
samp0  -2.6833  -0.4472   0.2236   1.8416
samp1  -2.1381   0.0000  -0.2673   2.1036
samp2  -0.8944   0.4472  -0.6708   1.8416

[[-2.6833 -0.4472  0.2236  1.8416]
 [-2.1381  0.     -0.2673  2.1036]
 [-0.8944  0.4472 -0.6708  1.8416]]

1import numpy as np
2
3def layer_norm(x, gamma, beta, eps=1e-5):
4    """Per-sample normalization across the feature axis."""
5    mu    = x.mean(axis=-1, keepdims=True)
6    var   = x.var(axis=-1, keepdims=True)
7    x_hat = (x - mu) / np.sqrt(var + eps)
8    return gamma * x_hat + beta
9
10x     = np.array([[2., 4., 6., 8.],
11                  [1., 3., 2., 6.],
12                  [5., 7., 3., 9.]], dtype=np.float32)
13gamma = np.array([2.0, 1.0, 0.5, 1.0], dtype=np.float32)
14beta  = np.array([0.0, 0.0, 0.0, 0.5], dtype=np.float32)
15
16y = layer_norm(x, gamma, beta)
17print(y)

LayerNorm in PyTorch

PyTorch ships $\texttt{nn.LayerNorm}$ with a fused CUDA kernel that computes mean, variance, standardization and affine in a single pass — roughly $3\times$ faster than the naive sequence of operations. Importantly, the math is byte-for-byte identical to our NumPy implementation, which we verify below.

        f0    f1    f2    f3
samp0  2.0   4.0   6.0   8.0
samp1  1.0   3.0   2.0   6.0
samp2  5.0   7.0   3.0   9.0

[[5.0],
 [3.0],
 [6.0]]  — same as NumPy.

[[5.00],
 [3.50],
 [5.00]]

         f0       f1       f2       f3
samp0  -1.3416  -0.4472   0.4472   1.3416
samp1  -1.0690   0.0000  -0.5345   1.6036
samp2  -0.4472   0.4472  -1.3416   1.3416

         f0       f1       f2       f3
samp0  -2.6833  -0.4472   0.2236   1.8416
samp1  -2.1381   0.0000  -0.2673   2.1036
samp2  -0.8944   0.4472  -0.6708   1.8416

tensor([[-2.6833, -0.4472,  0.2236,  1.8416],
        [-2.1381,  0.0000, -0.2673,  2.1036],
        [-0.8944,  0.4472, -0.6708,  1.8416]],
       grad_fn=<NativeLayerNormBackward0>)

1import torch
2import torch.nn as nn
3
4x = torch.tensor([[2., 4., 6., 8.],
5                  [1., 3., 2., 6.],
6                  [5., 7., 3., 9.]])
7
8ln = nn.LayerNorm(normalized_shape=4, eps=1e-5)
9with torch.no_grad():
10    ln.weight.copy_(torch.tensor([2.0, 1.0, 0.5, 1.0]))
11    ln.bias.copy_(torch.tensor([0.0, 0.0, 0.0, 0.5]))
12
13y = ln(x)
14print(y)

Pre-LN vs Post-LN in Transformers

The original Attention Is All You Need (Vaswani et al. 2017) placed LayerNorm afterthe residual addition, so one sub-block was $x \mapsto \text{LN}\bigl(x + \text{Attention}(x)\bigr)$.

This worked for the 6-layer original Transformer but turned out to be catastrophically unstable beyond ~12 layers. Xiong et al. (2020, Theorem 1, ICML On Layer Normalization in the Transformer Architecture) show that in Post-LN the expected gradient norm at the input of the last layer scales as $\mathcal{O}(\sqrt{L})$ at initialization, while the earliest layers receive disproportionately large gradients. The practical consequence is that Post-LN requires a small learning rate plus a long warmup schedule, whereas Pre-LN places LayerNorm inside the residual branch and achieves roughly uniform gradient magnitudes across depth. The visualizer below makes this depth-wise asymmetry concrete — sweep the layer count and watch the per-layer gradient norms diverge under Post-LN and stay flat under Pre-LN.

Pre-LN moves the normalization inside the residual branch — the rule becomes $x \mapsto x + \text{Attention}\bigl(\text{LN}(x)\bigr)$.

The residual stream is now an unnormalized highway; every sub-block reads a normalized copy but writes back to the raw stream. The consequence is a uniform gradient magnitude across depth — GPT-2, GPT-3, LLaMA, PaLM, and essentially every modern LLM use Pre-LN. The tradeoff is a slight quality drop at the same parameter count, which is overwhelmingly worth the ability to train stably at 70B+ parameters without warmup gymnastics.

RMSNorm: The Modern LLM Default

Zhang and Sennrich (2019) observed that LayerNorm spends half its operations on the mean-subtraction step — and that, empirically, removing that step doesn't hurt model quality. The result is RMSNorm, defined as $\text{RMSNorm}(x) = \gamma \cdot \frac{x}{\sqrt{\tfrac{1}{D}\sum_{d} x_d^{2} + \varepsilon}}$.

Three changes from LayerNorm: (i) no mean subtraction, (ii) variance is replaced by raw second moment (mean of $x^{2}$), (iii) no $\beta$. That removes one reduction, one subtraction, and one learnable parameter — about 7–15% fewer FLOPs and one fewer all-reduce when training across many GPUs.

Every flagship open-weights LLM from 2023 onward uses RMSNorm: LLaMA, LLaMA 2, LLaMA 3, Mistral, Mixtral, Qwen, Gemma, Phi-3. T5 used it earlier. The original LayerNorm survives mostly in older stacks (BERT family, GPT-2). Epsilons differ slightly: LLaMA-1/2/3 and Mistral set $\varepsilon = 10^{-5}$ in their published HuggingFace configs, while Gemma and T5 use $\varepsilon = 10^{-6}$. Either way, $\varepsilon$ only exists to keep the denominator finite for near-dead features — it has no measurable effect on loss for well-trained models.

1def rms_norm(x, gamma, eps=1e-5):
2    """No mean subtraction — divide by the root-mean-square only."""
3    rms = np.sqrt((x ** 2).mean(axis=-1, keepdims=True) + eps)
4    return gamma * (x / rms)
5
6x     = np.array([2., 4., 6., 8.], dtype=np.float32)
7gamma = np.array([1., 1., 1., 1.], dtype=np.float32)
8print(rms_norm(x, gamma))

Connections to Flash Attention, MHA, PE, KV-Cache, and Scaling

LayerNorm is the quiet infrastructure that every other Transformer component takes for granted. Pull it out and attention collapses, positional encodings drift, KV-caches become invalid, and scaling laws fail to hold. Below is how each modern system leans on it.

Multi-Head Attention: uniform input scale for every head

In MHA, the input token $x \in \mathbb{R}^{D}$ is projected into $H$ heads via $Q_h = x W_h^{Q}$, and scores are scaled by $1 / \sqrt{d_k}$. That scaling only gives unit-variance scores if $x$ itself has unit variance. LN is what enforces that precondition, so every head sees queries and keys at the same statistical scale and no single head dominates the softmax prematurely.

Positional Encodings: LN washes out arbitrary magnitudes

Sinusoidal PE has entries in $[-1, 1]$, but learned embeddings can have much larger magnitudes. When we compute $h = \text{Embed}(t) + \text{PE}(t)$ and pass it through LN, the layer strips out the overall scale and leaves only the direction. That's why empirical work shows model quality is remarkably robust to PE magnitude: LN re-standardizes whatever you hand it.

Rotary Position Embedding (RoPE) is applied after LN, inside the Q and K projections — and because it's a pure rotation in 2-D subspaces, it preserves the unit norm LN just imposed. That's not a coincidence: RoPE was designed to compose cleanly with RMSNorm.

Flash Attention: LN is the kernel boundary

FlashAttention (Dao et al. 2022) fuses $\text{softmax}(Q K^{\top} / \sqrt{d_k}) V$ into a single SRAM-resident CUDA kernel, eliminating the $O(T^{2})$ HBM traffic of materializing the attention matrix. But FlashAttention does not include the LayerNorm that precedes the Q/K/V projections — LN sits at the kernel boundary, and its output is the input stream that Flash consumes.

In practice modern stacks fuse LN+Linear (or RMSNorm+Linear) into their own kernel — the whole Pre-LN sub-block becomes:

• Kernel 1: RMSNorm + QKV projection (fused)
• Kernel 2: FlashAttention (fused softmax + matmul)
• Kernel 3: output projection + residual add (fused)

Three kernels instead of the naïve dozen, and LN is the hand-off point for each.

KV-Cache: LN is per-token, so caches stay valid

Autoregressive inference caches past keys and values to avoid recomputing them at every step. For that cache to stay valid, the operation that produces K and V must depend only on the current token's past. LayerNorm satisfies this perfectly: its statistics come from the token being processed, nothing else. BatchNorm would be catastrophic here — its statistics depend on the entire batch, so a cached$K_t$ would be inconsistent with a new batch of different size (see Section 3 for another reason BatchNorm breaks under gradient accumulation).

This is a hard architectural reason Transformers cannot use BatchNorm. The moment you want autoregressive generation with caching, the per-sample (per-token) nature of LN becomes non-negotiable.

Grouped-Query and Multi-Query Attention

GQA and MQA shrink the KV-cache by sharing keys/values across multiple query heads. Both still apply LN to the full $D$-dimensional input before projections, so the reduced-size K/V still has unit-variance inputs. The savings are in memory (the cache) and bandwidth, not in the normalization step — which is why LN's cost becomes relatively more visible in MQA models and motivates the switch to RMSNorm.

Transformer scaling: LN pays for itself

At LLaMA-3-70B scale a forward pass performs roughly $L \times 2$ LayerNorms per token (one before attention, one before FFN) across $L = 80$ layers. That's 160 normalizations per token, each touching $D = 8192$ floats. Moving from LayerNorm to RMSNorm removes one reduction per LN and one learnable parameter — at the 15-trillion-token training budget used for LLaMA-3, that's a measurable wall-clock win.

Scaling laws (Kaplan et al. 2020, Hoffmann et al. 2022) implicitly assume the network is well-conditioned at every scale. Without normalization, the compute-optimal curves break: at 70B+ parameters, a de-normalized network doesn't converge at any learning rate. Normalization is not a performance optimization — it is what makes scaling work at all.

The one-line summary. LayerNorm (and its RMSNorm descendant) is the precondition that lets every other piece of the modern Transformer — scaled dot-product attention, multi-head projection, positional encoding, KV-cache reuse, deep residual stacking, and trillion-token training — work as designed. Remove it and you lose depth, stability, and scale simultaneously.

References

• Ioffe, S., Szegedy, C. (2015). Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. ICML.
• Ba, J. L., Kiros, J. R., Hinton, G. E. (2016). Layer Normalization. arXiv:1607.06450.
• Santurkar, S. et al. (2018). How Does Batch Normalization Help Optimization? NeurIPS.
• Vaswani, A. et al. (2017). Attention Is All You Need. NeurIPS.
• Xiong, R. et al. (2020). On Layer Normalization in the Transformer Architecture. ICML.
• Zhang, B., Sennrich, R. (2019). Root Mean Square Layer Normalization. NeurIPS.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS.
• Wang, H. et al. (2022). DeepNet: Scaling Transformers to 1,000 Layers. arXiv:2203.00555.

Check Your Understanding

A short self-check to see that the key distinctions have stuck.