Part IV — Pretraining

Decoding & loading pretrained weights

Lesson 10 built the training loop and watched a tiny model overfit one short story — and its greedy decoder loops on the same phrase because it always grabs the single highest-scoring token. This lesson adds two knobs, temperature and top-k, that turn that rigid argmax into controllable sampling; then it does the move that makes everything before it pay off — loading OpenAI's published GPT-2 weights into the architecture we built, so we get a strong base model without paying for pretraining.

1 · Greedy decoding is deterministic — and it repeats

Every generation step in lesson 08 ended the same way: the model produced a logit vector over the vocabulary, and we took its argmax — the single highest-scoring token. That is greedy decoding. It is simple and it is defensible (each step picks the locally most probable token), but it has two properties that matter here.

First, it is deterministic. The same context always yields the same logits, so it always yields the same next token — run the model on one prompt a hundred times and you get one identical continuation every time. There is no way to ask for a second, different draft. Second, a greedy path is prone to degenerate loops: once the model is confident about a short high-probability cycle ("… and the and the and the …"), argmax keeps re-entering it, because at each step the loop's next token really is the most probable one. A tiny model that overfit its training text (lesson 10) is especially vulnerable — it has memorized a few high-probability continuations and returns to them.

The fix is to replace the hard argmax with sampling: convert the logits to a probability distribution and draw a token from it, so lower-ranked tokens sometimes win. Concretely, we swap the picker:

logits = model(tokens)              # length-V scores for the next token
# greedy (old): deterministic, repetitive
next = argmax(logits)
# sampling (new): draw proportional to probability
probs = softmax(logits)
next  = sample_from(probs)          # e.g. one multinomial draw

Sampling alone already breaks the determinism. But raw sampling from softmax(logits) is a fixed amount of randomness — we cannot make it bolder or more cautious. We want a knob. That knob is temperature.

2 · Temperature scaling

Temperature scaling is a one-line change with a large effect: before the softmax, divide every logit by a positive number T called the temperature.

Why does dividing the logits reshape the distribution? Softmax cares about the gaps between logits — the difference z_i − z_j sets the odds ratio e^{(z_i − z_j)} between tokens i and j. Dividing every logit by T divides every gap by T, which rescales all those odds ratios together:

• T = 1: no change — you get the original softmax probabilities. This is the neutral setting.
• T < 1 (sharpen): gaps grow, so the leading token pulls away and the distribution gets peakier and more confident. As T → 0 the biggest logit dominates completely and sampling collapses to argmax — temperature recovers greedy decoding in the limit.
• T > 1 (flatten): gaps shrink, so probabilities move toward each other and the distribution gets flatter and more adventurous. As T → ∞ every gap goes to 0 and the distribution approaches uniform over the vocabulary.

On our example, this is exactly the diversity/coherence trade-off in miniature. With T = 1, the model would pick "forward" roughly 60% of the time — mostly coherent, occasionally choosing a sensible alternative like "toward". Push to T = 0.1 and "forward" wins essentially 100% of the time — indistinguishable from greedy, safe but repetitive. Push to T = 5 and the bars nearly level out: "forward" drops to about a quarter of the mass, and low-probability tokens start getting drawn. That flattening is what buys variety — and also what lets nonsense through. At T = 5 the token "pizza" gets sampled roughly 4% of the time, producing continuations like "every effort moves you pizza". Temperature does not make the model smarter; it only decides how far down the ranked list you are willing to reach.

3 · Top-k sampling (and a word on top-p)

Temperature reshapes the whole distribution, which means even at a modest temperature there is always a little probability mass smeared across hundreds of terrible tokens — the long tail. Sampling can, rarely, draw from that tail and derail the sentence. Top-k sampling removes the tail outright: keep only the k most likely tokens and give every other token zero probability, so a bad token cannot be selected no matter how many times you sample.

The mechanism is the same masking trick we used for the causal attention mask (lesson 06): set the unwanted logits to −∞ before the softmax. Because e^{−∞} = 0, those positions get probability exactly 0, and the softmax automatically renormalizes the surviving k tokens so they sum to 1 — no separate renormalization step needed.

topk_vals, topk_idx = top_k(logits, k)      # the k largest logits and where they are
mask = fill(logits.shape, -inf)             # everything masked off by default
mask[topk_idx] = logits[topk_idx]           # let only the top-k logits through
probs = softmax(mask / T)                    # masked → 0 prob; survivors renormalize
next  = sample_from(probs)

Take our example with k = 3. The three largest logits are "forward" (6.75), "toward" (6.28), and "closer" (4.51); every other token — including "pizza" at −1.89 — is set to −∞. Softmax over the survivors gives roughly:

Now the model still samples — it might pick "toward" for variety — but it is structurally impossible to emit "pizza." Top-k gives you the diversity of sampling with a floor under the quality: it bounds how bad the chosen token can be. The knob k sets the width of the candidate set: k = 1 is greedy again (only the top token survives), while large k approaches plain temperature sampling. Temperature and top-k compose — a common recipe is a moderate temperature with a moderate k, flattening the head for variety while the mask deletes the tail.

4 · Loading pretrained weights — the payoff of building it faithfully

Here is the move this whole track was set up to make. Everything we built — the BPE tokenizer (01), the token + positional embeddings (02), causal multi-head attention (05–06), the pre-norm transformer block with its GELU feed-forward (07), the stack with a final LayerNorm and an LM head (08) — is the same architecture OpenAI used for GPT-2. We did not build a lookalike; we built the thing. That means we do not have to train it from scratch. OpenAI published the trained GPT-2 weights, and because our modules match theirs shape-for-shape, we can pour those weights straight into our model and inherit a capable base model — skipping the compute bill entirely.

GPT-2 was released in four sizes; all share the design, differing only in width d, number of heads h, and depth L:

All four use the same vocabulary V = 50{,}257 and context length 1024. The 124M row is exactly the GPT_CONFIG_124M we have been building against, which is why it drops in cleanly: pick the config, instantiate the model, and copy the published tensors into the matching parameters.

What "loading weights" actually means

Loading is a bookkeeping exercise, not a learning one. For every parameter tensor in our model there is a corresponding tensor in the published checkpoint; you walk the model layer by layer and assign. The only real work is matching the layout: our parameter names and tensor shapes must line up with the checkpoint's, and where a convention differs — for example, a weight matrix stored transposed, or GPT-2's habit of packing the query, key, and value projections into one combined tensor that we split into three — you reshape or transpose as you copy. A single assertion protects you: check that each source shape equals the destination shape before assigning, so a silent mismatch becomes a loud error.

def load_gpt2_weights(model, params):
    # embeddings
    assign(model.pos_emb,  params["wpe"])          # (context_length, d)
    assign(model.tok_emb,  params["wte"])          # (V, d)
    for b in range(L):                              # each transformer block
        # GPT-2 packs Q,K,V into one matrix — split it into our three
        q_w, k_w, v_w = split(params[b]["attn.c_attn.weight"], 3, axis=-1)
        assign(model.blocks[b].att.W_q, transpose(q_w))
        assign(model.blocks[b].att.W_k, transpose(k_w))
        assign(model.blocks[b].att.W_v, transpose(v_w))
        assign(model.blocks[b].att.W_o, transpose(params[b]["attn.c_proj.weight"]))
        # feed-forward, layer norms, biases … one assign per tensor
        ...
    assign(model.final_norm, params["ln_f"])
    # LM head is weight-tied to the token embedding in GPT-2
    assign(model.out_head,  params["wte"])

The specific keys and transposes depend on the checkpoint format you download, but the shape of the task never changes: every tensor in, every tensor matched, one assertion per copy. Note the last line — GPT-2 ties the output head to the token-embedding matrix (the same (V, d) weights are reused to embed inputs and to produce logits), so there is nothing extra to load there. Once the copy finishes, running our forward pass produces genuine GPT-2 quality text, and with the temperature/top-k decoder from this lesson you can steer it.

5 · Steer the decoder yourself

The widget holds the book's fixed nine-token logits — no randomness in the scores. The temperature slider divides those logits by T before softmax (watch the bars sharpen below 1 and flatten above 1); the k slider masks all but the k largest logits to −∞ and lets softmax renormalize the survivors. The bars are the resulting sampling distribution. The blue bar is the argmax ("forward"); the green marker shows which token a fixed-seed sampler would draw from this exact distribution — deterministic, so the same settings always highlight the same token. The KPIs quantify the trade-off: entropy and effective #tokens (e^{entropy}, roughly "how many tokens are genuinely in play") both rise as you flatten, and the nonsense risk is the probability mass sitting on tokens outside the sensible set — the chance of a "pizza"-style slip.

Failure modes & checklist

Checkpoint

Where this points next

With sampling knobs and real GPT-2 weights, we now have something genuinely useful: a base model that completes text fluently and that we can steer for variety or safety. But notice what it still will not do. Prompt it with "Classify this message as spam or not: 'You won a prize, click here'" and it does not answer — it just continues the passage, because completing text is the only behavior it has. A base model has knowledge and fluency but no notion of "produce a decision" or "follow this instruction." The next wall is turning a completer into something that acts. Lesson 12, Classification fine-tuning — a spam detector, takes the smallest possible step: swap the LM head for a tiny class head, read the last token, and fine-tune the base model into a spam classifier.

Interview prompts

• Why is greedy decoding deterministic, and why does that cause repetition? (§1 — the same context yields the same logits, so argmax always returns the same token; a model can re-enter a high-probability short cycle because at each step the loop's next token really is the most probable.)
• Derive how dividing logits by T reshapes the softmax. (§2 — softmax depends on logit gaps; dividing every logit by T divides every gap by T, scaling all odds ratios e^{(z_i−z_j)} together — smaller T widens gaps (peakier), larger T shrinks them (flatter).)
• What does temperature do in the limits T→0 and T→∞? (§2 — T→0 collapses to argmax/greedy; T→∞ approaches a uniform distribution over the vocabulary.)
• Why implement T=0 as a special case instead of dividing? (§2 trap — temperature is a divisor, so 0 is undefined and produces NaNs/Infs; "temperature 0" is a convention meaning argmax.)
• How does top-k work, and why mask to −∞ before softmax rather than zeroing after? (§3 — keep the k largest logits, set the rest to −∞; since e^{−∞}=0 those get probability 0 and softmax renormalizes the survivors automatically, so probabilities still sum to 1.)
• How does top-p differ from top-k? (§3 — top-p keeps the smallest set of top tokens whose cumulative probability exceeds p, adapting the candidate count to the distribution's shape; top-k always keeps exactly k.)
• Why can we load OpenAI's GPT-2 weights into our model at all, and what is the main practical hurdle? (§4 — our architecture matches GPT-2's shape-for-shape, so tensors correspond one-to-one; the work is matching layout — transposes and splitting GPT-2's packed QKV — with a shape assertion on every copy. The LM head is weight-tied to the token embedding.)

Companion reads: cs336 · 18 Inference covers decoding and serving at production scale (KV cache, batching, speculative decoding); gpt_mini · 02 Pretrain is the compact tour of pretraining the base model whose weights we load here.