Loss functions and optimizers

The idea

Training needs two more pieces:

• a loss function that turns predictions and targets into one number to minimise (you met MSE and log loss in the ML track);
• an optimizer that uses the gradients to update every parameter (gradient descent, with refinements).

PyTorch provides both, so the training step shrinks to a few standard lines.

Losses for each problem type

import torch
from torch import nn

pred = torch.tensor([52.0, 61.0, 44.0])
actual = torch.tensor([50.0, 65.0, 44.0])
print(nn.MSELoss()(pred, actual), nn.L1Loss()(pred, actual))

tensor(6.6667) tensor(2.)

Loss objects are created once, then called like functions: loss_fn(predictions, targets).

Logits and BCEWithLogitsLoss

For binary problems, the model outputs a raw score (logit), and the loss applies the sigmoid itself:

logits = torch.tensor([2.0, -1.0, 0.5])      # raw model outputs
targets = torch.tensor([1.0, 0.0, 0.0])
loss_fn = nn.BCEWithLogitsLoss()
print(loss_fn(logits, targets))

probs = torch.sigmoid(logits)
manual = -(targets * torch.log(probs) + (1 - targets) * torch.log(1 - probs)).mean()
print(manual)

tensor(0.4714)
tensor(0.4714)

The same log loss as the ML track. Why not apply sigmoid in the model and use nn.BCELoss on probabilities? Because for very confident predictions, log(sigmoid(z)) loses precision (a probability rounds to exactly 1.0, and log(0) appears). BCEWithLogitsLoss combines the two steps in a numerically stable way. Rule: models return logits; losses take logits.

Multi-class: CrossEntropyLoss

For k classes (say, which of 5 product categories a review is about), the model returns k logits per example, and the target is the index of the right class:

logits = torch.tensor([[2.0, 0.5, -1.0],      # example 1: class 0 looks likely
                       [0.1, 0.2, 3.0]])      # example 2: class 2 looks likely
targets = torch.tensor([0, 2])                # correct class indices (int64)
print(nn.CrossEntropyLoss()(logits, targets))
print(torch.softmax(logits, dim=1).round(decimals=3))

tensor(0.1755)
tensor([[0.7860, 0.1750, 0.0390],
        [0.0490, 0.0540, 0.8960]])

CrossEntropyLoss applies softmax (which turns k logits into k probabilities summing to 1) and log loss in one stable step. Again: no softmax in the model.

Optimizers

An optimizer holds a model's parameters and updates them from their .grad:

torch.manual_seed(0)
model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
x = torch.tensor([[1.0], [2.0], [3.0]])
y = torch.tensor([[3.0], [5.0], [7.0]])
loss_fn = nn.MSELoss()

for step in range(3):
    optimizer.zero_grad()            # 1. clear old gradients
    loss = loss_fn(model(x), y)      # 2. forward pass + loss
    loss.backward()                  # 3. compute gradients
    optimizer.step()                 # 4. update parameters
    print(step, round(loss.item(), 3))

0 22.744
1 4.525
2 0.911

Those four lines are the heart of every PyTorch training loop. Compare with the autograd module: zero_grad() replaces the manual .grad.zero_() calls, and step() replaces the update inside no_grad().

Which optimizer?

• SGD (torch.optim.SGD): plain gradient descent; with momentum=0.9 it keeps some of the previous direction, smoothing the path.
• Adam (torch.optim.Adam): adapts the step size for each parameter from its gradient history. Works well with little tuning, and is the usual starting point (lr=1e-3).
• AdamW: Adam with properly applied weight decay (L2 regularisation): the default for many modern models.

for name, opt_class in [("SGD", torch.optim.SGD), ("Adam", torch.optim.Adam)]:
    torch.manual_seed(0)
    model = nn.Linear(1, 1)
    optimizer = opt_class(model.parameters(), lr=0.05)
    for _ in range(200):
        optimizer.zero_grad()
        loss = loss_fn(model(x), y)
        loss.backward()
        optimizer.step()
    print(name, round(loss.item(), 5), round(model.weight.item(), 3), round(model.bias.item(), 3))

SGD 0.00014 1.986 1.031
Adam 0.04756 1.743 1.57

SGD reaches w = 2, b = 1 within 200 steps; Adam, with the same learning rate, is still on its way. Adam's per-parameter step sizes shine on big, messy problems with many parameters, not on a one-weight toy, and each optimizer has its own good range of learning rates. The learning rate is still the most important setting, whatever the optimizer.

Why this matters for clean code

Losses and optimizers are configuration, not logic: create them once, near the model, from clear settings (lr, weight_decay). Mismatches are the classic silent bugs: softmax in the model and in CrossEntropyLoss (training works, but worse), integer targets for BCE, float targets for CrossEntropy. Keeping the "logits in, loss out" convention everywhere prevents them.

Recap

• Regression: nn.MSELoss / nn.L1Loss. Binary: nn.BCEWithLogitsLoss on logits with float 0/1 targets. Multi-class: nn.CrossEntropyLoss on k logits with int64 class indices.
• Models return logits; losses apply sigmoid/softmax internally for numerical stability.
• Training step: optimizer.zero_grad(), loss = loss_fn(model(x), y), loss.backward(), optimizer.step().
• SGD (± momentum), Adam, AdamW; the learning rate matters most.