Milestone 2 of 7

Assemble and verify the decoder

Build the pre-norm causal decoder and verify every shape, parameter, attention row, and causality invariant.

Construct the complete forward path without training it. Every tensor and parameter should have one declared role before gradients are introduced.

Milestone goal

Implement a two-block pre-norm decoder in model.py. Produce a shape trace, a parameter ledger that sums to 368, a causal-mask intervention, and a decoded input-target example. Refuse training if any check fails.

Assemble the dependency chain

For integer token IDs $I:(B,T)$ :

I\to E[I]+P[:T]\to \underbrace{\text{Block}_1\to\text{Block}_2}_{\text{pre-norm}} \to\operatorname{LN}_{final}\to\text{tied vocabulary logits}.

Each block applies:

r'=r+\operatorname{MHA}(\operatorname{LN}_1(r)), \qquad r^{next}=r'+\operatorname{MLP}(\operatorname{LN}_2(r')).

Use two heads, no attention projection biases, a ReLU MLP with biases, no dropout, and a readout whose weight object is the token-embedding weight. A same-valued copy is not weight tying; verify object identity.

Keep a shape ledger

For the reference batch:

Value	Shape
token IDs and targets	$(B,T)$
token plus position records	$(B,T,4)$
attention weights when requested	$(B,2,T,T)$
each block output	$(B,T,4)$
vocabulary logits	$(B,T,8)$

Assert ranks and axes at their creation sites. Do not identify a head axis only because one dimension happens to have size two.

Reconstruct all 368 parameters

Component	Scalars
token embedding $8\times4$	32
learned positions $4\times4$	16
one block: two LayerNorms	16
one block: bias-free Q/K/V and output projections	64
one block: biased $4\to8\to4$ MLP	76
two complete blocks	$2(156)=312$
final LayerNorm	8
tied readout	0 additional
total	368

Compare this hand ledger with sum(parameter.numel() for parameter in model.parameters()). If the total is 400 rather than 368, the readout is probably untied. Other mismatches often come from projection biases or a missing normalization bias.

Verify the causal contract by intervention

For an input such as [<bos>, A, B, C], save all logits. Change only the final token to D and run again. Logit rows 0, 1, and 2 must remain unchanged within the declared float32 tolerance; row 3 may change.

This test is stronger than printing a triangular mask. It checks the behavior of the assembled computation. Also verify that every finite attention row sums to one and that no row assigns probability to a future key.

Verify loss alignment

Given sequences of length five, inputs are the first four IDs and targets are the last four. Flatten logits from $(B,T,V)$ to $(BT,V)$ and targets from $(B,T)$ to $(BT)$ in the same row-major order. Decode at least one aligned pair:

input:  <bos> A B C
target: A     B C <eos>

A valid loss scalar does not prove the alignment is correct.

Question. Recover the logits shape

For $B=3,T=4,V=8$ , what is the decoder logits shape?

Your answer

Answer it first, then check.

Hint

The final axis is vocabulary.

Solution

The shape is

(3,4,8)

Diagnose before continuing

Failure	First check
parameter count is too large	tied readout and projection biases
earlier logits change after a future-token edit	causal mask orientation and API convention
loss is near zero before training	target leakage or incorrect shift
attention contains NaN	rows with every key masked or invalid dtype
identical positions receive identical records	position addition and slicing

Acceptance gate

Continue only when:

the forward pass accepts at least two batch sizes and lengths 1 through 4;
all shape assertions and the 368-parameter ledger agree;
readout and embedding weights are the same parameter object;
the future-token intervention preserves all earlier logits within tolerance;
loss flattening preserves row and position alignment;
an overlength input is rejected instead of silently clipping positions.

Deliverable: model.py, the parameter and shape ledgers, causal-mask results, row-sum checks, one aligned decoded example, and the exact command that reruns them.