BatchNorm: Batch Normalization in 2D

Batch Normalization (BatchNorm2d) standardizes 2D feature maps channel-wise across batch and spatial dimensions to zero mean and unit variance. It then applies a learnable affine transformation with scale ($\gamma$) and shift ($\beta$) parameters to preserve representational expressiveness.

[!TIP] Theory & Curriculum Link: For the first-principles derivation of internal covariate shift, variance scaling, and the comparison with LayerNorm and RMSNorm, see Curriculum: Normalization & Regularization.

1. Quick Example: Training vs. Evaluation

import numpy as np
import netcl.autograd as ag
import netcl.nn as nn
from netcl.core.device import manager
from netcl.core.tensor import Tensor

q = manager.default("auto").queue

# 1. Initialize layer for 32 feature channels
bn = nn.BatchNorm2d(q, num_features=32)

# 2. 4D image batch: (Batch=4, Channels=32, Height=16, Width=16)
x = Tensor.from_host(q, np.random.randn(4, 32, 16, 16).astype(np.float32))

# 3. Training mode: updates running_mean and running_var
bn.train()
with ag.Tape() as tape:
    y = bn(ag.tensor(x))

print(f"Training output shape: {y.value.shape}")

# 4. Evaluation mode: uses frozen running statistics
bn.eval()
y_eval = bn(ag.tensor(x))
print("Inference step evaluated deterministically.")

2. Mathematical Formulation

For an input tensor $\mathbf{X} \in \mathbb{R}^{N \times C \times H \times W}$, normalization is computed independently for each channel $c \in \{1, \dots, C\}$ across all $M = N \cdot H \cdot W$ elements.

Training Mode

1. Channel Mean:

$$\mu_c = \frac{1}{M} \sum_{n=1}^N \sum_{h=1}^H \sum_{w=1}^W x_{n, c, h, w}$$

2. Channel Variance:

$$\sigma_c^2 = \frac{1}{M} \sum_{n=1}^N \sum_{h=1}^H \sum_{w=1}^W (x_{n, c, h, w} - \mu_c)^2$$

3. Standardization:

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

4. Learnable Affine Transformation:

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

Where $\gamma_c \in \mathbb{R}$ (scale) and $\beta_c \in \mathbb{R}$ (bias) are learnable parameters.

5. Running Statistics Tracking (Exponential Moving Average):

$$\mu_{\text{run}, c} \leftarrow (1 - m) \mu_{\text{run}, c} + m \cdot \mu_c$$

$$\sigma^2_{\text{run}, c} \leftarrow (1 - m) \sigma^2_{\text{run}, c} + m \cdot \sigma_c^2$$

Where $m$ is the momentum parameter (default: $0.1$).

Inference Mode (bn.eval())

During evaluation, batch statistics are bypassed because predictions for individual samples must be deterministic and independent of other samples in the batch:

$$y_{n, c, h, w} = \gamma_c \frac{x_{n, c, h, w} - \mu_{\text{run}, c}}{\sqrt{\sigma^2_{\text{run}, c} + \epsilon}} + \beta_c$$

3. Practical Usage: Conv-BN-ReLU Block

When a convolutional layer is directly followed by BatchNorm, disable bias in the convolution (bias=False). The shift parameter $\beta$ in BatchNorm already serves as the learnable offset, saving memory and compute.

import numpy as np
import netcl.autograd as ag
import netcl.nn as nn
from netcl.core.device import manager
from netcl.core.tensor import Tensor

q = manager.default("auto").queue

class ConvBlock(nn.Module):
    def __init__(self, queue, in_channels, out_channels):
        super().__init__()
        # No bias needed in Conv2d when followed by BatchNorm
        self.conv = nn.Conv2d(queue, in_channels, out_channels, kernel_size=3, padding=1, bias=False)
        self.bn = nn.BatchNorm2d(queue, num_features=out_channels)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.bn(self.conv(x)))

block = ConvBlock(q, in_channels=16, out_channels=32)
x_img = Tensor.from_host(q, np.random.randn(2, 16, 28, 28).astype(np.float32))

with ag.Tape() as tape:
    out = block(ag.tensor(x_img))

print(f"Block output shape: {out.value.shape}")  # (2, 32, 28, 28)

4. Parameter Reference

Related Documentation

• Curriculum: Normalization & Regularization: Deep dive into training dynamics.
• NN API Reference: Complete method signatures for all neural network layers.
• Concepts: ResNet: Using BatchNorm inside residual blocks.