Loss functions and optimisers in practice

CrossEntropyLoss and its logits expectation, BCEWithLogitsLoss, MSE variants, and SGD versus Adam versus AdamW with correct weight decay.

Choosing a loss

import torch
import torch.nn as nn

logits = torch.randn(8, 5)                 # batch of 8, 5 classes
targets = torch.randint(0, 5, (8,))        # int64 class indices

# CrossEntropyLoss takes raw logits and applies log-softmax internally
ce = nn.CrossEntropyLoss()
loss = ce(logits, targets)
manual = nn.functional.cross_entropy(logits, targets)
print(torch.allclose(loss, manual))        # True

# class weights and label smoothing
weights = torch.tensor([1.0, 2.0, 1.0, 1.0, 5.0])
ce_weighted = nn.CrossEntropyLoss(weight=weights, label_smoothing=0.05)
print(ce_weighted(logits, targets))

# binary / multi-label: always the logits variant
binary_logits = torch.randn(8, 1)
binary_targets = torch.randint(0, 2, (8, 1)).float()
bce = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([3.0]))
print(bce(binary_logits, binary_targets))

# regression
pred = torch.randn(8, 1)
y = torch.randn(8, 1)
print(nn.MSELoss()(pred, y), nn.L1Loss()(pred, y), nn.SmoothL1Loss()(pred, y))

CrossEntropyLoss expects (N, C) logits and (N,) integer targets. Passing a softmax output is the single most common mistake and produces a loss that still decreases, just to a worse optimum.
For multi-label problems use BCEWithLogitsLoss, never softmax: the classes are not mutually exclusive and their probabilities should not sum to one.
pos_weight reweights the positive class for imbalanced binary data. It changes the loss scale, so revisit the learning rate.
SmoothL1Loss is quadratic near zero and linear far out, which makes it robust to outliers in regression.

Optimisers and weight decay

from torch.optim import SGD, Adam, AdamW

# never apply weight decay to norms or biases
def split_params(model, decay=0.01):
    decay_params, no_decay = [], []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if param.ndim <= 1 or name.endswith(".bias"):
            no_decay.append(param)
        else:
            decay_params.append(param)
    return [{"params": decay_params, "weight_decay": decay},
            {"params": no_decay, "weight_decay": 0.0}]

groups = split_params(model)
opt = AdamW(groups, lr=3e-4, betas=(0.9, 0.95), eps=1e-8)

# the canonical step: zero_grad must come before backward, or gradients accumulate
opt.zero_grad(set_to_none=True)
loss = criterion(model(x), y)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
opt.step()

# inspect the current learning rate of each group
for g in opt.param_groups:
    print(g["lr"], g["weight_decay"])

Optimiser	Weight decay behaviour	Recommended use
`SGD(momentum=0.9)`	L2 added to the gradient	Vision models when chasing final accuracy
`Adam`	L2 added to the gradient: inconsistent per-parameter	Legacy code; prefer AdamW
`AdamW`	Decoupled, applied directly to the weights	The default for most modern training
`RMSprop`	L2 added to the gradient	Occasional recurrent workloads
`LBFGS`	Full-batch only	Small deterministic problems

💡

Weight decay is not L2 regularisation in Adam. Adam divides every gradient by an adaptive denominator, so a gradient-based penalty has a different effective strength on every parameter. AdamW applies decay separately and gives what you actually intended.

Loss and optimiser failure modes

Symptom	Cause	Check
Loss is `nan` immediately	Rate too high, or `log(0)` in a hand-written loss	Print logits before the loss; lower the rate 10x
Loss decreases then explodes	Rate too high late in training	Add a schedule or clip gradients
Loss barely moves	Frozen parameters, or `requires_grad=False`	Print `requires_grad` per parameter
Loss identical every step	Gradients are zero	Check that `backward` ran and grads are not `None`
Two identical losses from different batches	In-place op broke autograd, or the model is `eval()`	Compare inputs, check `model.training`

# gradient health check: run after backward, before step
def grad_report(model, top=5):
    rows = []
    for name, p in model.named_parameters():
        if p.grad is None:
            rows.append((name, "None", 0.0))
            continue
        rows.append((name, f"{p.grad.norm().item():.3e}", p.grad.abs().max().item()))
    rows.sort(key=lambda r: r[2], reverse=True)
    for name, norm, mx in rows[:top]:
        print(f"{name:40s} grad_norm={norm} max={mx:.3e}")
    total = sum(p.grad.norm().item() ** 2 for p in model.parameters() if p.grad is not None) ** 0.5
    print("global grad norm:", round(total, 6))

# a parameter with grad Norm 0 and also .grad None means it is unused in the graph
grad_report(model)

FAQ

Should I pass probabilities or logits to CrossEntropyLoss?

Logits, always. The layer computes log-softmax internally in a numerically stable way. If you apply softmax first the loss is still differentiable but mathematically wrong and less stable.

Why is my binary classification loss so large?

Usually a shape mismatch between predictions (N, 1) and targets (N,), which broadcasts into a (N, N) loss. Print both shapes and make them identical before the loss call.

The training loop Building networks: layers, containers and initialisation

Last refreshed 2026-09-18.

Loss functions and optimisers in practice

Choosing a loss

Optimisers and weight decay

Loss and optimiser failure modes

FAQ

Related