Batch Normalization Forward Pass

hardPRONeural NetworksDeep LearningCV: Deep Learning

Implement the forward pass of batch normalization using PyTorch tensors.

Batch normalization normalizes activations across the batch dimension:

$\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}$ $y_i = \gamma \hat{x}_i + \beta$

Where:

$\mu_B = \frac{1}{m}\sum_{i=1}^{m} x_i$ (batch mean)
$\sigma_B^2 = \frac{1}{m}\sum_{i=1}^{m} (x_i - \mu_B)^2$ (batch variance)
$\gamma, \beta$ are learnable parameters
$\epsilon$ is for numerical stability

In this simplified version, inference mode (training=False) also normalizes with the current batch's mean and variance; no running statistics need to be stored.

Return a float tensor with the same shape as x, with every value rounded to the nearest integer.

Example

Input

x = [[1, 2], [3, 4], [5, 6]]  # batch=3, features=2
gamma = [1, 1]
beta = [0, 0]
training = True

Output

[[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]]

Reasoning

Inputs: x is a batch of $m = 3$ samples with 2 features each, so each column is one feature measured across the batch. gamma = [1, 1] and beta = [0, 0] are the per-feature scale and shift; training is True.
Feature 1 (first column) holds $1, 3, 5$ : $\mu_B = 3$ and $\sigma_B^2 = \frac{(-2)^2 + 0^2 + 2^2}{3} = \frac{8}{3} \approx 2.667$ , so $\sqrt{\sigma_B^2 + \epsilon} \approx 1.633$ .
Normalizing gives $\frac{1-3}{1.633} \approx -1.2247$ , $\frac{3-3}{1.633} = 0$ and $\frac{5-3}{1.633} \approx 1.2247$ .
Feature 2 holds $2, 4, 6$ : every value is 1 higher, so $\mu_B = 4$ , the variance is still $\frac{8}{3}$ , and the normalized column is identical. Notice that batch norm erases a constant per-feature offset.
With $\gamma = 1$ and $\beta = 0$ , $y_i = \hat{x}_i$ . Rounding each value to the nearest integer turns $\pm 1.2247$ into $\pm 1$ , giving [[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]]: a $3 \times 2$ float tensor, the same shape as x, with each column centred on 0.

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

Learn this in

Deep LearningCV Study Plan Transformer ArchitectureLLM Study Plan

More from CV: Deep Learning

Implement Max Pooling with Indiceshard Depthwise Separable Convolutionmedium Softmax Cross-Entropy Gradientmedium Flatten Feature Mapeasy Max Poolingeasy

Concept primerPREVIEW

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 ( $\gamma$ and $\beta$ ) 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 $\mu_B$ and $\sigma_B^2$ 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.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

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.

solution.py

Python 3.13

Test results