Trace the Complete Tiny Decoder

Run an inspectable two-block decoder with real two-head causal attention, residual MLP updates, final normalization, tied logits, probabilities, and loss. Assertions verify shapes, 368 parameters, and future-token invariance.

The complete forward pass should expose its intermediate records before an optimizer changes any parameter. The implementation below uses the frozen vocabulary, embeddings, positions, dimensions, and tied readout. Its small projection values are chosen for auditing rather than model quality.

Run a two-block tied decoder

Inspect token records, two causal blocks, final records, logits, probabilities, loss, parameter count, and a future-token intervention.

What This Trace Proves

The assertions verify declared shapes, vocabulary width, probability normalization, parameter accounting, and one causal intervention. They do not prove that every weight matrix has the intended semantics or that training will succeed. More local checks remain useful inside attention, norms, and MLPs.

The two blocks use different MLP scales so their computations are visibly distinct. A trainable implementation would store separate projection and norm parameters for each block even if two arrays happened to share initial values.

Inspect Intermediate Values Before the Loss

If the loss is unexpected, locate the first changed contract in this order:

IDs and positions;
first block;
second block;
final normalization;
tied logits;
target alignment;
stable cross-entropy.

Q1. Explain the causality intervention

Why does the code compare only logits rows 0, 1, and 2 after changing input position 3?

Choose the expected dependence

Rows 0–2 cannot depend on the later token, but row 3 may depend on its own changed input.

All four rows must remain unchanged because the model is causal.

No row can be compared because changing one token invalidates the run.

Select one choice, then check.

Hint

Distinguish earlier query rows from the changed query row.

Solution

Rows 0–2 must be unchanged because position 3 is in their future. Row 3 may change because the inclusive causal mask lets it read its own input position.

Carry the Exact Model Forward

Chapter 7 will replace fixed illustrative weights with trainable parameters, but it will retain the vocabulary, dimensions, masks, target shift, tied readout, and parameter ledger unless a change is explicitly recorded.