Activation Functions & Gradient Problems

Without activation functions, a 100-layer network would collapse into one straight line, because stacking linear steps is still linear. Activations add the bends that let networks learn complex patterns. But a poor choice can make gradients vanish or explode, and training stops working. Let us see both, with real measurements.

01Why activations?

A layer computes x @ W + b, a straight-line (linear) operation. Two linear layers in a row equal one linear layer, so depth would be useless. An activation function applied after each layer adds a non-linear bend, and stacks of bends can fit almost any shape.

import numpy as np

z = np.array([-3.0, -1.0, 0.0, 1.0, 3.0])
acts = {
    "sigmoid":    1 / (1 + np.exp(-z)),
    "tanh":       np.tanh(z),
    "relu":       np.maximum(0, z),
    "leaky_relu": np.where(z > 0, z, 0.01 * z),
    "elu":        np.where(z > 0, z, np.exp(z) - 1),
}
print("z          ", z)
for name, a in acts.items():
    print(f"{name:11s}", a.round(3))

scores = np.array([2.0, 1.0, 0.1])                 # raw scores for 3 classes
soft = np.exp(scores) / np.exp(scores).sum()
print("\nsoftmax", scores, "->", soft.round(3), "sum =", soft.sum().round(3))

Softmax is special: it works on a whole vector, turning raw scores into probabilities that add up to 1. It is used on the output layer for multi-class problems, and it is how an LLM turns scores into next-token probabilities.

02Vanishing and exploding gradients

Backprop multiplies one local derivative per layer (Session 11). The derivative of sigmoid is at most 0.25. Multiply many numbers below 1 and the result shrinks towards zero: early layers get almost no gradient and stop learning. That is the vanishing gradient problem.

import numpy as np

# The derivative of sigmoid is at most 0.25 (at z = 0)
z = np.linspace(-6, 6, 1201); s = 1 / (1 + np.exp(-z))
print("max sigmoid'(z) =", (s * (1 - s)).max().round(3))

# Backprop multiplies one derivative per layer
for layers in [2, 5, 10, 20]:
    print(f"{layers:2d} layers: sigmoid gradient <= {0.25 ** layers:.1e}   relu gradient = {1.0 ** layers:.0f}")

# Exploding: weights a bit above 1 multiply up
print("\n1.5 multiplied 30 times =", f"{1.5 ** 30:,.0f}")

Now a real 20-layer PyTorch network. We measure the average gradient that reaches the first layer:

import torch, torch.nn as nn
torch.set_num_threads(1)                    # same results on every run

def first_layer_grad(act, he_init=False, depth=20):
    torch.manual_seed(0)
    layers = []
    for _ in range(depth):
        lin = nn.Linear(32, 32)
        if he_init: nn.init.kaiming_normal_(lin.weight, nonlinearity="relu"); nn.init.zeros_(lin.bias)
        layers += [lin, act()]
    net = nn.Sequential(*layers, nn.Linear(32, 1))
    net(torch.randn(64, 32)).pow(2).mean().backward()
    return net[0].weight.grad.abs().mean().item()

for name, act, he in [("Sigmoid", nn.Sigmoid, False), ("Tanh", nn.Tanh, False),
                      ("ReLU", nn.ReLU, False), ("ReLU + He init", nn.ReLU, True)]:
    print(f"{name:15s} 20 layers -> average gradient at layer 1: {first_layer_grad(act, he):.1e}")

Sigmoid’s gradient at layer 1 is around 10⁻²⁰: effectively zero, so the first layer never learns. Tanh and ReLU with default PyTorch initialisation also fade badly at this depth. Only ReLU with He initialisation keeps a healthy gradient (around 10⁻³).

The opposite problem, exploding gradients, happens when factors are above 1 (1.5³⁰ ≈ 191,751): the loss jumps to infinity or NaN. Common fixes are gradient clipping (torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)), a smaller learning rate and careful initialisation. RNNs suffer from both problems badly (Module 4).

03Choosing an activation

04Summary

• Activations add non-linearity; without them depth is useless.
• ReLU for hidden layers, sigmoid for yes/no outputs, softmax for multi-class outputs.
• Sigmoid’s derivative ≤ 0.25 → gradients vanish through many layers.
• ReLU + He initialisation keeps gradients alive (measured above).
• Exploding gradients: fix with clipping, a smaller learning rate and good initialisation.

05Quick quiz

06References

• Deep Sparse Rectifier Neural Networks (ReLU) (Glorot et al., 2011)
• Delving Deep into Rectifiers (He initialisation) (He et al., 2015)
• PyTorch: non-linear activations