Cross-Entropy From Logits

Compute mean multiclass cross-entropy from logits and integer labels.

Start with the concrete trace below. It shows the state the algorithm must carry as it runs.

Problem trace

Cross-Entropy From Logits: Compute stable log-normalizers per example, subtract the indexed correct logits, then average example losses.

Input and goalCompute mean multiclass cross-entropy from logits and integer labels.

Read logits and integer labelsUse logits [[2,1,0],[0,1,2]] with labels [0,1]. Each label selects one class in its own row.

example	class 0	class 1	class 2	label
0	2	1	0	0
1	0	1	2	1

Recognize it

Use it when multiclass logits and integer labels require negative log likelihood without explicitly forming potentially unstable probabilities.

Keep true

For each row, log_normalizer minus the shifted selected logit equals negative log probability of the label; max shifting cancels from both terms.

Reuse it

For stable classification losses, compute log-normalization in score space, gather the target score by aligned row indices, and reduce only after producing per-example losses.

Pattern: Stable log-softmax plus indexed selection.

Simple idea: Subtract each row maximum. Compute the log normalizer for each example. Subtract the correct-class logit, then average.

import numpy as np

def cross_entropy(logits: np.ndarray, labels: np.ndarray) -> float:
   shifted = logits - np.max(logits, axis=1, keepdims=True)
   log_normalizer = np.log(np.sum(np.exp(shifted), axis=1))
   correct_logits = shifted[np.arange(len(labels)), labels]
   return float(np.mean(log_normalizer - correct_logits))

Cost: $O(batch\times classes)$ time and output-sized temporary space.