3.2.9Training Deep Networks

Layer normalization

1,839 words8 min readdifficulty · medium3 backlinks

WHY does this exist?

Deep networks stack many layers. Layer \ell learns weights assuming a certain distribution of inputs. But those inputs are the outputs of layer 1\ell-1, whose weights are also changing. So the input distribution to layer \ell keeps shifting during training — informally called internal covariate shift. Big or lopsided activations also cause exploding/vanishing gradients and slow learning.

BatchNorm fixed this by normalizing across the batch dimension. But that ties every example's statistics to its batch-mates. This breaks when:

  • the batch is tiny (noisy statistics),
  • you process one token at a time (RNNs / autoregressive Transformers),
  • train vs. inference statistics differ.

HOW: derive it from scratch

We have one example's pre-normalization vector xRH\mathbf{x}\in\mathbb{R}^H.

Step 1 — mean over features. Why? To recenter to zero so the layer isn't biased high or low. μ=1Hi=1Hxi\mu = \frac{1}{H}\sum_{i=1}^{H} x_i

Step 2 — variance over features. Why? To measure spread so we can rescale to a controlled size. σ2=1Hi=1H(xiμ)2\sigma^2 = \frac{1}{H}\sum_{i=1}^{H} (x_i-\mu)^2

Step 3 — standardize. Why? Subtracting μ\mu centers it; dividing by σ2+ϵ\sqrt{\sigma^2+\epsilon} makes variance 1\approx 1. The ϵ\epsilon (e.g. 10510^{-5}) prevents division by zero when a vector is near-constant. x^i=xiμσ2+ϵ\hat{x}_i = \frac{x_i-\mu}{\sqrt{\sigma^2+\epsilon}}

Step 4 — affine restore. Why? Forcing mean 0 / var 1 might be too rigid (e.g. a sigmoid wants inputs in a wider range). So we give each feature a learnable scale γi\gamma_i and shift βi\beta_i. If the best thing is not to normalize, the network can learn γi=σ2+ϵ\gamma_i=\sqrt{\sigma^2+\epsilon}, βi=μ\beta_i=\mu and recover the original. yi=γix^i+βiy_i = \gamma_i \hat{x}_i + \beta_i

Figure — Layer normalization

BatchNorm vs LayerNorm — the crucial axis

Normalize over Depends on other examples? Train ≠ Inference?
BatchNorm the batch (same feature, all examples) Yes Yes (running stats)
LayerNorm the features (all features, one example) No No — identical

Worked examples


Common mistakes (steel-manned)


The 80/20

  • LayerNorm = normalize across features of one example → mean 0, var 1 → then learnable γ,β\gamma,\beta.
  • It's batch-independent ⇒ works for size-1 batches, RNNs, Transformers, and is identical at test time.
  • ϵ\epsilon prevents divide-by-zero; γ,β\gamma,\beta restore expressiveness.

Flashcards

LayerNorm normalizes across which dimension?
Across the features of a single example (a "row"), not across the batch.
BatchNorm normalizes across which dimension?
Across the batch for each fixed feature (a "column").
Formula for the normalized value x^i\hat{x}_i in LayerNorm
x^i=xiμσ2+ϵ\hat{x}_i=\dfrac{x_i-\mu}{\sqrt{\sigma^2+\epsilon}} with μ,σ2\mu,\sigma^2 over that example's features.
Why add learnable γ\gamma and β\beta?
To restore expressiveness — the net can rescale/shift or even fully undo normalization if that fits the next layer better.
Purpose of ϵ\epsilon in LayerNorm?
Numerical stability; avoids division by ~0 when a feature vector is nearly constant.
Why is LayerNorm preferred in Transformers/RNNs over BatchNorm?
It's batch-independent, so it works for batch size 1 and per-token processing, and behaves identically at train and test time.
After standardization (ε=0), what are the mean and variance over the features?
Mean 0 and variance 1 exactly.
Is LayerNorm invariant to scaling the whole input vector?
Yes — scaling all features by a constant leaves x^\hat{x} unchanged (before γ,β\gamma,\beta).
Does LayerNorm need running statistics at inference?
No; each example uses its own stats, same as training.

Recall Feynman: explain to a 12-year-old

Imagine each student (an example) answers a row of quiz questions (the features). One student's scores might all be huge, another's tiny — that makes the teacher's job messy. LayerNorm makes each student rescale their own row of scores so the average is 0 and the spread is 1. Now every row looks "normal" and the next teacher can grade fairly. It never mixes students together — each fixes only their own row. Then two knobs (γ,β\gamma,\beta) let the class adjust the final scale if the plain version was too strict.


Connections

  • Batch normalization — the batch-axis sibling; contrast the normalization dimension.
  • Internal covariate shift — the problem normalization mitigates.
  • Transformers — LayerNorm is applied around attention & feed-forward sublayers.
  • Residual connections — combined with LayerNorm as Add & Norm.
  • RMSNorm — a cheaper variant dropping mean-centering.
  • Vanishing and exploding gradients — why stable activation scales help.
  • Activation functions — why γ,β\gamma,\beta matter for the next nonlinearity.

Concept Map

motivates

fails on small batch and RNNs

Step 1

Step 2

centers

rescales

prevents div by zero

yields

Step 4

can undo normalization

Internal covariate shift

BatchNorm

Layer normalization

Mean over features mu

Variance over features

Standardize x-hat

Epsilon guard

Affine restore gamma beta

Output y

Zero mean unit variance

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Dekho, deep network mein har layer apne se pehle wali layer ka output leti hai. Jaise-jaise weights update hote hain, ye activations ka scale kabhi bahut bada, kabhi bahut chhota ho jaata hai — isse training slow ho jaati hai aur gradients phat ya mar jaate hain. Layer normalization ka kaam simple hai: har ek example ke andar jitne bhi features (neurons) hain, unka mean nikaal ke subtract karo aur standard deviation se divide karo, taaki har example ka row ban jaye "mean 0, variance 1". Isse next layer ko hamesha ek controlled scale ke numbers milte hain.

Sabse important baat: LayerNorm batch ke doosre examples ko nahi dekhta. Ye sirf ek example ke apne features par kaam karta hai (row-wise). Yahi wajah hai ki Transformers aur RNNs mein ye use hota hai — chahe batch size 1 ho ya aap ek-ek token process kar rahe ho, koi problem nahi. Aur training vs testing mein behaviour bilkul same rehta hai, BatchNorm ki tarah running-average ka jhanjhat nahi.

Normalize karne ke baad hum do learnable knobs lagate hain: γ\gamma (scale) aur β\beta (shift). Kyun? Kyunki forced mean-0 var-1 hamesha best nahi hota — kabhi next activation ko wider range chahiye. In knobs se network khud decide karta hai kitna scale/shift chahiye, aur zarurat pade to normalization ko poora undo bhi kar sakta hai. Ek chhota sa ϵ\epsilon bhi add karte hain taaki agar variance zero ke paas ho to divide-by-zero se NaN na aaye. Bas itna yaad rakho: Layer = line/row par normalize, Batch = column par.

Go deeper — visual, from zero

Test yourself — Training Deep Networks

Connections