Loss Functions - PyTorch Basics

Loss Functions

What you'll build: a numerically stable softmax and the three losses you will use most: cross-entropy (multi-class classification), mean squared error (regression) and binary cross-entropy (binary classification).

Background

The loss turns predictions into the single number that backward() differentiates. Picking the right one, and feeding it the right input (logits or probabilities, class indices or one-hot), is a constant source of silent bugs and a common interview topic.

Formulas

For logits z of shape (batch, num_classes):

Softmax (over the last dimension, shifted by the row max for stability):

softmax(z)[k] = exp(z[k] - max(z)) / sum_j exp(z[j] - max(z))

Cross-entropy with integer class labels t of shape (batch,), averaged over the batch:

CE = -(1 / batch) * sum_i log softmax(z_i)[t_i]
log softmax(z)[k] = z[k] - max(z) - log(sum_j exp(z[j] - max(z)))

Mean squared error for predictions and targets of the same shape:

MSE = mean((y_pred - y_true) ** 2)

Binary cross-entropy for probabilities p in (0, 1) and labels y in {0, 1}:

BCE = -mean(y * log(p) + (1 - y) * log(1 - p))

Inputs

• softmax(x): array of any shape; normalise over the last axis.
• cross_entropy_loss(logits, targets): logits (batch, num_classes) raw scores, targets (batch,) integers in [0, num_classes).
• mse_loss(y_pred, y_true): same-shaped float arrays.
• binary_cross_entropy(y_pred, y_true): y_pred probabilities, y_true 0/1 labels, same shape.

Expected Return Format

• softmax returns an array of the input's shape whose last axis sums to 1.
• The three losses return a Python float (mean over all elements or samples).

Edge Cases and Conventions

• Compute the cross-entropy through log-softmax, not log(softmax(...)): it stays exact even when one logit dominates (logits [[1000, 0]] with target 1 give a loss of exactly 1000).
• In the NumPy BCE, clip p to [1e-15, 1 - 1e-15] before taking logs so a prediction of exactly 0 or 1 doesn't produce inf.

Constraints

• PyTorch track: use the built-ins F.softmax, F.cross_entropy, F.mse_loss and F.binary_cross_entropy, and learn what each expects.
• NumPy track: implement each formula from scratch, vectorised.

Follow-up

• Why does PyTorch's F.cross_entropy take raw logits instead of probabilities?
• Why is BCEWithLogitsLoss preferred over sigmoid followed by BCELoss?

Input: cross_entropy_loss(np.array([[2.0, 1.0, 0.1]]), np.array([0]))
Output: 0.4170

Input: mse_loss(np.array([1.0, 2.0, 3.0]), np.array([1.0, 2.0, 4.0]))
Output: 0.3333

Input: binary_cross_entropy(np.array([0.9, 0.1, 0.8]), np.array([1.0, 0.0, 1.0]))
Output: 0.1446

import numpy as np

def softmax(x: np.ndarray) -> np.ndarray:
    """Numerically stable softmax over the last axis (same shape as x)."""
    # Your code here
    pass


def cross_entropy_loss(logits: np.ndarray, targets: np.ndarray) -> float:
    """
    Cross-entropy loss for classification, averaged over the batch.

    Args:
        logits: Raw scores (batch, num_classes)
        targets: Integer class labels (batch,)

    Returns:
        Scalar loss (float)
    """
    # Your code here
    pass


def mse_loss(y_pred: np.ndarray, y_true: np.ndarray) -> float:
    """Mean squared error (float)."""
    # Your code here
    pass


def binary_cross_entropy(y_pred: np.ndarray, y_true: np.ndarray) -> float:
    """
    Binary cross-entropy on probabilities (float).
    Clip y_pred to [1e-15, 1 - 1e-15] before taking logs.
    """
    # Your code here
    pass