↑ ↓ to move · enter to open

C07 · AI · tokens · probability · attention

how a language model writes

A language model writes one small piece at a time. It chops your words into tokens, turns them into numbers, scores every possible next token, picks one, and repeats. This page builds each step at toy size so you can watch every number: a tokenizer trained in your browser, a word-counting predictor, a four-number attention layer, and the full loop. Real models run the same loop with billions of learned parameters.

LAB 1

chopping text into tokens

A model can't read letters. It reads a list of whole numbers called tokens. A good token list is short, so common chunks like ·the should be one token, but rare words still need to work. Byte-pair encoding builds that list by counting: start from single bytes, find the pair of neighbors that appears most often, glue it into a new token, and repeat. Below it trains on a short story written for this page. Each dot · is a space.

training text
40

How to read this: each colored patch is one token. The colors only show where one token ends and the next begins. They carry no other meaning. In the chips below, the small number under each token is its id: 0 to 255 are single bytes, 256 and up are learned merges.

  • amber: the pair glued together most recently

Try this: press Play and watch the patches in the story grow longer as common chunks like ·the merge.

Vocabulary
—
Characters typed
—
Tokens
—
Characters per token
—
Vocabulary size

The merge rule

(a, b) = argmax count(a b)
replace every a b with new id 256 + k
vocab size = 256 + merges

Ids 0 to 255 are the 256 possible byte values, so any text in any language can be encoded, even if it was never seen. Each merge adds one id. Merges never cross a word boundary, because the text is first split into words that keep their leading space.

To tokenize new text, the model replays its merges in the order it learned them. Fewer tokens per sentence means the model sees more text for the same amount of work.

Now

CodePython 3


          

Save the training text above as corpus.txt. Ties go to the pair seen first, here and in the page.

LAB 2

guess the next word

A language model never stores sentences. It answers one question over and over: given the text so far, how likely is each possible next token? The simplest possible answer is to count. This model read the same story, split into words, and counted which word followed each pair of words. Its whole memory is a table of counts. Here it predicts whole words so you can read it; real models predict tokens from lab 1.

next-word odds
1.00
3 0.70

How to read this: each row is a word that followed this context in the story. The logit is its raw score, and the green bar is its chance of being picked next. Faded rows have been ruled out by the picking rule.

  • in the text below: your start
  • what the model wrote
  • underline: newest word

Try this: slide Temperature down to 0.1 and back up to 2.5 while watching the green bars.

—
Entropy
—
Effective choices 2H
—
Words allowed
—
Last pick

Scores to probabilities

logit zi = ln(counti)
pi = ezi/T / Σj ezj/T (softmax)
H = −Σ pi log2 pi bits

A network outputs one score, a logit, per token. Softmax turns scores into probabilities that are positive and add to 1. With logits equal to log-counts and T = 1, softmax gives exactly count ÷ total.

Temperature divides every logit before softmax. T below 1 sharpens the odds toward the favorite, and T above 1 flattens them. Entropy measures the spread: 0 bits means one certain word, and 2H is roughly how many words are really in play.

Picking a word

Greedy always takes the top word, so the text is repeatable and often loops. Top-k keeps the k likeliest words. Top-p, or nucleus sampling, keeps the smallest set of top words whose probabilities add up to at least p. Either way, the survivors are rescaled to add to 1 and one is drawn at random.

From counting to a network

Counting has two big flaws. A pair of words never seen in training has no row at all, so the model falls back to fewer words of context. And a word never seen after a context gets a count of 0, a logit of −∞, and probability exactly 0.

An LLM replaces the table with a neural network that reads the whole context, often thousands of tokens, and gives every token in its vocabulary a finite logit. Its numbers are learned by the same kind of training as C05 · How machines learn. The softmax, temperature and sampling steps stay exactly as here.

CodePython 3


          

Python's random generator differs from the page's, so the same seed gives a different but equally repeatable text.

LAB 3

attention: which words matter

To guess what comes after "it", a model has to know what "it" refers to. Attention lets each word look back at earlier words and pull in what it needs. Every word gets three vectors: a query (what am I looking for?), a key (what do I offer?), and a value (what do I pass along?). A query that lines up with a key gets a high score. The vectors below have 4 numbers each and were set by hand so that "it" looks for an animal. A trained model learns its vectors and uses hundreds or thousands of numbers. Click a word.

attention weights

How to read this: brighter green means a word pays more attention there, and each row adds up to 100%.

  • attention paid
  • striped: blocked, a word can't look at words after it
  • the strongest link in the row
  • the row you picked

Try this: click "it" and see which earlier word lights up. Then switch to Head 2 · previous word.

Row detail
—
Asking word
—
Strongest link
—
Row entropy
—
Row sum

Scaled dot-product attention

scoreij = qi · kj / √d
mask: scoreij = −∞ when j > i
outi = Σj softmax(scorei)j · vj

In matrix form that's softmax(QKᵀ/√d + M)·V, every word at once. The mask M blocks later words, because while writing, the model can't see words that don't exist yet.

Dividing by √d keeps scores from growing with the vector length. Without it, long vectors give huge scores and softmax snaps to a single word, which makes training unstable. Here d = 4, so √d = 2.

What the word carries now

Many heads

A real layer runs many heads side by side, each with its own Q, K and V. Head 2 here uses position instead of meaning: its key for word j is (1, j, j²), so its score is −6(j − (i−1))², highest on the word just before. The outputs of all heads are joined end to end and mixed by one more matrix, WO.

CodePython 3 + NumPy


        
LAB 4

the whole loop

Put the pieces together and you get one step of writing: tokens become vectors, layers of attention mix them, the last vector turns into a score for every token, softmax makes odds, one token is drawn, and it's added to the text. Then the whole thing runs again on the longer text. Training is the same pass in reverse: compare the odds with the real next token and nudge every number to make the right answer more likely.

one step of writing

How to read this: press Next stage to walk down the list. Amber is the step happening now, green steps are done, and gold is the token that got picked.

  • colored squares: numbers in a token's vector, green = positive
  • blue = negative (brighter = bigger)
  • small boxes: one per layer, outlined green once the text has passed through

Training · one gradient step at a time

Context "the cat". The story's answer is sat, but these starting logits favor "went". Each Gradient step pushes the gold bar up and the others down. The last column (p − y) shows how wrong each chance is, and each step moves every logit by −η(p − y).

1.00

How many numbers? · parameter counter

12
768
attention 4Ld²feed-forward 8Ld²embeddings Vd
—
≈ 12·L·d² + V·d
—
With biases, norms, positions
—
Weights at 16 bits each
—
Training loss now

Now

The math of one step

x = embed(tokens) + positions n × d
x ← x + attention(x); x ← x + FFN(x) × L
z = hlast · ET one logit per token, V of them
p = softmax(z / T) → sample → append

Embedding is a table lookup: token id → one row of a V × d table of learned numbers. C06 · Embeddings shows why nearby vectors mean similar things. Each layer adds its result back onto x, so information flows straight through.

Learning from mistakes

loss = −ln p(correct next token)
∂loss/∂zi = pi − yi
z ← z − η · ∂loss/∂z

If the model gave the right token probability 1, the loss is 0. At 0.5 it's 0.693. Guessing evenly over V tokens costs ln V, about 10.8 for V = 50,257. Training averages this loss over enormous amounts of text and lowers it by gradient descent, the method in C05 · How machines learn. Here only 5 logits move; a real step moves every parameter, with gradients found by backpropagation.

Counting the parameters

Each layer has four d × d attention matrices (Q, K, V, output): 4d². Its feed-forward part widens to 4d and back: 2 × d × 4d = 8d². That's 12d² per layer. One V × d table turns tokens into vectors and, shared, turns the last vector back into logits.

The ≈ formula leaves out biases and layer norms (13d per layer here), learned position vectors, a separate output table if it isn't shared, and design variants such as other feed-forward widths or shared keys across heads. For 12 × 768 with V = 50,257 and 1,024 positions the full count is 124,439,808, which matches the published size of GPT-2 small (2019).

CodePython 3