Deep Learning, Visually #20

Watch the Short: youtube.com/shorts/yYxpjlvW69s

Why the numbers drift

A deep network is a chain of layers. Each layer takes a weighted sum of the activations from the layer before it, passes that sum through an activation function, and hands the result to the next layer.

Training changes the weights at every step. So every layer keeps seeing inputs that have shifted since the last step: a new average, a new spread. In a deep network, those shifts pile up. With a big learning rate, the numbers can drift so far that they blow up, and training dies.

A plain network at a big learning rate: the numbers run off the chart and the weights flash red.

What batch norm does

Batch normalization adds a reset after each layer. For each neuron, over the mini-batch, it does three things:

  1. It finds the average (μ_B) and the spread (σ_B) of the neuron's weighted sums.
  2. It slides them so their average sits at zero, and squeezes them so their spread is one.
  3. It applies a learned scale (γ) and a learned shift (β), so each layer can still pick its own range.

In one line:

y = γ · (x − μ_B) / √(σ_B² + ε) + β

It sits after the weighted sum and before the activation function: Linear → BatchNorm → ReLU. Because β already gives that layer a shift, the layer needs no bias. In conv nets, there is one average and one spread per channel, taken over the mini-batch and every position.

One neuron's 60 weighted sums: centered, squeezed to a set spread, then scaled and shifted by the learned γ and β.

Training and testing

During training, batch norm uses each mini-batch's own average and spread, and keeps running averages of them. At test time it uses those saved running averages. In PyTorch, model.train() and model.eval() switch between the two. The defaults are eps = 1e-5 and momentum = 0.1.

What it bought

Sergey Ioffe and Christian Szegedy, at Google, tested it on Inception, an image network:

  • Inception: learning rate 0.0015. 31.0M steps to its best accuracy, 72.2%.
  • BN-x5: learning rate 0.0075 (5×). 2.1M steps to reach 72.2%, about 14 times fewer. Best 73.0%.
  • BN-x30: learning rate 0.045 (30×). 2.7M steps to reach 72.2%. Best 74.8%.

BN-x5 also used changes that batch norm made possible: no dropout, weight decay cut 5×, a learning rate lowered 6× faster, no local response normalization, better shuffling and fewer photo distortions.

With sigmoid as the activation function, plain Inception never beat 0.1% accuracy, while BN-x5-Sigmoid reached 69.8%.

The paper's result: a learning rate five times bigger, and the same accuracy in fourteen times fewer steps.

Our toy test

We tried it on scikit-learn's 8×8 digits, with 1,500 training images and 297 test images. The network had 8 hidden layers of 64 ReLU neurons with a He start, trained by plain SGD on mini-batches of 60.

First, a sweep: three seeds, 1,000 steps each.

  • Plain: it trained at learning rates 0.03 and 0.1 (96–99% test accuracy). At 0.3, two of three seeds blew up. At 0.5, all three did. At 1.0, two blew up and the third stayed at 7%. At 2.0, all three blew up.
  • Batch norm: every learning rate up to 2.0 trained (96–99%).

Then a race at learning rate 0.5, five times the plain net's biggest safe one, same start, same mini-batches.

  • Plain: the training loss fell from 3.2 to 0.67 by step 100, then jumped to 25 at step 130 and hit NaN at step 135. Test accuracy went from 83% at step 100 to 8% at step 132. Other seeds blew up at steps 128 and 348: the exact step is chaotic.
  • Batch norm: in layers 2, 4, 6 and 8, the middle stayed at 0.06 or less and the spread stayed between 0.99 and 1.06. Training loss fell to 0.06 by step 125. Test accuracy reached 96% by step 100 and 98% by step 135.

The race, same learning rate: the plain network blows up and its loss shoots off the chart; with batch norm the cloud stays centered and the loss stays low.

Why does it work?

The paper's idea is "internal covariate shift": each layer's inputs keep changing during training, and batch norm reduces that. Santurkar et al. (2018) tested this and argue that batch norm mostly makes the loss landscape smoother. Bjorck et al. (2018) find that it mainly allows larger learning rates. This is still debated.

Limits and cousins

Small mini-batches hurt it, because their averages are noisy. At a batch size of 2, Group Norm had 10.6% lower error than batch norm (ResNet-50, ImageNet). Transformers use Layer Norm, which normalizes one example at a time.

Try it in PyTorch 

nn.Sequential(nn.Linear(256, 256, bias=False), nn.BatchNorm1d(256), nn.ReLU()) nn.Sequential(nn.Conv2d(64, 64, 3, padding=1, bias=False), nn.BatchNorm2d(64), nn.ReLU())

References

This article is part of Deep Learning, Visually by Logit Solutions AI: short animated explainers that cover one idea at a time, each with a written version like this one.