Mini-GPT: Assemble a Model Whose Training You Can Explain

Assemble the interfaces

IDs -> token + position vectors -> causal blocks -> final LayerNorm -> bias-free vocabulary projection.

Shapes: (T,) -> (T,D) -> (T,V).

One logit row per next-token prediction.

Trace a tiny model

V=3, D=4, H=2, one block, maximum context 4.

For three IDs: embeddings (3,4), head scores (3,3), final logits (3,3).

Predict its unique parameter count.

Check the count

Token table 12 + position table 16 + block 244 + final norm 8 + output 12 = 292.

Registering the block twice can add a second optimizer update even when forward runs once.

Align the objective

Source [0,1,2,0] gives input [0,1,2], target [1,2,0].

Uniform logits give mean loss log(3) ≈ 1.098612.

Causality and target shift are separate invariants.

Prove the prefix is protected

Compare inputs [0,1,2] and [0,1,0].

Their first two logit rows must match under identical parameters.

Test the entire model; normalization can leak future information even with a correct attention mask.

Run a small training experiment

Fixed tiny stream; fresh graph each step. Zero leaves, backward mean loss, update unique parameters once.

Check finite differences and a decreasing loss trend before evaluating generated text.

Generate beyond the window

With a six-token seed and maximum context 4, use the last four IDs for the next prediction.

Positions reset to 0..3 inside the crop. Sample only the last logit row.

Return the full seed plus all samples.

Diagnose the failing contract

Wrong low loss: check shift and causality.

Plateau: check registration and graph paths. Extreme capped loss: check log-space computation.

Failure after the context fills: crop before position lookup. Explain each invariant before scaling up.