Mini-GPT: Integrate, Check, Then Train

This capstone combines the v1 components into a small, single-head causal language model. It is an integration exercise: generated text alone cannot tell you whether its loss, graph, and optimizer are correct. V2 offers the longer cumulative path with separate tokenization, bigram, MLP, embedding, and multi-head exercises; use that path if this jump feels too large.

You are ready to start when you can differentiate a shared Value, trace one attention row, explain both residual paths, and identify the vocabulary axis. Revisit attention and transformer blocks when any of those dependencies is unclear.

Start with one token transition

A next-token training window needs one more source token than input positions. For source IDs [0,1,2,0], use input [0,1,2] and targets [1,2,0]. Each output row predicts the token immediately after its input position. Do not manufacture a final target by repeating the last input token; that teaches a transition absent from the source.

The capstone receives IDs, not raw text. For a character demonstration, build a sorted character mapping from the training text, save it, and define how unknown characters will be handled. Tokenizer fitting and held-out split construction are not extra required APIs here. For the full data pipeline, go to v2 tokenization.

Follow the model stage by stage

For vocabulary size V=3, width D=4, maximum sequence length M=5, one block, and T=3 input positions:

Stage	Shape	First check
Token IDs	(3,)	Every ID is in [0,3)
Token and position lookup	each (3,4)	Repeated token IDs share trainable rows
Embedding sum	(3,4)	Add position 0,1,2 vectors elementwise
Pre-norm transformer block	(3,4)	Both residuals preserve width
Final LayerNorm	(3,4)	Normalize features of each row
Bias-free vocabulary projection	(3,3)	Each row has one logit per class

The module's num_heads argument is retained for compatibility but v1 uses one head. It does not imply multi-head behavior. Linear stores weights (out_features,in_features) and applies each row to one input vector, unlike the ML matrix exercise's (in_features,out_features) batch convention.

Implement incrementally: embedding identity/gradient checks, per-row LayerNorm and FFN, attention causality, one block, then the final projection. Embedding must return the original parameter scalars; copying their numeric values into fresh leaves detaches learning.

Account for parameters before optimizing

Count unique trainable scalar identities, not paths through attributes. A block may be reachable through both a list and an alias, but an SGD loop must update each scalar once.

For width D, the four biased attention projections contribute 4*(D*D+D). The ReLU FFN has widths D -> 4D -> D, contributing 8D*D+5D. Two LayerNorms contribute 4D. One block therefore has 12D*D+13D parameters. The whole model adds V*D token embeddings, M*D position embeddings, 2D final norm parameters, and D*V output weights.

For V=3, D=4, M=5, and one block, the count is 12+20+244+8+12=296. If parameters() returns 540 entries but only 296 unique identities, a block was traversed twice: its 244 parameters will receive double updates. A falling loss can hide this bug. The supplied capstone Module now enumerates each identity once.

Compute cross-entropy from logits

Stable softmax prevents exponent overflow, but -log(softmax(target)) can still lose the target probability to underflow. For logits [1000,-1000] and target index 1, the mathematical loss is approximately 2000, not a clipped constant such as 23.

Choose a numeric shift m=max(logit.data), then keep all subsequent arithmetic on Value objects:

shifted_j = logit_j - m
row_loss = log(sum(exp(shifted_j))) - shifted_target
loss = mean(row_loss over positions)

Here the shifted logits are [0,-2000], exponential sum is about 1, and row loss is 0-(-2000)=2000. Logit gradients are approximately [1,-1]: the wrong class must decrease and the true class increase. A clipped probability loss can instead conceal both the large error and the useful gradient.

Average over positions once. Duplicating a logits/target batch should preserve its mean loss and total parameter gradient. This logit-space objective is also the bridge from the ML classification exercises' explicitly clipped diagnostic losses to the v2 language-model objective.

Prove causality independently of training

With causal=True, attention normalizes only positions at or before the query; there is no finite-penalty assumption. Reuse the allowed-prefix rule from the attention exercise in this separate capstone implementation.

Compare [0,1,2] with [0,2,1]: first-position logits must match exactly up to numerical tolerance because token 0 and its position are unchanged. Position embeddings do not excuse a discrepancy. A second check differentiates an early output with respect to later inputs and expects zero future gradients. Shape checks alone prove neither property.

Complete a loss and update trace

There are two positions and three classes. All logits start at zero; targets are [1,2]. Predict the mean loss and the six logit gradients before revealing them.

Check the uniform-logit trace

Each row assigns probability 1/3 to every class, so the mean loss is log(3). Before mean reduction, row errors are [1/3,-2/3,1/3] and [1/3,1/3,-2/3]. Divide each by two: [1/6,-1/3,1/6] and [1/6,1/6,-1/3]. The target class has a negative gradient, so subtracting the gradient raises its logit.

Train and evaluate in separate stages

For each independent step, zero parameter gradients, build a fresh forward graph, compute loss, backpropagate once, and then update each unique parameter. Do not reuse an old computation graph after changing parameters. Fit a tiny repeated pattern first to test optimization, then use held-out text to measure generalization. Neither test guarantees convincing language from a tiny scalar model.

Generation copies the seed IDs, repeatedly crops to the last max_seq_len tokens, obtains logits, and samples only the last row after division by a positive temperature. It retains the whole generated output while bounding the model input. Local position indices restart at zero after cropping. Temperature below one sharpens relative probabilities; it is not the same as guaranteed greedy selection. Zero temperature is outside this API.

Independent capstone evidence

Before considering the implementation complete, show: a known loss and gradient; unique parameter accounting; unchanged prefix logits after changing the suffix; and a stub-model generation trace showing which contexts are passed on each step. Then demonstrate a small training-loss decrease and separately report held-out behavior. Keep dimensions tiny so each failed property can be diagnosed. The stronger v2 capstone extends these same checks to learned multi-head attention.