L8.3 · LESSON 3 OF 6

Top-k

Compute temperature and candidate filtering, and distinguish sampling from greedy and beam decoding.

THE IDEA

Top-k: only the k best candidates

Top-k sampling keeps the k most probable tokens, sets the others to zero and renormalises so the survivors add up to 1 before sampling. It removes the long tail of unlikely tokens.

The catch: a fixed k means different things for different distributions. When the model is confident, two tokens may already hold 98% of the mass; when it is unsure, two tokens may hold only half.

token	probability	kept?	after renormalising
1	0.5	yes	0.667
2	0.25	yes	0.333
3	0.15	no	0
4	0.07	no	0
5	0.03	no	0

probability mass kept by the top 2: 0.75 The same k keeps very different amounts of mass for different distributions: try both shapes with k = 2.

Switch between a peaked and a flat distribution at the same k.

Predict first. Does k = 2 always preserve 80% of the probability?

Chapter context

Decoding chooses tokens from the model's output distribution. It changes generated behavior without retraining the model. Sampling settings influence diversity and concentration, but they do not independently guarantee correctness or creativity.

TRY IT YOURSELF

Part 1. Do it hands-on

Work out the kept mass and renormalised probabilities by hand, find the smallest k that keeps 90%, then code top-k.

token	probability	kept?	after renormalising
1	0.5	yes	0.5
2	0.25	yes	0.25
3	0.15	yes	0.15
4	0.07	yes	0.07
5	0.03	yes	0.03

probability mass kept by the top 5: 1 The same k keeps very different amounts of mass for different distributions: try both shapes with k = 2.

mass kept by k = 2

renormalised top-2:

The slider above is for the third task (smallest k holding at least 90%).

To do: Probabilities [0.5, 0.25, 0.15, 0.07, 0.03]. With top-k, k = 2, how much probability mass do the kept tokens hold?
To do: Renormalise the two kept tokens so they sum to 1. Type both probabilities (3 decimals)
To do: Choose k so that the kept tokens hold at least 90% of the probability, with as few tokens as possible

Part 2. Do it in code

top_k_filter(probs, k): keep the k largest probabilities (ties broken by lower index), set the others to 0 and renormalise. k must be at least 1 (ValueError); k larger than the vocabulary keeps everything.
retained_mass(probs, k): the probability mass held by the k largest entries before renormalising.
smallest_k(probs, mass): the smallest k whose retained mass is at least `mass`.
top_k_sample(probs, k, u): filter with top_k_filter and pick the index for the uniform number u in [0, 1) by walking the cumulative sums.

import numpy as np

def top_k_filter(probs, k):
    pass

def retained_mass(probs, k):
    pass

def smallest_k(probs, mass):
    pass

def top_k_sample(probs, k, u):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Does k=2 always preserve 80% of probability?

AYesBNo

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

p_i(T)=\frac{\exp(z_i/T)}{\sum_j\exp(z_j/T)},\quad T>0.

Top-p retains a probability-sorted prefix with cumulative mass at least p, then renormalizes the retained entries. Specify tie handling and any minimum-candidate rule in implementation.

Worked example

For logits [2,1,0] at temperature one, probabilities are approximately [0.665241,0.244728,0.090031]. With top-p=0.8, keep the first two because the first alone is below the threshold. Their renormalized probabilities are approximately [0.731059,0.268941].

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import exp,isclose
weights=[exp(v-2) for v in [2,1,0]]
probs=[w/sum(weights) for w in weights]
assert probs[0] < .8 < sum(probs[:2])
retained=[p/sum(probs[:2]) for p in probs[:2]]
assert isclose(probs[0],.6652409557748218)
assert isclose(retained[0],.7310585786300049)
assert isclose(sum(retained),1)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Top-p=.8 does not mean keeping 80% of vocabulary items. A high temperature does not create new factual knowledge. Filter-and-renormalize order must be explicit.