Top-k and top-p: choose what can be sampled

A language model scores its vocabulary for the next token. Sampling draws from a probability distribution instead of always taking the highest score. Top-k and top-p restrict which candidates participate in that draw.

Top-k retains a fixed number of the highest-scoring tokens. With k = 2, only two candidates remain, whether the model is confident or uncertain.

Top-p, also called nucleus sampling, retains the smallest highest-probability set whose cumulative probability reaches a threshold. At p = 0.8, probabilities 0.85, 0.08, 0.04, 0.03 retain just the first token. Probabilities 0.4, 0.3, 0.2, 0.1 retain three, reaching 0.9. The boundary token must be included to cross the threshold.

Apply the filter to model scores

This NumPy function accepts a single next-token logits vector and returns a vocabulary ID. It applies top-k first when supplied, then top-p to the renormalized surviving distribution. Pass a NumPy random generator as rng; valid settings are top_k >= 1 and 0 < top_p <= 1.

import numpy as np

def sample_token(logits, rng, top_k=None, top_p=1.0):
    ids = np.argsort(logits)[::-1]
    if top_k is not None:
        ids = ids[:top_k]
    probabilities = np.exp(logits[ids] - logits[ids].max())
    probabilities /= probabilities.sum()
    before = np.r_[0.0, np.cumsum(probabilities)[:-1]]
    keep = before < top_p
    ids, probabilities = ids[keep], probabilities[keep]
    probabilities /= probabilities.sum()
    return int(rng.choice(ids, p=probabilities))

The retained probabilities are normalized again before drawing. Removed candidates get no chance on this step, even though the model originally assigned them some probability.

Temperature reshapes probabilities before filtering in many implementations. Top-k and top-p remove candidates. These settings influence diversity and repetition; none checks whether the selected continuation is factually correct.