Math for AI from First PrinciplesPost 8 of 10 · 10 of 10 written
- Overview: eight pillars, one model
- 00The language of mathematics
- 01Linear algebra
- 02Calculus and matrix calculus
- 03Probability
- 04Statistics and evaluation
- 05Information theory
- 06Optimization
- 07Numerical and systems math
- 08Capstone: one transformer training step, through every pillar
- 09Learning path, resources and question bank
8. Capstone: one transformer training step, through every pillar
Math for AI · 8 of 103 min readOct 2026
Math for AI · Post 8 of 103 min readOct 2026
The pillars are not separate subjects; a single training step uses all eight. Walk through this sequence until you can explain every line without notes. When you can, you are fluent.
- Text becomes tokens. A tokenizer is a function from strings to sequences in Vⁿ, where V is a finite vocabulary set. Ask whether it is invertible and what information it loses. (Pillar 0)
- Tokens become vectors. Each token ID selects a row of an embedding matrix E with shape |V| × d. That lookup equals multiplying a one-hot vector by E. (Pillar 1)
- Position is added. Rotary embeddings (RoPE) rotate pairs of coordinates by angles that depend on position. A rotation is an orthogonal matrix, so lengths are preserved and dot products encode relative distance. (Pillar 1)
- Activations are normalized. RMSNorm divides each vector by its root-mean-square size. That keeps values in a range where floating point is accurate and gradients are well scaled. (Pillars 3 and 7)
- Attention mixes information across tokens. Three linear maps produce queries, keys and values. QKᵀ scores every pair, dividing by √d_k fixes the variance, and softmax turns scores into weights. (Pillars 1 and 3)
- A causal mask hides the future. Each position may only attend to earlier positions. That is the chain rule of probability built into the architecture: each token is predicted from its past only. (Pillar 3)
- Softmax is computed stably. The maximum score is subtracted before exponentiating, and attention cost grows as n² in sequence length. (Pillar 7)
- An MLP transforms each token. Two linear maps with a nonlinearity between them. Without the nonlinearity, the whole network would collapse into a single matrix. (Pillars 1 and 2)
- Residual connections add the input back. Each block computes x + f(x), whose derivative contains a clean “1”. Gradients reach early layers through that highway. (Pillar 2)
- Logits become a distribution. A final linear map gives one score per vocabulary token, and softmax turns them into a categorical distribution over the next token. (Pillars 1 and 3)
- The loss is cross-entropy. The loss is −log of the probability given to the true next token, averaged over positions. That is maximum likelihood, minimizing KL to the data, and its exponential is perplexity. (Pillars 4 and 5)
- The backward pass computes every gradient. Reverse-mode autodiff applies the chain rule from the loss back to each weight. At the output it starts as p − y and flows back through transposed weight matrices. (Pillar 2)
- Memory is spent. The backward pass needs stored activations, and Adam keeps two extra numbers per parameter. Weights, gradients, optimizer state and activations must all fit across the GPUs. (Pillar 7)
- The optimizer updates the weights. AdamW takes a step scaled per parameter, after gradient clipping, with a learning rate set by warmup and cosine decay. Weight decay is the Gaussian prior from MAP. (Pillars 4 and 6)
- Numbers stay in range. Matrix multiplies run in bf16 while an fp32 master copy of the weights accumulates small updates that bf16 would round away. (Pillar 7)
- The model is evaluated honestly. Held-out loss and perplexity track progress, and benchmark scores come with confidence intervals and paired comparisons. Contamination is checked. (Pillars 4 and 5)
- At inference, tokens are sampled. Temperature and top-p reshape the next-token distribution before sampling. The KV cache grows with every token, and decoding speed is limited by memory bandwidth. (Pillars 3 and 7)
The best way to make this real is to build it. Write a tiny GPT from scratch, first the forward pass in NumPy, then training in PyTorch, and annotate every line with the pillar it uses. Each bug you hit will send you back to exactly the pillar you need.