Loss Functions and Task Alignment

A loss function turns a training task into a scalar number that PyTorch can differentiate. The model produces tensors, the target tensors describe the desired answer, and the loss measures the mismatch in a way that creates useful gradients. Task alignment means those three pieces agree: the model output representation, the target encoding, and the mathematical assumptions inside the loss.

In this training section, loss selection is the link between forward computation and parameter updates. An optimizer only follows gradients; it does not know whether the objective represents house-price error, species classification, multilabel tagging, language modeling, or ranking. By the end of this lesson, you should be able to inspect output shapes, target shapes, activations, reductions, and weights, then choose and debug a loss that matches the learning problem.

What a PyTorch Loss Actually Does

Every standard PyTorch loss module or functional loss returns a tensor computed from predictions and targets. During training, that tensor is usually reduced to a scalar by taking a mean or sum. Calling backward() on that scalar asks autograd to compute partial derivatives of the loss with respect to tensors that contributed to it and have requires_grad=True. The optimizer then uses those gradients to update parameters.

The most important detail is that a loss is not only a score. It defines the shape of the optimization landscape. Mean squared error makes large numeric errors grow quadratically, so outliers can dominate. Cross entropy combines a logarithmic penalty with a normalization over classes, so it strongly rewards assigning probability mass to the correct class. Binary cross entropy treats each output position as an independent yes-or-no decision, which is different from choosing exactly one class.

Loss Anatomy

A loss choice has five practical parts. First is the output tensor: raw values, probabilities, logits, or structured predictions. Second is the target tensor: floats for regression, integer class indices for single-label classification, or multi-hot floats for multilabel classification. Third is the reduction: mean, sum, or none. Fourth is optional weighting, used to change the importance of classes, samples, or positions. Fifth is the activation contract: some losses expect raw logits and apply a stable transformation internally, while others expect already transformed values.

For classification in PyTorch, logits are preferred during training. A logit is an unconstrained real number produced by the model before softmax or sigmoid. CrossEntropyLoss expects logits and internally applies a numerically stable log-softmax plus negative log-likelihood. BCEWithLogitsLoss expects logits and internally applies a stable sigmoid formulation. Adding softmax before CrossEntropyLoss or sigmoid before BCEWithLogitsLoss weakens or distorts gradients and can create numerical problems.

Example 1: Regression With Mean Squared Error

For a regression model, output and target usually have the same shape and comparable units. Mean squared error computes the average squared difference. In this example the prediction errors are 0.5 and 1.0, so the mean squared error is (0.25 + 1.0) / 2 = 0.625. The gradient on each prediction equals the error when the reduction is the mean over two elements.

import torch
from torch import nn

pred = torch.tensor([[2.0], [4.0]], requires_grad=True)
target = torch.tensor([[1.5], [3.0]])
loss = nn.MSELoss()(pred, target)
loss.backward()

print(round(loss.item(), 3))
print([round(x, 3) for x in pred.grad.flatten().tolist()])

The expected output is 0.625 followed by [0.5, 1.0]. This is aligned when a one-unit error has the same meaning throughout the target range. It can be poorly aligned when the target has extreme outliers, when relative error matters more than absolute error, or when the target scale changes between datasets. In those cases, consider target normalization, L1Loss, SmoothL1Loss, or a loss written in the business unit that actually matters.

Example 2: Single-Label Classification With Cross Entropy

Single-label classification means exactly one class is correct for each example. The model should return one logit per class. Targets should be integer class indices, not one-hot rows. CrossEntropyLoss compares the correct class logit against all other class logits through a softmax normalization.

import torch
from torch import nn

logits = torch.tensor([[2.0, 0.0, -1.0]])
target = torch.tensor([0])
loss = nn.CrossEntropyLoss()(logits, target)

print(round(loss.item(), 6))

The expected output is 0.169846. The loss is small because class 0 has the highest logit by a clear margin. If the target were 1, the loss would be much larger because the model assigned more evidence to class 0. This loss is aligned for tasks such as digit recognition, image category prediction, or intent classification where one answer excludes the others.

A common mistake is to put nn.Softmax(dim=1) at the end of the model and then pass probabilities into CrossEntropyLoss. The model may still train slowly, but the objective no longer matches the stable logit-based formula. Keep the final layer linear during training and apply softmax only when reporting probabilities for evaluation or inference.

Example 3: Class Imbalance With Weighted Cross Entropy

When rare classes matter, unweighted mean loss can encourage a model to perform well on common classes while ignoring rare ones. CrossEntropyLoss(weight=...) multiplies the per-example loss by the target class weight. With mean reduction, PyTorch normalizes by the sum of the selected weights rather than just the batch size.

import torch
from torch import nn

logits = torch.tensor([[2.0, -1.0], [0.2, 1.4], [1.2, 0.1]])
target = torch.tensor([0, 1, 1])
weights = torch.tensor([1.0, 3.0])
losses = nn.CrossEntropyLoss(weight=weights, reduction="none")(logits, target)
mean_loss = nn.CrossEntropyLoss(weight=weights)(logits, target)

print([round(x, 4) for x in losses.tolist()])
print(round(mean_loss.item(), 4))

The expected output is [0.0486, 0.7898, 4.162] and then 0.7143. The third example is costly because the target is class 1 but the logit for class 0 is higher. The class weight makes this rare-class error more influential. This is useful when missing the rare class is expensive, but it is not free: heavier weights can increase false positives, amplify noisy labels, and make training less stable.

Design Choices and Trade-Offs

Choose the loss from the task semantics before tuning the model. For numeric prediction, decide whether absolute error, squared error, relative error, or thresholded error best represents success. For single-label classification, use cross entropy with integer targets. For multilabel classification, use BCEWithLogitsLoss with one binary target per label because multiple labels can be true at the same time. For ordinal targets, plain cross entropy may ignore ordering; predicting class 1 when the answer is 2 is treated no differently from predicting class 9. For dense segmentation or language modeling, the same loss may be applied over many positions after reshaping tensors and masking ignored positions.

Reduction is another design choice. mean keeps gradient magnitudes roughly comparable across batch sizes, which is usually convenient. sum makes the gradient grow with the number of elements and can require learning-rate changes. none is best for debugging, masking padded tokens, applying custom sample weights, or logging which examples dominate training.

Reliability depends on numerical stability. Prefer CrossEntropyLoss over manually composing softmax, log, and negative indexing. Prefer BCEWithLogitsLoss over sigmoid followed by BCELoss. These combined losses avoid avoidable overflow and underflow when logits are large. Performance also depends on tensor layout: flattening token logits incorrectly can silently train against the wrong targets, while unnecessary probability tensors can add memory pressure for large vocabularies or images.

Failure Modes and Troubleshooting

• Symptom: CrossEntropyLoss raises a target type or shape error. Cause: targets are one-hot floats such as [batch, classes] instead of integer ids such as [batch]. Diagnose: print logits.shape, target.shape, and target.dtype. Correct: convert one-hot targets with argmax(dim=1) only if exactly one class is true, or switch to BCEWithLogitsLoss for multilabel data.
• Symptom: training loss decreases but validation accuracy stays near chance. Cause: output columns and target class ids use different class ordering. Diagnose: inspect the dataset class-to-index mapping and the final layer size, then overfit a tiny labeled batch and check predictions manually. Correct: persist one class mapping and reuse it for dataset creation, model interpretation, metrics, and inference.
• Symptom: loss becomes nan or suddenly explodes. Cause: excessive learning rate, invalid target values, manual log of zero, or unstable probability-based loss composition. Diagnose: check for nonfinite inputs, logits, loss, and gradients after each stage on a small batch. Correct: use logit-based built-in losses, lower the learning rate, normalize inputs or targets, and clip gradients when the model architecture justifies it.
• Symptom: a model predicts only the majority class. Cause: the unweighted loss is aligned with average accuracy but not with rare-class recall. Diagnose: compare per-class precision, recall, confusion matrix rows, and unreduced loss by class. Correct: add class weights, resample training data, adjust thresholds, or change the metric and model-selection criterion.

Hands-On Lab: Audit and Train a Tiny Classifier

Prerequisites: a Python environment with PyTorch installed, basic tensor familiarity, and a terminal or notebook where you can run short scripts. The lab checks the alignment between logits, integer targets, loss value, and gradients.

1. Create a two-class linear model with nn.Linear(2, 2). Do not add softmax to the model.
2. Build a small batch of four two-dimensional points and integer targets such as torch.tensor([0, 0, 1, 1]).
3. Run a forward pass and print logits shape and target shape. Verify that they are [4, 2] and [4].
4. Compute nn.CrossEntropyLoss(reduction="none") and print the four per-example losses. Identify which examples currently dominate the batch.
5. Switch to default mean reduction, call backward(), and verify that every trainable parameter has a non-None gradient.
6. Run 50 optimizer steps with a modest learning rate. Verify that loss on the same tiny batch decreases. This proves the wiring is trainable, not that the model generalizes.
7. Cleanup by deleting the experimental script or notebook cell outputs. If you changed a shared environment, remove temporary files and restore any dependency changes you made for the lab.

Verification should include the printed shapes, a finite initial loss, non-None gradients, and a lower final loss on the tiny batch. If loss does not decrease, inspect target dtype, remove accidental softmax, lower the learning rate, and confirm the labels are not shuffled independently from the inputs.

Assessment Exercises

1. A model returns logits of shape [32, 5], and the dataset returns targets of shape [32, 5] containing zeros and ones. Should you use CrossEntropyLoss? Explain what task assumption decides the answer.
2. You train a regression model with MSELoss, but a few mislabeled examples dominate the gradients. Propose two changes and explain the trade-off each introduces.
3. For token classification, logits have shape [batch, tokens, classes]. Describe how you would reshape logits and targets for CrossEntropyLoss, and where you would use an ignore index for padding.
4. A rare but important class has high recall after adding a large class weight, but precision collapses. What diagnostics would you inspect before deciding whether the weight is too large?
5. Explain why applying sigmoid before BCEWithLogitsLoss is a loss-alignment bug even if the tensor shapes are correct.

Summary

Loss functions encode the training task. In PyTorch, correct alignment means the output representation, target encoding, activation contract, reduction, and weighting all match the problem. Use mean squared or related losses for numeric targets, cross entropy for mutually exclusive classes, and binary cross entropy with logits for independent labels. Debug loss problems by checking shapes, dtypes, class mappings, unreduced losses, finite values, and gradients on a small batch before trusting full training runs.