🎓 A1Learner
🔍

PyTorch Why Normalization Layers

Normalization layers are layers that pin the scale of the values passing through a network. Without them, the scale drifts from layer to layer in a deep stack.

Measuring a set of values

Two numbers describe a set of activations. The mean says where the values sit, and the std says how spread out they are.

import torch

x = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
print(f"Mean: {x.mean():.4f}")
print(f"Std: {x.std():.4f}")

Scale passes through a layer

A Linear layer passes scale along. Ten times the input spread gives roughly ten times the output spread.

import torch

torch.manual_seed(0)
layer = torch.nn.Linear(4, 4)
x = torch.randn(8, 4)
print(f"Normal input std: {layer(x).std():.4f}")
print(f"Ten-times input std: {(layer(x * 10)).std():.4f}")

A deep stack, left alone

Each layer in a deep stack reads the last layer's output. Any shrink or growth compounds, layer after layer.

h = x
for layer in layers:
    h = torch.relu(layer(h))  # nothing resets the scale

Example

import torch

torch.manual_seed(0)
x = torch.randn(64, 16)
layers = [torch.nn.Linear(16, 16) for _ in range(10)]

print(f"Input: mean {x.mean():.4f}, std {x.std():.4f}")
h = x
for i, layer in enumerate(layers):
    h = torch.relu(layer(h))
    print(f"Layer {i + 1}: mean {h.mean():.4f}, std {h.std():.4f}")

Key points

  • Activation: a value a layer outputs, before the next layer reads it
  • Scale: how spread out a set of values is, measured here with std()
  • Std: the standard deviation, the typical distance of values from their mean
  • Use mean() and std() on activations to measure drift layer by layer
  • A plain stack can shrink activations toward zero, as the closing table shows
  • Normalization pins each layer's scale, so later layers receive values of a steady size

In this example, a batch of 64 rows enters a stack of 10 Linear layers with ReLU. The input starts with mean 0.0322 and std 1.0280. The first layer already cuts the std to 0.3545. By layer 3 the std is 0.0993, about a tenth of the input spread. The last layers sit near 0.08, and layer 9 dips to 0.0535. Nothing in the stack pulls the scale back. The activations shrink as they travel, and later layers work with tiny values.

Note

Drift is a scale problem, not a wrong-answer problem, and shallow networks may never show it.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does the closing example show happening to activation std through a plain deep stack?