Batch Normalization

Batch normalization is a trainable transformation that standardizes each feature using statistics from a mini-batch. For feature

j

in a batch of

m

examples,

\mu_j = \frac{1}{m}\sum_{i=1}^{m}x_{ij}, \qquad \sigma_j^2 = \frac{1}{m}\sum_{i=1}^{m}(x_{ij}-\mu_j)^2,

\hat{x}_{ij} = \frac{x_{ij}-\mu_j}{\sqrt{\sigma_j^2+\epsilon}}, \qquad y_{ij} = \gamma_j\hat{x}_{ij}+\beta_j.

The normalization step fixes the batch mean and variance of each feature. The learned parameters

\gamma_j

and

\beta_j

then restore the scale and offset needed by the model. During training, the layer uses mini-batch statistics and updates running estimates. During inference, it uses the stored running estimates, so the output of an example no longer depends on the other examples in its batch. Batch normalization was introduced by Ioffe and Szegedy (2015) to accelerate optimization and permit larger learning rates. The paper attributed these improvements to reduced internal covariate shift, meaning changes in a layer’s input distribution as earlier layers are updated. Santurkar et al. (2018) later tested this explanation and found that distributional stability was not the source of the optimization benefit. The derivation below first develops the original internal-covariate-shift argument. The final section then evaluates that argument using the evidence reported by Santurkar et al.

import torch
import numpy as np
import matplotlib.pyplot as plt

torch.manual_seed(0)

# Standard normal CDF (used for the dead-ReLU fraction); torch.erf avoids a scipy dependency
def Phi(t):
    t = torch.as_tensor(t, dtype=torch.float32)
    return 0.5 * (1.0 + torch.erf(t / 2.0 ** 0.5))

Backpropagation through a single neuron

Consider a sigmoid neuron with input vector

x

, weight vector

w

, and bias

b

z = w^\top x + b, \qquad a = \sigma(z) = \frac{1}{1+e^{-z}}.

For a scalar loss

L

, the chain rule gives

\frac{\partial L}{\partial w_i} = \frac{\partial L}{\partial a}\frac{\partial a}{\partial z}\frac{\partial z}{\partial w_i} = \frac{\partial L}{\partial a}\,\sigma'(z)x_i.

Define the local error signal

\delta = \frac{\partial L}{\partial a}\,\sigma'(z).

The parameter gradients are therefore

\nabla_w L = \delta x, \qquad \frac{\partial L}{\partial b}=\delta.

The weight gradient is proportional to the input vector. Its magnitude depends on both the input scale and the local error signal. For a sigmoid,

\sigma'(z)=a(1-a)

approaches zero as

|z|

grows, so saturation can suppress the gradient even when the input is large. The bias gradient contains no direct factor of

x

# Single neuron: verify d a / d w_i = sigma'(z) * x_i
n = 8
x = torch.randn(n)
w = torch.randn(n, requires_grad=True)
b = torch.zeros(1, requires_grad=True)

z = w @ x + b
a = torch.sigmoid(z)
a.backward()                      # gradient of the scalar output a

sigma_prime = (a * (1 - a)).item()
ratio = w.grad / x                # should be constant = sigma'(z) for every i
print(f"sigma'(z)      = {sigma_prime:.6f}")
print(f"grad_i / x_i   = {ratio.tolist()}")

assert torch.allclose(ratio, torch.full((n,), sigma_prime), atol=1e-5), "grad not proportional to input"
print("OK: gradient is proportional to the input, slope = sigma'(z)")

sigma'(z)      = 0.233472
grad_i / x_i   = [0.2334718108177185, 0.2334717959165573, 0.2334718108177185, 0.2334718108177185, 0.2334717959165573, 0.2334717959165573, 0.2334718108177185, 0.2334718108177185]
OK: gradient is proportional to the input, slope = sigma'(z)

Backpropagation through a dense layer

A dense layer followed by a ReLU computes

z = Wx+b, \qquad a=\operatorname{ReLU}(z).

Let

g=\partial L/\partial a

be the upstream gradient and let

r=\mathbb{1}[z>0]

be the ReLU gate. The error signal at the pre-activation is

\delta = g\odot r.

Applying the chain rule to each element

W_{jk}

gives

\frac{\partial L}{\partial W_{jk}}=\delta_jx_k,

and the complete parameter gradients are

\frac{\partial L}{\partial W}=\delta x^\top, \qquad \frac{\partial L}{\partial b}=\delta.

The weight gradient is an outer product. Every row is the layer input scaled by one component of the error signal. For a mini-batch, the update is the sum or mean of these outer products over the examples. Consequently, changes in the distribution and scale of the layer inputs change the distribution and scale of its weight gradients. The bias gradient again has no direct multiplicative dependence on the input.

# FC + ReLU: verify dL/dW = outer(delta, x) and dL/db = delta
d_in, d_out = 6, 4
x = torch.randn(d_in)
W = torch.randn(d_out, d_in, requires_grad=True)
b = torch.randn(d_out, requires_grad=True)

z = W @ x + b
a = torch.relu(z)
g = torch.randn(d_out)            # arbitrary upstream gradient dL/da
a.backward(g)

m = (z > 0).float()               # ReLU gate
delta = g * m
expected_dW = torch.outer(delta, x)
expected_db = delta

assert torch.allclose(W.grad, expected_dW, atol=1e-6), "dL/dW mismatch"
assert torch.allclose(b.grad, expected_db, atol=1e-6), "dL/db mismatch"
print("OK: dL/dW = outer(delta, x), every row is x scaled; dL/db = delta")

OK: dL/dW = outer(delta, x), every row is x scaled; dL/db = delta

The internal covariate shift hypothesis

For layer

\ell

z^{(\ell)}=W^{(\ell)}h^{(\ell-1)}+b^{(\ell)}, \qquad \frac{\partial L}{\partial W^{(\ell)}} =\delta^{(\ell)}{h^{(\ell-1)}}^\top.

An update to an earlier layer changes

h^{(\ell-1)}

. The next optimization step therefore presents layer

\ell

with a different input distribution and a different distribution of weight gradients. Ioffe and Szegedy called this moving target internal covariate shift. This term is specific to hidden activations during training and should not be confused with covariate shift between training and test data. The original argument proceeds as follows:

Earlier layers continually change the coordinates received by later layers.
Later layers must adapt to those changing coordinates.
Standardizing intermediate features should make their scale more stable, improve gradient flow, and allow larger learning rates.

In the common Linear -> BatchNorm -> ReLU arrangement, batch normalization standardizes the pre-activation

z^{(\ell)}

before the nonlinearity. This directly controls the coordinates delivered to the ReLU and indirectly controls the input delivered to the next layer. It does not directly normalize

h^{(\ell-1)}

in the gradient formula for the current weight matrix, an important limitation of the simple argument. If a wide layer has approximately Gaussian pre-activations, the learned affine transformation gives

y \approx \mathcal{N}(\beta,\gamma^2).

After a ReLU, the probability of an output being clipped to zero is approximately

P(y\leq 0)=\Phi\!\left(-\frac{\beta}{|\gamma|}\right).

This calculation is useful, but it is not a definition of batch normalization. The layer fixes only the first two moments, not the shape of the distribution, and

\gamma

and

\beta

are learned.

N = 200000
zhat = torch.randn(N)             # BN output before the affine
gamma, beta = 1.0, 0.0
zt = gamma * zhat + beta
out = torch.relu(zt)

dead_measured = (out == 0).float().mean().item()
dead_predicted = Phi(-beta / gamma).item()
print(f"dead fraction measured  = {dead_measured:.4f}")
print(f"dead fraction predicted = Phi(-beta/gamma) = {dead_predicted:.4f}")
assert abs(dead_measured - dead_predicted) < 0.01

dead fraction measured  = 0.5019
dead fraction predicted = Phi(-beta/gamma) = 0.5000

Evidence against the internal covariate shift explanation

The optimization benefit of batch normalization is well established. The claim that the benefit is caused by reduced internal covariate shift is not. Santurkar et al. (2018) tested the causal claim directly. Their experiments found that the distributional stability of layer inputs had little connection to successful training. They instead observed that batch normalization made the loss and gradients vary more smoothly with the parameters. A smoother objective makes a gradient step more predictive and supports larger learning rates. This result rejects a strong version of the original story. A moving hidden distribution may occur during training, but reducing that movement is neither a sufficient explanation nor the measured source of batch normalization’s benefit. Santurkar et al. instead support an optimization account based on a smoother loss landscape and more stable, predictive gradients.

References

Ioffe, S., and Szegedy, C. (2015). Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. ICML.
Santurkar, S., Tsipras, D., Ilyas, A., and Madry, A. (2018). How Does Batch Normalization Help Optimization?. NeurIPS.
LeCun, Y., Bottou, L., Orr, G. B., and Muller, K.-R. (1998). Efficient BackProp.
Data preprocessing and whitening.

PyTorch reference

PyTorch class	Description
`nn.BatchNorm1d`	Applies Batch Normalization over a 2D or 3D input.
`nn.BatchNorm2d`	Applies Batch Normalization over a 4D input.
`nn.BatchNorm3d`	Applies Batch Normalization over a 5D input.

Edit this page on GitHub or file an issue.

Neural Networks

Backpropagation

Whitening

Normalization

Regularization

Hyperparameter Optimization

Transfer Learning

Batch Normalization

Backpropagation through a single neuron

Backpropagation through a dense layer

The internal covariate shift hypothesis

Evidence against the internal covariate shift explanation

References

PyTorch reference

​Backpropagation through a single neuron

​Backpropagation through a dense layer

​The internal covariate shift hypothesis

​Evidence against the internal covariate shift explanation

​References

​PyTorch reference

Backpropagation through a single neuron

Backpropagation through a dense layer

The internal covariate shift hypothesis

Evidence against the internal covariate shift explanation

References

PyTorch reference