Skip to content

PigeonColumba compressaTheme

Lesson 5 of 10

The distillation loss and its gradient

Hinton's distillation loss compares the student's and teacher's softened distributions:

L = T^2 * mean_i KL(p_t || p_s), with p_t = softmax(z_t / T), p_s = softmax(z_s / T)

and KL(p || q) = sum_c p_c (log p_c - log q_c).

Why T^2? The gradient of the softened loss shrinks like 1/T^2 as T grows; multiplying back keeps its size comparable to the cross-entropy term when the two are mixed.

The gradient with respect to the student's logits works out to

dL/dz_s = T * (p_s - p_t) / N

Derive it yourself first: it's the cross-entropy gradient from lesson 1 with a soft target, with the chain rule through z / T.

Your code

import numpy as np


def kd_loss(zs, zt, T=4.0):
    """Return (loss, dloss/dzs) for student logits zs and teacher logits zt, both (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

Compute log-probabilities with a stable log-softmax; log(softmax(...)) can hit log(0).

Hint 2 of 2

The T^2 in the loss and the 1/T from the chain rule leave a single T in the gradient.

The reference solution

import numpy as np


def _log_softmax(z, T):
    z = z / T
    z = z - z.max(1, keepdims=True)
    return z - np.log(np.exp(z).sum(1, keepdims=True))


def kd_loss(zs, zt, T=4.0):
    ls, lt = _log_softmax(zs, T), _log_softmax(zt, T)
    ps, pt = np.exp(ls), np.exp(lt)
    N = len(zs)
    loss = T * T * (pt * (lt - ls)).sum() / N
    return float(loss), T * (ps - pt) / N

Where this is used