Skip to content

PigeonColumba compressaTheme

Lesson 1 of 10

Softmax, cross-entropy and its gradient

Every classifier in this lab ends the same way: a vector of logits z (one number per class) becomes probabilities with softmax, and the loss is cross-entropy against the true class.

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

Dividing the logits by a temperature T before the softmax flattens the distribution when T > 1. Distillation depends on that knob.

Two things to get right:

Numerical stability. exp(1000) overflows. Subtracting the row's maximum first changes nothing mathematically (it cancels in the ratio) and keeps every exponent at or below 0.
The gradient. For L = -log softmax(z)[y], the derivative with respect to the logits is beautifully simple: p - onehot(y). Averaged over a batch of N rows, divide by N.

Implement the three functions. All inputs are 2-D: one row per example.

Your code

import numpy as np


def softmax(z, T=1.0):
    """Row-wise softmax of z / T. z has shape (N, C)."""
    raise NotImplementedError


def cross_entropy(z, y):
    """Mean of -log softmax(z)[i, y[i]] over the N rows."""
    raise NotImplementedError


def cross_entropy_grad(z, y):
    """d cross_entropy / d z, shape (N, C)."""
    raise NotImplementedError

Checking your code and saving a draft need JavaScript, which is turned off in this browser. Turn it on and reload the page to check your code. You can still edit it here, and the hints and the reference solution are below.

Pigeon's hints (2)

Hint 1 of 2

z.max(axis=1, keepdims=True) keeps the shape (N, 1) so it broadcasts.

Hint 2 of 2

For the gradient, start from softmax(z) and subtract 1 at each row's true class, then divide by N because the loss is a mean.

The reference solution

import numpy as np


def softmax(z, T=1.0):
    z = z / T
    z = z - z.max(axis=1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)


def cross_entropy(z, y):
    p = softmax(z)
    return float(-np.log(p[np.arange(len(y)), y] + 1e-12).mean())


def cross_entropy_grad(z, y):
    p = softmax(z)
    p[np.arange(len(y)), y] -= 1
    return p / len(y)

Where this is used

Knowledge distillation (Hinton)Runs hereTrain a small model to match a big model's softened output probabilities instead of (or as well as) the one-hot labels.