01 · Loss Functions 🎯

4 min read

A model can't improve until it knows how wrong it is. The loss function is the number that defines "wrong" — and everything the model learns flows from it.

1. The 'Why'

Training a neural network is really an optimization problem, and every optimization problem needs something to optimize toward. The loss function (also called the cost or objective) is that target: it takes the model's predictions and the true answers and boils them down to a single number measuring how badly the model did. Low loss means good predictions; high loss means poor ones. This one number is the anchor for the entire learning process — it's what autograd differentiates, and it's what the optimizer tries to shrink. Without a well-chosen loss, the model has no compass, no notion of "better," and can't learn anything meaningful.

The critical insight is that different tasks need different notions of "wrong." Predicting a house price (a continuous number) calls for a loss that penalizes being numerically far off — that's mean squared error. Classifying an image into one of ten categories calls for a loss that rewards putting high probability on the correct class — that's cross-entropy. Choosing the wrong loss for your task is one of the most common ways to get a model that trains but never works. This sub-module teaches you the two workhorse losses, when each applies, and the subtle-but-important detail about logits that trips up nearly every beginner using cross-entropy.

2. Core Concepts

Loss = a single scalar. A loss function reduces (predictions, targets) to one number. Because it's scalar, you can call .backward() on it directly (recall the scalar rule from Module 02).

Regression → nn.MSELoss. Mean squared error averages the squared differences between predictions and targets. Use it when the output is a continuous quantity.

Classification → nn.CrossEntropyLoss. Measures the gap between the predicted class distribution and the true class. Use it for multi-class classification. It expects raw logits (unnormalized scores) as input and integer class indices as targets — it applies the softmax internally.

Reduction. By default losses return the mean over the batch (reduction='mean'). You can switch to 'sum' or 'none', but mean is the standard choice.

Loss is a module. You instantiate a loss (criterion = nn.CrossEntropyLoss()) once, then call it each batch (criterion(preds, targets)).

3. Code in Action

Regression with `MSELoss`

import torch
import torch.nn as nn

criterion = nn.MSELoss()                 # mean squared error

preds   = torch.tensor([2.5, 0.0, 2.1])  # model outputs (continuous)
targets = torch.tensor([3.0, -0.5, 2.0]) # ground-truth values

loss = criterion(preds, targets)         # single scalar
print(loss)                              # mean of (preds - targets)^2

Classification with `CrossEntropyLoss`

import torch
import torch.nn as nn

criterion = nn.CrossEntropyLoss()        # for multi-class classification

# RAW LOGITS: shape (batch=2, num_classes=3) — NOT softmaxed
logits = torch.tensor([[2.0, 0.5, 0.1],
                       [0.2, 1.5, 2.3]])
# Targets are INTEGER class indices, not one-hot vectors
targets = torch.tensor([0, 2])           # sample 0 is class 0, sample 1 is class 2

loss = criterion(logits, targets)        # softmax + negative log-likelihood, internally
print(loss)                              # a single scalar

Loss as the first link in the training chain

import torch
import torch.nn as nn

model = nn.Linear(4, 3)                  # a tiny classifier
criterion = nn.CrossEntropyLoss()

x = torch.randn(8, 4)                    # batch of 8
y = torch.randint(0, 3, (8,))            # 8 integer labels

logits = model(x)                        # forward pass -> raw logits (8, 3)
loss = criterion(logits, y)             # measure wrongness
loss.backward()                          # autograd fills every .grad — ready for the optimizer
print(loss.item())                       # .item() pulls the Python float out of the scalar tensor

4. Common Pitfalls

Pitfall 1 — Softmaxing before CrossEntropyLoss. nn.CrossEntropyLoss applies softmax internally, so adding a Softmax (or LogSoftmax) layer before it double-applies the operation and cripples training. Feed it raw logits. (If you truly need probabilities elsewhere, apply softmax separately for display only.)

Pitfall 2 — Wrong target format. CrossEntropyLoss wants integer class indices of shape (batch,), not one-hot vectors and not floats. Passing one-hot targets or the wrong dtype raises errors or gives nonsense. For regression with MSELoss, by contrast, predictions and targets must have matching float shapes.

Pitfall 3 — Shape mismatches. A frequent bug is predictions of shape (batch, 1) versus targets of shape (batch,) in regression, which broadcasts into a wrong loss silently. Print shapes and squeeze/reshape so they line up exactly.

5. Further Reading & Watch List

➡️ Next: 02 · Optimizers