LLM Workshop Lab
Chapter 1 · Build Your Own LLM

Next-word prediction & the art of sampling

An LLM does one thing: given the tokens so far, it outputs a probability distribution over the next token. Everything else — chat, code, essays — falls out of running that prediction in a loop and choosing how to pick from the distribution. This lab lets you turn the knobs yourself: temperature, top-K, and top-P.

§1The autoregressive loop

"Autoregressive" means the model's own output is fed back in as input. One forward pass predicts one token; generation is just this loop until the model emits the special <eos> (end-of-sequence) token or hits a length cap.

Step 1

Context in

The prompt plus every token generated so far, e.g. the cat sat on the

Step 2

Model forward pass

One pass through the network produces a score (logit) for every token in the vocabulary — ~50,000 for GPT-2.

Step 3

Softmax → probabilities

Logits become a probability distribution that sums to 1. Most mass sits on a few tokens; the rest is a long tail.

Step 4

Sample one token

Greedy takes the argmax. Sampling draws randomly — shaped by temperature, top-K, and top-P.

↻ Append the chosen token to the context and repeat — until the model picks <eos> or hits max_new_tokens

§2Sampling lab

This is the lecture's Excel demo, made live. The model just saw the cat sat on the and produced this distribution over what comes next. Dashed outlines are the original probabilities; solid bars are what's left after your settings rescale, filter, and renormalize. Then sample from it — a lot — and watch which tokens actually get picked.

<1 sharpens · 1 = as-is · >1 flattens toward uniform
Keep only the K most likely tokens (max = off)
Keep the smallest set covering ≥ P of the probability
No draws yet
Tokens kept
—
Top token prob
—
Entropy
—
Draws so far
0
original probability after temp + filters (renormalized) filtered out — can never be sampled last drawn

Try the classics: temperature 0.05 (≈ greedy — every draw is mat), temperature 3 (near-uniform chaos — spaceship starts showing up), top-K 3 (only three candidates survive no matter what), and top-P 0.9 (the long tail is amputated but the head keeps its shape).

§2.5The last step: rolling a weighted die

Temperature, top-K and top-P all just reshape the menu — none of them actually picks. The pick is a weighted random draw over whatever survived: lay the survivors end to end along a line from 0 to 1, each owning a slice as wide as its probability, then draw a random number and see where it lands. A token holding 47% of the line gets hit 47% of the time. That is all torch.multinomial does — and it's the step do_sample=True switches on.

This line is your §2 survivors — move a slider up there and watch it repartition.
0.00.250.50.751.0
0 rolls

There's no knob for this step because do_sample is the knob, and it's binary: False gives you argmax and ignores temperature, top-K and top-P entirely; True gives you this draw. Argmax is still reachable from inside sampling — top-K 1 leaves one survivor owning the whole line, and temperature near 0 stretches the leader across almost all of it. If you want repeatable output without giving up sampling, the knob you actually want is the random seed: set_seed(42) makes the same prompt and settings produce the same rolls every run.

§3Generation playground

A real (tiny!) language model running in your browser: a bigram model counted from a small corpus about cats, dogs, and mice. It predicts the next word from the current word only — but the autoregressive loop, EOS stopping, and sampling behavior are exactly what GPT-2 does at 50,000× the scale. Click any bar to force that token, or let it sample.

Set temperature to 0.05 (≈ greedy) and auto-run: identical every time — and it gets stuck repeating "the cat sat on the cat sat on…". That's greedy degeneration, live.

§4The same knobs in code

Everything you just did by hand maps one-to-one onto Hugging Face's model.generate() — the exact call from the Colab exercise. 

from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed

tok   = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

ids = tok("The cat sat on the", return_tensors="pt").input_ids

set_seed(42)  # same seed + same settings = same rolls, every run (§2.5)

out = model.generate(
    ids,
    max_new_tokens=40,        # cap on the autoregressive loop (§1)
    do_sample=True,           # False → greedy argmax every step
    temperature=0.8,          # <1 sharper, >1 flatter (§2 slider 1)
    top_k=50,                 # keep the 50 most likely tokens (slider 2)
    top_p=0.95,               # nucleus: smallest set covering 95% (slider 3)
    eos_token_id=tok.eos_token_id,  # the stop token
)
print(tok.decode(out[0]))

§5Test yourself

Eleven questions on everything in this chapter. Instant feedback, explanations included.

§6Key takeaways

One job, loopedAn LLM predicts a probability distribution over the next token. Generation = predict, sample, append, repeat until <eos>.
Greedy ≠ bestAlways taking the argmax produces repetitive, degenerate text. Randomness is a feature, not a bug.
Temperature reshapesIt divides the logits before softmax: low T exaggerates the leader, high T flattens toward uniform. It never removes tokens.
Top-K / top-P removeTop-K keeps a fixed count; top-P keeps a probability budget, so its count adapts to how confident the model is. Survivors are renormalized to sum to 1.
Order of operationsTemperature rescales → top-K filters → top-P filters → renormalize → sample. That's what model.generate does under the hood.
The filters don't pickThey only set the menu. The final token is a weighted random draw over the survivors — each owns a stretch of the 0–1 line as wide as its probability. do_sample=False replaces that draw with argmax and makes the other knobs no-ops; set_seed makes the draws repeatable.
The long tail is realOver 50k tokens, thousands of individually-unlikely tokens add up — that's why unfiltered high-temperature sampling goes off the rails.
Colab code exercise ↗ Excel sampling exercise ↗ Slide deck ↗
Next chapter: finish the next lecture, paste its summary to Claude, and this lab grows a new chapter — same idea: interactive concepts first, then a quiz.