Batch Normalization Forward Pass
Implement the forward pass of batch normalization using PyTorch tensors.
Batch normalization normalizes activations across the batch dimension:
x^i​=σB2​+ϵ​xi​−μB​​ yi​=γx^i​+β
Where:
- μB​=m1​∑i=1m​xi​ (batch mean)
- σB2​=m1​∑i=1m​(xi​−μB​)2 (batch variance)
- γ,β are learnable parameters
- ϵ is for numerical stability
Also track running mean/variance for inference mode.
Example:
x = [[1, 2], [3, 4], [5, 6]] # batch=3, features=2 gamma = [1, 1] beta = [0, 0] training = True
[[-1.22, -1.22], [0, 0], [1.22, 1.22]]
Feature 1: values [1,3,5], mean=3, std=1.63 Normalized: [-1.22, 0, 1.22]
Feature 2: values [2,4,6], mean=4, std=1.63 Normalized: [-1.22, 0, 1.22]
With gamma=1, beta=0: output equals normalized values.
Constraints:
- x: Input tensor of shape (batch_size, features)
- gamma, beta: Learnable parameters (features,)
- training: Boolean for train vs inference mode
- Return: Normalized output tensor
Batch Normalization Forward Pass: Background & Strategy
Background Knowledge
What is Batch Normalization?
Batch normalization (BN) is a technique that normalizes the inputs to each layer during training by centering and scaling activations across the batch dimension. The core insight is that by standardizing intermediate representations, you reduce internal covariate shift—the phenomenon where the distribution of layer inputs changes as earlier layers update their weights. This instability can slow training and require careful learning rate tuning. BN addresses this by forcing activations to have consistent statistics (mean near 0, variance near 1) within each mini-batch.
Why Does It Work?
The normalization step (dividing by standard deviation) acts as a form of adaptive learning rate scaling: features with larger magnitudes are scaled down more aggressively, while smaller features receive gentler updates. The learnable affine parameters (γ and β) allow the network to recover expressiveness after normalization—the model can learn to "undo" the normalization if beneficial. Additionally, BN introduces a regularization effect because statistics are computed per mini-batch rather than globally, adding noise that improves generalization.
Training vs. Inference Distinction
A critical aspect of BN is the dual-mode behavior: during training, you normalize using batch statistics (computed from the current mini-batch), but during inference, you must use running statistics (exponential moving averages accumulated during training). This separation exists because at inference time, you may have a single sample or a different batch composition, making batch statistics unreliable.
Algorithm/Approach
The forward pass has two distinct paths:
-
Training Mode: Compute batch mean and variance from the current mini-batch, normalize using these statistics, apply the affine transformation, and update running statistics for later use.
-
Inference Mode: Use the pre-computed running mean and variance (accumulated during training) to normalize, then apply the affine transformation.
Both paths share the same normalization and scaling formulas, but differ in which statistics they use.
Step-by-Step Strategy
Step 1: Compute Batch Statistics (Training Only)
Calculate μB​ and σB2​ by reducing across the batch dimension (dimension 0). For a 4D tensor (batch, channels, height, width), you typically normalize per-channel, so reduce over dimensions [0, 2, 3], keeping the channel dimension intact.
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.