Lesson 14 of 20Training and Sampling

Sampling Strategies: Temperature, Top-k & Top-p

Once trained, generating text means repeatedly asking "given everything so far, what's the distribution over the next word?" and picking one. How you pick matters a lot:

Temperature rescales the logits before softmax. Low temperature sharpens the distribution toward the top choice (more predictable, more repetitive); high temperature flattens it (more diverse, more prone to nonsense).
Top-k simply deletes every candidate outside the k most likely, then renormalizes, a safety rail that stops the model from ever picking something wildly unlikely, no matter how flat the distribution gets.
Top-p (nucleus sampling) does something subtly different: instead of a fixed count of candidates, it keeps the smallest set of top candidates whose probabilities add up to at least p, then renormalizes just those.

"The cat sat on the ___", adjust temperature and sampling cutoff, then sample

mat

0.65

floor

0.14

couch

0.09

chair

0.03

table

0.02

roof

0.00

bed

0.05

rug

0.01

Temperature: 1.00

Top-k: 8

8 of 8 candidates survive this cutoff.

Push temperature toward 0 and hit Sample repeatedly, you'll get "mat" almost every time, since the distribution collapses onto the single highest-scoring word. Push it up toward 2 with a high top-k and you'll start seeing genuinely surprising picks. This exact tradeoff, coherent-but-boring versus diverse-but-risky, is why every text-generation setting exposes knobs like these.

Switch the widget above to top-p and watch the "candidates survive" count as you adjust temperature. That count is the real difference from top-k: with top-k fixed at, say, 4, the model always considers exactly 4 words, whether it's extremely confident or genuinely torn between many options. Top-p adapts, when one word dominates the distribution, a single candidate alone can cross a high p threshold; when the distribution is flatter, top-p automatically opens up to more candidates to reach the same cumulative probability. That's why top-p and top-k are usually described as solving the same problem (cut off the unlikely tail) with different definitions of "cut off."

Both the training loop and the sampling function, in plain Python, using the exact same numbers as the widgets above:

Python

import math, random

words = ["mat", "floor", "couch", "chair", "table", "roof", "bed", "rug"]
base_logits = [4.0, 2.5, 2.0, 1.0, 0.5, -1.0, 1.5, 0.2]
target = 0  # "mat"

def softmax(logits):
    m = max(logits)
    exps = [math.exp(l - m) for l in logits]
    s = sum(exps)
    return [e / s for e in exps]

# training: the gradient of softmax + cross-entropy is exactly
# softmax(logits) - one_hot(target) — no calculus needed to use it
logits = [0.0] * len(words)
lr = 0.5
for step in range(30):
    probs = softmax(logits)
    loss = -math.log(probs[target])
    if step % 5 == 0 or step == 29:
        print(f"step {step:>2}: loss={loss:.3f}  P('mat')={probs[target]:.3f}")
    grad = [p - (1 if i == target else 0) for i, p in enumerate(probs)]
    logits = [l - lr * g for l, g in zip(logits, grad)]

print()

def temperature_sample(logits, temperature, k):
    scaled = [l / temperature for l in logits]
    probs = softmax(scaled)
    top_k = sorted(range(len(probs)), key=lambda i: probs[i], reverse=True)[:k]
    filtered = [p if i in top_k else 0 for i, p in enumerate(probs)]
    total = sum(filtered)
    probs = [p / total for p in filtered]
    r = random.random()
    for i, p in enumerate(probs):
        r -= p
        if r <= 0:
            return i
    return len(probs) - 1

print("samples at temperature=1.0, top_k=8:")
print([words[temperature_sample(base_logits, 1.0, 8)] for _ in range(8)])

def top_p_filter(probs, p):
    ranked = sorted(range(len(probs)), key=lambda i: probs[i], reverse=True)
    keep, cumulative = set(), 0.0
    for i in ranked:
        if cumulative >= p:
            break
        keep.add(i)
        cumulative += probs[i]
    filtered = [probs[i] if i in keep else 0 for i in range(len(probs))]
    total = sum(filtered)
    return [f / total for f in filtered], len(keep)

peaked_probs = softmax([6.0, 1.0, 0.5, 0.2, 0.1, 0.0, -0.5, -1.0])
flat_probs = softmax([1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3])

_, n_peaked = top_p_filter(peaked_probs, 0.9)
_, n_flat = top_p_filter(flat_probs, 0.9)
print()
print(f"top_p=0.9 keeps {n_peaked} candidates when one word dominates")
print(f"top_p=0.9 keeps {n_flat} candidates when the distribution is flat")

import math, random

words = ["mat", "floor", "couch", "chair", "table", "roof", "bed", "rug"]
base_logits = [4.0, 2.5, 2.0, 1.0, 0.5, -1.0, 1.5, 0.2]
target = 0  # "mat"

def softmax(logits):
    m = max(logits)
    exps = [math.exp(l - m) for l in logits]
    s = sum(exps)
    return [e / s for e in exps]

# training: the gradient of softmax + cross-entropy is exactly
# softmax(logits) - one_hot(target) — no calculus needed to use it
logits = [0.0] * len(words)
lr = 0.5
for step in range(30):
    probs = softmax(logits)
    loss = -math.log(probs[target])
    if step % 5 == 0 or step == 29:
        print(f"step {step:>2}: loss={loss:.3f}  P('mat')={probs[target]:.3f}")
    grad = [p - (1 if i == target else 0) for i, p in enumerate(probs)]
    logits = [l - lr * g for l, g in zip(logits, grad)]

print()

def temperature_sample(logits, temperature, k):
    scaled = [l / temperature for l in logits]
    probs = softmax(scaled)
    top_k = sorted(range(len(probs)), key=lambda i: probs[i], reverse=True)[:k]
    filtered = [p if i in top_k else 0 for i, p in enumerate(probs)]
    total = sum(filtered)
    probs = [p / total for p in filtered]
    r = random.random()
    for i, p in enumerate(probs):
        r -= p
        if r <= 0:
            return i
    return len(probs) - 1

print("samples at temperature=1.0, top_k=8:")
print([words[temperature_sample(base_logits, 1.0, 8)] for _ in range(8)])

def top_p_filter(probs, p):
    ranked = sorted(range(len(probs)), key=lambda i: probs[i], reverse=True)
    keep, cumulative = set(), 0.0
    for i in ranked:
        if cumulative >= p:
            break
        keep.add(i)
        cumulative += probs[i]
    filtered = [probs[i] if i in keep else 0 for i in range(len(probs))]
    total = sum(filtered)
    return [f / total for f in filtered], len(keep)

peaked_probs = softmax([6.0, 1.0, 0.5, 0.2, 0.1, 0.0, -0.5, -1.0])
flat_probs = softmax([1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3])

_, n_peaked = top_p_filter(peaked_probs, 0.9)
_, n_flat = top_p_filter(flat_probs, 0.9)
print()
print(f"top_p=0.9 keeps {n_peaked} candidates when one word dominates")
print(f"top_p=0.9 keeps {n_flat} candidates when the distribution is flat")

Your turn

Implement top_p_filter(probs, p): keep the smallest set of highest-probability candidates whose cumulative probability reaches at least p, zero out the rest, then renormalize so the kept probabilities sum to 1.

Python

Your turn

Implement greedy_decode(logits): return the index of the highest-scoring logit, exactly what temperature_sample converges to as temperature approaches 0.

Python

As temperature approaches 0, sampling becomes equivalent to…

Every sample above costs the same to compute no matter how long the conversation gets so far, but real generation has a performance wrinkle worth understanding before this arc wraps up: KV caching, next.

← Training: From Logits to Loss KV Caching: Making Generation Fast →