Lesson 11 of 15

Training and Sampling Playground

A stack of transformer blocks turns a token into a rich, context-aware vector — but that's still not a prediction. One more layer turns it into logits: one raw score per word in the vocabulary. Higher score means the model thinks that word is more likely to come next.

Right after initialization, those scores are close to random. Training nudges them toward reality using cross-entropy loss: given the correct next word, penalize the model for putting low probability on it. The gradient of this loss with respect to the logits has a beautifully clean form — softmax(logits) − one_hot(correct_word) — the exact same gradient-descent loop from the derivatives lesson, just with a vector-valued output instead of a single number.

Predict before you look

Before you scrub or hit play: the correct next word is "mat." Roughly what do you expect P('mat') to look like at step 0 versus step 30 — and will the loss curve fall smoothly, or in fits and starts?

Scrub through training and watch two things happen together: the loss curve falls, and the predicted probability on the correct word ("mat", highlighted) climbs toward 1 while every other candidate gets squeezed toward 0. That's all training is — millions of tiny nudges just like this one, repeated across every position in every sentence in the training set.

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.

"The cat sat on the ___" — adjust temperature and top-k, 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.00Top-k: 8

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 these two knobs.

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)])

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

One practical wrinkle in generation: producing each new token requires computing attention against every token so far — which means computing K and V vectors for the whole sequence. Doing that from scratch at every single step means position 1's K/V get recomputed at step 2, step 3, step 4... every step, forever, even though position 1 never changes. That's wasted work, and it grows fast: generating N tokens naively costs roughly 1+2+3+...+N K/V computations — quadratic in sequence length.

Generating token by token — which K/V vectors get (re)computed?

The

cat

sat

the

mat

Without cache

With cache

The fix is exactly what it sounds like: a KV cache. Compute each position's K and V vectors once, store them, and on every later step only compute K/V for the newest token, reusing the cached ones for everything before it. That turns the quadratic cost into a linear one — and it's why real chat models can hold long conversations without generation grinding to a halt as the conversation grows.

Python

What problem does a KV cache solve during text generation?

You now have every piece: tokenize, embed, attend, assemble into blocks, train with gradient descent, and sample. The capstone puts them all together into one small program that actually runs.

As temperature approaches 0, sampling becomes equivalent to…

← Assembling the Transformer Fine-Tuning Intuition (LoRA) →