Why Loss Functions Matter and How to Choose One

What a Loss Function Does

A loss function takes the model's predictions and the true labels, and returns a single number representing how wrong the model is. Training minimizes this number. Everything the model learns flows from this signal.

If your loss function does not align with what you actually care about, your model will optimize for something other than what you want - and do it very efficiently.

Cross-Entropy Loss: For Classification

Cross-entropy is the standard loss for classification tasks. It measures the difference between the predicted probability distribution and the true distribution.

Binary cross-entropy (two classes):

pythonCopy
import torch
import torch.nn as nn

criterion = nn.BCEWithLogitsLoss()
# Note: BCEWithLogitsLoss combines sigmoid + BCE for numerical stability
# Input: raw logits (before sigmoid), Target: 0.0 or 1.0

logits = model(X)      # Shape: (batch_size,)
loss = criterion(logits, y.float())

Categorical cross-entropy (multiple classes):

pythonCopy
criterion = nn.CrossEntropyLoss()
# Note: CrossEntropyLoss takes raw logits (before softmax)
# Target: integer class indices (not one-hot)

logits = model(X)      # Shape: (batch_size, num_classes)
loss = criterion(logits, y)  # y: shape (batch_size,), dtype long

Mean Squared Error: For Regression

MSE penalizes large errors heavily (squaring amplifies them). Use when large errors are disproportionately costly.

pythonCopy
criterion = nn.MSELoss()

# Equivalent formula:
mse = ((y_pred - y_true) ** 2).mean()

RMSE (root MSE) is the same metric in the original units, making it easier to interpret.

Mean Absolute Error: For Regression With Outliers

MAE penalizes all errors linearly. More robust to outliers than MSE - a prediction that is 100 units wrong contributes proportionally, not quadratically.

pythonCopy
criterion = nn.L1Loss()  # L1Loss = MAE

# Use when: outliers are legitimate data points, not measurement errors
#           and you want the model to be equally wrong across the full range

Focal Loss: For Severe Class Imbalance

Focal loss down-weights the contribution of easy examples (those the model already predicts correctly with high confidence), forcing the model to focus on the hard cases. Designed for object detection but useful for any severely imbalanced task.

pythonCopy
def focal_loss(logits, targets, gamma=2.0, alpha=0.25):
    bce = nn.functional.binary_cross_entropy_with_logits(logits, targets, reduction='none')
    p_t = torch.exp(-bce)
    focal_loss = alpha * (1 - p_t) ** gamma * bce
    return focal_loss.mean()

Custom Loss Functions: Aligning With Business Metrics

Business metrics often do not have differentiable closed forms. False negatives may be 10x more costly than false positives. A custom loss can express this:

pythonCopy
def asymmetric_loss(logits, targets, fn_weight=10.0):
    """Penalize false negatives fn_weight times more than false positives."""
    probs = torch.sigmoid(logits)
    # False negative: target=1, prediction low
    fn_loss = targets * (-torch.log(probs + 1e-8)) * fn_weight
    # False positive: target=0, prediction high
    fp_loss = (1 - targets) * (-torch.log(1 - probs + 1e-8))
    return (fn_loss + fp_loss).mean()

The Alignment Check

Before committing to a loss function, answer: "If my model perfectly minimizes this loss, will it do what the business needs?"

A model that minimizes cross-entropy on a 99%/1% imbalanced dataset will learn to always predict the majority class. A model that minimizes MAE where a 1% error on large values costs 100x more than a 1% error on small values is ignoring the asymmetry.

The loss function is your contract with the optimizer. Write it precisely.

Common Mistakes

Using MSE for classification problems. Mean squared error treats a predicted probability of 0.4 (when the true label is 1) as only slightly wrong. Cross-entropy loss penalizes confident wrong predictions exponentially, which is what you actually want - you need the model to be heavily discouraged from being confidently incorrect, not just slightly off.

Choosing a loss function before defining the business objective. MSE and MAE both measure regression error, but MAE is robust to outliers while RMSE magnifies them. Picking one before deciding whether outlier errors matter in your use case leads to a model that is optimized for the wrong thing, even though the training loop runs without errors.

Not verifying numerical stability of the loss computation. Computing log(p) directly when p is near zero produces -inf or NaN, silently corrupting gradients and making debugging extremely difficult. Always use framework-provided stable implementations like F.binary_cross_entropy_with_logits or apply the log-sum-exp trick when implementing custom losses.