PyTorch Activation Functions

Activation functions introduce non-linearity into a neural network. Without them, stacking ten linear layers gives you the same result as a single linear layer — the network can only learn straight-line relationships. Activation functions let networks learn curves, corners, and complex patterns.

Why Non-Linearity Matters

Without activation:
Layer1(x) = W1 * x + b1
Layer2(x) = W2 * (W1*x + b1) + b2
           = (W2*W1)*x + (W2*b1 + b2)
           = W_combined * x + b_combined  ← still just one linear layer

With activation:
Layer1(x) = relu(W1*x + b1)   ← non-linear output
Layer2(x) = W2 * relu_output + b2  ← can't collapse to one layer

ReLU – Rectified Linear Unit

ReLU is the most widely used activation function. It returns the input if positive, and zero if negative. This simplicity makes it fast to compute and it solves the vanishing gradient problem for most networks.

import torch
import torch.nn as nn

relu = nn.ReLU()
x = torch.tensor([-3.0, -1.0, 0.0, 1.0, 3.0])
print(relu(x))   # tensor([0., 0., 0., 1., 3.])

ReLU Diagram:
        │
    3 ─ │              /
    2 ─ │            /
    1 ─ │          /
    0 ─ │─────────/─────────
   -1 ─ │ (flat at 0 for negative inputs)
        └──────────────────→ input
           negative   positive

Leaky ReLU

Standard ReLU outputs exactly zero for negative inputs. When many neurons output zero, they stop contributing to gradient flow — this is called a "dying ReLU." Leaky ReLU allows a small non-zero output for negative inputs, keeping neurons alive.

leaky = nn.LeakyReLU(negative_slope=0.01)
x = torch.tensor([-3.0, -1.0, 0.0, 1.0, 3.0])
print(leaky(x))   # tensor([-0.0300, -0.0100,  0.0000,  1.0000,  3.0000])

Sigmoid

Sigmoid maps any input to a value between 0 and 1. It is used in the output layer for binary classification — where you need a probability between 0 and 1.

sigmoid = nn.Sigmoid()
x = torch.tensor([-5.0, -1.0, 0.0, 1.0, 5.0])
print(sigmoid(x))   # tensor([0.0067, 0.2689, 0.5000, 0.7311, 0.9933])

Sigmoid Diagram:
  1.0 ─ │          ─────────
  0.5 ─ │       /
  0.0 ─ │─────/──────────────
        └──────────────────→ input
                   Squeeze everything to [0, 1]

Avoid sigmoid in hidden layers of deep networks — its gradient is at most 0.25, causing vanishing gradients. Use it only in output layers for binary classification.

Tanh

Tanh maps inputs to values between -1 and 1 — similar to sigmoid but centered at zero. Zero-centered outputs help gradients flow more symmetrically and make training more stable in some cases.

tanh = nn.Tanh()
x = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0])
print(tanh(x))   # tensor([-0.9640, -0.7616,  0.0000,  0.7616,  0.9640])

Softmax

Softmax converts a vector of raw scores into a probability distribution — all values sum to 1.0, and each value is between 0 and 1. Use softmax in the output layer for multi-class classification.

softmax = nn.Softmax(dim=1)
scores = torch.tensor([[2.0, 1.0, 0.5]])   # 3 class scores for 1 sample
probs = softmax(scores)
print(probs)         # tensor([[0.6232, 0.2291, 0.1477]])
print(probs.sum())   # tensor(1.)  ← always sums to 1

Softmax converts raw scores to probabilities:
  [2.0, 1.0, 0.5]
       ↓ softmax
  [0.62, 0.23, 0.15]   ← sums to 1.0

GELU

GELU (Gaussian Error Linear Unit) is used in transformer models like BERT and GPT. It behaves similarly to ReLU but with a smooth curve around zero — which can help gradients flow more stably in very deep networks.

gelu = nn.GELU()
x = torch.tensor([-3.0, -1.0, 0.0, 1.0, 3.0])
print(gelu(x))   # tensor([-0.0036, -0.1587,  0.0000,  0.8413,  2.9964])

Choosing the Right Activation

Location                   Activation to Use
────────────────────────────────────────────────────
Hidden layers (general)    ReLU  (fast, simple)
Hidden layers (deep)       Leaky ReLU or GELU
Output: binary class       Sigmoid
Output: multi-class        Softmax (or none + CrossEntropyLoss)
Output: regression         None (or ReLU if output must be positive)
RNN hidden states          Tanh
Transformers               GELU

Applying Activations in a Model

import torch.nn as nn

model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 32),
    nn.ReLU(),
    nn.Linear(32, 1),
    nn.Sigmoid()   # for binary classification output
)

Summary

Activation functions add non-linearity so networks can learn complex patterns. ReLU is the go-to choice for hidden layers — fast and effective. Leaky ReLU prevents dying neurons. Sigmoid and Tanh are used in output layers or recurrent networks. Softmax turns class scores into probabilities. GELU is the preferred activation in transformers. Always place an activation after each hidden linear layer. The output layer's activation depends on the task: sigmoid for binary, softmax for multi-class, none for regression.