Careeroza

Top-k and Top-p

Cutting Off the Long Tail

  • Top-k: keep only the k most likely tokens, renormalize, then sample. k=3 means never pick outside the top 3.
  • Top-p (nucleus sampling): keep the smallest set of top tokens whose probabilities add up to at least p (e.g. 0.9). The set is small when the model is confident and larger when it is unsure.
  • Both remove rare, nonsense tokens while keeping some variety. They are often combined with temperature.

A Sampler with Temperature, Top-k and Top-p

Example

This is how a sampler in an inference engine (like the one inside Ollama or vLLM) chooses a token:

import math, random

def sample(tokens, logits, temperature=1.0, top_k=None, top_p=None, rng=None):
    m = max(logits)
    exps = [math.exp((s - m) / temperature) for s in logits]
    z = sum(exps)
    ranked = sorted(zip(tokens, [e / z for e in exps]), key=lambda x: -x[1])
    if top_k:
        ranked = ranked[:top_k]
    if top_p:
        kept, total = [], 0.0
        for tok, pr in ranked:
            kept.append((tok, pr)); total += pr
            if total >= top_p:
                break
        ranked = kept
    z2 = sum(pr for _, pr in ranked)
    ranked = [(t, pr / z2) for t, pr in ranked]
    choice = rng.choices([t for t, _ in ranked], weights=[pr for _, pr in ranked])[0]
    return choice, ranked

tokens = ["on", "delayed", "delivered", "cancelled", "banana"]
logits = [3.1, 2.5, 2.4, 2.2, -1.0]
for label, kw in [("top_k=2", {"top_k": 2}), ("top_p=0.8", {"top_p": 0.8}), ("no filter", {})]:
    rng = random.Random(7)
    _, kept = sample(tokens, logits, **kw, rng=rng)
    picks = [sample(tokens, logits, **kw, rng=rng)[0] for _ in range(10)]
    print(f"{label:10} candidates: {[(t, round(pr, 2)) for t, pr in kept]}")
    print(f"{'':10} 10 picks:   {' '.join(picks)}")

Output

top_k=2    candidates: [('on', 0.65), ('delayed', 0.35)]
           10 picks:   on delayed on on on on on on on on
top_p=0.8  candidates: [('on', 0.49), ('delayed', 0.27), ('delivered', 0.24)]
           10 picks:   on delayed on delayed on on delayed on on on
no filter  candidates: [('on', 0.41), ('delayed', 0.22), ('delivered', 0.2), ('cancelled', 0.16), ('banana', 0.01)]
           10 picks:   on delivered on delayed on on delayed on delayed on