Lesson 3 - Top-k Sampling

Welcome to Top-k Sampling#

In Lesson 2 you gave your MiniGPT a voice with real choice behind it. Instead of always taking the argmax, you divided the logits by a temperature and drew the next character with np.random.choice, and you watched a single prompt produce several different Lantern Bay continuations. Temperature is the master knob: low temperature sharpens the distribution toward the top token, high temperature flattens it so more characters become plausible.

But temperature reshapes the whole distribution, including the part you would rather not touch. Every character in the vocabulary keeps a nonzero probability, and the long tail of very unlikely tokens still has a small chance of being picked on any given step. Most of the time that is harmless. But push the temperature up for variety and the tail swells, and every so often the sampler reaches into it and draws a character that makes no sense in context, a k where a vowel belonged, and from there the text can spiral into gibberish. This lesson adds a second knob that closes that door directly: top-k sampling. Before you sample, you throw away everything except the k most likely tokens. The tail cannot bite you if it is not there to draw from.

By the end of this lesson, you will be able to:

• Explain why plain temperature sampling still risks drawing from the low-probability tail, and why that risk grows with temperature
• Implement top_k_filter(logits, k) in pure NumPy: keep the k largest logits and set the rest to $-\infty$ so softmax zeros them out
• Write a generate function that combines top-k filtering with temperature and samples with np.random.choice
• Read which tokens survive top-k for a real next-token distribution and how the surviving probabilities renormalize
• Reason about the trade-off: small k is safe but less diverse (with k = 1 reproducing greedy), large k is more varied but lets the tail back in

You should have Lesson 2’s temperature sampling fresh in mind, along with softmax, encode, and decode. Let’s cap that tail.

The problem: temperature never removes the tail#

Recall exactly what temperature does. Given the final-position logits, we divide by a temperature T and apply softmax:

Dividing by T > 1 shrinks the gaps between logits, so softmax comes out flatter and rare characters gain probability. That is the whole point when you want variety. The catch is that flattening is indiscriminate: it lifts the tail along with everything else. The 21 least-likely characters do not disappear, they collectively get a little more mass, and np.random.choice will eventually land on one of them.

On our low-entropy Lantern Bay corpus the model is very confident, so at temperature 1.0 the tail is negligible and you rarely notice. But the moment you raise the temperature to loosen things up, the tail becomes a real hazard: one bad character breaks a word, and because generation is autoregressive, that broken word becomes the context for the next prediction, nudging the model further off the corpus it knows. Temperature gives you no way to say “be more varied, but never pick something absurd.” That is precisely the gap top-k fills.

The idea: keep k, cut the rest, renormalize#

Top-k sampling is a single, blunt rule applied before you sample. Look at the next-token distribution, keep the k characters with the highest probability, and throw the rest away entirely. “Throw away” means set their logits to $-\infty$ so that after softmax they carry exactly zero probability. Then renormalize over the survivors (which softmax does for you) and sample as usual.

The guarantee is worth stating plainly: you can never sample a token outside the top k. The worst character the model could possibly emit is the k-th best one, which, if k is small, is still a reasonable choice. The unbounded downside of plain sampling, the chance of reaching into the deep tail, is gone. In exchange you accept a smaller menu of options at every step.

Notice the small but important detail in the figure: after filtering, the kept probabilities go up a little. The tail held about 0.012 of the total mass; once it is removed, that mass is spread across the survivors during renormalization, so w climbs from 0.4152 to 0.4201. Top-k does not just delete the tail, it hands the tail’s probability back to the tokens you kept.

Training the MiniGPT#

We need a trained model to sample from, so we train the same char-level MiniGPT on the Lantern Bay corpus with the canonical config, exactly as in the previous lessons. The run is short (1000 Adam steps at lr = 3e-3) but drives the loss well down.

import numpy as np, time, warnings
warnings.filterwarnings("ignore")

CORPUS = (
    "the lamp keeper walks the shore at dusk. "
    "the lamp keeper lights the lantern when the fog rolls in. "
    "the boats come home when the lantern glows. "
    "the boats wait outside the bay when the fog is thick. "
    "a small boat drifts near the rocks at dawn. "
    "the keeper rings the bell when a boat drifts near the rocks. "
    "the tide rises at dusk and falls at dawn. "
    "the gulls call over the bay when the tide is low. "
    "the keeper counts the boats and writes the count in a book. "
    "the lantern burns all night and rests at dawn. "
) * 20

chars = sorted(set(CORPUS)); vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)}; itos = {i: c for i, c in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)
data = np.array(encode(CORPUS), dtype=np.int64)

# MiniGPT, Adam, softmax, get_batch are the pieces you built in Modules 1-7.
from minigpt import MiniGPT, Adam, softmax, get_batch

block_size, n_embd, n_head, n_layer = 32, 64, 4, 3
np.random.seed(42)
rng = np.random.default_rng(42)

model = MiniGPT(vocab_size, block_size, n_embd, n_head, n_layer, seed=42)
opt = Adam(model.P, lr=3e-3)
train_data = data[: int(len(data) * 0.9)]

t0 = time.time()
for step in range(1000):
    xb, yb = get_batch(train_data, block_size, 32, rng)
    _, loss = model.forward(xb, yb)
    grads = model.backward()
    opt.step(model.P, grads)
    if step % 200 == 0 or step == 999:
        print(f"step {step:4d}  loss {loss:.4f}")
print("trained in", round(time.time() - t0, 1), "s")

step    0  loss 3.2026
step  200  loss 0.1894
step  400  loss 0.1510
step  600  loss 0.1522
step  800  loss 0.1272
step  999  loss 0.1170

The loss falls from 3.2026 (about ln(24), uniform guessing over 24 characters) to 0.1170 in around 20 seconds. Same trained model as the last two lessons; now we sample from it differently.

Implementing top_k_filter#

The filter takes a 1-D logit vector and a k, and returns a copy where every logit outside the top k is set to $-\infty$. We do it without sorting the whole array’s identities: np.sort(logits)[-k] is the k-th largest value (the threshold), and any logit strictly below it gets set to -np.inf. When softmax later exponentiates, exp(-inf) is 0, so those tokens receive exactly zero probability, and the surviving k renormalize among themselves.

def top_k_filter(logits, k):
    logits = logits.astype(np.float64).copy()   # don't mutate the caller's array
    if k >= logits.shape[-1]:
        return logits                            # k >= vocab: keep everything
    kth = np.sort(logits)[-k]                     # the k-th largest logit value
    logits[logits < kth] = -np.inf               # discard everything below it
    return logits

Two details matter. First, we copy and cast to float64 so we never corrupt the model’s logits and so -np.inf behaves cleanly. Second, the guard if k >= vocab_size means top_k_filter(logits, vocab_size) is a no-op, which makes “no filtering at all” (plain temperature sampling) just a special case of the same code. Let us apply it to a real distribution and look at the survivors.

context = np.array([encode("the boats ")[-block_size:]])  # (1, T)
logits, _ = model.forward(context)
last = logits[0, -1].astype(np.float64)                    # (vocab_size,)

def survivors(k):
    probs = softmax(top_k_filter(last, k))
    idx = np.where(probs > 0)[0]
    idx = idx[np.argsort(-probs[idx])]                     # sort by prob, high to low
    return [(itos[i], round(float(probs[i]), 4)) for i in idx]

print("nonzero tokens in full distribution:", int((softmax(last) > 1e-9).sum()))
print("k=3 survivors :", survivors(3))
print("k=10 survivors:", survivors(10))

nonzero tokens in full distribution: 24
k=3 survivors : [('w', 0.4201), ('c', 0.4129), ('a', 0.167)]
k=10 survivors: [('w', 0.4155), ('c', 0.4083), ('a', 0.1651), ('t', 0.0041), ('r', 0.0022), ('n', 0.0016), ('b', 0.0015), ('l', 0.0011), ('g', 0.0003), ('h', 0.0002)]

The full distribution has all 24 characters nonzero, but the model is enormously confident: w, c, and a (the starts of “wait”, “come”, “and”) hold about 0.99 of the mass between them. With k = 3 those three survive and renormalize to 0.4201, 0.4129, 0.1670, summing to 1. With k = 10 we also keep t, r, n, b, l, g, h, whose probabilities are all under 0.005. Those extra seven are exactly the tail that top-k is designed to gate: harmless here at temperature 1, but the first thing to cause trouble once the temperature climbs.

Generating with top-k and temperature#

Now fold the filter into a generation loop. It is the temperature loop from Lesson 2 with one line added: run top_k_filter on the final-position logits before dividing by the temperature and calling softmax. Everything else, the autoregressive append, the block_size context window, np.random.choice, is unchanged.

def generate(model, prompt, n_new, k, temp=1.0, rng=None):
    block_size = model.cfg[1]
    ids = encode(prompt)
    for _ in range(n_new):
        context = np.array([ids[-block_size:]])            # (1, T)
        logits, _ = model.forward(context)
        filtered = top_k_filter(logits[0, -1], k)          # keep top k, rest -inf
        probs = softmax(filtered / temp)                    # temper the survivors
        next_id = int(rng.choice(len(probs), p=probs))      # sample from the k
        ids.append(next_id)
    return decode(ids)

Because the discarded logits are -inf, filtered / temp is still -inf (dividing -inf by a positive temperature leaves it -inf), so temperature only ever reshapes the surviving k. The order matters conceptually but not numerically here: filter first, then temper the kept tokens.

k = 1 is exactly greedy#

The cleanest sanity check: with k = 1 only the single most likely token survives, softmax over one token gives probability 1.0, and np.random.choice has no choice at all. That must reproduce greedy decoding character for character, regardless of temperature or seed.

def generate_greedy(model, prompt, n_new):
    ids = encode(prompt)
    for _ in range(n_new):
        context = np.array([ids[-model.cfg[1]:]])
        logits, _ = model.forward(context)
        ids.append(int(logits[0, -1].argmax()))
    return decode(ids)

k1     = generate(model, "the ", 70, k=1, temp=1.0, rng=np.random.default_rng(42))
greedy = generate_greedy(model, "the ", 70)
print("k=1   :", repr(k1))
print("greedy:", repr(greedy))
print("equal :", k1 == greedy)

k=1   : 'the lantern when the fog rolls in. the boats come home when the lantern gl'
greedy: 'the lantern when the fog rolls in. the boats come home when the lantern gl'
equal : True

equal: True. Greedy is just top-k sampling with k = 1, the extreme end of the safety-versus-diversity dial. As you raise k, you open the door to more of the distribution.

The trade-off: small k vs large k#

Top-k’s single parameter is a dial between two failure modes. Small k is safe but can be repetitive, at k = 1 you are back to deterministic greedy and its loops. Large k restores diversity but, once k is large enough to include the tail, brings back exactly the risk you were trying to avoid. The right k depends on the vocabulary: here vocab_size is only 24, so k = 3 is quite tight (three candidates) while k = 10 already admits low-probability characters. On a real model with tens of thousands of tokens, typical k values are 40 to 50.

Let us feel the dial. We generate at k = 3 and k = 10 at the same temperature (1.5, warm enough to make the choice visible) with three seeds each.

print("k=3:")
for s in (42, 7, 123):
    print("  ", repr(generate(model, "the ", 70, k=3, temp=1.5, rng=np.random.default_rng(s))))
print("k=10:")
for s in (42, 7, 123):
    print("  ", repr(generate(model, "the ", 70, k=10, temp=1.5, rng=np.random.default_rng(s))))

k=3:
   'the tide rises at dusk and falls at dawn. the gulls call over the bay when'
   'the lantern when the fog rolls in. the boats come home when the lantern gl'
   'the lantern when the fog rolls in. the boats come home when the lantern gl'
k=10:
   'the rocks. the tide rises at dusk and falls at dawn. the gulls call over t'
   'the lantermps the gulls at dawn. the gulles call over the bay when the tid'
   'the lantern when the fog rolls in. the boats come home walll fover the cog'

Read the two blocks side by side. Every k = 3 line is clean Lantern Bay prose, the three survivors at each step are always real, corpus-consistent continuations, so the text stays on the rails. The k = 10 lines are more varied but start to fray: seed 7 gives "lantermps" and "gulles" (plausible-looking but not corpus words), and seed 123 ends in "walll fover the cog", which is the tail showing through, low-probability characters k = 10 admitted that k = 3 would have blocked. That is the whole trade-off in three lines of each: smaller k buys coherence at the cost of variety; larger k buys variety at the cost of the occasional derailment.

Why top-k matters most at high temperature#

The tail’s danger scales with temperature, which is where top-k earns its keep. To show it, we crank the temperature to 2.0 and generate 200 characters two ways: plain sampling (k = vocab_size, no filtering) versus k = 5. We measure how much of each output actually appears in the corpus by counting the fraction of its 4-character windows that are substrings of CORPUS, a simple “how on-corpus is this text” score.

def in_corpus_rate(text):
    windows = [text[i:i+4] for i in range(len(text) - 3)]
    good = sum(1 for w in windows if w in CORPUS)
    return round(good / len(windows), 3)

plain = generate(model, "the ", 200, k=vocab_size, temp=2.0, rng=np.random.default_rng(3))
topk  = generate(model, "the ", 200, k=5,          temp=2.0, rng=np.random.default_rng(3))
print("plain (k=24) in-corpus 4-gram rate:", in_corpus_rate(plain))
print("top-k (k=5)  in-corpus 4-gram rate:", in_corpus_rate(topk))
print("  plain:", repr(plain[:90]))
print("  topk :", repr(topk[:90]))

plain (k=24) in-corpus 4-gram rate: 0.662
top-k (k=5)  in-corpus 4-gram rate: 0.846
  plain: 'the bcome hll lantern gl when the fksear wall the rises at dusk at dusk canthe d count in '
  topk : 'the bay when the tiide i. ises thick. a small boat drifts near the rocks. tide rises at du'

At temperature 2.0 plain sampling scores 0.662, only two thirds of its 4-grams are real corpus fragments, and you can see why: "bcome hll", "fksear wall" are the tail leaking through a flattened distribution. Restricting to k = 5 lifts the score to 0.846 and the text reads far more like Lantern Bay, because at every step the five allowed characters are still sensible even when the temperature would happily have handed probability to nonsense. Same temperature, same seed, same model, the only difference is that top-k refused to let the tail play.

Practice Exercises#

Exercise 1. Confirm the renormalization arithmetic by hand. Take the last logit vector after "the boats ", compute the full softmax, and record the total probability of everything outside the top 3 (sum of all but the three largest). Then apply top_k_filter(last, 3), softmax it, and check that the three surviving probabilities sum to 1.0 and that each is larger than its pre-filter value. Explain, in one sentence, where the tail’s mass went.

Exercise 2. Sweep k and watch coherence change. Reuse the in_corpus_rate helper and, at a fixed temp=2.0 and fixed seed, generate 200 characters for each k in [1, 2, 3, 5, 10, 24] from the prompt "the ". Print k and its in-corpus 4-gram rate. Describe the shape of the curve: which end is safest, and where does letting more tokens in start to visibly cost you?

Exercise 3. Show that top-k caps the worst case regardless of temperature. Write a function that, for a given k, generates 300 characters and returns the set of distinct characters that appear. Compare the character set for k = 3 against k = vocab_size (plain sampling) at temp=2.5. Verify the k = 3 run never emits a character that was outside the top 3 of some step, while plain sampling reaches characters that only the deep tail could have produced. Why does a small k bound the set of reachable characters more tightly than any temperature setting can?

Summary#

Top-k sampling is a small addition to your generation loop with an outsized effect on quality. Plain temperature sampling reshapes the entire next-token distribution but always leaves the low-probability tail reachable, and that tail grows more dangerous as you raise the temperature for variety. Top-k closes the gap: before sampling, keep only the k most likely tokens, set the rest to $-\infty$ so softmax zeros them, renormalize over the survivors, and draw from those. You implemented top_k_filter and a top-k generate in a few NumPy lines, saw w, c, a survive k = 3 for a real distribution while k = 10 admitted a seven-character tail, confirmed k = 1 reproduces greedy exactly, and measured how k = 3 kept text on the Lantern Bay corpus where plain sampling at temperature 2.0 drifted into gibberish. The parameter k is a safety-versus-diversity dial, tight at the bottom, risky at the top, tuned relative to your vocabulary size.

Key Concepts#

• Top-k filtering — keep the k highest-probability tokens, set the rest to $-\infty$; top_k_filter(logits, vocab_size) is a no-op, so plain sampling is the k = vocab_size special case.
• Renormalization — after zeroing the tail, softmax redistributes its mass across the survivors, so kept probabilities rise slightly (w: 0.4152 → 0.4201 at k = 3).
• Worst-case cap — no token outside the top k can ever be drawn, so the deep tail cannot derail generation no matter how high the temperature.
• k = 1 is greedy — one survivor means probability 1.0 on the argmax, which is deterministic greedy decoding.
• Safety vs diversity — small k (like 3 on a 24-token vocab) stays coherent but less varied; large k restores variety but re-admits the tail; you saw the in-corpus rate rise from 0.662 to 0.846 by switching plain sampling for k = 5 at temperature 2.0.

Why This Matters#

Top-k sampling is one of the two decoding strategies (with top-p, next lesson) that essentially every deployed language model uses, and building it from scratch shows there is nothing mysterious about it: it is one np.sort, one comparison, and a -inf assignment slotted in front of the sampler you already had. The deeper lesson is that good generation is not about a single perfect setting but about shaping the distribution before you draw from it, temperature controls the shape, top-k controls the support, and together they let you dial coherence against surprise on purpose rather than by accident. Once you can see the tail as a concrete thing you are choosing to keep or cut, tuning a model’s output stops being guesswork and becomes a design decision you fully understand.

Continue Building Your Skills#

You now have two knobs that work together: temperature to reshape the kept distribution and top-k to bound how many tokens are kept. But top-k has one awkward edge, k is a fixed count that ignores how peaked or flat the distribution actually is. When the model is very confident, keeping 10 tokens needlessly admits junk; when it is genuinely uncertain, keeping only 3 may be too stingy. The next lesson fixes exactly that with top-p (nucleus) sampling: instead of a fixed number of tokens, you keep the smallest set of tokens whose probabilities sum to at least p, so the cutoff adapts to the shape of each distribution, tight when the model is sure and generous when it is not. You will implement it in pure NumPy by sorting the probabilities, taking a cumulative sum, and cutting where the total crosses p, and compare its behavior directly against the top-k you just built.