Loss Functions

Learning Objectives

What is a Loss Function?

A loss function (also called a cost function or objective function) is the mathematical compass that guides neural network training. It quantifies how wrong your model's predictions are compared to the ground truth, providing a single number that the optimizer works to minimize.

The Core Insight: Neural networks learn by adjusting their parameters to minimize the loss. The loss function defines what "good" means for your task—it translates the abstract goal of "make better predictions" into a concrete mathematical objective that gradient descent can optimize.

Think of the loss function as the feedback signal:

1. Measures prediction quality — Compares model output to ground truth
2. Provides gradients — Tells each parameter how to change to improve
3. Shapes learning dynamics — Different losses lead to different learned behaviors

The Learning Loop

Every training step follows this cycle:

Mean Squared Error (MSE)

Mean Squared Error is the workhorse of regression tasks. It measures the average of the squared differences between predictions and targets:

Where:

• $n$ is the number of samples in the batch
• $y_i$ is the true value (target) for sample $i$
• $\hat{y}_i$ is the predicted value for sample $i$

Why Square the Errors?

Squaring serves several important purposes:

The Gradient of MSE

The gradient of MSE with respect to a prediction is:

This elegant result means the gradient is proportional to the error—larger errors produce larger gradients, pushing the model harder to fix bigger mistakes.

Interactive: MSE Explorer

Explore how MSE works by adjusting the regression line. Watch how the squared errors (visualized as squares) change size, and how the total loss responds to the fit quality.

Cross-Entropy Loss

While MSE works well for regression, cross-entropy loss is the standard for classification tasks. It measures the difference between two probability distributions: the true label distribution and the model's predicted distribution.

Information-Theoretic Foundation

Cross-entropy comes from information theory. For probability distributions $p$ (true) and $q$ (predicted):

This measures the expected number of bits needed to encode data from distribution $p$ using a code optimized for distribution $q$. When $q = p$, cross-entropy equals the entropy $H(p)$—the optimal encoding.

Connection to KL Divergence

Cross-entropy decomposes as:

Where $H(p)$ is the entropy of the true distribution (constant during training) and $D_{\text{KL}}(p \| q)$ is the KL divergence. Thus, minimizing cross-entropy is equivalent to minimizing KL divergence—making our predictions match the true distribution.

Why Cross-Entropy for Classification? The true distribution $p$ is typically one-hot (all probability on the correct class). Cross-entropy then reduces to $-\log q(\text{true class})$—maximizing the log-probability of the correct class. This is equivalent to Maximum Likelihood Estimation!

Binary Cross-Entropy (BCE)

For binary classification (two classes), we use Binary Cross-Entropy:

Where:

• $y_i \in \{0, 1\}$ is the true binary label
• $\hat{p}_i \in (0, 1)$ is the predicted probability of class 1 (from sigmoid)

The Beautiful Gradient

When combined with sigmoid activation, BCE has an elegant gradient:

Where $z$ is the logit (pre-sigmoid value). This simplification occurs because the sigmoid derivative $\sigma(z)(1-\sigma(z))$ cancels with terms in the cross-entropy gradient!

Interactive: BCE Explorer

Explore how Binary Cross-Entropy works with sigmoid activation. Adjust the logit value and watch how the loss and gradients respond. Try the gradient descent animation to see the optimization in action.

Multi-Class Cross-Entropy

For classification with more than two classes, we extend to categorical cross-entropy:

Where:

• $C$ is the number of classes
• $y_c$ is 1 if class $c$ is correct, 0 otherwise (one-hot)
• $\hat{p}_c$ is the predicted probability for class $c$ (from softmax)

Simplification with One-Hot Targets

When targets are one-hot encoded, only the term for the true class contributes:

The loss is simply the negative log-probability assigned to the correct class. This connects directly to Maximum Likelihood Estimation—we're maximizing the likelihood of the data under our model.

Softmax + Cross-Entropy

In PyTorch, nn.CrossEntropyLoss combines log-softmax and negative log-likelihood:

Where $z_k$ is the logit for the correct class $k$. This form is numerically stable and efficient.

Interactive: Cross-Entropy Explorer

Explore multi-class cross-entropy loss. Adjust the predicted probabilities for each class and observe how the loss changes. Notice that only the probability of the true class matters!

MSE vs Cross-Entropy

Why does cross-entropy dominate classification while MSE rules regression? The answer lies in their gradient behaviors.

Key Differences

Interactive: 3D Loss Landscape

Visualizing loss functions in 3D helps build intuition about how gradient descent navigates the optimization landscape. The surface shows how loss varies with predicted probabilities, and the path traces gradient descent from random initialization to the optimum.

Key Observation: Notice how cross-entropy creates steep "cliffs" when predictions are confidently wrong, while MSE forms a gentler bowl. These steep gradients are exactly why cross-entropy trains classification models faster—the model is strongly pushed away from bad predictions.

Interactive: Gradient Flow

Watch how gradients flow through a neural network with different loss functions. Compare cross-entropy and MSE to see why cross-entropy provides better gradient flow, especially in early layers.

Numerical Stability

One of the most common training bugs involves numerical instability in loss computation. When logits are large, naive softmax computation can overflow, producing NaN losses. Understanding and avoiding these pitfalls is essential for robust training.

The Log-Sum-Exp Trick

The key insight is that we can subtract the maximum logit before exponentiating without changing the result:

This ensures all exponents are ≤ 0, preventing overflow. Try the interactive demo below with extreme logit values to see the difference:

Weighted Loss for Imbalanced Data

Real-world datasets are often imbalanced—some classes appear much more frequently than others. A naive model might achieve 95% accuracy by always predicting the majority class, while completely ignoring the minority class. Class weighting solves this problem.

The Class Weighting Formula

We modify the loss to weight each sample by its class:

Common weighting strategies:

Contrastive Losses

Contrastive learning has revolutionized self-supervised learning. Instead of predicting labels, we learn embeddings where similar items are close and dissimilar items are far apart. These losses power models like CLIP, SimCLR, and modern face recognition.

Triplet Loss

The foundation of metric learning, introduced in FaceNet for face recognition:

Where $a$ is the anchor, $p$ is a positive (same class), and $n$ is a negative (different class).

InfoNCE Loss

The loss behind CLIP and SimCLR, treating contrastive learning as classification:

Where $\tau$ is the temperature parameter controlling sharpness.

Other Loss Functions

Beyond MSE and cross-entropy, several specialized loss functions address specific challenges:

L1 Loss (Mean Absolute Error)

More robust to outliers than MSE since errors are not squared. However, the gradient is discontinuous at zero, which can cause training instability.

Huber Loss (Smooth L1)

The best of both worlds: quadratic for small errors (smooth gradients), linear for large errors (outlier robustness). Common in object detection for bounding box regression.

Focal Loss

Designed for imbalanced datasets. The $(1 - p_t)^\gamma$ factor down-weights easy examples (high $p_t$), focusing training on hard examples. Introduced in RetinaNet for object detection.

Label Smoothing

Instead of one-hot targets like $[1, 0, 0, 0]$, use soft targets like $[0.9, 0.033, 0.033, 0.033]$. This prevents overconfidence and improves generalization by encouraging the model to maintain some uncertainty.

Domain-Specific Losses

Different domains have developed specialized losses tailored to their unique challenges. Here's a reference of the most important ones:

Image Segmentation

Object Detection

Generative Models

Natural Language Processing

Calibration & Confidence

A well-calibrated model's predicted probabilities reflect true likelihoods. If a model predicts 70% confidence, it should be correct 70% of the time. Modern neural networks are often overconfident—they predict high confidence even when wrong.

Why Calibration Matters

• Decision making: Medical diagnosis, autonomous driving require knowing when the model is uncertain
• Model selection: Uncalibrated confidence makes it hard to compare models
• Uncertainty estimation: Downstream tasks depend on accurate uncertainty

Expected Calibration Error (ECE)

The standard metric for calibration quality:

Lower ECE is better. A perfectly calibrated model has ECE = 0.

Debugging Loss Functions

Training neural networks involves many potential pitfalls. Understanding common loss-related issues and their solutions can save hours of debugging time.

Essential Debugging Code

# Check for NaN/Inf in forward pass
def check_tensor(name, tensor):
    if torch.isnan(tensor).any():
        print(f"NaN detected in {name}!")
    if torch.isinf(tensor).any():
        print(f"Inf detected in {name}!")

# Monitor gradient magnitudes
for name, param in model.named_parameters():
    if param.grad is not None:
        grad_norm = param.grad.norm()
        if grad_norm > 100:
            print(f"Large gradient in {name}: {grad_norm:.2f}")

# Use anomaly detection mode
with torch.autograd.detect_anomaly():
    loss = model(x)
    loss.backward()  # Will print where NaN/Inf originated

PyTorch Implementation

PyTorch provides all common loss functions in torch.nn. Here's how to use them and implement custom losses.

Built-in Loss Functions

Logits [2.0, 1.0, 0.1] -> softmax [0.66, 0.24, 0.10]

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5# Sample predictions and targets
6predictions = torch.tensor([0.9, 0.2, 0.1])  # Model outputs
7targets = torch.tensor([1.0, 0.0, 0.0])       # Ground truth
8
9# --- Mean Squared Error ---
10mse_loss = nn.MSELoss()
11mse = mse_loss(predictions, targets)
12print(f"MSE Loss: {mse.item():.4f}")  # 0.0067
13
14# --- Binary Cross-Entropy ---
15bce_loss = nn.BCELoss()
16# Predictions must be in [0, 1]
17bce = bce_loss(predictions, targets)
18print(f"BCE Loss: {bce.item():.4f}")  # 0.1536
19
20# --- Cross-Entropy for Classification ---
21# For multi-class, use logits (before softmax)
22logits = torch.tensor([[2.0, 1.0, 0.1]])  # Raw scores
23labels = torch.tensor([0])                 # Class index
24ce_loss = nn.CrossEntropyLoss()
25ce = ce_loss(logits, labels)
26print(f"CE Loss: {ce.item():.4f}")  # 0.4170

Complete Classification Pipeline

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class ImageClassifier(nn.Module):
6    def __init__(self, num_classes=10):
7        super().__init__()
8        self.fc1 = nn.Linear(784, 256)
9        self.fc2 = nn.Linear(256, num_classes)
10        # Note: No softmax here! CrossEntropyLoss handles it
11
12    def forward(self, x):
13        x = x.flatten(1)  # Flatten images
14        x = F.relu(self.fc1(x))
15        x = self.fc2(x)   # Output raw logits
16        return x
17
18# Create model, loss, optimizer
19model = ImageClassifier(num_classes=10)
20criterion = nn.CrossEntropyLoss()
21optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
22
23# Training step
24images = torch.randn(32, 1, 28, 28)  # Batch of 32 MNIST images
25labels = torch.randint(0, 10, (32,))  # Random labels for demo
26
27# Forward pass
28logits = model(images)           # Shape: (32, 10)
29loss = criterion(logits, labels) # Cross-entropy loss
30
31# Backward pass
32optimizer.zero_grad()
33loss.backward()
34optimizer.step()
35
36print(f"Loss: {loss.item():.4f}")

Custom Loss Functions

You can create custom losses by subclassing nn.Module:

Easy (p=0.9): weight=(0.1)^2=0.01, Hard (p=0.1): weight=(0.9)^2=0.81

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5# --- Label Smoothing Loss ---
6class LabelSmoothingLoss(nn.Module):
7    def __init__(self, num_classes, smoothing=0.1):
8        super().__init__()
9        self.smoothing = smoothing
10        self.num_classes = num_classes
11
12    def forward(self, logits, targets):
13        # Create smoothed one-hot targets
14        with torch.no_grad():
15            smooth_targets = torch.zeros_like(logits)
16            smooth_targets.fill_(self.smoothing / (self.num_classes - 1))
17            smooth_targets.scatter_(1, targets.unsqueeze(1), 1.0 - self.smoothing)
18
19        # Cross-entropy with soft targets
20        log_probs = F.log_softmax(logits, dim=-1)
21        loss = -(smooth_targets * log_probs).sum(dim=-1).mean()
22        return loss
23
24# --- Focal Loss (for imbalanced data) ---
25class FocalLoss(nn.Module):
26    def __init__(self, alpha=1.0, gamma=2.0):
27        super().__init__()
28        self.alpha = alpha
29        self.gamma = gamma
30
31    def forward(self, logits, targets):
32        ce_loss = F.cross_entropy(logits, targets, reduction='none')
33        pt = torch.exp(-ce_loss)  # p_t = softmax probability of correct class
34        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
35        return focal_loss.mean()
36
37# Usage
38logits = torch.randn(32, 10)
39targets = torch.randint(0, 10, (32,))
40
41smooth_loss = LabelSmoothingLoss(num_classes=10, smoothing=0.1)
42focal_loss = FocalLoss(gamma=2.0)
43
44print(f"Label Smoothing: {smooth_loss(logits, targets):.4f}")
45print(f"Focal Loss: {focal_loss(logits, targets):.4f}")

Choosing the Right Loss

Selecting the appropriate loss function depends on your task and data characteristics:

Decision Flowchart

1. Classification or Regression? Classification → cross-entropy family; Regression → MSE/MAE family
2. Binary or multi-class? Binary → BCE; Multi-class → CrossEntropyLoss
3. Multi-label (can have multiple correct answers)? Use BCE per label
4. Imbalanced classes? Consider focal loss or class weights
5. Outliers in regression? Use Huber or MAE instead of MSE
6. Model overconfident? Try label smoothing

Real-World Case Studies

Understanding how famous models use loss functions helps build intuition for your own projects. Here are key examples across different domains:

Computer Vision

Natural Language Processing

Other Domains

The Pattern: Notice how most successful models don't just use vanilla cross-entropy. They carefully design or combine losses to match their task's unique requirements. This is where deep learning engineering meets science.

Test Your Understanding

Test your knowledge of loss functions with this quiz. Each question has a detailed explanation to reinforce the concepts.

Summary

Exercises

Conceptual Questions

1. Why does cross-entropy loss produce unbounded gradients as the predicted probability approaches 0, while MSE gradients stay bounded?
2. Explain why minimizing cross-entropy is equivalent to Maximum Likelihood Estimation for classification problems.
3. Under what circumstances would you choose MAE over MSE for a regression problem? What trade-off are you making?

Coding Exercises

1. Implement MSE from Scratch: Write a function that computes MSE loss and its gradient without using any PyTorch loss classes. Verify against nn.MSELoss.
2. Implement Cross-Entropy from Scratch: Write a function that computes cross-entropy loss for multi-class classification, including the softmax step. Compare numerical stability with and without the log-sum-exp trick.
3. Gradient Comparison: Train a binary classifier on a simple dataset using (a) MSE loss and (b) BCE loss. Plot training curves and compare convergence speed.
4. Focal Loss Implementation: Implement Focal Loss and train on an imbalanced dataset (e.g., 90% class 0, 10% class 1). Compare to standard cross-entropy.

Challenge: Multi-Task Loss

Create a neural network that simultaneously performs classification and regression (e.g., predicting both a class label and a continuous value). Design a combined loss function:

• How do you choose $\lambda$?
• What if the two losses are on very different scales?
• Implement and train on a suitable dataset

In the next section, we'll explore Normalization Layers—techniques like Batch Normalization and Layer Normalization that stabilize training and allow for deeper networks by controlling the distribution of activations.