5.2.5 Activation Functions

Why do we need activation functions?

In Tensors, Matrices, and Related Mathematical Functions, we saw that a layer of a network is mostly a matrix multiplication plus a bias: $z = Wx + b$. That is a linear operation. If we stack two such layers, we get $W_2(W_1x + b_1) + b_2$, which is just another linear function of $x$. However many layers we stack, the whole network could still only draw straight lines and flat planes. Multi-Layer Perceptrons shows this collapse happening in code.

An activation function fixes this. It is a small non-linear function applied to every value a layer produces, before the next layer sees it:

$$a = \phi(z) = \phi(Wx + b).$$

With a non-linear $\phi$ between the layers, stacking layers really does add power: the network can bend its decision boundaries and fit curved relationships. This page walks through the activation functions you will actually meet, how their gradients behave (which matters for training), and which one to pick in each part of a network.

You do not need to memorise every formula here. What you should take away is the shape of each function, what goes wrong with the older ones, and the simple rule for choosing.

What you need to know

By the end of this module, you should be comfortable with these ideas:

That is enough to read and design the networks in this guide.

Learn it properly

This page gives you a guided tour with small experiments. A proper reading adds the details and practice that make the ideas stick, so pair the two.

Pick one main course

The topic is small, so one good chapter is enough. Read it, run its code, and then come back here for the experiments.

Our pick: Dive into Deep Learning, section 5.1

Read the Activation Functions part of section 5.1 in Dive into Deep Learning. It gives the formula for ReLU, sigmoid and tanh, plots each function and its derivative with PyTorch, and is written for exactly this stage. The rest of section 5.1, on why hidden layers help, pairs well with Multi-Layer Perceptrons.

The exercises at the end of the section are worth trying, but the ones that ask for proofs are optional at this level.

If you prefer something else

• Course notes: Stanford CS231n's Commonly used activation functions covers sigmoid, tanh, ReLU and Leaky ReLU, with blunt practical advice on which to use.
• A gentle start: Google's Machine Learning Crash Course page on activation functions is visual and short.
• A video: StatQuest's ReLU In Action shows, step by step, how ReLU units bend a straight line into a curve.
• More depth on vanishing gradients: section 5.4 of Dive into Deep Learning, Vanishing Gradients. This goes further than you need for now.

A quick tour

The rest of this page shows each activation function in a line or two of PyTorch, and then two small experiments that explain why the modern choices won. Run the snippets yourself; they are short.

Where the activation goes

Let us start with one layer and one activation. The layer turns 3 input numbers into 4 values, and the activation is applied to each of those 4 values separately:

import torch
from torch import nn

torch.manual_seed(0)
layer = nn.Linear(3, 4)
x = torch.tensor([1.0, -2.0, 0.5])

z = layer(x)            # the linear part: Wx + b
a = torch.relu(z)       # the activation, value by value
print(z)
print(a)

tensor([-1.4129, -0.2851, -1.1906,  0.4664], grad_fn=<ViewBackward0>)
tensor([0.0000, 0.0000, 0.0000, 0.4664], grad_fn=<ReluBackward0>)

ReLU kept the one positive value and replaced the three negative ones with 0. The grad_fn part only tells us that PyTorch is tracking these values for backpropagation; you can ignore it for now.

In a whole network, we put an activation after every hidden layer. This is the pattern you will see in almost every small PyTorch model:

model = nn.Sequential(
    nn.Linear(3, 16),
    nn.ReLU(),          # activation after the first layer
    nn.Linear(16, 16),
    nn.ReLU(),          # activation after the second layer
    nn.Linear(16, 1),   # no activation after the last layer
)

Notice that the last layer has no activation. We will see why in Which activation where.

The five functions

Now we can meet the functions themselves. We will feed the same five inputs to each one:

z = torch.tensor([-2.0, -0.5, 0.0, 0.5, 2.0])

Sigmoid squashes any number into the range (0, 1):

$$\sigma(z) = \frac{1}{1 + e^{-z}}.$$

torch.sigmoid(z)

tensor([0.1192, 0.3775, 0.5000, 0.6225, 0.8808])

You met it already in logistic regression, where it turns a score into a probability. That is still its main job today.

Tanh (the hyperbolic tangent) has the same S-shape, but squashes into (−1, 1) and is centred on 0:

$$\tanh(z) = \frac{e^{z} - e^{-z}}{e^{z} + e^{-z}}.$$

torch.tanh(z)

tensor([-0.9640, -0.4621,  0.0000,  0.4621,  0.9640])

ReLU (rectified linear unit) keeps positive values and sets negative ones to zero:

$$\operatorname{ReLU}(z) = \max(0, z).$$

torch.relu(z)

tensor([0.0000, 0.0000, 0.0000, 0.5000, 2.0000])

It is almost free to compute, which matters when a network has millions of units.

Leaky ReLU multiplies negative values by a small slope $\alpha$ instead of setting them to 0. PyTorch's default slope is 0.01:

$$\operatorname{LeakyReLU}(z) = \begin{cases} z & z > 0 \\ \alpha z & z \le 0 \end{cases}$$

import torch.nn.functional as F

F.leaky_relu(z, negative_slope=0.01)

tensor([-0.0200, -0.0050,  0.0000,  0.5000,  2.0000])

GELU (Gaussian error linear unit) is a smooth version of ReLU. It multiplies $z$ by the probability that a standard normal variable is below $z$, written $\Phi(z)$:

$$\operatorname{GELU}(z) = z\,\Phi(z).$$

F.gelu(z)

tensor([-0.0455, -0.1543,  0.0000,  0.3457,  1.9545])

For large positive inputs it behaves like ReLU, and for large negative ones it goes to 0, but the corner at 0 is rounded off. Transformer models, such as BERT and the GPT family, use it. You do not need the exact formula of $\Phi$ at this level.

Here are all five together:

Their gradients

Training works by following gradients backwards through every layer (Backpropagation explains how). Each time the gradient passes back through an activation, it gets multiplied by that activation's derivative. So the derivative decides whether a useful signal survives the trip:

Two facts in this picture explain almost everything about choosing activations. Sigmoid's derivative is small everywhere, and sigmoid and tanh both go flat (i) for large inputs. ReLU's derivative, on the other hand, is exactly 1 whenever the unit is active. The next two experiments show what that means in practice.

Vanishing gradients

Let us see what a small derivative does over many layers. We pass one number through ten sigmoids in a row, then ask PyTorch how much the final output changes when the starting number changes:

start = torch.tensor(0.5, requires_grad=True)
value = start
for layer in range(10):
    value = torch.sigmoid(value)

value.backward()   # work out the gradient
print(start.grad)

tensor(3.4831e-07)

The gradient is about 0.00000035. Each sigmoid multiplied it by at most 0.25, and $0.25^{10}$ is already less than one in a million. Now the same with ReLU:

start = torch.tensor(0.5, requires_grad=True)
value = start
for layer in range(10):
    value = torch.relu(value)

value.backward()
print(start.grad)

tensor(1.)

With ReLU the gradient arrives untouched. In a real network the weights also scale the gradient, but the effect of the activation is the same: with sigmoid in every hidden layer, the early layers of a deep network receive almost no signal and barely learn. This is the vanishing gradient problem, and it is the main reason ReLU replaced sigmoid and tanh in hidden layers. (Good weight initialisation also matters here; it is a later, Silver-level topic.)

Dead ReLUs

ReLU has its own weakness. If a unit's input is negative for every example, it always outputs 0, its derivative is always 0, and its weights never change again. The unit is "dead". Let us build one on purpose, with a large negative bias:

weight = torch.tensor(1.0, requires_grad=True)
bias = torch.tensor(-5.0, requires_grad=True)
inputs = torch.tensor([-1.0, 0.0, 1.0, 2.0])

outputs = torch.relu(inputs * weight + bias)
outputs.sum().backward()
print(outputs)
print(weight.grad, bias.grad)

tensor([0., 0., 0., 0.], grad_fn=<ReluBackward0>)
tensor(0.) tensor(0.)

Every output is 0 and both gradients are 0, so gradient descent has nothing to work with. Leaky ReLU keeps a small slope for negative inputs, so the same unit still gets a gradient and can recover:

weight = torch.tensor(1.0, requires_grad=True)
bias = torch.tensor(-5.0, requires_grad=True)

outputs = F.leaky_relu(inputs * weight + bias, negative_slope=0.01)
outputs.sum().backward()
print(weight.grad, bias.grad)

tensor(0.0200) tensor(0.0400)

In practice, a few dead units are normal and harmless. It becomes a problem when many die at once, usually because the learning rate is too high. That is when Leaky ReLU or GELU are worth trying.

Which activation where

Now we can put it together. A network has two kinds of layers, and they follow different rules.

Hidden layers (every layer except the last) need a non-linearity so the network can learn curved patterns:

• Use ReLU by default. It is fast, and its gradient does not shrink.
• Try GELU or Leaky ReLU if training is unstable or many units seem to die. GELU is also the standard choice inside transformers.
• Avoid sigmoid and tanh in the hidden layers of ordinary networks, because of vanishing gradients. You will still see them inside some specialised layers, such as the gates of recurrent networks; those come much later.

The output layer is different. Its job is to produce the kind of answer the task needs, and in PyTorch that job is shared with the loss function:

These raw scores are called logits. Leaving the output activation to the loss is not laziness: combining the two in one function is more numerically stable, as Loss Functions shows. When you want actual probabilities after training, apply torch.sigmoid or torch.softmax yourself.

Softmax deserves a mention here because it is the multiclass version of sigmoid: it turns $K$ scores into $K$ probabilities that add up to 1,

$$\operatorname{softmax}(z)_k = \frac{e^{z_k}}{\sum_{j} e^{z_j}}.$$

It is not used between layers. We look at it as a function in Tensors, Matrices, and Related Mathematical Functions, and at how the loss uses it in Loss Functions.

Activations in PyTorch

PyTorch gives you each activation in two forms, and they compute exactly the same thing:

• as a module such as nn.ReLU(), which you put inside nn.Sequential or store as an attribute of your model;
• as a function such as torch.relu(z) or F.relu(z) (with import torch.nn.functional as F), which you call inside your model's forward method.

relu_module = nn.ReLU()
values = torch.tensor([-1.0, 3.0])

print(relu_module(values))
print(torch.relu(values))

tensor([0., 3.])
tensor([0., 3.])

Use whichever reads better in your code. The module names are the ones you would guess: nn.Sigmoid, nn.Tanh, nn.ReLU, nn.LeakyReLU, nn.GELU. The torch.nn documentation lists the rest; most of them are variations you do not need yet.

Common mistakes

The last three give no error message at all, which is why they are worth remembering: the code runs, the model just learns badly.

How to practise

• Work through the Dive into Deep Learning section with its code open, and plot each function and its derivative yourself, as we did above.
• Change the experiments on this page. How small does the gradient get after 20 sigmoids? What happens to the dead ReLU if you change the bias to −0.5?
• Once you reach A Simple MLP in PyTorch, train the same network twice, once with nn.Sigmoid() and once with nn.ReLU() in the hidden layers, and compare how fast the training loss falls.

Check yourself

What comes next

• Loss Functions: how the network's raw scores are turned into a single number that training can reduce, and why the output activation lives inside the loss.