Architecture

The model is a decoder-only transformer with pre-LayerNorm blocks, written with NumPy and hand-written backward functions. This page shows its structure. The pages it links to hold the equations.

The forward pass

transformer.forward takes one sequence of ids and returns logits and a cache for the backward pass. It first checks the length (transformer.py:156). The longest allowed sequence is 32 tokens, and training and generation use windows of 16.

Each box below is one function call or inline expression, in code order. The name in code font is what the code runs there. The boxes in the frame are the body of block_forward, which forward calls once per layer.

Embeddings

1. Token plus position embeddingsembed_forward(T, d)(16, 32)transformer.py:158

Blocks, repeated num_layers = 2 times transformer.py:162

2. LayerNorm 1layer_norm_forward(T, d)(16, 32)transformer.py:133
3. Attentionattention_forward(T, d)(16, 32)transformer.py:134
4. Residual addX + att(T, d)(16, 32)transformer.py:135
5. LayerNorm 2layer_norm_forward(T, d)(16, 32)transformer.py:136
6. Feed-forwardffn_forward(T, d)(16, 32)transformer.py:137
7. Residual addX1 + ff(T, d)(16, 32)transformer.py:139

Output

8. Final LayerNormlayer_norm_forward(T, d)(16, 32)transformer.py:165
9. Linear headXf @ params["head_w"] + params["head_b"](T, V)(16, V)transformer.py:166
10. Logitsreturn logits(T, V)(16, V)transformer.py:168

Plain names have no page yet. The link under each name opens the real lines at the pinned commit.

Each block reads normalised copies, and each residual adds to the value that was not normalised. T is at most 16. V depends on the corpus, so it stays a symbol.

What is not inside forward

forward stops at the logits (transformer.py:168). It applies no softmax to them and computes no loss.

Step	Where it runs	Line
Cross-entropy, in training	`loss_and_grads`	transformer.py:173
Cross-entropy, in validation	`validation_loss`	train.py:31
Softmax on logits, in generation	`sample_next`	generate.py:11

The only softmax inside forward turns attention scores into attention weights (transformer.py:83). Attention covers it.

Parameters

init_params returns a flat dictionary of arrays. The names below come from calling it. b{l} is the block index, starting at 0, so b0.wq is the query weights of the first block. The "With config" column uses config.py. V depends on the corpus, so it stays a symbol. head_w is a separate matrix (transformer.py:40).

Name	Shape	With config	Scalars	Role
`tok_emb`transformer.py:21	`(V, d)`	`(V, 32)`	`V × d`32V	Token table: one row per vocabulary id.Starts as Xavier uniform. Weight decay.
`pos_emb`transformer.py:22	`(T_max, d)`	`(32, 32)`	`T_max × d`1,024	Position table: one row per position.Starts as Xavier uniform. Weight decay.
`b{l}.ln1_g`transformer.py:29	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 1 scale, before attention.Starts as ones. No decay.
`b{l}.ln1_b`transformer.py:29	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 1 shift, before attention.Starts as zeros. No decay.
`b{l}.wq`transformer.py:32	`(d, d)`	`(32, 32)`	`d × d`1,024 per block	Query projection weights.Starts as Xavier uniform. Weight decay.
`b{l}.bq`transformer.py:33	`(d,)`	`(32,)`	`d`32 per block	Query projection bias.Starts as zeros. No decay.
`b{l}.wk`transformer.py:32	`(d, d)`	`(32, 32)`	`d × d`1,024 per block	Key projection weights.Starts as Xavier uniform. Weight decay.
`b{l}.bk`transformer.py:33	`(d,)`	`(32,)`	`d`32 per block	Key projection bias.Starts as zeros. No decay.
`b{l}.wv`transformer.py:32	`(d, d)`	`(32, 32)`	`d × d`1,024 per block	Value projection weights.Starts as Xavier uniform. Weight decay.
`b{l}.bv`transformer.py:33	`(d,)`	`(32,)`	`d`32 per block	Value projection bias.Starts as zeros. No decay.
`b{l}.wo`transformer.py:32	`(d, d)`	`(32, 32)`	`d × d`1,024 per block	Attention output projection weights.Starts as Xavier uniform. Weight decay.
`b{l}.bo`transformer.py:33	`(d,)`	`(32,)`	`d`32 per block	Attention output projection bias.Starts as zeros. No decay.
`b{l}.ln2_g`transformer.py:35	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 2 scale, before feed-forward.Starts as ones. No decay.
`b{l}.ln2_b`transformer.py:35	`(d,)`	`(32,)`	`d`32 per block	LayerNorm 2 shift, before feed-forward.Starts as zeros. No decay.
`b{l}.ff1_w`transformer.py:36	`(d, F)`	`(32, 128)`	`d × F`4,096 per block	First feed-forward layer weights.Starts as Xavier uniform. Weight decay.
`b{l}.ff1_b`transformer.py:36	`(F,)`	`(128,)`	`F`128 per block	First feed-forward layer bias.Starts as zeros. No decay.
`b{l}.ff2_w`transformer.py:37	`(F, d)`	`(128, 32)`	`d × F`4,096 per block	Second feed-forward layer weights.Starts as Xavier uniform. Weight decay.
`b{l}.ff2_b`transformer.py:37	`(d,)`	`(32,)`	`d`32 per block	Second feed-forward layer bias.Starts as zeros. No decay.
`lnf_g`transformer.py:39	`(d,)`	`(32,)`	`d`32	Final LayerNorm scale.Starts as ones. No decay.
`lnf_b`transformer.py:39	`(d,)`	`(32,)`	`d`32	Final LayerNorm shift.Starts as zeros. No decay.
`head_w`transformer.py:40	`(d, V)`	`(32, V)`	`V × d`32V	Linear head weights. A separate matrix, not tied to tok_emb.Starts as Xavier uniform. Weight decay.
`head_b`transformer.py:40	`(V,)`	`(V,)`	`V`V	Linear head bias.Starts as zeros. No decay.

Symbols: T_max is context_len, d is emb_dim, F is ff_dim, L is num_layers and V is the vocabulary size.

Summing the "Scalars" column gives the total. The first line is collected from the shapes above and the second puts in the values from config.py.

Implemented

N = 2V d + T_{\max} d + 2d + V + L\left(4d^{2} + 2d F + 9d + F\right)

N = 65\,V + 26{,}496

count_params in transformer.py sums the array sizes, but only tests call it. Training does not print this number.

Repository map

Every Python file at the pinned commit. The links open the file in the Code Explorer.

File	Lines	What it holds
checkpoint.py	46	Saves and loads weights, model shape and tokenizer.
config.py	44	Hyperparameters as module-level constants, plus two import-time assertions.
data.py	43	Corpus loading, encoding, the split, sliding windows and batches.
generate.py	35	Sampling and the generation loop.
gradcheck.py	33	Finite-difference gradient checker. Only tests use it.
ops.py	60	Operations with hand-written backward functions.
optim.py	50	Adam with decoupled weight decay, the learning-rate schedule and gradient clipping.
talk.py	30	Interactive inference entry point.
tests/test_checkpoint.py	76	Checkpoint tests.
tests/test_data.py	83	Data pipeline tests.
tests/test_generate.py	82	Sampling tests.
tests/test_gradcheck.py	39	Gradient-check tests.
tests/test_ops.py	88	Operation tests.
tests/test_optim.py	49	Optimizer, schedule and clipping tests.
tests/test_tokenizer.py	50	Tokenizer tests.
tests/test_transformer.py	110	Model tests.
tokenizer.py	119	BPE tokenizer that works on characters: pre-tokens, merges, encode and decode.
train.py	85	Training entry point.
transformer.py	190	Parameters, the forward pass and the backward pass.

data/corpus.csv and checkpoints/ are git-ignored, so they are not part of the commit. A fresh clone needs a corpus before it can train.