Sampling from scratch in NumPy

You already know what temperature and top-p do to the output of an API. This section writes them, because the from-scratch version is a common warm-up question and because implementing nucleus sampling is the shortest way to understand the argument for it. Everything here operates on one vector of logits — the raw scores the model produces for every token in the vocabulary, before anything turns them into probabilities.

The softmax and greedy decoding#

Logits are unnormalised and can be any real number. The softmax turns them into a distribution; greedy decoding skips it entirely, because the largest logit is the largest probability and argmax does not care about the monotonic transformation in between.

import numpy as np


def softmax(logits):
    z = logits - logits.max()        # stability: the largest exponent becomes exp(0)
    e = np.exp(z)
    return e / e.sum()


def greedy(logits):
    return int(np.argmax(logits))    # no softmax needed: argmax is unchanged by it

Temperature#

Temperature divides the logits before the softmax. Below one, the gaps between logits widen and the distribution sharpens; above one they narrow and it flattens. In the limit as temperature goes to zero the distribution becomes one-hot on the argmax, which is why temperature zero and greedy are the same operation — and why a real implementation special-cases it rather than dividing by zero.

def sample_temperature(logits, temperature, rng):
    if temperature == 0.0:
        return greedy(logits)        # the limit, not a division
    p = softmax(logits / temperature)
    return int(rng.choice(len(p), p=p))

Note what temperature cannot do. It rescales the gaps between logits, so where one token is already far ahead of the rest, even a high temperature leaves it far ahead. Variety appears only where the model was genuinely undecided.

Top-k#

Top-k keeps the k highest-scoring tokens, renormalises over them and samples. The point is to cut off the long tail of tokens with tiny but nonzero probability, which collectively hold enough mass to be chosen more often than you would like over a few hundred steps. Use argpartition rather than a full sort: it finds the k largest in linear time without ordering them, and order does not matter here.

def sample_top_k(logits, k, temperature, rng):
    logits = logits / temperature
    # the k largest, unordered - cheaper than argsort over a 100k vocabulary
    keep = np.argpartition(logits, -k)[-k:]
    p = softmax(logits[keep])        # renormalise over the survivors only
    return int(rng.choice(keep, p=p))

Nucleus sampling#

Top-p keeps the smallest set of tokens whose probabilities add up to at least p, then renormalises over that set. The size of the set is not fixed: it is whatever the model's confidence makes it. Sort descending, take the cumulative sum, and cut at the first index where the total crosses the threshold — including that index, so the set always has at least one member even when the top token alone exceeds p.

def sample_top_p(logits, top_p, temperature, rng):
    logits = logits / temperature
    order = np.argsort(logits)[::-1]          # descending by probability
    p = softmax(logits[order])
    cum = np.cumsum(p)
    # first index where the running total reaches the threshold, kept inclusively
    cut = int(np.searchsorted(cum, top_p)) + 1
    keep, p = order[:cut], p[:cut]
    return int(rng.choice(keep, p=p / p.sum()))

Why top-p adapts and top-k does not#

The argument is easier to run than to read. Take two logit vectors over a vocabulary of fifty: one peaked, where a single token holds 95% of the mass, and one flat, where the top token holds under 4%.

def nucleus_size(logits, top_p, temperature=1.0):
    p = np.sort(softmax(logits / temperature))[::-1]
    return int(np.searchsorted(np.cumsum(p), top_p)) + 1


rng = np.random.default_rng(0)

peaked = np.full(50, -6.0)
peaked[7], peaked[3] = 6.0, 3.0           # one clear winner, one runner-up
flat = rng.normal(size=50) * 0.4          # no token much ahead of any other

print(nucleus_size(peaked, 0.9))          # 1
print(nucleus_size(flat, 0.9))            # 42

Now set k to ten and ask what it does on each. On the peaked step it admits nine tokens the model had all but ruled out, one of which will eventually be drawn; that is the mechanism behind a model producing a fluent sentence with one word that does not belong. On the flat step it discards thirty-two tokens the model considered roughly as plausible as the ones it kept, which flattens variety precisely where variety was warranted.

with k = 10

peaked step:  top-k keeps 10  |  top-p(0.9) keeps  1
flat   step:  top-k keeps 10  |  top-p(0.9) keeps 42

Two checks worth running#

Sampling code is easy to get subtly wrong and hard to notice, because wrong output still looks like text. Both of these are cheap.

# temperature approaching zero must agree with greedy
assert all(sample_temperature(flat, 1e-6, rng) == greedy(flat) for _ in range(20))

# top-k with k = vocabulary size must reproduce the full distribution
draws = [sample_top_k(flat, 50, 1.0, rng) for _ in range(5000)]
empirical = np.bincount(draws, minlength=50) / 5000
assert np.abs(empirical - softmax(flat)).max() < 0.02