Skip to content

← Loss Functions step 3 of 8

Medium Primitives

Implement Cross-Entropy Loss

For classification with several classes, cross-entropy rewards assigning high probability to the correct class. For each example, select the probability of its target class and take its negative natural logarithm. Then average over the examples.

Suppose the probability rows are [[0.7, 0.2, 0.1], [0.1, 0.8, 0.1]] and the target class indices are [0, 1]. Select 0.7 from the first row and 0.8 from the second. The loss is (-log(0.7) - log(0.8)) / 2, approximately 0.289909.

Each example contributes one selected probability. Do not average over all classes or choose the most likely class instead of the target class.

Your task

Implement cross_entropy_loss(probs, targets):

  • probs is a floating-point tensor of shape (N, C), with N > 0 and C > 0. Each row contains probabilities in [0, 1] summing to 1.
  • targets is an integer tensor of shape (N,). Each value is a zero-based class index from 0 through C - 1.
  • Return a scalar tensor containing the mean negative log probability of the target classes.

Before taking logarithms, clamp probabilities to [1e-7, 1.0]. Do not add epsilon to all values or renormalize after clamping. A target-class probability of zero therefore produces a finite loss of -log(1e-7).

Inputs are already probabilities. In contrast, PyTorch's built-in cross_entropy expects logits, so passing these inputs to that function would compute a different objective. Use basic tensor operations here.

With its default target of 1, the plot shows -log(p) for the probability assigned to a correct class, away from the clamped endpoint. The exercise averages this quantity across rows.

Loading visualization…