Choosing an Activation Function

Activation Functions: Introducing Non-linearity

An activation function is the non-linear transformation applied after each layer's weighted sum. Without it, stacking layers produces only linear transformations—the network cannot learn non-linear patterns. Choosing the right activation function affects convergence speed, training stability, and final performance.

The Expressiveness of Activation Functions

A powerful result in deep learning theory: a single hidden layer with a non-linear activation can approximate any continuous function, given enough neurons. However, deep networks with multiple layers can achieve the same expressiveness with exponentially fewer neurons. The activation function is what makes this depth-efficiency possible.

Without activation functions:

# Just linear transformations
output = X @ W1 @ W2 @ W3
# This simplifies to X @ (W1 @ W2 @ W3), equivalent to a single matrix

With activation functions:

# Non-linear depth enables exponential expressiveness
h1 = relu(X @ W1)
h2 = relu(h1 @ W2)
output = h2 @ W3
# Cannot simplify; depth is essential

This is why deep learning works: depth + non-linearity = exponential expressiveness with reasonable parameter counts.

The diagram shows four common activation functions plotted side by side. Notice sigmoid and tanh saturate (approach constant slopes) at extreme values, while ReLU remains linear and unbounded in the positive region.

Sigmoid: The Classic Choice (Now Outdated for Hidden Layers)

Sigmoid maps inputs to (0, 1):

σ(z) = 1 / (1 + e^(-z))

It was the default choice for decades. However, it has a critical drawback: vanishing gradients.

The derivative of sigmoid:

d σ(z) / dz = σ(z) * (1 - σ(z))

This derivative is always ≤ 0.25 (peaks at z=0 where σ(z)=0.5, so 0.5 * 0.5 = 0.25). In a deep network, these small gradients multiply together across layers via the chain rule. In a 10-layer network, the gradient can shrink by a factor of (0.25)^10 ≈ 10^-7, making weight updates negligible.

Quantitative analysis: If each of 20 sigmoid layers has gradient ≈ 0.2 (typical for mixed inputs), the backpropagated gradient becomes 0.2^20 ≈ 10^-14—effectively zero. Early-layer weights barely update.

import numpy as np
import matplotlib.pyplot as plt

z = np.linspace(-5, 5, 100)
sigmoid = 1 / (1 + np.exp(-z))
sigmoid_grad = sigmoid * (1 - sigmoid)

plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.plot(z, sigmoid, 'b-', linewidth=2, label='sigmoid(z)')
plt.grid(True, alpha=0.3)
plt.xlabel('z')
plt.ylabel('sigmoid(z)')
plt.title('Sigmoid Function')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(z, sigmoid_grad, 'r-', linewidth=2, label="sigmoid'(z)")
plt.axhline(y=0.25, color='k', linestyle='--', alpha=0.5, label='max gradient ≈ 0.25')
plt.grid(True, alpha=0.3)
plt.xlabel('z')
plt.ylabel("sigmoid'(z)")
plt.title('Sigmoid Gradient')
plt.legend()
plt.tight_layout()
plt.show()

When to use sigmoid:

• Output layer for binary classification (maps logits to probabilities in [0, 1]).
• Specific problem domains where you need bounded outputs.

Avoid for hidden layers in deep networks due to vanishing gradients.

Tanh: Symmetric Sigmoid

Tanh is similar to sigmoid but outputs (-1, 1):

tanh(z) = (e^z - e^(-z)) / (e^z + e^(-z))

The gradient of tanh is 1 - tanh(z)^2, which peaks at 1 (when z=0) and approaches 0 at extremes. This is larger than sigmoid's maximum of 0.25, making tanh slightly better for deep networks.

z = np.linspace(-5, 5, 100)
tanh = np.tanh(z)
tanh_grad = 1 - np.tanh(z) ** 2

plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.plot(z, tanh, 'g-', linewidth=2, label='tanh(z)')
plt.grid(True, alpha=0.3)
plt.xlabel('z')
plt.ylabel('tanh(z)')
plt.title('Tanh Function')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(z, tanh_grad, 'purple', linewidth=2, label="tanh'(z)")
plt.axhline(y=1, color='k', linestyle='--', alpha=0.5, label='max gradient = 1')
plt.grid(True, alpha=0.3)
plt.xlabel('z')
plt.ylabel("tanh'(z)")
plt.title('Tanh Gradient')
plt.legend()
plt.tight_layout()
plt.show()

When to use tanh:

• Legacy code or specific architectures (RNNs without gates sometimes prefer tanh).
• When outputs need to be centered around 0 (helps with optimization).

Modern practice: ReLU-based functions are preferred for hidden layers.

ReLU: The Modern Default

ReLU (Rectified Linear Unit) is simply:

ReLU(z) = max(0, z) = {
  z,     if z ≥ 0
  0,     if z < 0
}

It is astonishingly simple, yet highly effective. For positive inputs, the gradient is 1; for negative inputs, the gradient is 0.

Derivative of ReLU:

d ReLU(z) / dz = {
  1,     if z > 0
  0,     if z < 0
  undefined (usually treated as 0 or ignored), if z = 0
}

This constant gradient of 1 for positive inputs means gradients do not shrink in backpropagation—early-layer weights receive meaningful updates even in very deep networks.

z = np.linspace(-5, 5, 100)
relu = np.maximum(0, z)
relu_grad = (z > 0).astype(float)

plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.plot(z, relu, 'orange', linewidth=2, label='ReLU(z)')
plt.grid(True, alpha=0.3)
plt.xlabel('z')
plt.ylabel('ReLU(z)')
plt.title('ReLU Function')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(z, relu_grad, 'r-', linewidth=2, label="ReLU'(z)")
plt.grid(True, alpha=0.3)
plt.xlabel('z')
plt.ylabel("ReLU'(z)")
plt.axvline(x=0, color='k', linestyle='--', alpha=0.5)
plt.title('ReLU Gradient (1 if z > 0, else 0)')
plt.legend()
plt.tight_layout()
plt.show()

Why ReLU is so good:

• No vanishing gradient: Gradient is 1 for positive inputs, enabling deep networks.
• Computational efficiency: max(0, z) is a single comparison, much faster than sigmoid or tanh.
• Sparse activations: Negative inputs are zeroed, creating sparsity that can reduce overfitting.

The dying ReLU problem: If a neuron's weights update such that it always outputs negative values, it stays "dead"—the gradient is 0 forever, and it never recovers. This is rare in practice but can happen with aggressive learning rates or poor initialization.

Here is ReLU in PyTorch:

import torch
import torch.nn as nn

# Create a simple network with ReLU
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 32),
    nn.ReLU(),
    nn.Linear(32, 2)
)

# Sample input
x = torch.randn(5, 10)
output = model(x)
print(f"Output shape: {output.shape}")

Variants: Leaky ReLU and GELU

To address the dying ReLU problem, variants have been proposed.

Leaky ReLU:

Leaky ReLU(z) = z if z > 0, else alpha * z (typically alpha = 0.01)

Instead of 0 for negative inputs, use a small slope. This ensures the gradient is never exactly 0.

import torch.nn.functional as F

z = torch.randn(5, 10)
leaky_relu_out = F.leaky_relu(z, negative_slope=0.01)

GELU (Gaussian Error Linear Unit):

GELU(z) ≈ z * Φ(z)

where Φ(z) is the cumulative distribution function of the standard normal. GELU is smoother than ReLU and has become popular in transformer-based models.

import torch.nn.functional as F

z = torch.randn(5, 10)
gelu_out = F.gelu(z)

Practical guidance:

• Start with ReLU for hidden layers. It is fast, well-understood, and effective.
• Use Leaky ReLU if you suspect dying ReLU (monitor for dead neurons during training).
• Use GELU in transformer models (it is the standard) or if you want smoother gradients.
• Avoid sigmoid/tanh for hidden layers in deep networks due to vanishing gradients.

Comprehensive activation function comparison:

| Function | Formula | Output Range | Max Gradient | Gradient Behavior | Vanishing Gradient Risk | Deep Network Viability | |----------|---------|------|---|------|------|------| | Sigmoid | 1/(1+e^(-z)) | (0, 1) | 0.25 @ z=0 | Decays at extremes | Very high (shrinks 0.25^n) | Poor for >5 layers | | Tanh | (e^z - e^(-z))/(e^z + e^(-z)) | (-1, 1) | 1 @ z=0 | Decays at extremes | High (shrinks 1^n ≈ 1) | Mediocre, ReLU better | | ReLU | max(0, z) | [0, ∞) | 1 for z>0 | Constant for positive | None (gradient = 1) | Excellent | | Leaky ReLU | z if z>0, else αz (α=0.01) | (-∞, ∞) | 1 or α | Constant | None | Excellent | | ELU | z if z>0, else α(e^z-1) | (-α, ∞) | 1 or varies | Smooth | Low | Very good | | GELU | z * Φ(z) | ~(-0.17 to 1) | Smooth | Smooth, continuous | None | Excellent (modern) |

Output Layer Activations: Problem-Dependent

The activation function for the output layer depends on your task:

| Task | Output Activation | Loss Function | Output Range | |------|------------------|---------------|---------------| | Binary classification | Sigmoid | Binary Cross-Entropy | (0, 1) | | Multi-class classification | Softmax | Cross-Entropy | Probability distribution | | Regression | None (linear) | MSE or MAE | (-∞, ∞) | | Regression (bounded) | Sigmoid or Tanh | MSE or MAE | (0, 1) or (-1, 1) |

Binary Classification

model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 1)
)
loss_fn = nn.BCEWithLogitsLoss()  # Combines sigmoid + BCE loss

# BCEWithLogitsLoss applies sigmoid internally, so don't add it to the model

Multi-Class Classification

model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 10)  # 10 classes
)
loss_fn = nn.CrossEntropyLoss()  # Applies softmax internally

# CrossEntropyLoss expects raw logits, not softmax

Regression

model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 1)  # Raw output, no activation
)
loss_fn = nn.MSELoss()

The Problem with Sigmoid: Saturation in Deep Networks

Let's concretely demonstrate why sigmoid causes problems in deep networks:

# Simulate gradient flow through sigmoid layers
np.random.seed(42)

def sigmoid(z):
    return 1 / (1 + np.exp(-np.clip(z, -500, 500)))

def sigmoid_grad(z):
    s = sigmoid(z)
    return s * (1 - s)

# Network with 20 sigmoid layers
num_layers = 20
upstream_grad = 1.0
gradients = [upstream_grad]

for layer in range(num_layers):
    # Simulated pre-activation: uniform in [-3, 3] (reasonable inputs)
    z = np.random.uniform(-3, 3)
    # Gradient through this layer's sigmoid
    local_grad = sigmoid_grad(z)
    upstream_grad *= local_grad
    gradients.append(upstream_grad)

# Plot gradient magnitude through layers
plt.figure(figsize=(10, 5))
plt.plot(gradients, 'b-', linewidth=2, marker='o')
plt.ylabel('Gradient magnitude')
plt.xlabel('Layer (backward from output)')
plt.title('Gradient Flow Through 20 Sigmoid Layers: Vanishing Problem')
plt.yscale('log')
plt.grid(True, alpha=0.3)
plt.show()

print(f"Input gradient: {gradients[0]}")
print(f"Gradient at layer 10: {gradients[10]:.2e}")
print(f"Gradient at layer 20: {gradients[20]:.2e}")

Expected output: gradient shrinks by orders of magnitude. By layer 20, it is ~10^-10—effectively zero. Early-layer parameters barely update.

Compare with ReLU: if all pre-activations are positive, the gradient is 1 at every layer, so upstream_grad remains constant. This is why ReLU enabled training of 100+ layer networks.

Initialization Interacts with Activation Function

The choice of activation function affects how you should initialize weights:

• Sigmoid/Tanh: Xavier initialization is ideal; it balances the variance of pre-activations to avoid saturation.
• ReLU: He initialization; accounts for ReLU zeroing half the outputs, preventing early saturation.

from torch.nn import init

# Xavier for sigmoid/tanh
for layer in model:
    if isinstance(layer, nn.Linear):
        init.xavier_uniform_(layer.weight)

# He for ReLU
for layer in model:
    if isinstance(layer, nn.Linear):
        init.kaiming_uniform_(layer.weight, nonlinearity='relu')

Using the wrong initialization with the wrong activation can slow training significantly. PyTorch defaults to He, which is safe for ReLU. If you switch to tanh, consider re-initializing.

Complete Example: Building and Training a Network

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import make_classification
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

# Create synthetic data
X, y = make_classification(n_samples=200, n_features=10, n_classes=2, random_state=42)
scaler = StandardScaler()
X = scaler.fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.long)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.long)

# Define model with appropriate activations
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),  # Hidden layer activation
    nn.Linear(64, 32),
    nn.ReLU(),  # Hidden layer activation
    nn.Linear(32, 2)  # Output: raw logits, no activation
)

loss_fn = nn.CrossEntropyLoss()  # Applies softmax + CE loss
optimizer = optim.Adam(model.parameters(), lr=0.01)

# Training
for epoch in range(50):
    # Forward pass
    logits = model(X_train)
    loss = loss_fn(logits, y_train)

    # Backward pass
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 10 == 0:
        with torch.no_grad():
            train_pred = model(X_train).argmax(dim=1)
            test_pred = model(X_test).argmax(dim=1)
            train_acc = (train_pred == y_train).float().mean()
            test_acc = (test_pred == y_test).float().mean()
            print(f"Epoch {epoch+1}: Loss={loss.item():.4f}, Train acc={train_acc:.2%}, Test acc={test_acc:.2%}")

This network uses ReLU in hidden layers (efficient, prevents vanishing gradients) and no activation in the output layer (raw logits go into CrossEntropyLoss).

Softmax: The Output Activation for Multi-Class Classification

For multi-class classification, the output layer should produce a probability distribution over classes. Softmax converts raw logits to probabilities:

softmax(z_i) = exp(z_i) / Σ_j(exp(z_j))

Each output is between 0 and 1, and they sum to 1—a valid probability distribution.

import torch.nn.functional as F

logits = torch.tensor([[2.0, 1.0, 0.1],
                       [1.0, 3.0, 0.5]])
probabilities = F.softmax(logits, dim=1)
print(probabilities)
# Output: [[0.6524, 0.2388, 0.1088],
#          [0.0900, 0.6652, 0.2448]]
# Each row sums to 1

In practice: Use CrossEntropyLoss, which applies softmax internally. Do not manually apply softmax before the loss function—it expects raw logits.

# RIGHT
logits = model(X)
loss = nn.CrossEntropyLoss()(logits, y_true)  # Softmax is built-in

# WRONG
logits = model(X)
probabilities = F.softmax(logits, dim=1)
loss = nn.CrossEntropyLoss()(probabilities, y_true)  # Softmax applied twice!

Activation Functions for Regression

Regression tasks (predicting continuous values) usually need no output activation. The model outputs raw, unbounded values.

Exception: If your target is bounded, use a matching activation:

• Target in [0, 1] (e.g., percentage): Use sigmoid on output.
• Target in [-1, 1]: Use tanh on output.
• Target in [0, ∞) (e.g., positive magnitudes): Use ReLU or softplus (log(1 + exp(z))) on output.

# Regression with bounded output
class BoundedRegressionModel(nn.Module):
    def __init__(self):
        super(BoundedRegressionModel, self).__init__()
        self.fc1 = nn.Linear(10, 64)
        self.fc2 = nn.Linear(64, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.sigmoid(self.fc2(x))  # Output in [0, 1]
        return x

Activation Functions Across Architectures

Different architectures have different conventions:

| Architecture | Hidden Activations | Output Activation | Notes | |--------------|-------------------|------------------|-------| | MLP (dense) | ReLU (standard) | None / Sigmoid / Softmax | Task-dependent | | CNN | ReLU | None / Sigmoid / Softmax | ReLU universal in CNNs | | RNN / LSTM | Tanh (RNN) or None (LSTM/GRU) | Task-dependent | LSTMs have internal gates; no standard activation | | Transformer | GELU or ReLU | Task-dependent | GELU is standard in BERT/GPT |

For your first project: use ReLU for hidden layers, no activation for regression output, softmax + cross-entropy for classification output.

Activation Functions and Batch Normalization

Modern networks almost always include batch normalization between layers, which rescales activations to have mean 0 and std 1 (per batch). This dramatically stabilizes training and makes the choice of activation function less critical.

With batch norm, tanh and sigmoid recover some of their appeal because activations are normalized, preventing saturation. However, ReLU + batch norm remains the standard.

# Modern convention: ReLU + BatchNorm
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.BatchNorm1d(64),
    nn.ReLU(),
    nn.Linear(64, 32),
    nn.BatchNorm1d(32),
    nn.ReLU(),
    nn.Linear(32, 2)
)

Batch norm is covered in depth in later lessons. For now, recognize that it is a companion to activation functions—together they enable stable, fast training of deep networks.

Common Mistake: Using ReLU in the Output Layer for Classification

A frequent error is applying ReLU to the final output layer:

# WRONG: ReLU in output layer
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 10),
    nn.ReLU()  # Wrong! Zeros out negative logits
)
loss_fn = nn.CrossEntropyLoss()

ReLU clamps negative logits to 0, destroying the model's ability to express low probabilities. This breaks the loss function's assumptions.

Correct approach:

# RIGHT: No activation in output layer
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 10)  # Raw logits
)
loss_fn = nn.CrossEntropyLoss()  # Handles softmax internally

Remember: Activation functions are for hidden layers. Output layer activations are task-specific and often not needed (the loss function handles them).