Batch Normalization: Forward Pass Calculation

As discussed, internal covariate shift describes the changing distribution of layer inputs during training, which can hinder the learning process. Batch Normalization (BN) aims to mitigate this by normalizing the inputs to a layer for each mini-batch. Let's examine how this normalization is calculated during the forward pass.

The core idea is to take the activations arriving at the Batch Normalization layer for a given mini-batch and transform them so they have approximately zero mean and unit variance. This standardization happens independently for each feature or channel. However, simply forcing a zero mean and unit variance might limit the layer's representational power. Therefore, BN introduces two learnable parameters per feature, $\gamma$ (gamma) and $\beta$ (beta), that allow the network to scale and shift the normalized values. This means the network can learn the optimal scale and mean for the inputs to the next layer.

Consider a mini-batch $\mathcal{B} = \{x_1, x_2, ..., x_m\}$ of activations for a specific feature (e.g., the output of a single neuron for $m$ different examples in the mini-batch). The Batch Normalization forward pass involves these steps:

1. Calculate Mini-Batch Mean ($\mu_{\mathcal{B}}$): Compute the average value of the activations for this feature across the mini-batch.

   $$\mu_{\mathcal{B}} = \frac{1}{m} \sum_{i=1}^{m} x_i$$

2. Calculate Mini-Batch Variance ($\sigma^2_{\mathcal{B}}$): Compute the variance of the activations for this feature across the mini-batch.

   $$\sigma^2_{\mathcal{B}} = \frac{1}{m} \sum_{i=1}^{m} (x_i - \mu_{\mathcal{B}})^2$$

3. Normalize ($\hat{x}_i$): Normalize each activation $x_i$ in the mini-batch using the calculated mean and variance. A small constant $\epsilon$ (epsilon, e.g., $1e-5$) is added to the variance inside the square root for numerical stability, preventing division by zero if the variance is very small.

   $$\hat{x}_i = \frac{x_i - \mu_{\mathcal{B}}}{\sqrt{\sigma^2_{\mathcal{B}} + \epsilon}}$$

   After this step, the normalized activations $\hat{x}_i$ for the mini-batch will have a mean close to 0 and a variance close to 1.

4. Scale and Shift ($y_i$): Transform the normalized activation $\hat{x}_i$ using the learnable parameters $\gamma$ and $\beta$. These parameters are initialized (often to 1 and 0, respectively) and updated during backpropagation just like other network weights.

   $$y_i = \gamma \hat{x}_i + \beta$$

   The output $y_i$ is the final result of the Batch Normalization layer for the input activation $x_i$, and it's passed on to the subsequent layer (typically followed by a non-linear activation function).

Flow of the Batch Normalization forward pass for a single feature across a mini-batch. Inputs $x_i$ are used to compute mini-batch statistics ($\mu_{\mathcal{B}}, \sigma^2_{\mathcal{B}}$), which then normalize each $x_i$ to $\hat{x}_i$. Finally, learnable parameters $\gamma$ and $\beta$ scale and shift $\hat{x}_i$ to produce the output $y_i$.

It's important to remember that $\gamma$ and $\beta$ are learned per feature. If a fully connected layer has $N$ output neurons, there will be $N$ pairs of $(\gamma, \beta)$. If a convolutional layer has $C$ output channels, there will be $C$ pairs of $(\gamma, \beta)$ used to normalize across the batch, height, and width dimensions for each channel.

This process ensures that the inputs to the next layer have a stable distribution during training, controlled by the learned $\gamma$ and $\beta$ parameters, which significantly helps in stabilizing and accelerating the training process.