Normalization Techniques - Stabilizing Training Across Architectures

Reading time: ~40 min | Interview relevance: High | Roles: MLE, AI Eng, Research Engineer, LLM Engineer

The Real Interview Moment

You are in a Google MLE interview. The interviewer asks: "You are training a ResNet-50 with batch size 2 on a limited-memory GPU. Your model diverges during training. Why might this happen, and what normalization technique would you switch to?"

You answer "batch normalization with small batch size," and she follows up: "Derive what happens to BatchNorm statistics when batch size is 2. Then explain why Layer Normalization or Group Normalization would fix the problem. Write the formulas for all three."

This question tests a chain of knowledge: the mathematical definition of each normalization, how the statistics depend on the batch, why that creates problems, and what alternatives exist. It is a favorite at Google, Meta, and Amazon because normalization choices have real impact on production models - and because most candidates only know the surface level.

Candidates who say "BatchNorm normalizes across the batch" without being able to write the formula or explain the training-vs-inference discrepancy get a "lean no-hire." Candidates who can derive BatchNorm, explain the running statistics, compare it to LayerNorm with precise axis specifications, and discuss RMSNorm in the context of LLMs get a "strong hire."

What You Will Master

• Derive Batch Normalization from first principles with exact axes of normalization
• Explain BatchNorm's training vs inference behavior (running statistics)
• Derive Layer Normalization and explain why it works for Transformers
• Explain Group Normalization and Instance Normalization with precise definitions
• Derive RMSNorm and explain why it replaced LayerNorm in modern LLMs
• Diagnose when BatchNorm fails (small batches, sequence models, distributed training)
• Choose the right normalization for any architecture (CNNs, Transformers, GANs, diffusion)
• Solve interview problems on normalization design and debugging

Self-Assessment: Where Are You Now?

Skill	1 -- Cannot	2 -- Vaguely	3 -- Can Explain	4 -- Can Derive	5 -- Can Teach	Your Score
Write the BatchNorm formula						___
Explain BatchNorm train vs inference						___
Write the LayerNorm formula						___
Explain why LayerNorm for Transformers						___
Define GroupNorm and InstanceNorm						___
Derive RMSNorm from LayerNorm						___
Diagnose BatchNorm failures						___
Choose normalization for a given architecture						___

Target: All 4s and 5s before your interview.

Part 1 - The Core Problem: Internal Covariate Shift

Why Normalize?

During training, the distribution of inputs to each layer changes as the parameters of the preceding layers are updated. This phenomenon - called internal covariate shift - makes training difficult because each layer must constantly adapt to a moving input distribution.

Consequences of not normalizing:

1. Gradients can vanish or explode as activations grow or shrink through layers
2. Optimization becomes sensitive to learning rate and initialization
3. Training requires smaller learning rates, leading to slower convergence

What normalization does:

1. Centers activations to have zero mean
2. Scales activations to have unit variance
3. Applies learnable scale ($\gamma$) and shift ($\beta$) to restore representational power

The general normalization formula is:

$$\hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta$$

Where $\mu$ and $\sigma^2$ are computed over some subset of the activations (which subset defines the normalization type), and $\gamma$, $\beta$ are learned parameters.

60-Second Answer

"Normalization stabilizes training by keeping activations in a well-conditioned range. The key difference between normalization types is WHICH dimensions they average over. BatchNorm averages over the batch dimension - it normalizes each feature across all examples. LayerNorm averages over the feature dimension - it normalizes each example across all features. GroupNorm divides features into groups and normalizes within each group. RMSNorm simplifies LayerNorm by removing the mean centering, using only the root mean square. The choice depends on the architecture: BatchNorm for CNNs, LayerNorm for Transformers, RMSNorm for modern LLMs, GroupNorm when batch sizes are small."

Part 2 - Batch Normalization

The Formula

Given a mini-batch of activations $x \in \mathbb{R}^{B \times C \times H \times W}$ (for a CNN with batch size $B$, channels $C$, height $H$, width $W$):

Step 1: Compute per-channel statistics across the batch and spatial dimensions.

$$\mu_c = \frac{1}{B \cdot H \cdot W} \sum_{b=1}^{B} \sum_{h=1}^{H} \sum_{w=1}^{W} x_{b,c,h,w}$$ $$\sigma_c^2 = \frac{1}{B \cdot H \cdot W} \sum_{b=1}^{B} \sum_{h=1}^{H} \sum_{w=1}^{W} (x_{b,c,h,w} - \mu_c)^2$$

Step 2: Normalize.

$$\hat{x}_{b,c,h,w} = \frac{x_{b,c,h,w} - \mu_c}{\sqrt{\sigma_c^2 + \epsilon}}$$

Step 3: Scale and shift with learned parameters (per channel).

$$y_{b,c,h,w} = \gamma_c \hat{x}_{b,c,h,w} + \beta_c$$

Learned parameters: $\gamma \in \mathbb{R}^C$ and $\beta \in \mathbb{R}^C$ (one scale and one shift per channel).

Training vs Inference Behavior

This is one of the most important and most-tested aspects of BatchNorm.

During training:

• $\mu$ and $\sigma^2$ are computed from the CURRENT mini-batch
• A running average is maintained for inference:

$$\mu_{\text{running}} \leftarrow (1 - \alpha) \cdot \mu_{\text{running}} + \alpha \cdot \mu_{\text{batch}}$$ $$\sigma^2_{\text{running}} \leftarrow (1 - \alpha) \cdot \sigma^2_{\text{running}} + \alpha \cdot \sigma^2_{\text{batch}}$$

where $\alpha$ is the momentum (typically 0.1).

During inference:

• $\mu_{\text{running}}$ and $\sigma^2_{\text{running}}$ are used (NOT batch statistics)
• This is because at inference time, you may have a single example (batch size 1), and computing statistics from one example is meaningless

The critical implication: BatchNorm behaves DIFFERENTLY during training and inference. Forgetting to switch to eval mode (model.eval() in PyTorch) is a classic production bug.

# PyTorch: BatchNorm training vs inference
import torch.nn as nn

bn = nn.BatchNorm2d(num_features=64)

# Training: uses batch statistics, updates running stats
bn.train()
output_train = bn(input_tensor)

# Inference: uses running statistics
bn.eval()
output_eval = bn(input_tensor)
# These produce DIFFERENT outputs for the same input!

Instant Rejection

If asked "what happens when you deploy a model with BatchNorm but forget to call model.eval()?" and you cannot answer, that is a very bad sign. The answer: the model uses batch statistics from whatever data happens to be in the current batch, leading to inconsistent and often degraded predictions. For batch size 1, the mean equals the single example and the variance is zero, causing a division-by-zero (or near-zero) catastrophe.

Why BatchNorm Works

The original paper (Ioffe and Szegedy, 2015) attributed BatchNorm's success to reducing internal covariate shift. However, subsequent research (Santurkar et al., 2018) showed that the real reason is:

1. Smoothing the loss landscape. BatchNorm makes the loss surface smoother (more Lipschitz continuous), which allows larger learning rates.
2. Decoupling layer interactions. Each layer's input distribution is more stable, reducing the dependence between layers during optimization.
3. Implicit regularization. The noise from batch statistics acts as a regularizer (similar to dropout).

When BatchNorm Fails

Problem 1: Small batch sizes.

When $B$ is very small (e.g., 1-4), the batch statistics $\mu_{\text{batch}}$ and $\sigma^2_{\text{batch}}$ are noisy estimates of the true statistics. The normalization is dominated by noise, destabilizing training.

Quantifying the issue: the standard error of $\mu_{\text{batch}}$ is $\sigma / \sqrt{B \cdot H \cdot W}$. For $B = 2$ and a 1x1 feature map (as in FC layers), this is $\sigma / \sqrt{2}$ - the estimate has 70% of the true standard deviation as its own uncertainty.

Problem 2: Sequence models (RNNs, Transformers).

In sequence models, the batch dimension and sequence dimension interact in complex ways. Applying BatchNorm across the batch dimension means different time steps within the same sequence are normalized using the same batch statistics, which does not account for the temporal structure. Additionally, sequence lengths often vary, making batch statistics inconsistent.

Problem 3: Distributed training.

In data-parallel training across multiple GPUs, each GPU computes batch statistics from its local mini-batch. If the local batch size per GPU is small (even if the global batch is large), the statistics are noisy. SyncBatchNorm addresses this but adds communication overhead.

Problem 4: GANs.

The batch statistics create information leakage between examples in the batch, which can destabilize GAN training. Spectral normalization or instance normalization are preferred for GANs.

Common Trap

Do NOT say "BatchNorm fails with small batches because there isn't enough data." Be precise: the issue is that the sample mean and variance become high-variance estimates of the true mean and variance. A batch of 2 gives statistics with very high uncertainty, making the normalization noisy and unreliable. The noise does not cancel out - it amplifies through the network.

Part 3 - Layer Normalization

The Formula

Given activations $x \in \mathbb{R}^{B \times T \times D}$ (for a Transformer with batch size $B$, sequence length $T$, dimension $D$):

Normalize across the feature dimension for each individual token:

$$\mu_{b,t} = \frac{1}{D} \sum_{d=1}^{D} x_{b,t,d}$$ $$\sigma_{b,t}^2 = \frac{1}{D} \sum_{d=1}^{D} (x_{b,t,d} - \mu_{b,t})^2$$ $$\hat{x}_{b,t,d} = \frac{x_{b,t,d} - \mu_{b,t}}{\sqrt{\sigma_{b,t}^2 + \epsilon}} \cdot \gamma_d + \beta_d$$

Key difference from BatchNorm: Statistics are computed over the feature dimension $D$, independently for each example and each position. There is no dependence on other examples in the batch.

Axes of Normalization: The Critical Distinction

Normalization	Input Shape	Statistics Computed Over	Independent Per
BatchNorm	B x C x H x W	B, H, W	Channel C
LayerNorm	B x T x D	D	Example B, Position T
InstanceNorm	B x C x H x W	H, W	Example B, Channel C
GroupNorm	B x C x H x W	(C/G), H, W	Example B, Group

Why LayerNorm for Transformers

1. Batch-independent. Each token is normalized independently. This means:

   • Works with batch size 1
   • No running statistics needed (same behavior at train and test time)
   • No cross-example information leakage

2. Position-independent. Each position is normalized independently, which is appropriate for variable-length sequences.

3. Consistent behavior. No train/test discrepancy (unlike BatchNorm). The computation is identical during training and inference.

4. Parallelizable. Since each token is normalized independently, LayerNorm is trivially parallelizable across the sequence and batch dimensions.

Interviewer's Perspective

The strongest answer to "why LayerNorm for Transformers?" is: "BatchNorm computes statistics across the batch dimension, which means it depends on what other examples are in the batch. In Transformers processing variable-length sequences with padding, this creates two problems: the statistics are contaminated by padding tokens, and they vary depending on the batch composition. LayerNorm computes statistics per token across features, which is completely independent of batch composition, padding, and sequence length."

LayerNorm Parameters

Learned parameters: $\gamma \in \mathbb{R}^D$ and $\beta \in \mathbb{R}^D$ (one scale and one shift per feature dimension).

For $D = 4096$ (typical LLM): 8,192 parameters per LayerNorm layer. For a model with 2 LayerNorm layers per Transformer block and 32 blocks: $32 \times 2 \times 8192 = 524,288$ total - negligible compared to the attention and FFN parameters.

Part 4 - RMSNorm

The Simplification

RMSNorm (Zhang and Sennrich, 2019) removes the mean centering from LayerNorm, using only the root mean square for normalization:

$$\text{RMS}(x) = \sqrt{\frac{1}{D} \sum_{d=1}^{D} x_d^2}$$ $$\hat{x}_d = \frac{x_d}{\text{RMS}(x)} \cdot \gamma_d$$

Note: no $\beta$ (shift) parameter and no mean subtraction.

Why RMSNorm Replaced LayerNorm in LLMs

1. Computational savings.

LayerNorm requires:

• Computing the mean: $O(D)$ operations
• Subtracting the mean: $O(D)$ operations
• Computing the variance: $O(D)$ operations
• Total: ~3 passes over the data, plus the affine transform

RMSNorm requires:

• Computing the RMS: $O(D)$ operations (one pass)
• Dividing by RMS: $O(D)$ operations
• Total: ~2 passes over the data

RMSNorm is approximately 10-15% faster than LayerNorm on GPU.

2. Empirical equivalence.

Multiple studies have shown that mean centering provides negligible benefit for Transformer training. The dominant effect is the scaling (normalization by magnitude), not the centering.

3. Theoretical justification.

The re-centering (mean subtraction) in LayerNorm forces all activations to sum to zero, which constrains the representational capacity. RMSNorm avoids this constraint while still preventing activation magnitudes from growing uncontrollably.

Which Models Use RMSNorm

Model	Normalization
BERT	LayerNorm (post-norm)
GPT-2	LayerNorm (pre-norm)
GPT-3	LayerNorm (pre-norm)
T5	LayerNorm (pre-norm)
LLaMA	RMSNorm (pre-norm)
LLaMA 2	RMSNorm (pre-norm)
Mistral	RMSNorm (pre-norm)
PaLM	LayerNorm (pre-norm)
Gemma	RMSNorm (pre-norm)

The trend is clear: RMSNorm is becoming the default for new LLMs.

Company Variation

At Google, you might be asked about the computational difference between LayerNorm and RMSNorm in the context of TPU/GPU efficiency. At Anthropic and OpenAI, the focus is more on "why did LLaMA switch to RMSNorm and did it affect model quality?" (Answer: no measurable quality loss, measurable speed improvement.) At startups, the question is more practical: "should we use RMSNorm for our new model?" (Answer: yes, unless you have a specific reason to need mean centering.)

Part 5 - Group Normalization and Instance Normalization

Group Normalization

GroupNorm (Wu and He, 2018) divides the $C$ channels into $G$ groups and normalizes within each group:

$$\mu_{b,g} = \frac{1}{(C/G) \cdot H \cdot W} \sum_{c \in \text{group } g} \sum_{h,w} x_{b,c,h,w}$$ $$\sigma_{b,g}^2 = \frac{1}{(C/G) \cdot H \cdot W} \sum_{c \in \text{group } g} \sum_{h,w} (x_{b,c,h,w} - \mu_{b,g})^2$$

Key insight: GroupNorm is a generalization that includes LayerNorm and InstanceNorm as special cases:

• $G = 1$: GroupNorm = LayerNorm (all channels in one group)
• $G = C$: GroupNorm = InstanceNorm (each channel is its own group)

Why GroupNorm for Small-Batch CNNs

GroupNorm was designed specifically for the case where BatchNorm fails: small batch sizes.

Batch Size	BatchNorm	GroupNorm (G=32)
32	Excellent (stable statistics)	Good
8	Good	Good
4	Degraded (noisy statistics)	Good
2	Poor (very noisy)	Good
1	Fails (zero variance)	Good

GroupNorm is batch-independent - its statistics are computed per example. Performance is identical regardless of batch size. The default $G = 32$ works well for most CNN architectures.

Instance Normalization

InstanceNorm normalizes each channel of each example independently, using only spatial dimensions:

$$\mu_{b,c} = \frac{1}{H \cdot W} \sum_{h,w} x_{b,c,h,w}$$ $$\sigma_{b,c}^2 = \frac{1}{H \cdot W} \sum_{h,w} (x_{b,c,h,w} - \mu_{b,c})^2$$

Primary use case: style transfer and image generation.

InstanceNorm normalizes out instance-specific contrast and style information from each channel. This is exactly what you want for style transfer - you want to strip the content image of its style characteristics.

Normalization	Best For	Why
BatchNorm	CNNs with batch >= 16	Batch statistics provide regularization
GroupNorm	CNNs with small batch	Batch-independent, stable with any batch size
InstanceNorm	Style transfer, image gen	Strips instance-specific style info
LayerNorm	Transformers, RNNs	Batch-independent, position-independent
RMSNorm	LLMs	Faster LayerNorm, no quality loss

Part 6 - Normalization in Practice: Choosing the Right One

Decision Framework

Common Combinations in Production Models

Model	Architecture	Normalization	Why
ResNet-50	CNN	BatchNorm	Large batch training, strong regularization
YOLO	CNN	BatchNorm	Object detection, large batch
Mask R-CNN	CNN	GroupNorm	Small batch per GPU (images are large)
BERT	Transformer	LayerNorm (post-norm)	2018 design, pre-norm not yet standard
GPT-3	Transformer	LayerNorm (pre-norm)	Stable training at scale
LLaMA 2	Transformer	RMSNorm (pre-norm)	Faster, no quality loss
Stable Diffusion U-Net	CNN + cross-attn	GroupNorm	Varying batch sizes, image generation
StyleGAN	CNN (GAN)	InstanceNorm (modulated)	Style control
ViT	Transformer	LayerNorm (pre-norm)	Vision Transformer, standard Transformer norm

Part 7 - Mathematical Properties and Gradient Analysis

BatchNorm Gradient

The gradient of BatchNorm with respect to its input is:

$$\frac{\partial L}{\partial x_i} = \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} \left( \frac{\partial L}{\partial \hat{x}_i} - \frac{1}{m}\sum_{j=1}^{m}\frac{\partial L}{\partial \hat{x}_j} - \frac{\hat{x}_i}{m}\sum_{j=1}^{m}\frac{\partial L}{\partial \hat{x}_j}\hat{x}_j \right)$$

where $m = B \cdot H \cdot W$.

Key observation: The gradient depends on the entire batch through the sums. This means:

1. Each example's gradient is influenced by all other examples in the batch
2. This creates an implicit coupling between examples
3. It provides a regularization effect (similar to dropout noise)

LayerNorm Gradient

The LayerNorm gradient has the same structure but the sums are over the feature dimension $D$ instead of the batch:

$$\frac{\partial L}{\partial x_d} = \frac{\gamma_d}{\sqrt{\sigma^2 + \epsilon}} \left( \frac{\partial L}{\partial \hat{x}_d} - \frac{1}{D}\sum_{j=1}^{D}\frac{\partial L}{\partial \hat{x}_j} - \frac{\hat{x}_d}{D}\sum_{j=1}^{D}\frac{\partial L}{\partial \hat{x}_j}\hat{x}_j \right)$$

Key observation: The gradient depends only on the current example - there is no cross-example coupling. This is why LayerNorm does not provide the same implicit regularization as BatchNorm.

Why $\gamma$ and $\beta$ Are Necessary

Without $\gamma$ and $\beta$, the normalization forces the activations to have zero mean and unit variance. This constrains the representational power of the layer - the following layer can only work with standardized inputs.

With $\gamma$ and $\beta$, the network can learn to undo the normalization if that is optimal. In particular, if $\gamma = \sigma$ and $\beta = \mu$, then:

$$\hat{x} \cdot \gamma + \beta = \frac{x - \mu}{\sigma} \cdot \sigma + \mu = x$$

The normalization becomes an identity. This means the normalization layer can never reduce the network's representational power - it can only make optimization easier.

Common Trap

Do NOT say "we can remove the bias term from the preceding linear layer when using BatchNorm." While this is technically true (the bias is absorbed into $\beta$), the reason is subtle: the bias shifts all activations by a constant, but BatchNorm subtracts the mean, removing any constant shift. So the bias has no effect on the normalized output. This DOES matter in practice - frameworks like PyTorch allow bias=False in Conv2d before BatchNorm to save a few parameters.

Part 8 - Advanced Topics

Weight Normalization

An alternative approach: instead of normalizing activations, normalize the weights:

$$w = g \cdot \frac{v}{\|v\|}$$

Where $v$ is the weight vector, $g$ is a learned scalar magnitude. This decouples the direction and magnitude of the weight vector.

Advantages: No dependence on batch size, no running statistics. Disadvantages: Does not normalize activations, less effective than BatchNorm for CNNs.

Spectral Normalization

Used primarily in GANs, spectral normalization constrains the spectral norm (largest singular value) of each weight matrix:

$$W_{\text{SN}} = \frac{W}{\sigma_1(W)}$$

This enforces a Lipschitz constraint on each layer, stabilizing GAN training. Unlike BatchNorm, it normalizes weights rather than activations.

Adaptive Normalization (AdaIN, AdaLN)

Used in style transfer and diffusion models:

$$\text{AdaIN}(x, y) = \sigma(y) \cdot \frac{x - \mu(x)}{\sigma(x)} + \mu(y)$$

The scale and shift parameters come from a conditioning input $y$ (style image or diffusion timestep), not from learned parameters. This allows the normalization to be conditioned on external information.

AdaLN-Zero (used in DiT for diffusion): Initialize $\gamma$ and $\beta$ to produce zero output, so the initial model computes the identity through residual connections. This is a normalization variant of zero-initialization, which helps training stability.

Normalization-Free Networks

Recent work (Brock et al., 2021, NF-Nets) shows that carefully initialized networks with adaptive gradient clipping can train without any normalization:

$$\text{AGC}: g_i \leftarrow \lambda \frac{\|W_i\|}{\|g_i\|} g_i \quad \text{if} \quad \frac{\|g_i\|}{\|W_i\|} > \lambda$$

NF-Nets achieved state-of-the-art ImageNet accuracy without BatchNorm, but the technique requires careful tuning and has not been widely adopted.

Part 9 - Common Interview Questions

Q1: "Why not just use BatchNorm everywhere?"

BatchNorm fails in four scenarios:

1. Small batch sizes (< 8): noisy statistics destabilize training
2. Variable-length sequences (NLP): padding contamination, inconsistent statistics
3. Inference with batch size 1: must use running statistics, which may differ from actual inference distribution
4. Distributed training with small local batches: each GPU has noisy statistics unless using SyncBatchNorm

Q2: "Can you use BatchNorm in a Transformer?"

Technically yes, but it is a bad idea for multiple reasons:

1. NLP batches often have variable sequence lengths with padding - BatchNorm statistics would be contaminated by padding tokens
2. The batch dimension is the wrong axis for sequence data - you want per-token normalization
3. Train/test discrepancy: inference may process single sequences
4. PowerNorm (2003) and other attempts to use batch statistics in Transformers have shown limited success

Q3: "What is the relationship between BatchNorm and dropout?"

Both provide regularization:

• BatchNorm: noise from mini-batch statistics
• Dropout: noise from randomly zeroing activations

Using both together can be harmful - the noise from one can interfere with the other. In modern practice:

• CNNs: BatchNorm (no dropout in conv layers, sometimes dropout in FC layers)
• Transformers: LayerNorm + dropout (dropout on attention weights and after FFN)

Q4: "If RMSNorm removes mean centering, doesn't that lose information?"

The mean information is not lost - it is simply not removed from the activations. LayerNorm forces zero mean; RMSNorm does not. The downstream layers can still use mean information because it is preserved in the activations. The empirical finding is that the mean-centering step in LayerNorm provides negligible benefit - what matters is the magnitude normalization.

Practice Problems

Problem 1: Normalization Axis Challenge

You have a tensor of shape [B=4, C=64, H=32, W=32]. For each normalization type, specify:

(a) The dimensions averaged over to compute statistics (b) The number of mean/variance statistics computed (c) The shape of the learned $\gamma$ and $\beta$ parameters

Hint 1 -- Direction

For each normalization type, think about which dimensions define a "group" that shares the same statistics. BatchNorm groups by channel (averages over batch and spatial). LayerNorm groups by example (averages over features). GroupNorm groups by example and channel group.

Hint 2 -- Insight

The number of statistics equals the product of dimensions NOT averaged over. BatchNorm averages over B, H, W, leaving C = 64 statistics. LayerNorm averages over C, H, W, leaving B = 4 statistics. GroupNorm (G=8) averages over C/G, H, W within each group, leaving B x G = 32 statistics.

Hint 3 -- Full Solution + Rubric

Norm Type	Axes Averaged	Num Statistics (mean + var)	$\gamma$, $\beta$ shape
BatchNorm	B, H, W	64 (one per C)	[64]
LayerNorm	C, H, W	4 (one per B)	[64, 32, 32] or [64]
InstanceNorm	H, W	4 x 64 = 256	[64]
GroupNorm (G=8)	C/G, H, W	4 x 8 = 32	[64]

Important subtleties:

1. LayerNorm in CNNs is rare. When applied, it typically normalizes over C, H, W (or just C in Transformers where input is [B, T, D]).

2. GroupNorm with G=8 means 8 groups of 64/8 = 8 channels each. Statistics are shared within each group for each example.

3. InstanceNorm computes separate statistics for each (example, channel) pair using only spatial dimensions.

4. BatchNorm's $\gamma$, $\beta$ are per-channel [64], same as GroupNorm and InstanceNorm. LayerNorm's $\gamma$, $\beta$ can be per-element if normalizing over all of C, H, W.

Scoring Rubric:

• Strong Hire: Correct axes and statistics counts for all four, explains the subtleties of LayerNorm shape in CNN context, understands GroupNorm as interpolation between LN and IN.
• Lean Hire: Gets BatchNorm and LayerNorm correct but makes errors on GroupNorm or InstanceNorm.
• No Hire: Cannot specify the axes for BatchNorm correctly or confuses per-channel with per-example statistics.

Problem 2: BatchNorm Debugging

Your ResNet-50 achieves 92% validation accuracy during training. When you deploy it (with model.eval() correctly called), accuracy drops to 78%. What could cause this?

Hint 1 -- Direction

The 14% drop between training-time validation and deployment suggests a discrepancy between training and inference behavior. The most likely culprit is the running statistics in BatchNorm - if they do not match the actual inference data distribution, the normalization will be wrong.

Hint 2 -- Insight

Several scenarios cause running statistics to diverge from actual inference statistics: (1) training data distribution differs from inference distribution, (2) the momentum for running statistics was too high, meaning the final running stats over-weight recent batches, (3) data augmentation during training changes the activation distribution, but the running stats reflect augmented data while inference sees unaugmented data.

Hint 3 -- Full Solution + Rubric

Most likely causes (ranked):

1. Distribution shift between training and deployment data. The running statistics ($\mu_{\text{running}}$, $\sigma^2_{\text{running}}$) reflect the training data distribution. If the deployment data has a different distribution (lighting conditions, image resolution, preprocessing pipeline differences), the normalization will be incorrect.

2. Data augmentation mismatch. Training uses augmentation (random crops, color jitter, etc.), which changes activation distributions. Running statistics reflect augmented data. At inference, data is not augmented, causing a systematic shift in activation distributions.

3. Running statistics momentum issue. With momentum $\alpha = 0.1$ (default in PyTorch), the running mean is an exponential moving average that overweights recent batches. If the last few training batches are not representative, the running stats are biased.

4. Dropout interaction. If dropout is used alongside BatchNorm, the activation distributions during training (with dropout) differ from inference (without dropout). BatchNorm's running stats are computed with dropout active, but used without dropout at inference.

Diagnostic steps:

1. Compute BatchNorm statistics on the deployment data and compare with running statistics
2. Try recalibrating running statistics by passing the deployment data through the model in training mode (without updating weights)
3. Replace BatchNorm with GroupNorm or use torch.nn.SyncBatchNorm if the issue is distributed training-related
4. Check if the preprocessing pipeline is identical between training and deployment

Fix: Recalibrate running statistics.

# Recalibrate BatchNorm statistics on deployment data
model.train()  # Switch to training mode (use batch stats)
with torch.no_grad():
    for batch in deployment_calibration_loader:
        model(batch)  # Updates running stats
model.eval()  # Switch back to eval mode
# Now running stats match deployment distribution

Scoring Rubric:

• Strong Hire: Identifies distribution shift and augmentation mismatch as top causes, proposes recalibration as the fix, explains the running statistics mechanism, provides diagnostic steps.
• Lean Hire: Identifies the running statistics discrepancy but cannot explain multiple causes or propose a systematic fix.
• No Hire: Suggests the issue is model overfitting or proposes retraining as the only solution.

Problem 3: Normalization Design for a New Architecture

You are designing a model that processes 3D point clouds (unordered sets of 3D coordinates). Batch size varies from 1 to 32. Point clouds have 1024-10000 points each. You need to choose a normalization strategy.

(a) Why is BatchNorm problematic here? (b) Why is standard LayerNorm insufficient? (c) Design an appropriate normalization approach.

Hint 1 -- Direction

Point clouds have two key properties: they are unordered sets (permutation-invariant), and they have varying sizes. BatchNorm is problematic because of variable batch sizes. LayerNorm normalizes over all features, but in a point cloud, different points may have very different feature distributions (e.g., points on a flat surface vs. an edge).

Hint 2 -- Insight

The best approach is GroupNorm or a per-point normalization. PointNet and its variants use a per-point MLP with BatchNorm (which works when batches are large) or GroupNorm. For varying batch sizes, GroupNorm is robust. The normalization should be applied per-point across the feature channels, similar to how LayerNorm works per-token in Transformers.

Hint 3 -- Full Solution + Rubric

(a) BatchNorm problems:

1. Batch size varies from 1 to 32 - BatchNorm statistics are unreliable at batch size 1-4
2. Point clouds have varying numbers of points - padding/masking complicates batch statistics
3. Train/test discrepancy when deployment uses batch size 1

(b) Standard LayerNorm limitations: LayerNorm normalizes over all features of a single example. For a point cloud processed by a per-point MLP (like PointNet), the "features" at each point are a $D$-dimensional vector. LayerNorm over all points and features would mix statistics across very different regions of the point cloud, washing out local structure. A flat wall and a sharp edge have different feature distributions - normalizing them together is inappropriate.

(c) Recommended approach: Per-point GroupNorm (or per-point LayerNorm)

Treat each point like a "token" in a Transformer. For point features $x \in \mathbb{R}^{B \times N \times D}$ where $N$ is the number of points:

$$\mu_{b,n} = \frac{1}{D} \sum_{d=1}^{D} x_{b,n,d}, \quad \sigma^2_{b,n} = \frac{1}{D} \sum_{d=1}^{D} (x_{b,n,d} - \mu_{b,n})^2$$

This is exactly LayerNorm applied per-point across the feature dimension. It is:

• Independent of batch size (works with 1-32)
• Independent of point cloud size (same computation for 1024 or 10000 points)
• Preserves per-point feature structure

If the feature dimension $D$ is small (e.g., 64), GroupNorm with $G = 4$ or $G = 8$ provides more stable statistics by grouping features.

Alternative: Adaptive instance normalization for style-conditioned point generation.

Scoring Rubric:

• Strong Hire: Correctly identifies both BatchNorm and LayerNorm limitations for this domain, proposes per-point LayerNorm or GroupNorm with clear justification, addresses variable batch size and variable point count, draws the analogy to Transformer token processing.
• Lean Hire: Proposes GroupNorm as a batch-independent alternative but does not address the per-point structure.
• No Hire: Proposes BatchNorm with large batch sizes as the solution, or does not recognize the variable-size challenge.

Problem 4: RMSNorm vs LayerNorm Ablation

Your team is deciding between RMSNorm and LayerNorm for a new 7B parameter LLM. Design an ablation study to compare them.

(a) What metrics would you track? (b) What should be held constant? (c) What is your expected outcome, and under what conditions would you choose LayerNorm?

Hint 1 -- Direction

A proper ablation changes only one variable (normalization type) and measures multiple aspects: training loss, wall-clock time, downstream task performance, and training stability. The expected outcome based on prior work (LLaMA, Mistral) is that RMSNorm matches or slightly outperforms LayerNorm while being faster.

Hint 2 -- Insight

The critical insight is that the comparison should be at equal training compute, not equal steps. Since RMSNorm is faster per step, running both for the same number of steps gives RMSNorm a wall-clock advantage and LayerNorm a FLOP advantage. The fairest comparison is equal wall-clock time OR equal FLOPs. Also, the difference may only manifest at scale - a 7B model ablation is meaningful because small models may not show the efficiency difference clearly.

Hint 3 -- Full Solution + Rubric

(a) Metrics to track:

Category	Metric	Why
Loss	Training loss at fixed token count	Core comparison
Loss	Validation perplexity at checkpoints	Generalization
Efficiency	Wall-clock time per training step	RMSNorm advantage
Efficiency	Tokens per second throughput	Real-world impact
Stability	Max gradient norm over training	Training stability
Stability	Loss spikes (frequency and magnitude)	Robustness
Quality	Downstream evals (MMLU, HumanEval, etc.)	Final quality
Memory	Peak GPU memory usage	Resource impact

(b) Held constant:

• Model architecture (layers, dimensions, heads, FFN, vocabulary)
• Training data and data ordering (use the same data seed)
• Optimizer and learning rate schedule
• Batch size and context length
• Hardware (same GPUs, same parallelism strategy)
• Random seed for initialization
• Total training tokens (NOT steps - since per-step cost differs)

(c) Expected outcome:

Based on prior work:

• Training loss: Near-identical at equal token count
• Throughput: RMSNorm 5-15% faster (measured in tokens/second)
• Training stability: Both stable with pre-norm and proper initialization
• Downstream quality: No statistically significant difference

When to choose LayerNorm instead:

1. If your model has unusually asymmetric activation distributions (mean centering provides benefit)
2. If you observe training instability with RMSNorm that does not occur with LayerNorm (rare but possible in specific architectures)
3. If you are fine-tuning a pretrained model that used LayerNorm (switching mid-training is not recommended)
4. If your deployment framework does not have an optimized RMSNorm kernel (in which case the speed benefit is lost)

Scoring Rubric:

• Strong Hire: Specifies comprehensive metrics including both quality and efficiency, identifies the "equal tokens vs equal steps" fairness issue, predicts the correct outcome based on prior work, provides nuanced conditions for choosing LayerNorm.
• Lean Hire: Proposes a reasonable ablation but misses the fairness issue or the efficiency metrics.
• No Hire: Cannot design a proper ablation (e.g., proposes comparing on different datasets or without controlling for compute).

Problem 5: Normalization from Scratch

Implement BatchNorm (training mode) and LayerNorm in NumPy. Show how their behaviors differ for a simple example.

Hint 1 -- Direction

Both follow the same formula: normalize to zero mean and unit variance, then apply learned scale and shift. The difference is which axes you compute the mean and variance over. BatchNorm: over B (and spatial if CNN). LayerNorm: over D (the feature dimension).

Hint 2 -- Insight

For a 2D input [B, D], BatchNorm computes D means (one per feature, across the batch) while LayerNorm computes B means (one per example, across features). Make sure your implementation handles the axes correctly with keepdims=True for proper broadcasting.

Hint 3 -- Full Solution + Rubric

import numpy as np

def batchnorm_train(x, gamma, beta, eps=1e-5):
    """
    x: [B, D] input activations
    gamma: [D] scale parameter
    beta: [D] shift parameter
    Returns: normalized output, (mean, var) for running stats
    """
    # Compute statistics across batch dimension (axis=0)
    mu = np.mean(x, axis=0, keepdims=True)       # [1, D]
    var = np.var(x, axis=0, keepdims=True)        # [1, D]

    # Normalize
    x_hat = (x - mu) / np.sqrt(var + eps)         # [B, D]

    # Scale and shift
    out = gamma * x_hat + beta                     # [B, D]

    return out, mu.squeeze(), var.squeeze()


def layernorm(x, gamma, beta, eps=1e-5):
    """
    x: [B, D] input activations
    gamma: [D] scale parameter
    beta: [D] shift parameter
    Returns: normalized output
    """
    # Compute statistics across feature dimension (axis=1)
    mu = np.mean(x, axis=1, keepdims=True)        # [B, 1]
    var = np.var(x, axis=1, keepdims=True)         # [B, 1]

    # Normalize
    x_hat = (x - mu) / np.sqrt(var + eps)         # [B, D]

    # Scale and shift
    out = gamma * x_hat + beta                     # [B, D]

    return out


def rmsnorm(x, gamma, eps=1e-5):
    """
    x: [B, D] input activations
    gamma: [D] scale parameter (no beta)
    Returns: normalized output
    """
    # Compute RMS across feature dimension
    rms = np.sqrt(np.mean(x ** 2, axis=1, keepdims=True) + eps)  # [B, 1]

    # Normalize and scale
    out = (x / rms) * gamma  # [B, D]

    return out


# Demonstration: how they differ
np.random.seed(42)
x = np.array([
    [1.0, 2.0, 3.0, 4.0],   # example 1: mean=2.5
    [10.0, 20.0, 30.0, 40.0] # example 2: mean=25
])  # Shape: [2, 4]

gamma = np.ones(4)
beta = np.zeros(4)

bn_out, _, _ = batchnorm_train(x, gamma, beta)
ln_out = layernorm(x, gamma, beta)

# BatchNorm: per-feature normalization across batch
# Feature 0: mean of [1, 10] = 5.5, so both examples shift relative to 5.5
# LayerNorm: per-example normalization across features
# Example 0: mean of [1, 2, 3, 4] = 2.5, normalized to zero mean

The key difference in output:

• BatchNorm: Features that are consistently large across the batch (e.g., feature 3 is always the largest) remain relatively large after normalization. But the batch-level variation is removed.
• LayerNorm: Each example is independently normalized. Example 2 (values 10-40) produces the same normalized output as example 1 (values 1-4) because the relative structure is the same.

Scoring Rubric:

• Strong Hire: Correct implementation of all three with proper axis handling, demonstrates the behavioral difference with a clear example, uses keepdims=True for broadcasting.
• Lean Hire: Correct BatchNorm and LayerNorm but wrong axis or missing keepdims.
• No Hire: Cannot implement either correctly or uses the wrong normalization axis.

Interview Cheat Sheet

Concept	Key Formula	One-Liner	Red Flag
BatchNorm	$\hat{x} = (x - \mu_B)/\sqrt{\sigma_B^2 + \epsilon}$	Normalize per feature across batch	"BN works with batch size 1"
BN train vs eval	Train: batch stats; Eval: running stats	Different behavior at test time	Not knowing about model.eval()
LayerNorm	$\hat{x} = (x - \mu_D)/\sqrt{\sigma_D^2 + \epsilon}$	Normalize per example across features	"LN depends on other batch examples"
RMSNorm	$\hat{x} = x / \text{RMS}(x)$	LayerNorm without mean centering	"RMSNorm loses important information"
GroupNorm	Normalize per group of channels	Between LN and IN	"GroupNorm is always better than BN"
InstanceNorm	Normalize per channel per example	Strips instance-level style	"IN is good for classification"
BN fails when	Small batch, sequences, variable length	Noisy statistics	"BN always works"
LN for Transformers	Batch-independent, position-independent	No cross-example dependency	"BN also works for Transformers"
Pre-norm placement	$X + \text{Sub}(\text{LN}(X))$	Clean residual gradient path	"Post-norm and pre-norm are equivalent"
$\gamma$, $\beta$ purpose	Restore representational power	Can learn identity transform	"We can remove them to save params"

Spaced Repetition Checkpoints

Day 0 -- Initial Learning

• Read this entire page
• Write the BatchNorm, LayerNorm, and RMSNorm formulas from memory
• Draw the normalization axes diagram for all four types (BN, LN, GN, IN)
• Complete the self-assessment

Day 3 -- First Recall

• Without notes, explain BatchNorm training vs inference behavior
• Give the "60-Second Answer" covering all normalization types, out loud, timed
• Write the decision framework: which norm for which architecture

Day 7 -- Connections

• Explain why BatchNorm fails for small batches AND for Transformers (two separate arguments)
• Do Practice Problem 1 (axis challenge) on paper without hints
• Compare RMSNorm vs LayerNorm: formula difference, speed difference, quality difference

Day 14 -- Application

• Do Practice Problem 2 (BatchNorm debugging) under timed conditions (10 minutes)
• Implement BatchNorm and LayerNorm from scratch in NumPy
• Explain to an imaginary interviewer the normalization choices in LLaMA 2

Day 21 -- Mock Interview

• Have someone ask: "Explain all normalization types, when to use each, and why LLMs use RMSNorm"
• Time yourself: full explanation in under 8 minutes
• Do all 5 practice problems in sequence under timed conditions (50 minutes total)

Key Takeaways

1. The choice of normalization is an architecture decision, not a hyperparameter. BatchNorm for CNNs with large batches, LayerNorm for Transformers, RMSNorm for modern LLMs, GroupNorm for small-batch vision - these are established best practices backed by both theory and massive-scale experiments.

2. BatchNorm's train/test discrepancy is a source of real production bugs. Forgetting model.eval(), distribution shift between training and deployment data, and augmentation mismatch are all common failure modes that every ML engineer should be able to diagnose.

3. RMSNorm is the new default for LLMs because mean centering adds cost without adding value. The speed improvement (10-15%) is meaningful at scale, and no quality loss has been observed. This is a case where simplifying a technique improved it.

4. Understanding which axes are normalized is the key to understanding all normalization types. BatchNorm averages over the batch. LayerNorm averages over features. GroupNorm averages over channel groups. InstanceNorm averages over spatial dimensions. Every interview answer about normalization should start from this axis specification.