Normalization Techniques: Batch Norm and Layer Norm

A decoder that never sees a batch worth normalizing

Picture the inference server behind a Hindi-English chat model of the kind Sarvam AI or Krutrim would run in production. A request arrives, the model starts generating tokens one at a time, and at every decoding step the server is processing a ragged, constantly-changing pile of sequences: some requests are three tokens into their answer, others three hundred, and the batch that reaches the GPU this millisecond might contain a single lonely sequence if traffic is light. Now suppose someone tried to bolt a BatchNorm layer into that decoder the way you would into an image classifier. Two things break immediately. First, whenever the batch happens to contain just one sequence, BatchNorm's batch statistics degenerate — you'll see exactly why, numerically, in a moment. Second, even when the batch is large, the "population" BatchNorm is drawing from at decode step 47 is not stable across requests: it depends on which sequences from which users happen to be co-resident on the GPU at that instant, which has nothing to do with the content of any individual sequence. That's not internal covariate shift; that's noise from an accident of scheduling.

Transformers were never designed to tolerate this, which is why every major transformer architecture — the original 2017 paper, GPT-2 through GPT-4-class models, LLaMA, Mistral — uses a normalization layer that never looks at the batch axis at all. This chapter is about that layer, LayerNorm, and about the further simplification the field converged on for large-scale training, RMSNorm. You've already seen how BatchNorm tames training with per-batch mean/variance and exponential-moving-average buffers carried into inference. (It's worth knowing, as an aside, that even BatchNorm's original justification has been revisited: Santurkar, Tsipras, Ilyas & Madry, NeurIPS 2018, showed its benefit correlates more with smoothing the optimization landscape than with reducing "internal covariate shift" as Ioffe & Szegedy originally argued in 2015 — a reminder that a technique can work for a different reason than the one that motivated it.) Here we ask the sharper question that sequence models force on you: what do you do when the batch axis itself is not a trustworthy thing to average over?

Same statistics, a different reduction axis

Every normalization layer in deep learning does the same three things to some group of numbers: compute a mean, compute a variance, and rescale to zero-mean unit-variance before applying a learned affine transform. The entire family — BatchNorm, LayerNorm, InstanceNorm, GroupNorm — differs in exactly one design choice: which axis of the activation tensor defines "the group."

Take a transformer's hidden activations at one layer, shaped (N, T, d): N sequences in the batch, T positions per sequence, d features per position. For a fixed feature j, BatchNorm's statistics are computed by pooling across the batch (and, in the sequence setting, typically across positions too):

mean_j = (1/N) * sum over i of x[i, j]
var_j  = (1/N) * sum over i of (x[i, j] - mean_j)^2

LayerNorm inverts which axis is "pooled" and which is "kept separate." For a single token's feature vector x ∈ R^d (one specific i, one specific t), it computes:

mean_i = (1/d) * sum over j of x[i, j]
var_i  = (1/d) * sum over j of (x[i, j] - mean_i)^2
x_norm[i, j] = (x[i, j] - mean_i) / sqrt(var_i + eps)
y[i, j] = gamma[j] * x_norm[i, j] + beta[j]

Notice what this buys you: the statistics for token i depend only on token i's own d features. Nothing about any other token, any other sequence, or how many sequences happen to be in the batch enters the computation. Ba, Kiros & Hinton introduced this in "Layer Normalization" (arXiv, 2016) specifically to remove the batch-size dependence that made BatchNorm awkward for recurrent and sequence models — the exact failure mode from the decoder scenario above. gamma and beta are still learned, still one value per feature, still exactly as many parameters (2d) as BatchNorm's affine transform — the parameter count is identical; only the reduction axis moved.

Reading the two axes off the same tensor

The diagram below freezes one timestep and shows a batch of four token vectors (B1–B4), four features each (F1–F4). LayerNorm's statistics run horizontally, across the four features of a single token. BatchNorm's statistics run vertically, across the four tokens, for a single feature. The two groups genuinely intersect at one cell (B2, F3, shown in purple) — that scalar activation belongs to both a LayerNorm group and a BatchNorm group — but each normalizer only ever reduces along its own axis; it never sees the other group's members.

The collapse nobody warns you about

Here is the misconception worth killing early. Most students assume BatchNorm at batch size 1 is merely "numerically unstable" — a division that gets shaky as the denominator approaches zero, patched by the epsilon term. That is not what happens, and the real failure is worse because it is silent rather than a crash. With a batch of exactly one example, the batch mean for every feature equals that example's own value: mean_j = x[0, j]. Subtracting the mean from the value therefore gives exactly zero, every time, regardless of epsilon:

import numpy as np

x_single = np.array([[2.0, 4.0, 4.0, 6.0]])  # a "batch" of 1 token
mean = x_single.mean(axis=0)
print(mean)              # [2. 4. 4. 6.]
print(x_single - mean)   # [[0. 0. 0. 0.]]

The variance is also exactly zero, so the normalized output is 0 / sqrt(0 + eps) ≈ 0 for every feature — the entire activation vector is annihilated before the affine transform even runs, no matter what that token's actual values were. This is not an edge case you can dodge by picking a slightly larger epsilon; it is a structural property of computing a mean over a single-element set. It's why "just disable BatchNorm's running-stats mode at inference" doesn't save you when your production batch size genuinely is 1 (a single active request, or the last partial batch of a stream) — using training-mode batch statistics zeroes the input, and falling back to the stored running mean/variance from training data introduces a mismatch between what the model was trained to expect at that layer and what it now receives at inference for an out-of-distribution single input. LayerNorm never encounters this failure mode because its reduction axis is the d features of one token, and a single token vector with d > 1 has perfectly well-defined non-degenerate statistics regardless of batch size.

Worked example: normalizing one token vector

Take a toy token embedding after some sublayer, x = [2.0, 4.0, 4.0, 6.0], d = 4, eps = 1e-5.

mean = (2 + 4 + 4 + 6) / 4 = 4.0
var  = ((2-4)^2 + (4-4)^2 + (4-4)^2 + (6-4)^2) / 4
     = (4 + 0 + 0 + 4) / 4 = 2.0
std  = sqrt(var + eps) = sqrt(2.00001) ≈ 1.414217

x_norm = (x - mean) / std
       ≈ [-1.4142, 0.0000, 0.0000, 1.4142]

With identity parameters gamma = [1,1,1,1], beta = [0,0,0,0], the output equals x_norm exactly. Now apply non-trivial learned parameters gamma = [0.5, 0.5, 2.0, 2.0], beta = [1.0, 1.0, 0.0, 0.0]:

y = gamma * x_norm + beta
  = [0.5*(-1.4142)+1, 0.5*0+1, 2*0+0, 2*1.4142+0]
  ≈ [0.2929, 1.0000, 0.0000, 2.8284]

Verify both with code — this uses NumPy's biased variance (ddof=0, dividing by d, not d−1), which matches the formula above and is what every deep-learning framework's LayerNorm uses:

import numpy as np

def layer_norm(x, gamma, beta, eps=1e-5):
    mean = x.mean(axis=-1, keepdims=True)
    var = x.var(axis=-1, keepdims=True)  # ddof=0, divides by d
    x_norm = (x - mean) / np.sqrt(var + eps)
    return gamma * x_norm + beta

x = np.array([2.0, 4.0, 4.0, 6.0])

g_id = np.array([1.0, 1.0, 1.0, 1.0])
b_id = np.array([0.0, 0.0, 0.0, 0.0])
print(np.round(layer_norm(x, g_id, b_id), 4))
# [-1.4142  0.      0.      1.4142]

gamma = np.array([0.5, 0.5, 2.0, 2.0])
beta  = np.array([1.0, 1.0, 0.0, 0.0])
print(np.round(layer_norm(x, gamma, beta), 4))
# [0.2929 1.     0.     2.8284]

Every step traces exactly: layer_norm is fully defined above, called with concrete arrays, and the two printed lines match the hand computation to four decimal places (the tiny 1e-5 epsilon shifts the true value by less than 0.0001, invisible at this rounding).

Placement matters as much as the formula: Pre-LN vs Post-LN

Where you put the normalization inside a transformer block changes training dynamics as much as which normalizer you pick. The original transformer (Vaswani et al., 2017) applies LayerNorm after the residual add — "Post-LN":

out = LayerNorm(x + Sublayer(x))

Xiong et al. (ICML 2020), "On Layer Normalization in the Transformer Architecture," showed why Post-LN transformers need a learning-rate warmup phase to avoid diverging: at initialization, the expected gradient norm near the output layer of a Post-LN transformer stays large and roughly constant regardless of depth, while gradients reaching the early layers vanish quickly — warmup exists to tame those large output-layer gradients, not to prevent an early-layer blowup. Pre-LN, already used in GPT-2, is the placement Xiong et al. later showed removes the need for warmup, because its gradient norm decreases with depth instead: normalize the sublayer's input, and let the residual stream carry the raw, un-normalized sum forward:

out = x + Sublayer(LayerNorm(x))

This keeps an unimpeded identity path (the plain "+ x") all the way from the embedding layer to the output, which is what actually stabilizes gradients at depth and removes the need for careful warmup scheduling. The tradeoff: because Pre-LN never renormalizes the residual stream itself, its magnitude can drift upward as more sublayer outputs accumulate into it across dozens of layers. That's precisely why GPT-2-style architectures add one extra LayerNorm — often called ln_f — immediately before the final output projection, to rescale the accumulated residual stream one last time before it hits the vocabulary head. Drop that final LayerNorm and the logits are computed from an increasingly large, poorly calibrated activation, degrading the softmax's effective temperature in a depth-dependent way.

RMSNorm: dropping the mean

Zhang & Sennrich (NeurIPS 2019), "Root Mean Square Layer Normalization," asked whether LayerNorm's re-centering step (subtracting the mean) was actually doing useful work, or whether the rescaling step (dividing by a spread statistic) was carrying the benefit. Empirically, rescaling alone matched LayerNorm's performance in their experiments. RMSNorm keeps only the rescaling: no mean subtraction, and consequently no beta shift parameter either — just a per-feature learned gain gamma:

RMS(x) = sqrt( (1/d) * sum(x_j^2) + eps )
y_j = gamma_j * x_j / RMS(x)

Continue the same token vector, x = [2.0, 4.0, 4.0, 6.0]:

mean of squares = (4 + 16 + 16 + 36) / 4 = 72 / 4 = 18.0
RMS = sqrt(18.00001) ≈ 4.242642

y (gamma = 1) = x / RMS
  ≈ [0.4714, 0.9428, 0.9428, 1.4142]

import numpy as np

def rms_norm(x, gamma, eps=1e-5):
    rms = np.sqrt(np.mean(x**2) + eps)
    return gamma * (x / rms)

x = np.array([2.0, 4.0, 4.0, 6.0])
g_id = np.array([1.0, 1.0, 1.0, 1.0])
print(np.round(rms_norm(x, g_id), 4))
# [0.4714 0.9428 0.9428 1.4142]

Compare this to LayerNorm's [-1.4142, 0, 0, 1.4142] on the same input: RMSNorm's output is different because it never subtracted the mean, so the two zero-deviation entries (the two 4.0's) don't collapse to zero here — they retain a nonzero value reflecting their distance from zero, not their distance from the vector's mean. That difference is the entire point: skipping mean-centering removes one reduction (computing the mean) and one subtraction per element, which is a meaningful savings when this operation runs on every token, every layer, every forward and backward pass, at the scale of a multi-billion-parameter model — normalization layers are memory-bandwidth-bound, not compute-bound, so cutting the number of passes over each activation tensor has a real, measurable effect on training throughput. This is why LLaMA, Mistral, and PaLM use RMSNorm in place of full LayerNorm in their transformer blocks.

Active recall

Attempt each question before reading its answer.

1. Given the token vector [2.0, 4.0, 4.0, 6.0] from the worked example, why does applying BatchNorm in training mode with a batch size of 1 produce a zero vector — not just an unstable one — regardless of epsilon?
2. Compute the RMSNorm output for x = [2.0, 4.0, 4.0, 6.0] with gamma = [2.0, 2.0, 2.0, 2.0] and eps = 1e-5.
3. A GPT-2-style Pre-LN transformer has its final LayerNorm (ln_f) removed before the output head. What happens to the scale of the residual stream reaching the head as depth grows, and why does that specifically hurt the softmax?
4. True or false: "LayerNorm needs fewer learnable parameters than BatchNorm because it doesn't need a running average." Justify your answer by parameter count.
5. Reusing the LayerNorm worked example with gamma = [1,1,1,1], beta = [0,0,0,0]: if eps is changed from 1e-5 to an unusually large 1.0, recompute the normalized output for x = [2.0, 4.0, 4.0, 6.0]. What does this reveal about treating eps as always negligible?
6. Why can LayerNorm be applied cleanly during single-token autoregressive decoding (batch = 1, one new token per step), while BatchNorm cannot, even setting aside the batch-size-1 collapse from Q1?

Answers

1. With N = 1, the batch mean for every feature equals that single example's own value (mean_j = x[0,j]), so x[0,j] − mean_j = 0 exactly, for every feature, before epsilon is even applied. The numerator is deterministically zero; epsilon only affects the denominator, which was never the problem. The output is the zero vector regardless of how epsilon is chosen — a silent information loss, not a numerical-stability symptom a bigger epsilon could fix.

2. RMS is unchanged by gamma: RMS = sqrt(18.00001) ≈ 4.242642. The unscaled normalized vector is [0.4714, 0.9428, 0.9428, 1.4142] (from the worked example). Multiplying by gamma = [2,2,2,2] gives [0.9428, 1.8856, 1.8856, 2.8284].

3. In Pre-LN, each block adds Sublayer(LayerNorm(x)) directly into the residual stream without ever renormalizing the stream itself. As more layers add their (roughly independent, non-zero-variance) contributions, the stream's overall magnitude grows with depth — intuitively like accumulating variance across independent additions, so it grows roughly with the square root of the number of layers. Without ln_f rescaling it back down immediately before the output projection, the final hidden vector feeding the vocabulary head has a depth-dependent, uncontrolled magnitude. Since softmax logits are that vector dotted with the output embedding matrix, an inflated magnitude sharpens (or distorts) the softmax the way dividing by an uncalibrated temperature would, independent of what the model actually "believes" — a purely architectural artifact corrupting the output distribution.

4. False. Both LayerNorm and BatchNorm learn a per-feature gamma and beta — 2d learnable parameters either way, identical count. The distinction is that BatchNorm additionally carries non-learnable running-mean and running-variance buffers (updated by an exponential moving average during training, read back at inference) that LayerNorm never needs, because LayerNorm recomputes its statistics fresh from the current example at both training and inference time. Fewer buffers, not fewer parameters.

5. var + eps = 2.0 + 1.0 = 3.0, so std = sqrt(3.0) ≈ 1.7321 instead of sqrt(2.00001) ≈ 1.4142. The normalized vector becomes (x − mean)/1.7321 ≈ [−1.1547, 0, 0, 1.1547] — visibly smaller in magnitude than the eps = 1e-5 result of [−1.4142, 0, 0, 1.4142]. This shows eps is not always a harmless numerical-safety constant: at an unusually large value it actively shrinks the normalized range, behaving like a regularizer. It also matters in exactly the low-variance regime real models hit — near-duplicate token embeddings or padding positions — where a poorly chosen eps can meaningfully distort supposedly "normalized" activations rather than just guard against division by zero.

6. LayerNorm's statistics are computed entirely from the current token's own d features, so they're well-defined and unaffected by how many other sequences share the batch, how long any of them are, or how many tokens have already been generated. BatchNorm's statistics are defined by the batch's composition, which in a decoding server changes every step as sequences finish, new requests arrive, and dynamic batching reshuffles who's co-resident on the GPU — even ignoring the batch = 1 collapse, that composition dependence means the "population" a running-average BatchNorm was calibrated on during training bears no reliable relationship to any specific inference-time batch, making it structurally unsuited to autoregressive serving in a way LayerNorm simply is not.

Think About It

Think about this: How would you explain normalization techniques: batch norm and layer norm to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where normalization techniques: batch norm and layer norm is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting normalization techniques: batch norm and layer norm to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind normalization techniques: batch norm and layer norm, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.