Project tutorials · 6 steps

Cross-entropy loss and KL divergence in Python

Write the loss that language models train on, from scratch with NumPy, along with KL divergence, perplexity and the gradient.

What you'll build

Every language model you've used was trained by making one number smaller: cross-entropy. It measures how surprised the model was by the right answer. This tutorial writes it and its relatives in a few lines of NumPy, so you can see exactly what they compute.

What you'll learn

Compute cross-entropy between a true distribution and a model's prediction
Show that with one correct answer, cross-entropy is −log of the probability given to it
Compute KL divergence and explain it as the extra cost of using the wrong distribution
Turn an average loss into perplexity

What you need

Python, NumPy arrays and what a logarithm is. Everything here runs anywhere with pip install numpy.

Step 1: Write cross-entropy

p is the true distribution and q is the model's prediction. Cross-entropy adds up each true probability times the log of the predicted one, with a minus sign in front. Clipping q to a tiny value stops log(0), which would give infinity.

import numpy as np

def cross_entropy(p, q, eps=1e-12):
    p = np.asarray(p, dtype=float)
    q = np.clip(np.asarray(q, dtype=float), eps, None)
    return float(-np.sum(p * np.log(q)))

Step 2: The loss for one token

When training a language model there's one right next token, so p is all zeros and a single one. Every term in the sum vanishes except that one, which leaves minus the log of the probability given to the right answer. That's why you'll see it called negative log likelihood.

def token_loss(q, target):
    one_hot = np.zeros(len(q))
    one_hot[target] = 1
    return cross_entropy(one_hot, q)

Rather write it yourself?

In the lesson you write this from a starter, in your browser with nothing to install, and checks tell you what's wrong as you go. A free account opens it and all 244 lessons.

Create a free account

Step 3: KL divergence: the cost of being wrong

KL divergence measures the extra cost of using the distribution q when the truth is p. It's zero when the two match, and it's never negative. Entries where p is zero contribute nothing, so the mask skips them and avoids log(0).

def kl_divergence(p, q):
    p, q = np.asarray(p, dtype=float), np.asarray(q, dtype=float)
    mask = p > 0
    return float(np.sum(p[mask] * np.log(p[mask] / q[mask])))

Step 4: Perplexity: a loss you can picture

Perplexity is the exponential of the average loss. It reads as a count: a perplexity of 6 means the model is as unsure as if it were choosing evenly among six options.

def perplexity(probs_of_correct):
    return float(np.exp(np.mean(-np.log(probs_of_correct))))

Step 5: The gradient that training follows

For a model that produces logits and applies softmax, the gradient of this loss has a simple form: the predicted probabilities minus a one for the right token. It pushes the right token's logit up and the others down, each in proportion to the probability it was given.

def logit_gradient(q, target):
    grad = np.array(q, dtype=float)
    grad[target] -= 1
    return grad

Step 6: Try it

The first line is the loss when the right token was given 60%, then 5%: less probability means much more loss. The second is the KL divergence from a skewed distribution to a uniform one. The third shows a perplexity of 6 for a model that gives every right token 1 in 6, then the gradient, which is negative only for the right token.

q = [0.6, 0.25, 0.1, 0.05]
print(round(token_loss(q, 0), 3), round(token_loss(q, 3), 3))
print(round(kl_divergence([0.7, 0.2, 0.08, 0.02], [0.25] * 4), 4))
print(round(perplexity([1 / 6] * 50), 3), logit_gradient(q, 0))

What it prints

0.511 2.996
0.5344
6.0 [-0.4   0.25  0.1   0.05]

Where to take it next

Replace the softmax probabilities with your own model's, and watch the loss as it trains.
Compare your numbers with torch.nn.functional.cross_entropy, which takes logits rather than probabilities.

Mistakes to watch for

Thinking cross-entropy is symmetric in p and q.
Letting a zero probability reach log and produce infinity.
Comparing losses computed in bits with ones in nats.
Believing a lower loss always means better answers on every task (it measures next-token prediction).

The complete program

import numpy as np

def cross_entropy(p, q, eps=1e-12):
    p = np.asarray(p, dtype=float)
    q = np.clip(np.asarray(q, dtype=float), eps, None)
    return float(-np.sum(p * np.log(q)))

def token_loss(q, target):
    one_hot = np.zeros(len(q))
    one_hot[target] = 1
    return cross_entropy(one_hot, q)

def kl_divergence(p, q):
    p, q = np.asarray(p, dtype=float), np.asarray(q, dtype=float)
    mask = p > 0
    return float(np.sum(p[mask] * np.log(p[mask] / q[mask])))

def perplexity(probs_of_correct):
    return float(np.exp(np.mean(-np.log(probs_of_correct))))

def logit_gradient(q, target):
    grad = np.array(q, dtype=float)
    grad[target] -= 1
    return grad

q = [0.6, 0.25, 0.1, 0.05]
print(round(token_loss(q, 0), 3), round(token_loss(q, 3), 3))
print(round(kl_divergence([0.7, 0.2, 0.08, 0.02], [0.25] * 4), 4))
print(round(perplexity([1 / 6] * 50), 3), logit_gradient(q, 0))

Write it yourself

In the lesson you write this from a starter, in your browser with nothing to install, and checks tell you what's wrong as you go. A free account opens it and all 244 lessons.

Create a free account