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

The capstone combines the components into a decoder-only language model. Its accomplishment is not convincing generated prose: it is a complete path from token IDs through a causal forward pass to a finite loss, correct gradients, an SGD update, and autoregressive generation.

Use tiny dimensions and short sequences. Scalar Value graphs are intentionally inspectable, but much slower and more memory-intensive than tensor kernels. Implement the supplied mini-gpt interface in stages rather than debugging all components at once.

Write the shape contract first

For vocabulary V, embedding width D, maximum context M, and current sequence length T:

Stage	Shape	Shared across positions?
Input token IDs	(T,)	Not parameters
Token lookup	(T,D) from (V,D) table	Yes
Position lookup	(T,D) from (M,D) table	Rows selected by position
Elementwise embedding sum	(T,D)	—
Each causal transformer block	(T,D)	Yes
Final LayerNorm	(T,D)	Yes
Bias-free vocabulary projection	(T,V)	Yes

Require 1 <= T <= M for training and generation calls. Positions are 0..T-1 within this input window. The final LayerNorm is additional to each block's two normalizations. The vocabulary projection is Linear(D,V,bias=False); its weights are independent of the input embedding table in this exercise.

Use the full-input Q/K/V projections described in the transformer lesson. Register the list of blocks once; traversal must return each parameter identity once even if a child also has an alias.

Work a complete shape and parameter trace

Choose V=3, D=4, H=2, one block, and M=4. For IDs [0,1,2], T=3:

token vectors + position vectors:  (3,4)
head projections:                 (3,2) per head
attention matrices:               (3,3) per head
block result and final norm:      (3,4)
vocabulary logits:                (3,3)

Count unique trainable scalars before training:

token table       3*4 = 12
position table    4*4 = 16
one block             = 244
final norm        2*4 = 8
output projection 3*4 = 12
total                 = 292

The block count is 12D*D + 13D with this pack's biased Q/K/V/output projections and biased FFN. Seeing 536 parameters for this model would suggest the 244-parameter block has been registered twice. That would make a simple optimizer update block parameters twice per step.

The objective is one prediction per position

Take a source stream [0,1,2,0]. Training input is [0,1,2], targets are [1,2,0]. Logit row 0 predicts token 1 after seeing input token 0. Row 1 predicts token 2 after seeing [0,1]. Row 2 predicts token 0 after seeing [0,1,2].

A single forward call computes all three predictions, but each row is constrained by the causal mask. If row 0 could see input token 1, it could read its own target and make the training objective misleadingly easy.

For uniform logits with V=3, each loss is log(3) ≈ 1.098612 and their mean is the same. With B equal-length sequences, average over B*T target positions. The exercise helper handles one sequence; a batch loop may average its per-sequence losses. Unequal-length sequences require weighting by the number of positions if the intended objective is a per-token mean.

Use shifted log-sum-exp directly for each target. Logits [1000,-1000] with target 1 must produce a finite loss near 2000, not a capped value from clamping a probability. Return the mean as a Value so every layer remains connected.

A minimal training experiment

Fix the random seed, a tiny model, and a small repeated token stream. Sample a valid shifted window, then:

logits = model(x)
loss = cross_entropy_loss(logits, y)
for p in model.parameters():
    p.grad = 0.0
loss.backward()
for p in model.parameters():
    p.data -= learning_rate * p.grad

Rebuild the graph on each iteration. Update only after all gradients for the step are ready. Record the pre-update loss and inspect several gradients from different parts of the model. A single zero gradient can be expected from a masked path, unused embedding row, or inactive ReLU; a whole disconnected parameter group needs investigation.

Before a longer run, finite-difference a selected output weight and compare it to autograd on a fixed example. Check that unique parameter identities match the count above. Then verify that a suitably small learning rate reduces loss on the tiny stream over several steps. One sample of generated text is not a correctness test.

Prove the model cannot look ahead

Run [0,1,2] and [0,1,0] through identical parameters. The first two logit rows must match because their visible prefixes are identical. The third row is allowed to differ. This exercises masking in every block and catches accidental normalization across positions.

This test is different from a target-shift check. A correctly causal model trained against unshifted targets still learns the wrong task. Check both independently.

Generate with a bounded context

Start with a nonempty seed and copy it. For each new token:

1. Crop a model input to the most recent M IDs.
2. Recompute logits with local position IDs beginning at zero.
3. Take only the last row and divide by a finite positive temperature.
4. Apply stable softmax, sample one ID, and append it to the copied sequence.

Return the full original seed plus all generated IDs. Cropping limits what the model sees; it does not delete earlier IDs from the returned sequence. With M=4 and a six-token seed, the first prediction uses the last four seed tokens, then shifts the window for the next prediction.

This implementation recomputes every prefix and has no key/value cache or gradient-disabled execution mode. Generation does not call backward or update weights. Cropped windows reset learned absolute positions, so this is the explicit inference convention, not unlimited context memory.

Diagnose the first failing invariant

Symptom	Targeted check
Excellent training loss, poor next-token behavior	Inspect target shift and run prefix causality test
Loss plateaus near log(V)	Verify parameter registration and graph connectivity
Finite probabilities but wrong extreme loss	Compute cross-entropy in log space
Training changes when only aliases change	Count unique parameter identities
Failure once generation exceeds M	Crop context before position lookup

Checkpoint. Should generation sample from every returned logit row?

Check your reasoning

no. Earlier rows predict tokens already supplied in the context. Only the last row predicts the next unseen token.

Completing the capstone means you can explain the contracts and validate their composition. Larger datasets and faster tensor implementations are subsequent engineering steps; they do not replace these invariants.