Loss Functions › Classification Losses

Classification tasks require predicting one of $K$ discrete classes. The structure of the problem is fundamentally different from regression: instead of predicting a continuous value, the model produces scores (logits) for each class, which are then converted to probabilities.

The foundation of all classification losses is maximum likelihood estimation. We want to find model parameters that maximize the probability of observing the training data. If the model outputs a probability distribution $p$ over classes, the likelihood of the true class $y$ is $p(y)$. We want this to be high.

Maximizing likelihood is equivalent to minimizing negative log-likelihood:

$$\mathcal{L}_{\text{NLL}} = -\log p(y)$$

where:

• $p(y)$: the model's predicted probability for the true class $y$
• $\log$: the natural logarithm

The log transformation serves two purposes. First, it converts a product of probabilities (over a dataset) into a sum, which is numerically more stable. Second, the log is a monotone transformation, so maximizing $\log p$ is equivalent to maximizing $p$. Because log of a probability between 0 and 1 is negative, we negate it to get a positive loss.

Negative log-likelihood has an intuitive interpretation: it measures how surprised the model is to see the true label. If the model assigns probability 0.99 to the correct class, the loss is $-\log(0.99) \approx 0.01$ (barely surprised). If the model assigns probability 0.01 to the correct class, the loss is $-\log(0.01) \approx 4.6$ (very surprised). The loss grows without bound as the predicted probability approaches zero, which strongly penalizes confident wrong predictions.

This unbounded growth is intentional. A loss that caps at some maximum value would allow the model to be confidently wrong without incurring a proportional penalty. The logarithmic penalty means there is no "good enough" point where the model can stop improving its confidence in correct predictions, which drives the model toward sharper, more accurate probability estimates.

Cross-entropy loss is the most widely used loss for classification. It combines softmax (to convert logits to probabilities) with negative log-likelihood into a numerically stable combined operation.

Given the model's raw scores (logits) $\mathbf{z} = [z_1, z_2, \ldots, z_K]$ for $K$ classes, the softmax function converts them to a probability distribution:

$$p_k = \frac{e^{z_k}}{\sum_{j=1}^{K} e^{z_j}}$$

where:

• $z_k$: the raw score (logit) for class $k$
• $e^{z_k}$: the exponential of the logit, which ensures positivity
• $\sum_{j=1}^{K} e^{z_j}$: the normalizing constant, which ensures probabilities sum to 1
• $p_k$: the predicted probability for class $k$

Cross-entropy loss then applies negative log-likelihood to the true class:

$$\mathcal{L}_{\text{CE}} = -\sum_{k=1}^{K} y_k \log p_k$$

where:

• $y_k$: a one-hot indicator, equal to 1 for the true class and 0 elsewhere
• $p_k$: the predicted probability for class $k$

Because $y_k$ is one-hot, this sum collapses to just $-\log p_{y^*}$ where $y^*$ is the index of the true class. The sum notation is more general, covering soft labels (non-binary $y_k$ values), but the one-hot version is standard for supervised classification.

Numerical stability is a critical implementation detail. Computing softmax followed by log involves exponentials that can overflow (for large logits) or underflow (for very negative logits). The standard trick is the log-sum-exp stabilization:

$$\log \sum_{j=1}^{K} e^{z_j} = c + \log \sum_{j=1}^{K} e^{z_j - c}$$

where $c = \max_j z_j$ is subtracted before exponentiation to prevent overflow. PyTorch's F.cross_entropy() implements this automatically when you pass raw logits. Always pass logits directly to F.cross_entropy() rather than computing softmax separately and then applying NLL, because the combined implementation is both more numerically stable and more computationally efficient.

The gradient of cross-entropy with respect to the logits has a clean form that reveals why the loss is so well-suited to gradient descent. For the logit $z_k$:

$$\frac{\partial \mathcal{L}_{\text{CE}}}{\partial z_k} = p_k - y_k$$

This is simply the difference between the predicted probability and the target. For the true class ($y_k = 1$), the gradient is $p_k - 1$, which is negative (pushing the logit up) unless $p_k = 1$. For incorrect classes ($y_k = 0$), the gradient is $p_k$, which is positive (pushing the logit down). The optimizer directly reduces the predicted probability of incorrect classes and increases the probability of the correct class. This clean interpretation is one reason cross-entropy is so widely used: the gradients have an immediate intuitive meaning.

Binary classification is a special case where $K = 2$ (positive or negative). Rather than using a 2-class softmax, it is standard to use a single output neuron with a sigmoid activation, which directly outputs the probability of the positive class.

Given a model output $\hat{y} = \sigma(z)$ where $\sigma$ is the sigmoid function and $z$ is the logit, binary cross-entropy (BCE) is:

$$\mathcal{L}_{\text{BCE}} = -\left[ y \log \hat{y} + (1 - y) \log(1 - \hat{y}) \right]$$

where:

• $y \in \{0, 1\}$: the true binary label
• $\hat{y} \in (0, 1)$: the model's predicted probability of the positive class
• $y \log \hat{y}$: the loss contribution when $y = 1$ (penalizes predicting low probability for a positive example)
• $(1 - y) \log(1 - \hat{y})$: the loss contribution when $y = 0$ (penalizes predicting high probability for a negative example)

The two terms are mutually exclusive: when $y = 1$, the second term vanishes, and when $y = 0$, the first term vanishes. This loss strongly penalizes confident wrong predictions: if the model outputs $\hat{y} = 0.001$ for a positive example ($y = 1$), the loss is $-\log(0.001) \approx 6.9$.

The gradient of BCE with respect to the logit $z$ (before sigmoid) has a particularly clean form:

$$\frac{\partial \mathcal{L}_{\text{BCE}}}{\partial z} = \hat{y} - y$$

where:

• $\hat{y} = \sigma(z)$: the sigmoid of the logit
• $y$: the true label

This is the simplest possible gradient: just the prediction error. The sigmoid's derivative cancels with the derivative of the log, producing a numerically stable and interpretable gradient signal.

Binary cross-entropy is also the right loss for multi-label classification, where each example can belong to multiple classes simultaneously. In multi-label settings, you have $K$ output neurons, each representing an independent binary prediction. BCE is applied independently to each output, and the total loss is the sum over all $K$ binary losses. This is different from multiclass (categorical) classification, where exactly one class is correct and the outputs must sum to 1.

The choice between categorical and binary cross-entropy depends on the problem structure:

• Binary cross-entropy applies when each example belongs to one of two classes or when predicting multiple independent binary attributes simultaneously (multi-label classification).
• Categorical cross-entropy applies when each example belongs to exactly one of $K > 2$ classes (multi-class classification).

In multi-label classification, each example can have multiple positive labels (e.g., a news article tagged as both "politics" and "international"). Binary cross-entropy applies independently to each label, computing $K$ separate losses and summing them. The key distinction is that the $K$ outputs do not need to sum to 1, and each output is trained independently.

Class imbalance is one of the most pervasive problems in NLP. Named entity recognition datasets might have 95% non-entity tokens. Sentiment datasets might have 80% neutral examples. Toxicity classifiers deal with much more benign text than harmful text. In these settings, an unweighted cross-entropy loss will cause the model to learn to predict the majority class almost all the time, achieving low loss while being useless for the rare classes of interest.

Weighted cross-entropy assigns a higher loss multiplier to underrepresented classes:

$$\mathcal{L}_{\text{weighted}} = -\sum_{k=1}^{K} w_k \cdot y_k \log p_k$$

where:

• $w_k$: the weight for class $k$, typically set to the inverse frequency of class $k$ in the training set
• $y_k$: the one-hot target (1 for true class, 0 otherwise)
• $p_k$: the predicted probability for class $k$

A standard heuristic for setting weights is:

$$w_k = \frac{N}{K \cdot n_k}$$

where $N$ is the total number of training examples and $n_k$ is the number of examples in class $k$. This ensures that each class contributes equally to the total loss in expectation, regardless of its frequency. In PyTorch, you can pass a weight tensor directly to F.cross_entropy().