Where does a language model’s next token come from?

Build the context from tiles

Start with a marker, two color labels, and first or last. Every tile has an exact ID in a seven-token vocabulary. The model predicts a next token from those IDs; it does not receive a hidden English interpretation.

Add a learned position

Each ID selects an eight-number embedding row. A learned vector for that input position is added. Replacing a token changes its embedding; moving it also changes the position contribution.

Calculate attention

Query and key vectors produce scaled dot-product scores. Causal masking forbids later positions. Softmax normalizes the allowed scores, and those weights mix value vectors rather than token IDs.

Keep following the calculation

The mixed value is projected, added through a residual path, transformed by a ReLU feed-forward network, added through another residual and projected into seven output logits. A second softmax normalizes candidate tokens.

Investigate a surprising arrow

For red–blue–first, the next-token winner is red although the final attention row mostly weights first. The display retains those actual values. Inspect the full computation instead of replacing the attention pattern with a convenient story.

Draw and append

A categorical sample chooses a token from the output distribution. Append it, then recompute from the extended context. Stop at the end marker or when another inference would exceed five input positions; no tokens are silently discarded.

Probe an unsupported format

Give the model first without any color. It still places high probability on red. The output exists, but that format has no established color-selection answer in the training task.

A deliberately reduced trained Transformer

This original decoder has one block, one head, width 8, feed-forward width 16, 695 parameter slots, learned positions, separate input/output matrices and no normalization or dropout. It is a symbolic color-selection model, not an English chatbot or a replica of a modern large model.

Exact vocabulary and tokenization

IDs 0–6 mean <bos>, <eos>, red, blue, green, first and last. The editor canonicalizes whitespace and rejects unknown labels. An ID is a row selector; its numeric magnitude is not its meaning or probability. The checkpoint and vocabulary must stay matched.

A separate subword arithmetic card

The optional red redder exercise begins with characters and retains the space. Four chosen merges produce red, space, red, der. This demonstrates that tokens need not equal words. Its hand-constructed vocabulary and ranks are not used by the neural checkpoint.

Queries, keys and values

For row vectors x, the model forms q=xWq, k=xWk and v=xWv. Each allowed attention score is q·k/√8. A stable softmax subtracts the maximum allowed score before exponentiating. Future positions have score −∞ and weight exactly zero.

Two different softmax operations

Attention softmax distributes mass over allowed context positions. Output softmax distributes mass over seven candidate next-token IDs. A 0.9 attention weight on first does not mean a 0.9 probability of emitting first.

The entire downstream path matters

Attention’s weighted value sum passes through Wo and adds to x. A position-wise ReLU network and a second residual produce z. Output logits are zU+b. The exploded stack shows actual vectors at one selected position; cell height is a declared compressed visual scale, not a physical dimension.

A worked value coordinate

For <bos> red blue first, the final attention weights are approximately 0.000858485, 0.003125844, 0.096708236 and 0.899307436. Multiplying them by value-coordinate-0 entries and summing gives approximately −1.957841407. Full-precision values drive the visualization.

A defined artificial intervention

Connection removal sets one allowed attention entry to zero, renormalizes the remaining row, and recomputes downstream layers with the original queries, keys, values and parameters. Removing the only remaining allowed connection is undefined. Replacing a token changes its vectors and normalization, so it tests a different intervention.

Attention is evidence, not a complete causal account

Published work debates which explanatory claims attention supports and how to test them. A heatmap alone, or one ablation, does not establish a complete mechanism. This workbench shows measured intermediates and specific changed calculations without inventing a natural-language inner monologue.

Causal prediction and shifted targets

All 18 six-token sequences supply five input–next-target pairs. At input position i, the current token is known, so attention to itself is allowed. Inputs at positions later than i cannot influence its logits. Showing a full training worksheet does not grant the model access to those future answers.

Teacher-forced training and autoregressive inference

Training uses the ground-truth preceding tokens while computing 90 next-token targets per full batch. Generation instead appends its selected output before the next forward pass. Ordinary inference neither computes a gradient update nor changes the checkpoint.

Temperature operates after logits

For positive temperature τ, output probabilities are proportional to exp((logit−maximum)/τ). Changing τ leaves current logits and attention fixed. Greedy choice is a separate operation, with exact ties resolved by lowest token ID. A different sampled token can change later context and therefore later attention.

Reproducible categorical draws

The educational 32-bit LCG advances state once for each random draw, then u=state/2³². Choose the first cumulative probability strictly greater than u. Seed 49 begins with 0.2550580417737365 and 0.7230514015536755. It is a classroom generator, not cryptography.

A sampling calculation you can do on paper

For the separate logits [ln 6, ln 3, 0], temperature 1 gives [0.6,0.3,0.1]. Temperature 0.5 gives [36/46,9/46,1/46]. This card is intentionally distinct from the trained model’s often much sharper distributions.

All training data are disclosed

The corpus enumerates three first colors, three second colors and two queries: 18 sequences. All are used in training, without a held-out split. The 18/18 answer-position result is a training-task check, not evidence of generalization or a reasoning benchmark.

Why loss has a nonzero floor

Earlier identical prefixes genuinely admit several next targets: three first colors, three second colors and two queries. The empirical conditional loss floor is (ln 3+ln 3+ln 2)/5≈0.578074352 nats per target. Near-perfect answer selection does not make mean loss over all five targets zero.

Recorded training, fixed inference

The downloadable original code initializes from NumPy PCG64 seed 49 and performs 2000 full-batch Adam updates in float64. The final loss is approximately 0.578083150. The ledger’s points are recorded values, not live training. The exact checkpoint and corpus have separately recorded SHA-256 hashes.

Where this analogy ends

Larger systems can have different architectures, richer pretraining, post-training, reasoning behavior, tools and context handling. A next-token interface does not settle what algorithms a learned computation implements. This tiny experiment neither demonstrates nor rules out those larger capabilities.

Reading an AI visualization critically

Ask whether arrows show measured attention, a proposed explanation or actual intervention results. Follow the values and later computations; do not assume the prettiest arrow identifies a cause.

Comparing prompts and decoding settings

A prompt edit changes context computation; temperature changes selection from fixed logits; training changes parameters. Keeping these separate helps design an interpretable comparison.

Checking a confident answer

A categorical model probability is not a verified probability that an independent claim is true. Use appropriate external evidence and task evaluation, especially for formats outside the model’s demonstrated data.

After the four specified merges, how many tokens encode red redder in the separate arithmetic card?

• Two words, so two tokens
• Four: red, space, red, der
• Ten unchanged characters

Change a later input while keeping positions 0 and 1 fixed. Their logits should…

• Remain unchanged
• Improve by reading the new future
• Stop attending to their own known tokens

The worked attention/value sum at coordinate 0 is approximately…

• −1.957841407
• 0.899307436
• The average token ID

A small attention weight on red proves…

• The animation must be wrong
• Red is causally irrelevant
• Neither; inspect the full path and defined interventions

Temperature 0.5 on logits [ln 6, ln 3, 0] gives…

• [36/46,9/46,1/46] with unchanged parameters
• New training evidence that red is true
• Doubled attention weights

u=0.7230514 on the [0.6,0.3,0.1] arithmetic card selects…

• Red because it is most probable
• Blue because u lies in [0.6,0.9)
• Green because the seed is large

Why can 18/18 training answer selections coexist with loss near 0.57808?

• Earlier prefixes have multiple possible next targets
• Every prompt secretly retrains the weights
• A perfect answer count forces zero total loss

A high red probability for <bos> first, without a supplied color, establishes…

• A verified red fact
• A confident continuation outside the trained grammar
• That all larger models cannot reason

Scaled attention, causal masking, residual paths and feed-forward transformations

Vaswani and colleagues, original Transformer paper. Our reduced decoder differs substantially; all product diagrams are original, not redistributed paper figures.

Subword units and tokenizer-specific segmentation

Sennrich, Haddow and Birch, ACL 2016. Our optional ordered-merge arithmetic card is hand constructed and separate from the checkpoint vocabulary.

Tokenization is a specified encoding procedure

Kudo and Richardson, EMNLP 2018. No SentencePiece model or external vocabulary is used in this seven-token lesson.

Limits of interpreting attention as explanation

Jain and Wallace, NAACL 2019. Supports testing explanatory claims; not a theorem that all attention displays are useless.

A complementary view of attention interpretation

Wiegreffe and Pinter, EMNLP 2019. Their proposed tests motivate explicit interventions and qualified interpretation.

Adaptive training updates in the original reproducer

Kingma and Ba, Adam. The corpus, seed, learning rate and finite run here are explicitly authored choices.

Sampling and output-distribution transformations

Holtzman and colleagues: autoregressive decoding and sampling. No paper-specific performance claims are transferred to the miniature.

Output likelihood and truth require separate evaluation

Lin, Hilton and Evans, ACL 2022. Historical TruthfulQA results do not rate current systems or calibrate this color model.

Larger learned systems and reasoning need their own evidence

Wei and colleagues study intermediate-step prompting in larger models. This miniature does not reproduce their training or results.

Recorded seed and checkpoint are different reproducibility artifacts

NumPy qualifies cross-version random-stream guarantees. The exact checkpoint and complete corpus are supplied with hashes rather than relying on a seed alone.