Activation Functions and Loss Functions

A football player's training regime (the loss function) determines what they get better at. Train them to minimise goals conceded — they become a defender. Train them to maximise goals scored — they become a striker. The same player, the same training intensity, but the objective determines the skill.

The activation function is the player's physical capability — how much they can bend, how fast they can turn. A player with no flexibility (linear activation) cannot do anything a simple regression cannot. A player with full agility (ReLU, GELU) can learn arbitrarily complex patterns.

import torch
import torch.nn as nn
import numpy as np

# ── Implement and visualise all six activations ────────────────────────
z = torch.linspace(-3, 3, 100)

activations = {
    'ReLU':       nn.ReLU(),
    'LeakyReLU':  nn.LeakyReLU(0.01),
    'GELU':       nn.GELU(),
    'Sigmoid':    nn.Sigmoid(),
    'Tanh':       nn.Tanh(),
}

print("Activation function outputs at key points:")
print(f"{'Activation':<14} {'z=−2':>8} {'z=−1':>8} {'z=0':>8} {'z=1':>8} {'z=2':>8}")
print("─" * 55)

for name, fn in activations.items():
    vals = [fn(torch.tensor(float(z))).item() for z in [-2, -1, 0, 1, 2]]
    print(f"  {name:<12}  " + "  ".join(f"{v:>8.4f}" for v in vals))

# ── Derivative magnitudes — vanishing gradient check ──────────────────
print("
Derivative magnitude at z=−3 (vanishing gradient test):")
for name, fn in activations.items():
    z_test = torch.tensor(-3.0, requires_grad=True)
    out    = fn(z_test)
    out.backward()
    deriv = z_test.grad.item()
    flag  = '← near zero — vanishing!' if abs(deriv) < 0.01 else '← gradient flows'
    print(f"  {name:<14}: f′(−3) = {deriv:+.6f}  {flag}")

# ── GELU vs ReLU in practice ──────────────────────────────────────────
# GELU is smooth everywhere — ReLU has a kink at 0
# This matters for second-order optimisers and some architectures
z_compare = torch.linspace(-2, 2, 7)
relu_out  = nn.ReLU()(z_compare)
gelu_out  = nn.GELU()(z_compare)

print("
GELU vs ReLU — the smooth difference:")
print(f"{'z':>8} {'ReLU':>10} {'GELU':>10} {'diff':>10}")
for z_val, r, g in zip(z_compare, relu_out, gelu_out):
    print(f"  {z_val.item():>6.2f}  {r.item():>10.4f}  {g.item():>10.4f}  {(g-r).item():>+10.4f}")
print("GELU gates the input smoothly — ReLU hard-zeros it.")

import torch
import torch.nn as nn
import numpy as np

# ── All six loss functions on concrete examples ────────────────────────

# ── Regression losses ──────────────────────────────────────────────────
y_reg  = torch.tensor([41.0, 29.0, 55.0, 38.0])
y_pred = torch.tensor([35.0, 31.0, 48.0, 42.0])

mse   = nn.MSELoss()(y_pred, y_reg)
mae   = nn.L1Loss()(y_pred, y_reg)
huber = nn.HuberLoss(delta=5.0)(y_pred, y_reg)

print("Regression losses (DoorDash delivery time):")
print(f"  MSE:   {mse.item():.4f}  (sensitive to large errors)")
print(f"  MAE:   {mae.item():.4f}  (proportional to error)")
print(f"  Huber: {huber.item():.4f}  (MSE for small, MAE for large)")

# Demonstrate outlier sensitivity
y_with_outlier = torch.tensor([41.0, 29.0, 55.0, 97.0])   # last is outlier
mse_out   = nn.MSELoss()(y_pred, y_with_outlier)
mae_out   = nn.L1Loss()(y_pred, y_with_outlier)
huber_out = nn.HuberLoss(delta=5.0)(y_pred, y_with_outlier)
print(f"
With outlier (97 vs predicted 42):")
print(f"  MSE:   {mse_out.item():.1f}  ← dominated by outlier")
print(f"  MAE:   {mae_out.item():.1f}  ← less affected")
print(f"  Huber: {huber_out.item():.1f}  ← best balance")

# ── Classification losses ──────────────────────────────────────────────
# Binary — raw logits (not sigmoid!) into BCEWithLogitsLoss
logits_binary = torch.tensor([2.1, -0.5, 1.8, -1.2])   # raw scores
y_binary      = torch.tensor([1.0, 0.0, 1.0, 0.0])

# WRONG: nn.BCELoss()(torch.sigmoid(logits), y) — numerically unstable
# RIGHT: BCEWithLogitsLoss applies sigmoid internally with log-sum-exp trick
bce = nn.BCEWithLogitsLoss()(logits_binary, y_binary)
print(f"
Binary cross-entropy (BCEWithLogitsLoss): {bce.item():.4f}")

# Multi-class — raw logits (not softmax!) into CrossEntropyLoss
# Shape: (batch_size, n_classes)
logits_multi = torch.tensor([
    [2.1, 0.3, -0.5],   # sample 1 — class 0 most likely
    [0.1, 1.9, 0.2],    # sample 2 — class 1 most likely
    [-0.3, 0.4, 2.2],   # sample 3 — class 2 most likely
])
y_multi = torch.tensor([0, 1, 2])   # true class indices

# WRONG: torch.softmax then nn.NLLLoss — unstable
# RIGHT: CrossEntropyLoss takes raw logits and applies log-softmax internally
ce = nn.CrossEntropyLoss()(logits_multi, y_multi)
print(f"Cross-entropy (CrossEntropyLoss):         {ce.item():.4f}")

# ── Class weighting for imbalanced datasets ───────────────────────────
# If 95% class 0 and 5% class 1, upweight class 1 by 19×
weights = torch.tensor([1.0, 19.0, 1.0])
ce_weighted = nn.CrossEntropyLoss(weight=weights)(logits_multi, y_multi)
print(f"Weighted cross-entropy (19× class 1):     {ce_weighted.item():.4f}")

Computing log(sigmoid(z)) directly overflows for large |z|. The numerically stable version uses the log-sum-exp trick:

PyTorch's BCEWithLogitsLoss and CrossEntropyLoss implement this trick internally. Using nn.Sigmoid() + nn.BCELoss() skips it — leading to NaN at training time when logits are large.

import torch
import torch.nn as nn
import numpy as np

# ── Demonstrate numerical instability ─────────────────────────────────
large_logit = torch.tensor([100.0])   # extreme but not impossible after bad init
y_true      = torch.tensor([1.0])

# WRONG — sigmoid first, then BCE
try:
    prob_unstable = torch.sigmoid(large_logit)
    loss_unstable = nn.BCELoss()(prob_unstable, y_true)
    print(f"BCELoss after sigmoid:          {loss_unstable.item()}")
except Exception as e:
    print(f"BCELoss after sigmoid:          ERROR — {e}")

# RIGHT — BCEWithLogitsLoss handles it stably
loss_stable = nn.BCEWithLogitsLoss()(large_logit, y_true)
print(f"BCEWithLogitsLoss (raw logits): {loss_stable.item():.6f}  ← stable")

# ── Same issue with softmax + NLLLoss vs CrossEntropyLoss ─────────────
large_logits = torch.tensor([[100.0, 0.0, 0.0]])
y_cls        = torch.tensor([0])

# WRONG — softmax first
try:
    probs = torch.softmax(large_logits, dim=1)
    loss_wrong = nn.NLLLoss()(torch.log(probs), y_cls)
    print(f"
Softmax + NLLLoss:              {loss_wrong.item()}")
except Exception as e:
    print(f"
Softmax + NLLLoss:              {loss_wrong.item() if 'loss_wrong' in dir() else 'NaN or inf'}")

# RIGHT — CrossEntropyLoss uses log-softmax internally
loss_right = nn.CrossEntropyLoss()(large_logits, y_cls)
print(f"CrossEntropyLoss (raw logits):  {loss_right.item():.6f}  ← stable")

# ── Choosing the right final layer + loss combination ─────────────────
print("
Correct output layer + loss combinations:")
combos = [
    ('Binary classification',  'No activation (raw logits)', 'nn.BCEWithLogitsLoss()'),
    ('Multi-class',            'No activation (raw logits)', 'nn.CrossEntropyLoss()'),
    ('Regression',             'No activation (linear out)', 'nn.MSELoss() or nn.L1Loss()'),
    ('Multi-label binary',     'No activation (raw logits)', 'nn.BCEWithLogitsLoss()'),
    ('Probabilistic output',   'nn.Softmax() at inference',  'nn.CrossEntropyLoss() in training'),
]
for task, output, loss in combos:
    print(f"  {task:<25}: output={output:<30} loss={loss}")

import torch
import torch.nn as nn
from sklearn.metrics import accuracy_score, mean_absolute_error
import numpy as np
import warnings
warnings.filterwarnings('ignore')

torch.manual_seed(42)
np.random.seed(42)
n = 1000

# ── Task 1: Regression — delivery time prediction ─────────────────────
X_reg = torch.randn(n, 4)
y_reg = (X_reg[:, 0] * 7 + X_reg[:, 2] * 3 + torch.randn(n) * 2 + 35).unsqueeze(1)

model_reg = nn.Sequential(
    nn.Linear(4, 32), nn.ReLU(),
    nn.Linear(32, 1),              # linear output — no activation
)
opt_reg  = torch.optim.Adam(model_reg.parameters(), lr=0.01)
loss_reg = nn.MSELoss()

for _ in range(200):
    opt_reg.zero_grad()
    loss_reg(model_reg(X_reg), y_reg).backward()
    opt_reg.step()

with torch.no_grad():
    pred_reg = model_reg(X_reg)
    mae = mean_absolute_error(y_reg.numpy(), pred_reg.numpy())
print(f"Regression (MSELoss): MAE = {mae:.4f}")

# ── Task 2: Binary classification — fraud detection ───────────────────
X_cls = torch.randn(n, 4)
y_cls = ((X_cls[:, 0] + X_cls[:, 1] > 0.5).float())

model_bin = nn.Sequential(
    nn.Linear(4, 32), nn.ReLU(),
    nn.Linear(32, 1),              # raw logits — NO sigmoid
)
opt_bin  = torch.optim.Adam(model_bin.parameters(), lr=0.01)
loss_bin = nn.BCEWithLogitsLoss()  # combines sigmoid + BCE stably

for _ in range(200):
    opt_bin.zero_grad()
    loss_bin(model_bin(X_cls).squeeze(), y_cls).backward()
    opt_bin.step()

with torch.no_grad():
    logits  = model_bin(X_cls).squeeze()
    preds   = (logits > 0).float()   # threshold at 0 = threshold at 0.5 prob
    acc_bin = accuracy_score(y_cls.numpy(), preds.numpy())
print(f"Binary classification (BCEWithLogitsLoss): accuracy = {acc_bin:.4f}")

# ── Task 3: Multi-class — support ticket routing ──────────────────────
X_mc  = torch.randn(n, 8)
y_mc  = torch.randint(0, 4, (n,))
# Add signal
for c in range(4):
    X_mc[y_mc == c, c] += 2.0

model_mc = nn.Sequential(
    nn.Linear(8, 32), nn.ReLU(),
    nn.Linear(32, 4),              # 4 raw logits — NO softmax
)
opt_mc  = torch.optim.Adam(model_mc.parameters(), lr=0.01)
loss_mc = nn.CrossEntropyLoss()   # log-softmax + NLL internally

for _ in range(200):
    opt_mc.zero_grad()
    loss_mc(model_mc(X_mc), y_mc).backward()
    opt_mc.step()

with torch.no_grad():
    logits_mc = model_mc(X_mc)
    preds_mc  = logits_mc.argmax(dim=1)
    acc_mc    = accuracy_score(y_mc.numpy(), preds_mc.numpy())
    # At inference: apply softmax to get probabilities
    probs_mc  = torch.softmax(logits_mc, dim=1)
print(f"Multi-class (CrossEntropyLoss): accuracy = {acc_mc:.4f}")

Three weeks into training a fraud-scoring MLP, validation AUC plateaus well below where a similar model performed on a previous dataset. Loss is still decreasing, slowly, and nothing in the training logs looks obviously broken. An engineer adds forward hooks that log the fraction of each ReLU layer's outputs that are exactly zero, on every batch, and plots it over training.

Layer 2's dead-unit fraction starts around 8 percent at initialisation and climbs steadily to 46 percent by epoch 20 — and it is still rising. Nearly half the neurons in that layer produce zero output, and therefore zero gradient, for every single input. The root cause: a learning rate of 0.01 combined with a large negative bias initialisation pushed many pre-activations permanently negative in the first few hundred steps, and ReLU's zero gradient for negative inputs meant those neurons could never recover on their own.

The fix was two changes, not one: switch that layer to LeakyReLU so a small gradient always flows even for negative pre-activations, and lower the learning rate from 0.01 to 0.003. Dead-unit tracking stayed in the training pipeline permanently afterward — logged as a metric alongside loss and accuracy, because by the time AUC visibly plateaus, the dead units have usually been accumulating for many epochs already.

import torch
import torch.nn as nn

# ── Dead ReLU monitoring — a hook real teams bolt onto training ────────
# Attach to every ReLU layer, log the fraction of zero outputs per batch.
# A silently rising dead-unit fraction is the earliest warning sign of a
# training run that will plateau well below its real ceiling.

dead_fraction_log = {}

def make_relu_hook(name):
    def hook(module, layer_input, layer_output):
        dead = (layer_output == 0).float().mean().item()
        dead_fraction_log.setdefault(name, []).append(dead)
    return hook

model = nn.Sequential(
    nn.Linear(20, 128), nn.ReLU(),
    nn.Linear(128, 64), nn.ReLU(),
    nn.Linear(64, 1),
)

for name, layer in model.named_modules():
    if isinstance(layer, nn.ReLU):
        layer.register_forward_hook(make_relu_hook(name))

torch.manual_seed(42)
X = torch.randn(256, 20)

# Simulate several epochs of forward passes (training loop omitted)
for epoch in range(5):
    _ = model(X)

print(f"{'ReLU layer':<12} {'dead units (last batch)':>24}")
print("─" * 38)
for name, history in dead_fraction_log.items():
    print(f"  {name:<10} {history[-1]*100:>22.1f}%")

print("
Rule of thumb used on real teams: more than roughly 30-40% dead")
print("units in any hidden layer, especially if the fraction is still")
print("climbing epoch over epoch, is worth investigating before training")
print("finishes — it rarely self-corrects, and it means the effective")
print("network capacity is much smaller than the parameter count suggests.")

ReLU is a strong default, not a universal winner. Transformers — BERT, GPT, and effectively every modern large language model — use GELU by default because its smoothness (it is differentiable everywhere, unlike ReLU's hard kink at zero) measurably helps optimisation at that scale. LeakyReLU exists specifically because plain ReLU can permanently zero out neurons. RNNs and LSTMs still lean on tanh internally for zero-centred activations. "Default choice for hidden layers" is exactly right, but defaults get overridden the moment the architecture or task gives a reason to.

Softmax guarantees the outputs are non-negative and sum to 1 — that is a mathematical property of the formula, not a guarantee that a 0.97 output means the model is right 97% of the time. Modern, overparameterised networks are frequently overconfident: cross-entropy training pushes logits toward extreme values even after the network has already learned to classify correctly, inflating softmax outputs toward 0 or 1 well beyond what the true error rate justifies. Getting genuinely calibrated probabilities out of a classifier usually needs an explicit extra step, such as temperature scaling or Platt scaling, applied after training.

MSE measures squared distance from a target value, which is the right question for regression and the wrong question for classification. As this module's errors section shows directly: MSE on a binary problem is minimised by outputting the class mean (say 0.3 for a 30% positive rate) for every single input, because that constant genuinely minimises average squared error — the loss trains successfully by every metric except the one that actually matters. Cross-entropy is not an arbitrary alternative; it is the loss whose gradient actually pushes the network to separate classes.

In exact real-number arithmetic they are mathematically identical — but neural networks run in 32-bit floating point, where sigmoid(large logit) rounds to exactly 1.0 or 0.0 before log() ever sees it, producing log(0), which is negative infinity, which becomes NaN the instant it is combined with anything else. BCEWithLogitsLoss avoids this entirely by using the log-sum-exp identity to compute the same mathematical quantity without ever taking log(0) along the way. This is a correctness fix for a real failure mode, not a performance tweak — the naive version can silently produce wrong or NaN losses in production.

The two choices are coupled, not independent — this module's entire decision-guide table exists because getting one right without the other still breaks training. Using GELU in every hidden layer does nothing to fix a classifier trained with MSELoss; using the correct CrossEntropyLoss does nothing to fix a network that saturates because its output layer applies softmax before the loss instead of passing raw logits. The output layer's activation (or lack of one) and the loss function have to be chosen as a matched pair for the specific task — regression, binary, multi-class, or multi-label — not picked separately.

Sigmoid squashes a single value independently to (0, 1) — appropriate for binary classification (one logit, one probability) or multi-label problems where each of several outputs is an independent yes/no decision. Softmax takes a whole vector of scores and converts them jointly into a probability distribution that sums to exactly 1 across all classes — appropriate when exactly one class is correct out of several. A binary problem can technically be modelled as 2-class softmax instead of 1-unit sigmoid, and the two are mathematically equivalent in that special case, but sigmoid with a single output is simpler and is what BCEWithLogitsLoss expects.

If a logit is large (say 100), sigmoid(100) rounds to exactly 1.0 in float32 due to limited precision. BCELoss then needs log(1 − 1.0) = log(0), which is negative infinity, and that propagates as NaN through the rest of training. BCEWithLogitsLoss avoids ever computing sigmoid and log as separate operations — it uses the log-sum-exp trick to compute the mathematically equivalent quantity in a form that never requires taking the log of exactly zero, regardless of how large or small the logit is. That is why PyTorch's own documentation recommends it over the two-step version unconditionally.

MSE is being used for a classification task. MSE is minimised, for a fixed input distribution, by outputting something close to the mean of the targets — for a class- imbalanced binary problem that constant is close to the positive rate itself, and the model can get quite low average loss without ever learning to separate the classes. The fix is to switch to BCEWithLogitsLoss for binary or CrossEntropyLoss for multi-class — losses whose gradients specifically reward pushing predicted probabilities toward the correct class, rather than just minimising average squared distance.

When the data has real outliers you do not want to ignore but also do not want to dominate training. MSE squares the error, so one wildly off prediction contributes disproportionately to the total loss and can drag the whole model toward accommodating that single point. MAE treats every error proportionally regardless of size, which is robust to outliers but has a constant-magnitude gradient that can cause oscillation right at convergence. Huber loss behaves like MSE for small errors — smooth gradients near the optimum — and like MAE for large errors beyond the delta threshold, giving outlier robustness without sacrificing smooth convergence.

Start with weighted BCEWithLogitsLoss using pos_weight set to roughly the negative-to- positive ratio (about 99 here) — this makes a missed positive example cost proportionally more than a missed negative, counteracting the model's natural tendency to default toward the majority class. If that is not enough, focal loss goes further by explicitly downweighting easy, already-well-classified examples so the gradient signal concentrates on the hard, informative ones. Either way, accuracy stops being a meaningful metric at this imbalance — a model predicting "negative" for everything scores 99% accuracy while being useless, so evaluation should shift to precision, recall, and PR-AUC instead.