Turn logits into probabilities

Learn it

Words to know

logit

A raw score a model gives to one choice, before the scores are turned into probabilities.

softmax

The function that turns a list of scores into probabilities: all positive, adding up to 1.

How it works

1. Raise e to each score. This makes every number positive and stretches the gaps between them.
2. Divide each result by the total. Now they add up to 1.
3. One safety step first: subtract the largest score from every score. The probabilities do not change, because every power of e is divided by the same number, but e is never raised to a large number.

Worked example

Input
logits = [2, 1, 0]

Output
[0.665241, 0.244728, 0.0900306]

The three powers add up to 1.503, and the three results add up to 1. Each score is 1 apart, so each probability is about 2.7 times the next one.

Why it matters

The last step of a classifier or a GPT is a softmax over its scores. Attention in chapter 6 uses the same function to decide how much each word counts.

Your turn

Turn a list of scores into probabilities with softmax, in a way that cannot overflow.

logits

a list of one or more scores; any size, positive or negative

Returns

A list of probabilities, one per score, in the same order. They add up to 1.

Input
logits = [2, 1, 0]

Output
[0.665241, 0.244728, 0.0900306]

Input
logits = [1000, 1000]

Output
[0.5, 0.5]

Input
logits = [-50]

Output
[1]

Input
logits = [0, 0, 0]

Output
[0.333333, 0.333333, 0.333333]

Input
logits = [0, 2000]

Output
[0, 1]

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Where to start

   Two passes over the list: one to raise e to each score, one to divide each result by the total.

2. Hint 2: The key step

   Before raising e to anything, find the largest score and subtract it from every score. The largest becomes 0 and no score is left above 0.

3. Hint 3: The outline

   Set top to max(logits). Build exps as math.exp(x - top) for each x. Add them up into total. Return each item of exps divided by total.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

1. The largest score, to subtract from every score. After the shift the largest score is 0 and no score is above 0, so math.exp returns values from 0 to 1 and cannot overflow.
2. Raise e to each shifted score. The power makes every value positive and keeps the order: a higher score still gives a larger value.
3. Divide by the total so the results add up to 1. Each value becomes its share of the whole, which is what a probability is.

The same numbered steps are the comments in the solution.

import math


def solve(logits):
    # 1. The largest score, to subtract from every score
    top = max(logits)
    # 2. Raise e to each shifted score
    exps = [math.exp(x - top) for x in logits]
    # 3. Divide by the total so the results add up to 1
    total = sum(exps)
    return [e / total for e in exps]

Watch out

• Do not divide the raw scores by their sum. A negative score among positive ones would give a negative result, and scores that add up to 0 would divide by zero.
• Work out the largest score and the total once, outside the loops. Calling max or sum inside a loop repeats the same work for every item.

Check yourself