Why does the transformer block use LayerNorm or RMSNorm instead of BatchNorm?

**Option B is right.** The reason transformers use LayerNorm or RMSNorm is structural, not about cost or residuals or positions. ### What each normalization actually normalizes over **BatchNorm** computes a mean and variance for each channel ACROSS the batch dimension. For a CNN with input shape `(batch, channels, H, W)`, it averages over batch and spatial positions to get one mean and variance per channel. **LayerNorm** computes mean and variance for each example ACROSS the feature dimension. For a transformer with input shape `(batch, seq_len, d_model)`, it averages over d_model to get one mean and variance per token position per example. Each token is normalized using only its own features. **RMSNorm** is the same as LayerNorm but drops the mean subtraction and the learned bias. Divide by sqrt of mean-square along d_model, scale by a learned per-channel weight. ### Why BatchNorm fails for transformers 1. **Variable-length sequences need padding.** Real text sequences have different lengths, so batches are padded to the longest in the batch. Those padded positions are noise. If BatchNorm averages over padded positions, the padding pollutes the mean and variance for every real token. Masking them out helps but is fiddly and slow. 2. **Autoregressive inference uses batch size 1.** When generating one new token at a time for a single user, batch size is 1. BatchNorm with batch=1 has no batch statistic to compute; it has to fall back to running averages from training, which generalize poorly. 3. **Batch composition leaks information across examples.** Token i in example A is normalized using statistics that include tokens from examples B, C, and D. That is correlation across the batch, which is exactly what you do not want when each user's generation should be independent. ### Why the other options are wrong **Option A (cost)** is wrong. The kernel cost difference between BatchNorm and LayerNorm is small and not what drives the choice. **Option C (residuals)** is wrong. ResNets use BatchNorm with residual connections every day. BatchNorm and residuals are perfectly compatible. **Option D (positional information)** is wrong. Neither LayerNorm nor BatchNorm carries positional information. Position comes from the positional encoding (sinusoidal, learned absolute, or RoPE), never from normalization.

**Option B is the answer.** BatchNorm's statistics are computed across the batch axis. LayerNorm and RMSNorm compute their statistics within a single token along `d_model`. The axis of normalization is the load-bearing distinction. ### Three concrete failure modes of BatchNorm in a transformer **1. Padding contaminates batch statistics.** Variable-length text batches are padded to a common max length. BatchNorm averages over batch and sequence positions, so padded positions shift the mean and inflate the variance for every real token. Manual masking helps but is fiddly and slow. LayerNorm sidesteps this because it normalizes per-token along d_model, never crossing the batch or sequence dimension. **2. Autoregressive inference at batch size 1 breaks BatchNorm.** Serving one user's chat completion runs at batch=1. BatchNorm has no meaningful batch statistic there, so it falls back to training-time running averages that reflect a different distribution. LayerNorm is invariant to batch size; the same statistic is computed at batch=1 or batch=8192. **3. Batch composition leaks across examples.** In BatchNorm, token i in example A is normalized using statistics that include tokens from examples B, C, D. Two inputs that were independent become correlated through normalization, unwanted for generative serving where each user's output should be statistically independent. ### Why each wrong option is wrong **A (cost)**: small and not the deciding factor; RMSNorm shows even cheaper LayerNorm variants exist. **C (residuals)**: false; ResNet uses BatchNorm with residual connections in every block. **D (positional info)**: false; position lives in the positional encoding (sinusoidal, learned absolute, ALiBi, RoPE), never in normalization. The principle: normalize along an axis you can trust. CNNs on fixed-size images can trust the batch axis; transformers on variable-length text with batch=1 inference cannot.