Softmax & Cross-Entropy Loss

Once your network produces raw output numbers ("logits") for a classification task, you need two more things: a way to turn those arbitrary real numbers into a probability distribution, and a way to measure how wrong that distribution is compared to the true label. Those are softmax and cross-entropy loss, and together they're the standard head of essentially every classifier — image nets, language model next-token prediction, everything.

Softmax: logits to probabilities

Given a vector of logits z = [z_1, ..., z_n] (one per class), softmax produces:

softmax(z)_i = exp(z_i) / sum_j exp(z_j)

This guarantees every output is positive and the whole vector sums to 1 — a valid probability distribution — while preserving the relative ordering of the logits (the largest logit still gets the largest probability).

Numerical stability — the max-subtract trick. exp(z_i) overflows for even moderately large logits (e.g. exp(1000) is inf in floating point). Since softmax is invariant to adding a constant to every logit (the constant cancels in numerator and denominator), subtract the max first:

This shifts the largest logit to 0 before exponentiating, so the largest term is exp(0) = 1 and nothing overflows — every real ML framework does this, and skipping it is a classic source of NaN losses during training.

Cross-entropy loss: measuring the miss

Given the true class index label and the predicted distribution p = softmax(z), cross-entropy loss is:

loss = -log(p[label])

Intuitively: if the model assigns high probability to the correct class, p[label] is close to 1 and -log(p[label]) is close to 0 (low loss). If the model assigns the correct class near-zero probability, -log(p[label]) blows up toward infinity — the loss punishes confident wrongness much more harshly than a small miss, which is exactly the gradient signal you want during training: being badly wrong should produce a large gradient.

Because you already have p from softmax, this is a one-liner once you have the probabilities — the two concepts are almost always implemented together for exactly this reason (and combined implementations skip computing the full softmax vector when only the loss is needed, for efficiency — though this exercise asks you to expose both PROBS and LOSS separately).

Why natural log, and what "nats" means

Using math.log (base e) rather than log2 or log10 measures loss in nats instead of bits — this is just a convention (a constant scaling factor), but it's the one every ML framework (PyTorch's nn.CrossEntropyLoss, TensorFlow) uses by default, so matching it keeps your numbers comparable to real tooling.

Worked check

For logits = [2.0, 0.0], label = 0:

softmax([2, 0]) = [e^2/(e^2+e^0), e^0/(e^2+e^0)] ≈ [0.8808, 0.1192]
loss = -log(0.8808) ≈ 0.1269

Where this plugs into training

During backprop, the gradient of cross-entropy-after-softmax with respect to the logits has a famously clean closed form (p - one_hot(label)), which is part of why this pairing is so ubiquitous — but for this exercise, your job is just the forward computation: correct, numerically-stable softmax, and the loss built on top of it. The gradient will come later once this is wired into your autograd engine.