Activation Functions

Learning Objectives

The Big Picture

In the previous section, we learned that linear layers compute $y = Wx + b$. This is a powerful operation for combining features, but there's a fundamental limitation: stacking linear transformations just gives you another linear transformation.

The Core Problem: If you stack 100 linear layers without activation functions, the entire network is mathematically equivalent to a single linear layer. You gain nothing from depth! Activation functions break this linearity and give neural networks their power.

Activation functions are the non-linear transformations applied after each linear layer. They are the secret sauce that allows neural networks to approximate arbitrarily complex functions, learn intricate patterns, and solve real-world problems that linear models cannot.

Historical Context

The importance of non-linearity was recognized from the earliest days of neural networks. The original perceptron (1958) used a step function. For decades, sigmoid and tanh dominated. Then in 2012, the ReLU (Rectified Linear Unit) was used in AlexNet, winning ImageNet by a large margin. This simple change—using $\max(0, x)$ instead of sigmoid—was one of the key innovations that launched the deep learning revolution.

Why Non-linearity Matters

Let's prove mathematically why linear networks collapse into a single layer:

The Linear Network Collapse

Consider a network with two linear layers (no activation):

Substituting the first equation into the second:

The two-layer network is equivalent to a single linear transformation with weight $W' = W_2 W_1$ and bias $\mathbf{b}' = W_2\mathbf{b}_1 + \mathbf{b}_2$. No matter how many layers you add, you cannot escape this limitation.

Interactive: Activation Function Explorer

Explore the most important activation functions interactively. See their shapes, derivatives (which determine gradient flow), and trade-offs. Toggle the derivative to understand how gradients behave in different regions.

Classic Activation Functions

Let's examine the foundational activation functions that dominated neural networks for decades before ReLU.

Sigmoid (Logistic Function)

The sigmoid function squashes any input to the range $(0, 1)$. Its derivative is:

Key insight: The maximum gradient of sigmoid is only 0.25. This means gradients shrink by at least 4× at each layer, leading to the infamous vanishing gradient problem in deep networks.

Tanh (Hyperbolic Tangent)

Tanh is a rescaled sigmoid with output range $(-1, 1)$. Its derivative is:

The Problem: Saturation

Both sigmoid and tanh suffer from saturation: when |x| is large, the gradient becomes essentially zero:

• For sigmoid: $\sigma'(x) \approx 0$ when $|x| > 4$
• For tanh: $\tanh'(x) \approx 0$ when $|x| > 3$

In deep networks, this means gradients can shrink exponentially as they backpropagate through layers, making it nearly impossible to train the early layers.

The ReLU Revolution

The Rectified Linear Unit (ReLU) changed everything:

Its derivative is remarkably simple:

Why ReLU Works So Well

1. No gradient saturation (for positive x): Unlike sigmoid/tanh, ReLU's gradient is exactly 1 for all positive inputs, no matter how large. Gradients flow unchanged through many layers.
2. Computational efficiency: ReLU is just a threshold comparison—no exponentials, no divisions. This makes it 6× faster than sigmoid in practice.
3. Sparse activation: ReLU outputs exactly zero for negative inputs. This creates sparse representations, which can improve generalization and reduce computation.
4. Biological plausibility: Real neurons have a threshold below which they don't fire. ReLU mimics this behavior.

The Dying ReLU Problem

ReLU Variants

Several variants address the dying ReLU problem:

Interactive: Vanishing Gradients

Watch the vanishing gradient problem in action. See how gradients shrink (or stay stable) as they flow backward through layers. Compare different activation functions and understand why ReLU enabled training of deep networks.

Modern Activation Functions

Modern architectures, especially transformers, use more sophisticated activation functions that combine the benefits of ReLU with smoothness.

GELU (Gaussian Error Linear Unit)

Where $\Phi(x)$ is the standard Gaussian CDF. GELU can be thought of as:

• Probabilistic interpretation: GELU multiplies x by the probability that it's greater than a random Gaussian variable
• Smooth approximation of ReLU: Like ReLU but with a soft transition around zero
• Standard in transformers: BERT, GPT, ViT, and most modern transformers use GELU

Swish/SiLU (Sigmoid Linear Unit)

Swish was discovered through neural architecture search and has become popular in vision models. It's self-gated: the sigmoid acts as a gate that controls how much of x passes through.

Mish

Mish is smooth, non-monotonic, and self-regularizing. It's used in YOLOv4 and other vision architectures.

Hardswish

Hardswish is a hardware-efficient approximation of Swish, designed for mobile and embedded devices. Instead of computing a sigmoid (which requires expensive exponentials), it uses a piecewise linear function.

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5# PyTorch provides Hardswish natively
6hardswish = nn.Hardswish()
7
8# Manual implementation for understanding
9def hardswish_manual(x):
10    return x * F.relu6(x + 3) / 6
11
12# Usage in MobileNetV3
13class MobileNetV3Block(nn.Module):
14    def __init__(self, in_ch, out_ch, kernel_size, use_se=True):
15        super().__init__()
16        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, padding=kernel_size//2)
17        self.bn = nn.BatchNorm2d(out_ch)
18        self.act = nn.Hardswish()  # Key activation for MobileNetV3
19
20    def forward(self, x):
21        return self.act(self.bn(self.conv(x)))
22
23# Comparison: Swish vs Hardswish
24x = torch.randn(1000, 256)
25%timeit F.silu(x)        # Swish: ~0.5ms
26%timeit F.hardswish(x)   # Hardswish: ~0.3ms (40% faster!)

GLU (Gated Linear Unit)

GLU uses one linear projection as a gate for another. Variants like SwiGLU and GeGLU are used in modern language models like LLaMA and PaLM.

Interactive: Distribution Effects

See how different activation functions reshape the distribution of neural network activations. This matters for training stability and how information flows through the network.

PyTorch Implementation

PyTorch provides activation functions in three ways: as nn.Module classes, as functions in torch.nn.functional, and as tensor methods.

Basic Usage

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5# Method 1: Using nn.Module activation classes
6relu = nn.ReLU()
7sigmoid = nn.Sigmoid()
8tanh = nn.Tanh()
9leaky_relu = nn.LeakyReLU(negative_slope=0.01)
10gelu = nn.GELU()
11
12# Apply to a tensor
13x = torch.randn(4, 3)
14print("Input:", x)
15print("ReLU output:", relu(x))
16print("Sigmoid output:", sigmoid(x))
17
18# Method 2: Using torch.nn.functional (stateless)
19y = F.relu(x)
20y = F.sigmoid(x)
21y = F.gelu(x)
22
23# Method 3: Using torch tensor methods (some activations)
24y = x.relu()
25y = x.sigmoid()
26y = x.tanh()

Using Activations in Networks

1import torch.nn as nn
2
3class MLPWithActivations(nn.Module):
4    def __init__(self, input_dim, hidden_dim, num_classes):
5        super().__init__()
6
7        # Layer definitions
8        self.fc1 = nn.Linear(input_dim, hidden_dim)
9        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
10        self.fc3 = nn.Linear(hidden_dim, num_classes)
11
12        # Activation functions
13        self.relu = nn.ReLU()
14        self.dropout = nn.Dropout(0.5)
15
16    def forward(self, x):
17        # Layer 1 + ReLU
18        x = self.fc1(x)
19        x = self.relu(x)  # Non-linearity is crucial!
20        x = self.dropout(x)
21
22        # Layer 2 + ReLU
23        x = self.fc2(x)
24        x = self.relu(x)
25        x = self.dropout(x)
26
27        # Output layer - NO activation here!
28        # Apply softmax in loss function or after
29        x = self.fc3(x)
30        return x
31
32# Alternative: Using nn.Sequential
33model = nn.Sequential(
34    nn.Linear(784, 256),
35    nn.ReLU(),
36    nn.Dropout(0.5),
37    nn.Linear(256, 256),
38    nn.ReLU(),
39    nn.Dropout(0.5),
40    nn.Linear(256, 10),
41)

Modern Activation Functions

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5# GELU - Used in transformers (BERT, GPT, ViT)
6class TransformerMLP(nn.Module):
7    def __init__(self, d_model, d_ff):
8        super().__init__()
9        self.fc1 = nn.Linear(d_model, d_ff)
10        self.fc2 = nn.Linear(d_ff, d_model)
11        self.gelu = nn.GELU()
12
13    def forward(self, x):
14        return self.fc2(self.gelu(self.fc1(x)))
15
16# Swish/SiLU - Self-gated activation
17class SwishMLP(nn.Module):
18    def __init__(self, input_dim, hidden_dim):
19        super().__init__()
20        self.fc1 = nn.Linear(input_dim, hidden_dim)
21        self.fc2 = nn.Linear(hidden_dim, input_dim)
22        self.silu = nn.SiLU()  # Same as Swish
23
24    def forward(self, x):
25        return self.fc2(self.silu(self.fc1(x)))
26
27# Mish - Smooth self-regularizing
28mish = nn.Mish()
29
30# GLU - Gated Linear Unit (splits channels)
31class GLU(nn.Module):
32    def __init__(self, dim):
33        super().__init__()
34        self.fc = nn.Linear(dim, dim * 2)  # Double channels
35
36    def forward(self, x):
37        x, gate = self.fc(x).chunk(2, dim=-1)
38        return x * torch.sigmoid(gate)
39
40# PReLU - Learnable negative slope
41prelu = nn.PReLU(num_parameters=256)  # One slope per channel

Choosing Activation Functions

Here's a practical guide for choosing activation functions:

Hidden Layers

Simple Decision Tree

1. Transformers or modern NLP? Use GELU. It's the standard.
2. Computer vision with CNNs? Start with ReLU. Consider Swish or Mish for state-of-the-art.
3. MLPs for tabular data? ReLU or Leaky ReLU works well.
4. Worried about dying neurons? Use Leaky ReLU, ELU, or GELU.
5. Need bounded outputs in hidden layers? (rare) Use Tanh.

Output Layer Activations

The output layer requires special consideration because the activation determines how the network's output is interpreted:

1# Binary classification
2class BinaryClassifier(nn.Module):
3    def forward(self, x):
4        ...
5        return torch.sigmoid(self.fc_out(x))  # ✅ Sigmoid for binary
6
7# Multi-class classification
8class MultiClassifier(nn.Module):
9    def forward(self, x):
10        ...
11        return self.fc_out(x)  # ✅ Raw logits - let loss handle softmax
12
13# Regression
14class Regressor(nn.Module):
15    def forward(self, x):
16        ...
17        return self.fc_out(x)  # ✅ No activation for regression

Common Pitfalls

1# ❌ WRONG: No activations between layers
2def forward_wrong(self, x):
3    x = self.fc1(x)
4    x = self.fc2(x)  # Bug! This is just a linear network
5    return x
6
7# ✅ CORRECT: Activation between layers
8def forward_correct(self, x):
9    x = self.fc1(x)
10    x = F.relu(x)  # Non-linearity!
11    x = self.fc2(x)
12    return x
13
14# ❌ WRONG: Softmax before CrossEntropyLoss
15output = F.softmax(model(x), dim=-1)  # Bug!
16loss = F.cross_entropy(output, targets)  # Wrong!
17
18# ✅ CORRECT: Raw logits to CrossEntropyLoss
19output = model(x)  # Raw logits
20loss = F.cross_entropy(output, targets)  # Correct!

Test Your Understanding

Test your knowledge of activation functions with this comprehensive quiz. Each question has a detailed explanation.

Summary

Exercises

Conceptual Questions

1. Explain why the gradient of sigmoid is $\sigma(x)(1-\sigma(x))$. Derive this from the sigmoid formula.
2. Why is zero-centering (tanh) beneficial for training? Hint: think about gradient updates for the previous layer's weights.
3. In what sense is GELU "probabilistic"? What does $x \cdot \Phi(x)$represent?

Coding Exercises

1. Implement activations from scratch: Write ReLU, Leaky ReLU, and GELU using only basic tensor operations. Verify your implementations match PyTorch's.
2. Gradient flow experiment: Create a 20-layer network. Train it with sigmoid, tanh, and ReLU. Plot the gradient magnitude at each layer. Which activation maintains the healthiest gradients?
3. Dying ReLU detection: Train a network with ReLU and high learning rate. After training, count how many neurons have "died" (always output zero on the training set). Try fixing with Leaky ReLU.
4. GELU approximation: Implement the fast GELU approximation used in practice:
   $$\text{GELU}(x) \approx 0.5x\left(1 + \tanh\left(\sqrt{\frac{2}{\pi}}(x + 0.044715x^3)\right)\right)$$
   Compare accuracy and speed with the exact version.

Challenge: Custom Activation

Design your own activation function with these requirements:

• Non-zero gradient everywhere (no dying neurons)
• Bounded or slowly growing output (no exploding activations)
• Smooth (differentiable everywhere)

Implement it as a custom nn.Module and test it on MNIST. Can you beat ReLU?

In the next section, we'll explore Loss Functions—the functions that measure how wrong our network's predictions are and guide the learning process. You'll see how the choice of loss function is deeply connected to the probabilistic interpretation of neural networks.