Reach one complete prediction before revisiting the machinery

The full article is a multi-session lesson. For a first pass, follow the route below and work one example at each stop. The links jump to the relevant section; they do not remove slides from the complete presentation.

1. Set the prediction task, then use the last-token baseline to show why context matters. Treat the fixed-window recap in Section 3 as optional if Part 1 is fresh.
2. Motivate unequal weights and use one search query to distinguish a request, matching features, and retrieved information.
3. Return to the sentence. Follow matching and value mixing, then add the update. Keep the receiving position visible.
4. Exclude future tokens, then finish the vocabulary prediction. End this pass with a probability for the next token, not an unexplained attention vector.

In the next sitting, revisit the two bank contexts, the learned projections and score scaling (Sections 10–12), then consolidate the full walkthrough and parallel matrices (Sections 15–16). Preview the scaling factor on the first pass; derive it on the second.

After the matrix calculation, investigate position and order. Part III adds multiple heads and closes the text series with the TinyStories walkthrough, measured results and live generation. The detailed MLP/single-head lab remains in Notebook 5. The comparison of concatenation, averaging and attention costs is in optional Part 2B.

We want to predict the next word. As we add each part of the model, we will ask how it helps that prediction. Switch between the two contexts below and watch the probability bars move. The last three words are identical, so the model needs information about river or cheque from earlier in the sentence. Before opening that model, we’ll separate two ideas: a word’s vocabulary vector and the position where that word occurs.

All following bank calculations use these same four-coordinate sums. The named features and the small position offsets are hand-chosen teaching values. Attention and prediction can use the position information because it shares coordinates with the word features.

Start with a model that sends only the last token to the output head. Its entire input is one row of $d_{\text{model}}$ numbers. Switch the context and the bars stay put. The table below explains why: both sentences end with the at position 10, so they supply exactly the same starting row. The same row gives the same probabilities. In this toy, eight candidates also happen to tie. A last-token-only model need not produce a uniform distribution; it must produce the same distribution whenever it receives the same row.

Give the model more context by concatenating the last $w$ representations into one long row. Both network views use the MLP layout from Part I. Here we choose eight hidden units and twenty vocabulary scores. First identify the available inputs, then join their coordinates, then examine the network's equations and parameters. At $w=3$, river is outside. At $w=10$, every token fits because this sentence is short. Changing the window changes the required parameter shapes. These are architecture sketches, not trained predictions for every window size. The article also derives the optional simpler linear head separately.

A simpler linear head

This alternative connects the concatenated input directly to the vocabulary scores. It removes the MLP's hidden layer, so its single matrix W has a different shape from either W₁ or W₂ above.

In this direct linear head, each extra token adds 4 inputs and 4 × 20 = 80 weights. The 20 outputs stay fixed.

The input slots stay fixed

The network can respond to words inside those slots. The limit is which input rows reach it, not an inability to use the content they contain.

Counting the direct linear head

This count belongs to the single-matrix alternative above. It differs from the two-layer MLP's count $wdH+H+HV+V$.

More slots cost more parameters. A still longer prefix can put the clue outside again.

We now build context for a token already in the input. Begin at position 7, bank, using the seven known token rows. A pooled summary $m_i$ combines their corresponding coordinates, while the earlier $c_t$ concatenated rows. Try averaging first, then choose unequal contributions yourself. The summary stays four numbers wide. Later, we will build a summary at every input position before prediction. These arithmetic operations do not update the stored vocabulary table.

We will first express several questions as query vectors. Each coordinate has a named topic in this example, and a larger number represents a stronger request for that topic. After comparing the questions, we will describe videos on the same axes, then calculate which descriptions best match the request. These numerical examples are chosen by hand to make the roles visible. Learned representations do not generally have one human-readable topic per coordinate.

Hard retrieval returns the value of the highest-scoring item. Attention blends values according to how well their keys match the query. Follow that calculation in the tables: softmax converts the previous section's scores into the $\va{\alpha_j}$ weights, each weight scales its value row, and the bottom row adds the results. Start in hard mode, where one item gets all the weight. Switch to soft mode and lower the temperature to see it approach hard retrieval. The returned row is a mixture of values.

The query describes what the searcher needs. Keys help score the matches; values carry the information to return. Softmax gives us the weights for mixing that information. We can use the same roles for tokens in a sentence.

The video example keeps eight content coordinates throughout retrieval. For the sentence arithmetic, we return to our original smaller model: three query/key coordinates and two value coordinates. The matching and mixing rules stay the same.

First take the numerical queries, keys and values as supplied examples. Read their named axes and use them to collect context. After we have followed the message through the residual addition and into prediction, Section 11 explains how these exact same rows are produced from the token representations.

These are hand-designed teaching coordinates, not measured human concepts or axes that real models must learn. The same numerical model supplies every table; nothing is re-sampled between slides.

We will work through attention in two phases. Phase A uses queries and keys to score the available tokens, then turns those scores into weights that sum to one. Phase B multiplies each source's value by its weight and adds the results. The keys help choose where to read; the values carry the information.

Run phase A and click a key in the diagram to inspect its dot product, coordinate by coordinate. Then run phase B and follow each coordinate into the sum. Switch between the original value rows and the $\va{\alpha}\vv{v}$ products to see what the weights changed.

The message $m_7$ is a mixture of values, so it has $d_v$ value coordinates. To add it to bank's representation, we need $d_{\text{model}}$ representation coordinates. $W_O$ performs that mapping and gives us $\vd{\Delta e_7}$. The residual connection adds this update to the original $\ve{e_7}$, combining bank's existing representation with context.

Press Play or advance one step at a time to follow the diagram. The next table shows the numerical result; hover a cell in the final row to inspect the addition. Watch the axis names change at $W_O$: it maps the message's value coordinates into representation coordinates. Attention computes the update, and the residual connection adds it.

Run the same calculation for the river sentence and the cheque sentence, following bank in both. Its starting row is identical because its token and position are the same. Compare the $\va{\alpha_{7j}}$ columns: river gets most of the weight in sentence A, while cheque and deposited get most in sentence B. Follow those contributions into the message at the bottom and the $\vd{\Delta e}$ row. The update adds water information in A and finance information in B. Drag the slider to add it gradually. The same starting representation becomes context-dependent because it receives different information.

Queries, keys and values come from projecting the current representation. They are not separate learned lookup tables. Pick a token and compare the four rows: three fixed matrices turn the top $\ve{e}$ row into the query, key and value below it, without changing that input row. Their axis names describe different roles: what the token asks for, when it is a useful match, and what it sends. The matrices below are shared by every token. Step through the lifetime diagram to see when the rows are computed and used. These working vectors help produce $\vd{\Delta e}$; the row passed to the next sublayer is $\ve{e} + \vd{\Delta e}$.

The source's row becomes its key: when this source could help.

Same source row, different weights: the value is what this source sends.

We first inspect random query/key entries, then repeat the calculation over many pairs to measure the spread. A two-key example connects that spread to softmax saturation. These demonstrations use separate illustrative numbers. The final table returns to our actual bank query. Extra reading below includes the variance derivation, a fixed-score divisor slider and a standard-normal simulation.

The variance argument

Assume all query and key coordinates are mutually independent, mean 0, variance 1. Each product $q_lk_l$ has variance 1. Adding $d_k$ independent products gives

Dividing by $\sqrt{d_k}$ makes the variance 1 under those assumptions. Learned coordinates need not obey them exactly. Scaling controls the expected size of the scores, but attention can still be sharp.

Stable softmax uses $\exp(s_j-\max_l s_l)$. Subtracting one constant preserves every score difference and probability. Scaling divides those differences, so it changes the probabilities. These operations solve different problems.

scores = (Q @ K.T) / math.sqrt(Q.shape[-1])
alpha = scores.softmax(dim=-1)

scores = (q_bank @ K.T) / math.sqrt(K.shape[-1])
alpha = scores.softmax(dim=-1)

During training, the answer for position $i$ is already in the sentence at position $i+1$. We need to stop the model from reading it. Toggle the causal mask and watch row 5. With the mask on, token 5 (the) reads only positions 1 to 5. With it off, much of its weight moves to future tokens: river is the immediate target, and bank provides further future information. Send the values in both states and compare the messages. The mask enforces $j \le i$ by setting forbidden scores to $-\infty$ before softmax.

The same rule for four positions

Rows are receiving queries $i$. Columns are source keys and values $j$. A check allows that source, including the diagonal. A cross blocks a future source.

scores = scores.masked_fill(future, float('-inf'))
weights = scores.softmax(dim=-1)

Sections 09 to 13 built an updated vector for bank. Now extend the prefix through “and watched the” and predict the token after that final the. Its query, not bank’s query, determines this attention row. In this single-layer toy, earlier words reach the prediction through the update to that last position. Switch the context between the river and cheque sentences: the attention row, update and probabilities change. Then bypass attention and switch again: the probability row stays fixed, just as in section 02.

q_last = E[-1] @ W_Q  # the current last known position

scores = (q_last @ K.T) / math.sqrt(K.shape[-1])
alpha = scores.softmax(dim=-1)  # ten source weights; sum = 1

delta_last = (alpha @ V) @ W_O
e_last_new = E[-1] + delta_last

h = (e_last_new @ W_hidden + b_hidden).relu()  # W1, b1: 4 → 8
logits = h @ W_vocab + b_vocab  # W2, b2: 8 → 20
p_next = logits.softmax(dim=-1)

Forward computation changes the token states for this input. It does not change the stored projection matrices or embedding table.

Put the stages together for the ten-token sentence. Each table has one row per token and one column per coordinate, with the result in the last column or bottom row. The numbers come from the same calculation we have worked through by hand. Click a score or press Show the arithmetic to inspect it. At step 3, choose a different query token and advance again. The scores, mask and weights change, while the operations stay the same: score the sources for one query, mix their values, then project and add the update.

The walkthrough followed bank through attention. Every token applies the same operations to its own row, so we can compute them together with matrix products. Here we use the short seven-token prefix: each product carries out the corresponding walkthrough step for all seven tokens.

Follow row 7 as you advance through the operations. It contains bank's numbers at every stage. Each matrix row belongs to a token, and $\vp{E'} = \ve{E} + \vd{\Delta E}$ adds all their updates at once.

PyTorch · matrices are row-major; mask contains 0 or −∞.

A = (Q @ K.T / d_k**0.5 + mask).softmax(dim=-1)
E_new = E + (A @ V) @ W_O

We count the work of concatenation, averaging and attention in optional Part 2B: the cost of a longer context, after studying training and cached generation.

The fixed window and plain mean use fixed inclusion or pooling rules, although their outputs still depend on the words they receive. Compare them with distance-based weights and attention in one table.

Switch between the river and cheque sentences. The first three rows stay fixed because their inclusion or weighting rules do not read the content. The attention row gives more weight to different clues for the blank. Of these four schemes, only attention computes its mixing weights from the words.

Position is a design choice, not an extra word

Visual companions: Luis Serrano: positional encoding as a geometric displacement and Jia-Bin Huang: rotary position embedding and relative offsets. Our figures are original redrawings with separate, labelled toy numbers. The videos motivated the visual sequence; the mathematical claims are checked against the papers below.

The four-coordinate teaching model uses explicit additive position rows. The two-coordinate swap experiment above is independent and uses identity Q/K/V projections solely to make every number inspectable. Its results show that positional inputs can distinguish two token orders, not that an untrained model understands either sentence.

Appending a scalar index or a small position vector is valid. Scaling the index, choosing its range, and combining it with content are architectural decisions. Dividing indices by the current prefix length changes earlier features whenever the prefix grows; ordinary append-only KV caching relies on keeping those earlier representations unchanged.

For a sine/cosine pair at rate $\omega$, the angle-addition formulas express the features at $i+\delta$ as a linear transformation of those at $i$, using coefficients $\sin(\delta\omega)$ and $\cos(\delta\omega)$. For RoPE, the orthogonal rotations satisfy $R_i^\top R_j=R_{j-i}$, so $\langle R_iq,R_jk\rangle=q^\top R_{j-i}k$. This relative-offset structure coexists with content-dependent queries and keys.

Why can additive positions change a same-gap match? With identity projections, $(u+p_i)\cdot(v+p_j)=u\cdot v+u\cdot p_j+p_i\cdot v+p_i\cdot p_j$. For a sine/cosine pair, the last term depends on the gap, but the two content–position cross terms can still depend on the absolute indices. Learned projections can use these features; there is simply no built-in common-shift identity for the full additive score. Absolute location can itself be useful, for example for a start-of-sequence cue. The shifted sentence example preserves a local distance, not the full prediction: extra source tokens can change softmax weights, and contextual Q/K can change too.

Sources: Vaswani et al., Attention Is All You Need, §3.5; Shaw et al., relative position representations; Su et al., RoFormer / RoPE; Press et al., ALiBi; Chen et al., position interpolation.

We have used these objects several times. Can you now distinguish a key from a value, explain why a query is not the updated representation, and say what our $\vd{\Delta e}$ notation means?

Say your answer out loud before opening the reveal. Each answer includes a small table from the toy model with named columns, so you can check it against the arithmetic. Queries, keys and values are working vectors. The token row passed onward is $\vp{e_i'} = \ve{e_i} + \vd{\Delta e_i}$.

Try explaining attention without looking at the page. You can use a sentence in conversation, draw the sequence of operations on a whiteboard, or write the four matrix equations in an exam. The summaries below describe the same calculation in those three forms.

Play the chain to follow each stage in the diagram. Then switch the sentence in the last card: the highest-probability word changes, but the equations stay the same. Each token reads itself and earlier tokens, mixes their values, and receives the projected message through a residual addition.

PyTorch · one vocabulary distribution from the final updated row.

h = (E_new[-1] @ W_hidden + b_hidden).relu()
logits = h @ W_vocab + b_vocab
probs = logits.softmax(dim=-1)

The TinyStories classroom walkthrough now closes Part III: data setup, four-head attention, training, generation and the measured three-model comparison. The longer MLP/single-head calculation remains available as the illustrated Notebook 5.

The position vectors introduced at the start share the four word-feature coordinates. Moving an earlier word can therefore change its key and value, and the final prediction. Our numerical regression swaps “The” and “river” while keeping the final token fixed: the prediction changes with position vectors, but stays the same when those vectors are zeroed. This is a test of this single-layer example, not a claim that the hand-designed toy understands grammar.