Vanishing gradients, ReLU vs ELU and BatchNorm

Arjun’s first pastry-photo model was ambitious. If two hidden layers had learned rush hour, surely twenty layers could learn the difference between an éclair and a cinnamon roll. He stacked nineteen layers of 128 neurons with sigmoid activations between them, the activation he remembered best from Module 27’s probability lesson, and pressed run.

epoch 1  loss 2.307  accuracy 9.8%
epoch 2  loss 2.309  accuracy 11.3%
epoch 3  loss 2.307  accuracy 11.3%
epoch 5  loss 2.304  accuracy 9.8%

Ten categories of pastry. 2.303 is exactly $\ln (10)$, the loss you get from saying “10% each” to every photo. Accuracy hovered around one in ten. The network had learned nothing at all.

He checked the usual suspects. The learning rate was sensible. zero_grad() was there. The data loaded correctly. He doubled the epochs. Nothing.

“Don’t burn GPU hours on photos to debug this,” Priya said. “Run the same network on last week’s MNIST digits. Ten classes, tiny images, trains in seconds.”

He did. Loss 2.304, accuracy 9.8%. Same disease, smaller patient.

A message passed down a line

“Ever played Chinese whispers?” Priya asked. “Twenty people in a line. The first person whispers a sentence, each person passes on what they heard.”

“By the end it’s nonsense.”

“Now play it backwards, with a twist. Each person passes on only a quarter of what they heard, at a quarter of the volume. What does person one hear?”

Arjun did the arithmetic. “A quarter, of a quarter, of a quarter… nineteen times. Basically silence.”

“That’s your network. The loss is shouting at the last layer, and by the time the message gets back to the first layer, it’s a whisper so faint that the first layer’s weights don’t move. Nothing at the start of the network learns, so nothing useful reaches the end.”

Where the analogy breaks: in a real network, the “volume” passed back at each layer is not fixed. It depends on the activation function’s slope, on the size of the weights and on the input. That is good news, because all three are things you can change.

Following the gradient back through twenty layers

Remember the chain rule from Module 27: to find how a weight in layer 1 affected the loss, backpropagation multiplies the local slopes of every layer between it and the output. If each layer multiplies the signal by something less than 1, the product shrinks exponentially.

Look at sigmoid’s slope. Sigmoid squashes any input into 0 to 1, and its derivative is $\sigma (x)(1-\sigma (x))$. That is largest when $\sigma (x)=0.5$, where it equals $0.5\times 0.5=0.25$. Everywhere else it is smaller. So even in the best case, every sigmoid layer cuts the gradient to a quarter:

$$0.25^{10}\approx 0.000000950.25^{19}\approx 0.0000000000036$$

Arjun measured the real thing: the size of the gradient on each layer’s weights, after one batch, before any training:

Layer	Gradient size (sigmoid network)
Layer 20 (output)	0.78
Layer 10	0.0000000018
Layer 1 (input)	0.00000000000000024

From the output to the input, the gradient shrinks by a factor of about three thousand million million. With a learning rate of 0.01, layer 1’s weights move by about $10^{-18}$ per step. A float32 number cannot even represent a change that small relative to a weight of size 0.05; the update is literally rounded away. The first layers are frozen at their random starting values, so everything downstream receives random noise, and the network outputs “10% each” forever.

This is the vanishing gradient problem. Its twin, where the multipliers are bigger than 1 and the product grows exponentially, is exploding gradients. Both come from the same multiplication.

Fix one: activations that do not shrink the signal

The ReLU activation, $\max (0,x)$, has slope 1 for positive inputs and 0 for negative ones. Unlike sigmoid, it never squashes a positive signal. That is why ReLU replaced sigmoid inside networks around 2012, and why it remains the default.

ReLU has one flaw, the dead ReLU. If a neuron’s input is negative for every example, its output is always 0 and its slope is always 0, so it receives no gradient and can never recover. A big learning-rate step can kill a large share of a network’s neurons at once. You met this in Module 27 when the learning rate was 5.

ELU (exponential linear unit) keeps ReLU’s behaviour for positive inputs but replaces the hard zero with a smooth curve for negatives:

$$\text{ELU}(x)=\{\begin{array}{ll}x& x>0\\ \alpha (e^x-1)& x\le 0\end{array}$$

With the usual $\alpha =1$, an input of −1 gives −0.63, an input of −3 gives −0.95, and very negative inputs level off at −1. Two benefits follow. The slope for negative inputs is $e^x$, small but never exactly zero, so neurons cannot die completely. And because outputs can be negative, the average activation sits closer to zero, which tends to make the next layer’s job easier. The cost: computing an exponential is slower than a comparison, which matters a little on big models.

	ReLU	ELU	Sigmoid
Output for −2	0	−0.86	0.12
Output for 2	2	2	0.88
Slope for positives	1	1	at most 0.25
Slope for negatives	0 (can die)	$e^x$, never 0	small
Speed	fastest	slower	slower

You will also meet Leaky ReLU (a small fixed slope, like 0.01, for negatives) and GELU (a smooth ReLU used in transformers, which you will build in Module 30). They are all answers to the same question: how do you bend the signal without killing the gradient?

Fix two: start the weights at the right size

The pause above showed that activations are only half the story. The other half is how big the weights are when training begins: weight initialisation.

Priya’s analogy: a chain of twenty amplifiers between a microphone and a speaker. If each amplifier is set to 0.9× volume, the speaker is silent. If each is set to 1.1×, it screeches with feedback. You want each one at almost exactly 1×.

For ReLU layers, the standard recipe is Kaiming (or He) initialisation: draw weights with standard deviation $\sqrt{2/n_{\text{in}}}$, where $n_{\text{in}}$ is the number of inputs to the layer. The 2 compensates for ReLU zeroing roughly half the values. With 128 inputs, that is a standard deviation of $\sqrt{2/128}=0.125$.

from torch import nn

def init_kaiming(model: nn.Module) -> None:
    for layer in model.modules():
        if isinstance(layer, nn.Linear):
            nn.init.kaiming_normal_(layer.weight, nonlinearity="relu")
            nn.init.zeros_(layer.bias)

With Kaiming initialisation, the same 20-layer ReLU network’s activations stayed between 0.86 and 0.58 all the way through, and layer 1’s gradient was 0.71 instead of $10^{-7}$. It trained.

How sensitive is this? Arjun tried a standard deviation of 0.15 instead of 0.125 in every hidden layer, just 20% too big. Each layer now multiplied the signal by about 1.2. Over 18 layers, $1.2^{18}\approx 27$. The activations grew from 2.5 at layer 1 to 45 at layer 19, the starting loss was 38 instead of 2.3, the gradients on the last layer were nearly 500, and after the first few updates every number in the network was nan. Exploding gradients, from a 20% mistake.

Fix three: re-standardise between layers

Getting the initialisation exactly right is fragile, and the balance drifts once training starts changing the weights. Batch normalisation (BatchNorm) takes a more robust approach: after a layer, standardise its outputs, for every batch, so the next layer always receives values with mean 0 and standard deviation 1.

You already know this move. In Module 27 you standardised the inputs so no feature drowned out the others. BatchNorm does the same thing between every pair of layers. It is like recalibrating the scales at every station on the production line, not just at the start.

One neuron, one batch

Take one neuron’s outputs for a batch of four photos: $[2.0,4.0,6.0,8.0]$.

1. Batch mean: $\mu =5.0$.
2. Batch variance: $\frac{1}{4}\left((-3{)}^2+(-1{)}^2+1^2+3^2\right)=5.0$, so the standard deviation is $\sqrt{5}\approx 2.24$.
3. Standardise: $(x-5)/2.24$ gives $[-1.34,-0.45,0.45,1.34]$.
4. Scale and shift: multiply by a learned $\gamma$ and add a learned $\beta$. They start at 1 and 0, so nothing changes at first, but training can adjust them if this neuron works better with a different spread or centre.

$$\hat{x}=\frac{x-{\mu }_{\text{batch}}}{\sqrt{{\sigma }_{\text{batch}}^2+ϵ}}y=\gamma \hat{x}+\beta$$

The tiny $ϵ$ (epsilon) prevents division by zero. PyTorch’s nn.BatchNorm1d(1) on that batch prints [-1.3416, -0.4472, 0.4472, 1.3416].

Because each layer’s output is rescaled to a standard size, the signal cannot steadily shrink or grow from layer to layer, in either direction. That is why BatchNorm rescued even the badly initialised networks:

def deep_mlp(activation, depth=20, width=128, batchnorm=False):
    layers, d = [], 784                         # MNIST: 28 x 28 = 784 pixels
    for _ in range(depth - 1):
        layers.append(nn.Linear(d, width))
        if batchnorm:
            layers.append(nn.BatchNorm1d(width))   # standardise each neuron over the batch
        layers.append(activation())
        d = width
    layers.append(nn.Linear(d, 10))
    return nn.Sequential(*layers)

Those running averages update gradually. After the example batch above, PyTorch’s running mean had moved from its starting value of 0 to only 0.5 (it blends in 10% of each new batch by default). It takes hundreds of batches to settle, which is one more reason to evaluate only after real training.

Fix four: put a ceiling on the step

Sometimes you cannot prevent a large gradient: a rare, unusual batch, or a recurrent network multiplying through hundreds of time steps (Module 29). Gradient clipping sets a maximum: if the total size of the gradient exceeds a limit, scale it down to the limit, keeping its direction.

loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)   # between backward() and step()
optimizer.step()

Arjun’s exploding network, with the 20%-too-big initialisation, produced nan in the first epoch. With that single line added, it trained to 86% accuracy in five epochs. Clipping does not fix the root cause (the initialisation was still wrong, and it learned more slowly than the healthy networks), but it stops one bad step from destroying a run.

Everything, side by side

Arjun ran every combination on MNIST: 20 layers of 128 neurons, SGD with momentum, learning rate 0.01, 20,000 training images, five epochs, scored on the 10,000 test images.

Network	Gradient at layer 1, before training	Test accuracy after 5 epochs
Sigmoid	0.00000000000000024	9.8% (stuck)
ReLU, default init	0.000000095	11.3% (stuck)
ReLU, init 20% too large	40, then nan	9.8% (diverged)
ReLU, init 20% too large, clipped		86.3%
ReLU, Kaiming init	0.71	94.3%
ELU, Kaiming init	5.4	95.6%
Sigmoid + BatchNorm	0.97	90.9%
ReLU, default init + BatchNorm	38	95.6%

Three lessons in one table. The activation matters, but on its own it is not enough. The initialisation matters, and a small error compounds through depth. And BatchNorm makes the network forgiving: even sigmoid trains, and even the default initialisation works.

Where these fixes break

BatchNorm needs decent batches. With a batch of 2 to 8, common when each image is huge, batch statistics are too noisy. Alternatives normalise within each example instead: LayerNorm (used in every transformer, Module 30) and GroupNorm (common in vision with small batches).

BatchNorm’s two modes bite. Training uses batch statistics, evaluation uses running averages. Forgetting model.eval() gives batch-dependent predictions; evaluating a model that has only trained for a few batches uses running averages that have not settled. And if training and production images differ (brighter lighting, a new camera), the running averages are wrong for production.

Clipping can hide a real problem. If gradients are clipped on nearly every step, the learning rate or initialisation is wrong. Log how often clipping triggers.

Depth beyond about twenty layers needs more. Even with all four fixes, plain stacks get harder to train past a few dozen layers. The breakthrough that made 100-plus-layer networks routine was the residual connection: each block adds its input to its output, $y=x+f(x)$, giving gradients a direct path backwards that skips the multiplication. Every transformer uses them; you will build them in Module 30.

The deepest question is whether you need depth at all. Twenty dense layers on flattened pixels is the wrong shape of network for images, however well it trains. That is the next lesson.

Try it yourself

1. Build the 20-layer sigmoid network on MNIST and print each linear layer’s gradient size after one backward pass. Plot it on a log scale. Is the decay roughly a straight line? What does a straight line on a log scale mean?
2. Count dead ReLUs: for the ReLU + Kaiming network, pass 1,000 images through and, for each hidden layer, count neurons whose output is zero for every image. Repeat after training with a learning rate of 0.5. What changes?
3. Move nn.BatchNorm1d to after the activation instead of before it. Does it still train? (Both placements are used in practice.)

Show the answer

1. Yes, roughly straight: a straight line on a log scale means the gradient shrinks by about the same factor at every layer, which is exactly the repeated multiplication behind vanishing gradients.
2. With a sensible learning rate, few or no neurons are dead. With a learning rate of 0.5, large updates push many neurons permanently negative, and you will typically find a substantial share dead in some layers, sometimes whole layers. ELU or a lower learning rate prevents it.
3. Yes. The original BatchNorm paper placed it before the activation; many later models place it after. Differences are usually small. Pick one and be consistent.

Common mistakes

• Using sigmoid or tanh between hidden layers of a deep network. Keep them for outputs that need a range (probabilities, values between −1 and 1). Use ReLU, ELU or GELU inside.
• Assuming the default initialisation suits a very deep stack. For deep ReLU networks without normalisation, use Kaiming initialisation explicitly.
• Forgetting model.eval() with BatchNorm. Production predictions become dependent on whatever else is in the batch.
• Clipping as a substitute for fixing the cause. Clip to survive rare spikes, then fix the learning rate or initialisation.
• Debugging on the big dataset. Reproduce training problems on a small, fast dataset first, as Priya insisted.

Recap

• Backpropagation multiplies slopes and weights layer by layer, so gradients can vanish or explode exponentially with depth.
• Sigmoid’s slope is at most 0.25, so deep sigmoid stacks starve early layers; ReLU and ELU pass positive signals at full strength, and ELU never fully dies.
• Kaiming initialisation, standard deviation $\sqrt{2/n_{\text{in}}}$, keeps signals steady through ReLU layers; 20% too big can explode.
• BatchNorm re-standardises each layer’s outputs per batch, uses running averages in eval() mode, and makes deep networks forgiving.
• Gradient clipping caps the step size and rescues runs from rare spikes, but does not fix root causes.

Arjun switched the pastry network to ReLU with BatchNorm, and the loss finally fell: 2.30, 1.9, 1.6. He showed Meera the curve going down.

“So the first layers were deaf,” she said, “and now they can hear.”

“Exactly. The correction from the end of the network was too faint to reach the beginning. Now each layer re-levels the signal, so the message gets all the way back.”

Priya let him enjoy it for a minute. Then she pulled up one of his training photos, a croissant in the top-left corner of the frame, and the same croissant shifted to the bottom-right. “Your network sees these as two completely different lists of 12,288 numbers. It has to learn ‘croissant’ separately for every position in the picture. Tomorrow we use a network that already knows what a picture is.”

Go deeper

• PyTorch tutorials for torch.nn.init, nn.BatchNorm1d and nn.BatchNorm2d, and gradient clipping.
• Karpathy’s video above builds backpropagation by hand, which makes the multiplication behind vanishing gradients very concrete.