Part 2 used fixed weights to compute a contextual update. How does the model learn those weights?
Where Part 2 left us
Part 2 used one causal attention head and a 4 → 8 → 20 ReLU prediction MLP. The bank row gained water-related information in the river sentence and finance-related information in the cheque sentence. Those named axes and hand-chosen parameters carry into this lesson unchanged. We will compare the model’s prediction with the observed word water, calculate the loss, and use its gradients to update the parameters.
We ended with an attention update added to the input row, followed by a prediction MLP.
The same predictor: 4 inputs → 8 hidden activations → 20 scores
$$h=\operatorname{ReLU}(\vp{e'_{10}}W_1+b_1),\qquad \ell=hW_2+b_2,\qquad p=\operatorname{softmax}(\ell)$$Position 10 is the final the. Its updated row predicts position 11. We now learn these parameters from an observed continuation.
Notation carried into Part 3
The token bank starts from the same row in both sentences. Its contextual update depends on the rows it can read.
We have seen what the fixed parameters do. Learning gives us a way to find useful values for them.
To learn from a prediction, we need to know which token actually came next.
The observed next token tells us the error
The river prefix ends with the word the. The model assigns a probability to every vocabulary token before we reveal the observed continuation: water. Cross-entropy takes the negative logarithm of the probability assigned to water. If that probability is high, the loss is small; if it is low, the loss is large. The other logits still matter because they share the softmax denominator. You can inspect that denominator in the worksheet, then follow the water row into the loss calculation. Both use the same forward pass as Part 2.
From fixed examples to learning
Use the observed next token to improve the prediction.
Keep the attention head and prediction MLP from Part 2.
Predict → measure loss → compute gradients → update parameters.
Observed target:
Water is the continuation in this training example. Other continuations could also be grammatical.
The same final-position prediction
Full softmax worksheet for the final position
Loss worksheet
The model keeps its original prediction. We use the observed token as the target for scoring that prediction.
Some numbers are stored for the next example. Others are temporary results of this forward pass.
Which numbers will the optimizer change?
The token and position tables, projection matrices, and vocabulary head are parameters. Training changes their values and saves them. The rows $\ve{e}$, $\vq{q}$, $\vk{k}$, $\vv{v}$, scores, attention weights, messages, updates, logits, and probabilities are intermediates: we compute them again for each new sequence. Autograd keeps enough of these temporary values to calculate gradients, and the optimizer uses the gradients to update the parameters. These are two different kinds of update. $\vd{\Delta e}$ changes a representation during the forward pass; an optimizer step changes the stored weights used by future forward passes.
Stored parameters survive the example
Parameter and gradient shapes
Output projection and vocabulary head
A contextual update belongs to one input; an optimizer update changes the model used by future inputs.
The computation branches at E and joins again at the value mixture and residual sum.
The forward pass is a graph
The learned token table $E_{\text{tok}}$ and position table $P$ give us the current row stack $\ve{E}$. We project that same stack three ways: $Q=EW_Q$, $K=EW_K$, and $V=EW_V$. Query-key scores determine the attention weights $A$. The value branch supplies the vectors that $A$ mixes. The message matrix $H=AV$ becomes $\Delta E=HW_O$, then the residual sum gives $E'=E+\Delta E$. The vocabulary head maps the final row to logits, which we compare with the observed target to calculate the loss. This is still the simplified one-head model from Part 2. We will add the full pre-norm Transformer later.
Autograd follows these dependencies backward using the chain rule. Where paths meet, their gradient contributions add. The input $E$ receives a contribution through the direct residual path and others through attention. Both routes start from the input $E$, not the final $E'$. Gradients also reach every projection and the lookup rows used by this example. We do not optimize the target token or the fixed causal mask; they are data.
The same $E$ feeds all three projections and the residual path.
Each stored parameter $\theta$ receives a same-shaped gradient $\partial L/\partial\theta$.
optimizer.zero_grad() loss.backward() optimizer.step()
Backward computes slopes. The optimizer updates parameters.
Gradient sizes from the river example
Autograd differentiates the calculation that produced the loss. The optimizer then changes stored parameters, not the cached queries or attention weights.
The gradient points toward increasing loss. SGD steps in the opposite direction.
One update
The gradient gives a local slope for every parameter. Stochastic gradient descent multiplies each slope by the learning rate $\eta$, then subtracts it from the parameter. Try the three presets below. Each uses the same gradient from the river example and runs a fresh forward pass after the update. Compare the target probability and loss before and after the step. Larger steps help this one example more, but may not help other sentences or later updates. The final contextual row also changes because we updated all the contributing parameters together.
Choose the learning rate
The contextual row after the same update
The lower loss tells us that the update fit this example better. We still need other examples to test generalisation.
One causal forward pass predicts at every position. The mask keeps each prediction from reading its answer.
Train every position at once
Append the observed continuation water, then pair each input position with the token that follows it. The first position predicts fisherman, the next predicts sat, and the final position predicts water. Row $i$ can read only positions at or before $i$. We can still compute all rows in parallel: the causal mask blocks forbidden score cells before softmax. We then average the position losses, run one backward pass, and update the shared parameters once. That step may help some positions more than others. The optimizer follows their mean loss.
logits = model(tokens, causal_mask=True)
targets = tokens_shifted_left
loss = cross_entropy(logits, targets, reduction="mean")
optimizer.zero_grad()
loss.backward()
optimizer.step()All position probabilities and losses
Computing the rows in parallel saves work. Each row still predicts from its own prefix because the mask blocks later positions.
One head gives us one attention pattern. A token may need to read several kinds of relation at once.
Multiple heads
Take the sentence "Priya opened the notebook; she smiled." The row for she may need several pieces of context. One attention pattern could link she to Priya, another could retrieve the earlier event, and a third could use nearby structure. Each head has its own query, key, and value projections, and every projection reads the full input representation. We do not give each head a separate slice of the embedding. Each head returns a smaller projected message; we concatenate those messages by columns, then combine them with $W_O$. The rows below illustrate possible attention patterns. They were not learned by the river toy.
We have trained the one-head toy
What does a full block add?
Multiple heads can retrieve different information for the same token.
A feed-forward network processes each row; residual connections and normalization let us stack blocks.
The next examples illustrate these components. They are not the result of training the toy into a full Transformer.
Three possible needs for one query
Project, run each head, then combine the outputs
Illustrative head rows Illustrative
Head dimensions
The same causal mask $M$ blocks future sources in every head.
Each head returns one message row per input token.
We divide the width of the projected outputs, not the input. Every head reads the full input row through its own matrices.
Attention gathers information from other token rows. The feed-forward network then works on the features inside each row.
Attention mixes rows; the FFN mixes features
For one receiving token, attention combines value rows from several source positions. The feed-forward network (FFN) works within that token's row, combining its features. We apply the same FFN weights independently to every position. In the example below, a four-number representation becomes eight hidden numbers. ReLU replaces negative hidden values with zero, and a second projection returns to four numbers so we can add the result to the residual stream. We chose one-decimal weights to keep the arithmetic readable.
Across token rows
Attention combines $\vv{v_j}$ from positions $j\le i$ using $\va{\alpha_{ij}}$.
Within one row
The FFN transforms the coordinates of one token with shared weights.
The prediction MLP was 4 → 8 → 20, ending in vocabulary scores. This block FFN is 4 → 8 → 4, ending in a residual update.
Separate weights: W₁ (4 × 8), then W₂ (8 × 4), with a bias at each layer
hidden = torch.relu(row @ W1 + b1)
ffn_output = hidden @ W2 + b2First projection and all eight ReLU coordinates
Second projection
The FFN uses the same weights at every position, but processes each row separately. Attention is what carries information between positions.
Residual paths carry the rows forward. LayerNorm controls each row's scale before attention or the FFN reads it.
The residual stream and LayerNorm
The residual stream is the model-width row carried through the stack. Attention and the FFN read it, compute updates, and add those updates back. LayerNorm works within one row: subtract the row mean, then divide by $\sqrt{\operatorname{var}+\epsilon}$. The small $\epsilon$ keeps the denominator nonzero even when every value in the row is the same. Our pre-norm decoder applies LayerNorm before attention and before the FFN. Move the sliders to see the mean, centered values, variance, and normalized values change. The calculator uses $\epsilon=10^{-5}$ and fixes the scale and shift at $\gamma=1$, $\beta=0$. A trainable LayerNorm normally learns its scale and shift. See the PyTorch LayerNorm definition.
Residual stream with two branches
Live LayerNorm calculator
$\epsilon=10^{-5}$ prevents division by zero. Here $\gamma=1$ and $\beta=0$.
Choose the arrangement
LayerNorm prepares the row a branch will read. The residual stream still carries the model-width row forward.
Combine pre-norm attention and a pre-norm FFN to make one block. Then stack blocks.
One block, then many
One pre-norm decoder block adds two updates to the residual stream. LayerNorm first prepares the current rows for causal multi-head attention. Attention gathers information across positions, and we add its result back. A second LayerNorm prepares the resulting rows for the FFN, which mixes features at each position and returns the second update. Later blocks can use representations that earlier blocks have already changed. Within one block, however, every query, key, and value comes from that block's input rows. Other tokens cannot read a newly updated bank row until the next block. The multi-hop example below shows what this extra depth allows.
Representations across depth
They walked by the river bank and watched the ___
Block two reads the rows produced by block one. Its attention can therefore use relations that block one has already combined.
We can now connect embeddings, causal blocks, and the vocabulary head in one language model.
The complete decoder-only model
Our pre-norm decoder-only Transformer starts with token embeddings plus positions. Causal blocks update those rows, and a final LayerNorm prepares them for a linear vocabulary readout. Each block preserves model width. Training uses the distributions at all valid positions; generation chooses the next token from the last position’s distribution. This arrangement follows the GPT-2 implementation, which ties its linear readout to the token embedding table and has no separate output bias. The schematic permits an optional bias. Its nonlinear FFNs sit inside the blocks. This is a different architecture from our numerical one-head toy, whose final predictor remains a 4 → 8 → 20 ReLU MLP. The worksheet below runs that toy, not the full stack.
Tokens to next-token probabilities
The hidden layers are inside the blocks. This full-model diagram uses a linear final readout. Our Part 2 toy uses a separate prediction MLP.
Input and output shapes
Part 2 toy: one attention head and a 4 → 8 → 20 prediction MLP. No LayerNorm, block FFN, or stacked blocks.
The toy’s complete prediction MLP
Read the first product, add b₁, apply ReLU, then multiply by W₂ and add b₂. The source calls these parameters W_hidden, b_hidden, W_vocab, and b_vocab.
Our pre-norm decoder keeps one model-width row per token through causal blocks, applies final LayerNorm, then maps each row to vocabulary logits.
The model is larger, but its next-token cross-entropy is the same one we used for the character MLP.
Training: one pass, many guesses
One pass through the decoder produces a vocabulary distribution at every position. Shift the observed sequence by one position to get the targets. Because the causal mask limits each row to its prefix, we can use every row for next-token training. We average those losses into one scalar for the backward pass. The Part 1 aabid character model learned the same statistical task. What differs is how the models build their context representation: the MLP concatenates a fixed window, while the Transformer uses causal attention and stacked blocks over the available prefix.
Parallel next-token objective
The full Transformer uses the same next-token objective as our toy. Its forward calculation is larger; the loss rule is unchanged.
Each probability is for the observed next token, using only the prefix through position $i$. Average over the $T$ valid target positions to get one loss L.
aabid character MLP
Fixed window $\to$ concatenated embeddings $\to$ hidden layer $\to$ logits.
$$L=-\log p(\text{next character})$$
Decoder-only Transformer
Available prefix $\to$ causal blocks $\to$ final row $\to$ logits.
$$L=-\log p(\text{next token})$$
The next-token objective is unchanged; the Transformer builds its context representation with causal attention and depth.
Training scores every position at once. Generation chooses one token, appends it, and runs the forward pass again.
Generation: one token, then again
Read the probability distribution at the last position of the prefix. A decoding rule chooses one token; we append it and run the model for the new last position. Repeat until the stopping rule is met. Earlier keys and values cannot change when we append a later token under causal attention, so we can reuse them. A KV cache stores those projections and adds one new key and value per layer at each step. This saves projection work, but the new query still has to score all the cached keys. The cost display compares the size of the full score matrix with the growing cache.
From training to using the model
The model now has to choose its own next token.
Read the final row’s distribution, choose a token, append it, and predict again.
Parameters stay fixed. The prefix and its activations change; a cache lets us reuse earlier keys and values.
All cached rows at the selected length
Score cells and cached coordinates
Generation reuses cached keys and values, but each new query still compares with every cached key.
The toy has twenty position vectors, and this demonstration uses positions 1–11. It never fills in a missing position with a zero vector. Real models also have a supported context range; we need to decide what to do when that range is reached.
Encoders and decoders use the same attention calculation, with different sources and masks.
Three ways to arrange blocks
Transformer families mainly differ in their blocks and attention masks. In an encoder-only model, every token can read the full input, which suits tasks where the complete sequence is available. A decoder-only model uses causal self-attention to generate one token after another. An encoder-decoder model first builds representations of the source, then lets decoder queries read encoder keys and values through cross-attention. We still compare queries with keys, apply softmax, and mix values. The arithmetic is unchanged; the projected rows now come from different sequences. The 2017 Transformer paired an encoder stack with a decoder stack.
Encoder-only
Bidirectional self-attention over a complete input.
Decoder-only
Causal self-attention over the generated prefix.
Encoder-decoder
An encoder builds source rows; a causal decoder also reads them.
Reuse the arithmetic, change the sources
These are the bank-query and river-key numbers from Part 2, reused to review the dot product. They are not outputs of a trained encoder-decoder. Part 4 supplies a cross-attention example.
Cross-attention changes the source of keys and values, while keeping the query-key score and value mixture unchanged.
Self-attention connects allowed tokens directly. What does that gain, and what does it cost?
Sequence models side by side
A fixed-window MLP sees only the tokens in its window. An RNN can carry information over longer sequences, but distant tokens communicate through many recurrent updates. Temporal convolutions reach more positions as we add layers. Self-attention connects any allowed pair directly within one layer and computes positions in parallel during training. It also needs a score for every allowed query-key pair, so the number of score cells grows quadratically with sequence length. These are trade-offs, not a ranking. Memory, latency, data, and the task all affect which model we choose.
Fixed windows and recurrence
The recurring attention path
Self-attention connects tokens according to their content. Its direct paths come with a score count that grows as sequence length squared.
We know the training and generation paths. Now count their work, then examine the limits.
Context cost and attention limits
The cost of a longer context
Concatenation keeps slots. Averaging compresses them. Attention chooses what to read.
How much work does each approach do?
| Input | $n$ tokens, each with $d=d_{\rm model}$ numbers |
|---|---|
| Attention | $d_k$ numbers per query/key; $d_v$ per value |
| Prediction MLP | $h$ hidden units, then $C$ vocabulary scores |
| Concatenation | $w$ token slots per prediction; $w$ is fixed for a trained MLP |
The vocabulary size $C$ is unrelated to the value matrix $\vv{V}$. We count one sequence. A batch of $B$ sequences costs roughly $B$ times as much arithmetic.
One output cell uses $k$ multiply-accumulates (MACs).
There are $ab$ output cells. Total: $akb$ MACs.
Example: $[4\times3][3\times2]$ has 8 outputs, each using 3 products: 24 MACs.
Three ordered rows → 12 separate inputs. Earlier words are outside the window. Circles are scalar units; 4 of 20 outputs are drawn.
| Operation | Shapes | Work |
|---|---|---|
| Join $w$ embeddings | $w\times d$ into $1\times wd$ | Up to $O(wd)$ data movement |
| Hidden layer | $[1\times wd][wd\times h]$ | $wdh$ MACs |
| Vocabulary layer | $[1\times h][h\times C]$ | $hC$ MACs |
Add biases and ReLU between the two products. Softmax reads the $C$ logits.
For one prediction: $O(wdh+hC)$. Doubling $w$ doubles the hidden-layer work and its weights.
Ten rows → 4 coordinate-wise means → the hidden-layer predictor. Every token has equal weight in the summary.
| Operation | Work |
|---|---|
| Sum and divide $n$ rows | $O(nd)$ additions/scaling |
| Hidden layer, then vocabulary layer | $dh+hC$ MACs |
For one prediction from scratch: $O(nd+dh+hC)$. The MLP input stays $d$ numbers wide.
Query-dependent mixture → 2-number message → $W_O$ → 4-number update. Add $e_{10}$ before the MLP.
Ten context tokens → ten rows of $E$, each with four input numbers.
Queries and keys have the same width (3) for matching. Values carry 2 numbers to mix. One square = one scalar entry.
Full matrix view: compute all ten rows; the prediction head below reads only the final row.
$E'=E+\Delta E$ has shape $10\times4$. Take $e'_{10}$ (4 numbers) → 8 hidden units → 20 scores → softmax.
| Product | Shapes | MACs |
|---|---|---|
| $\vq{Q}=\ve{E}W_Q$ | $[n\times d][d\times d_k]$ | $ndd_k$ |
| $\vk{K}=\ve{E}W_K$ | $[n\times d][d\times d_k]$ | $ndd_k$ |
| $\vv{V}=\ve{E}W_V$ | $[n\times d][d\times d_v]$ | $ndd_v$ |
These costs grow linearly with $n$. The matrices $W_Q,W_K,W_V$ keep the same shapes.
Rows: receiving queries. Columns: source keys. One cell: one match score.
4 queries × 4 keys
16 scores
8 queries × 8 keys
64 scores
Twice as many queries, each comparing with twice as many keys: four times as many scores.
$n^2$ output cells, each with $d_k$ products: $n^2d_k$ MACs.
Scale by $\sqrt{d_k}$, add the causal mask, then apply row softmax: $O(n^2)$ scalar work.
$S$ stores scores. $\va{A}$ stores weights. Both have $n\times n$ entries in the straightforward implementation.
4 tokens: $1+2+3+4=10$
8 tokens: $1+\cdots+8=36$
Allowed pairs: $\dfrac{n(n+1)}2$. For large $n$, doubling $n$ gives approximately four times as many allowed pairs.
| Product | Shapes | MACs |
|---|---|---|
| $H=\va{A}\vv{V}$ | $[n\times n][n\times d_v]$ | $n^2d_v$ |
| $\vd{\Delta E}=HW_O$ | $[n\times d_v][d_v\times d]$ | $nd_vd$ |
Residual addition $\vp{E'}=\ve{E}+\vd{\Delta E}$ uses $nd$ additions.
The two pairwise products are $QK^\top$ and $AV$. The output projection is linear in $n$.
| Product | Shapes | MACs |
|---|---|---|
| Hidden activations | $[n\times d][d\times h]$ | $ndh$ |
| Vocabulary logits | $[n\times h][h\times C]$ | $nhC$ |
Bias, ReLU and vocabulary softmax add $O(nh+nC)$ work.
A large vocabulary can make this head expensive. “Quadratic attention” describes the pairwise part, not every operation in the model.
Fixed widths: $d=4,\ d_k=3,\ d_v=2,\ h=8,\ C=20$. All $n$ rows. Dense MAC counts.
| Operation group | MACs | If $n$ doubles |
|---|
These are calculated operation counts, not timings. GPU parallelism, memory traffic and kernel choices affect elapsed time.
Training example: The fisherman sat beside …
| Input row | Predict | Allowed context |
|---|---|---|
| The | fisherman | The |
| fisherman | sat | The fisherman |
| sat | beside | The fisherman sat |
There are $n$ prediction targets. We compare the work needed to produce all $n$ predictions, not one MLP prediction against a whole attention matrix.
| Method | All $n$ forward predictions |
|---|---|
| Concatenation, fixed $w$ | $O(nwdh+nhC)$ |
| Cumulative prefix average | $O(nd+ndh+nhC)$ |
With fixed $w$ and fixed widths, both are linear in $n$.
If we choose $w=n$ for a full-length concatenation model, its hidden product costs $O(n^2dh)$ across all targets. Its first matrix also grows to $nd\times h$.
$u_i$ is the sum through position $i$. Add one $d$-number row for each new position.
sums = E.cumsum(dim=0) # [n, d]
counts = torch.arange(1, n + 1, device=E.device)[:, None]
means = sums / counts # [n, d]All prefix summaries take $O(nd)$ work. Choosing an algorithm matters even before choosing a model.
Attention projections and mixing:
Add the prediction MLP: $ndh+nhC$.
When $d_k,d_v$ are proportional to $d$:
The $n^2d$ term compares and mixes token pairs. Position addition and embedding lookup add linear work.
| Method | All-position forward work |
|---|---|
| Concatenation | $O(nwdh+nhC)$ |
| Cumulative average | $O(nd+ndh+nhC)$ |
| Attention | $O(nd^2+n^2d+ndh+nhC)$ |
With fixed widths and a fixed window $w$: concatenation and averaging grow linearly; full attention includes quadratic token-pair work.
Lower cost does not mean equal information: a fixed window misses distant tokens; a plain mean loses order and query-dependent selection.
For $Z=XW$, let $G=\partial L/\partial Z$ be the gradient arriving at $Z$.
| Forward | $Z=XW$ |
|---|---|
| Gradient to input | $\partial L/\partial X=GW^\top$ |
| Gradient to weights | $\partial L/\partial W=X^\top G$ |
Each product has the same MAC count. Backward changes the multiplier, not linear versus quadratic growth in $n$.
Held-out scoring uses the forward pass only. Training also computes gradients and updates parameters.
All rows: $n^2$ scores
Final row only: $n$ scores
For our one-layer model, a cold prompt can use $n$ keys/values and only its final query. Its unused query rows need not be computed.
Full-sequence forward / usual stacked-model prefill: build contextual rows across the prompt, including the pairwise attention work.
Part II, one layer, final output only:
Read left to right: all K/V projections, one query, one score-and-mix row, output projection, prediction MLP.
With fixed widths, this special final-row calculation is linear in prompt length. Our full-matrix demonstration deliberately computes more rows.
| Already stored | New work at position $n$ |
|---|---|
| $k_1,\ldots,k_{n-1}$ $v_1,\ldots,v_{n-1}$ | Compute $q_n,k_n,v_n$ from the new row. |
| Earlier rows stay unchanged | Compare $q_n$ with $n$ keys. Mix $n$ values. |
New projections: $2dd_k+dd_v$. New score/mix row: $n(d_k+d_v)$. Then $d_vd+dh+hC$ for projection and prediction.
One cached step has linear context work. The cache stores $O(n(d_k+d_v))$ numbers per head per layer.
| Method and saved state | Work per new prediction |
|---|---|
| Concatenation: last $w$ rows | $O(wdh+hC)$ |
| Mean: running sum + count | $O(d+dh+hC)$ |
| Attention: KV cache | $O(d^2+nd+dh+hC)$ |
In the attention row, $d_k,d_v$ scale with $d$. Without saved sums or K/V, rebuilding past information adds work.
For $m$ next-token predictions after a prompt of length $n_0$, the query rows visit $mn_0+m(m-1)/2$ source positions in total.
| Pairwise arithmetic | $O(n^2d)$ when computing all attention rows |
|---|---|
| Stored attention matrix | $O(n^2)$ numbers in the straightforward implementation |
| Learned projections | Independent of $n$; weights are shared across tokens |
FlashAttention avoids storing the whole attention matrix. Dense pairwise arithmetic is still quadratic.
A learned absolute-position table has its own configured capacity. Longer input does not automatically create new trained rows.
Scope of the comparison
These counts use conventional dense matrix multiplication, one sequence and one attention head. Bias additions, nonlinearities, lookup/data movement and optimizer state have separate costs. Training on all positions, scoring all positions, and generating from one final row are different workloads. Concatenation fixes a window width; a running mean discards slot assignments; attention retains source rows and uses a query-dependent mixture. None of these arithmetic counts proves that one method is faster on a particular device or more accurate on a task.
Sources: Attention Is All You Need (attention architecture and pairwise cost); FlashAttention (exact attention without materializing the full score matrix). The baseline and predictor counts follow directly from the shapes above.
Each receiving token uses attention to combine value rows from allowed positions. Learned projections compute the weights from the current representations, so the pattern can adapt to the input and shorten paths between distant tokens. An attention map alone does not guarantee correct facts, sound reasoning, a complete explanation, or unlimited memory. One head may handle several relations, and several heads may share a relation. The causal mask is needed during training as well as generation. Context length is finite, and computing the scores has a cost. Use the checks below to distinguish what the calculation shows from what would need more evidence.
What it gives
Learned routing that depends on the current representations and respects the mask.
A direct within-layer path between allowed token positions.
What it does not ensure
Correct facts, reliable reasoning, complete explanations, or memory beyond the available context.
Test the claims
Attention weights tell us how the model mixed value information. We need other evidence to explain its broader behavior.
Compare the Transformer with our character MLP. How has the way we represent context changed?
Back to aabid
Both models embed discrete symbols, produce logits, apply softmax, and learn from next-token cross-entropy. The main difference is between the embeddings and the output head. Our aabid MLP concatenates a fixed number of previous embedding rows. We choose its context size in advance, and each window slot has separate input weights. The Transformer keeps one row per token and uses causal self-attention to read from any allowed position in the prefix. It then repeats attention and FFN updates through a stack. We have changed how the model builds context, while keeping the next-token learning objective.
aabid fixed-window MLP
Decoder-only Transformer
Both models optimize next-token loss; the Transformer replaces a fixed concatenated window with causal contextual rows and depth.
Eight questions to check what each part of the model does.
Pause and think
Answer each question before opening its panel. Then find an equation or diagram that supports your answer. Check that you can distinguish parameters from activations, full input rows from projected head outputs, and attention across positions from the FFN within a row. Explain why the causal mask is needed in both training and generation. The last question returns to the layer example: if a representation moves farther from its initial embedding, all we know is that the vector changed. Distance alone cannot tell us whether context caused the change or whether the result is more useful. To isolate the effect of context, change the surrounding tokens while keeping the model and receiving token fixed.
For each answer, explain what the model does, where the information flows, and what the calculation does or does not establish.
Summarize one training step, one block, and the generation loop.
Three summaries
We trained the Part 2 toy by comparing its prediction with an observed token, computing gradients, and taking an SGD step. We then built a larger architecture: pre-norm attention and an FFN each add a residual update inside a block. A stack of those blocks, followed by final LayerNorm and a vocabulary readout, produces next-token distributions. Generation appends a chosen token and repeats with fixed parameters, reusing cached keys and values.
Learning
Score the observed target, compute gradients, and update stored parameters.
One block
Attention mixes positions; the FFN mixes features. Each normalized branch adds a residual update.
The model
Causal blocks → final LayerNorm → vocabulary scores. Choose a token and repeat with fixed parameters.
The familiar attention branch
The attention calculation we worked through in Part 2 is one branch in this trainable residual stack. The whole model predicts and generates tokens.