The maths of machine learning

Drag the points (or focus one and use the arrow keys). The line, the red residuals and their squares — whose total area is the SSE — update as you drag. Squared error fits by the normal equations, shown with the four sums Σx, Σy, Σx², Σxy and the solve; absolute error fits by iteratively reweighted least squares, with the squared-error line dashed for comparison. Make an outlier (or drag one point far away) shows squared error chasing it while absolute error barely moves. The points slider goes from 2 to 20; the first five are the book’s study-hours data.

Underneath, gradient descent walks downhill on the MSE from θ = 0 over a contour picture of the loss in (b, w), with the excess loss per step on a log scale and the book’s table of steps 0, 10, 100, 1000 and 2000 next to the normal-equations answer. Raise η until it diverges. SGD runs mini-batch stochastic gradient descent on 1000 noisy points from y = 3x + 2 (batch size, epochs and η are sliders) and shows how it jitters around the exact answer.

The book’s ten students (hours studied, passed or not) and the curve σ(wx + b), with sliders for w and b, the live cross-entropy and its gradient, the decision boundary −b/w, and the predicted pass probabilities. The second panel is the cross-entropy over (w, b) as a contour picture with a dot for the current setting; click to move it. Train runs gradient descent from the dot and draws its path down the convex bowl. Edit the data as hours,passed pairs.

Build a computation graph either from an expression (sin(x*x), exp(2*x)/(1+x*x) — parsed by the site’s algebra parser) or step by step, one node per line: z = mul(w, x). The palette buttons append a node (+, ×, −, ÷, x², sin, exp, tanh, σ, ReLU, ln) fed by the current output. Names that are never defined are inputs; give them values below.

Forward ▶ animates the values flowing through the graph; ◀ Backward starts from 1 at the output and multiplies by the local derivative written on each edge, adding the contributions at nodes that feed several others — the result at each node is ∂output/∂node. Nodes can be dragged. The panels compare the reverse-mode gradient with forward-mode dual numbers (the table lists every node’s a + bε) and with a central finite difference.

Finding the sweet spot plots the error of forward and central differences against the step h on log–log axes for the graph’s output, with the theoretical best h (√ε and ∛ε) marked. Switch to 32-bit floats — every operation rounded to single precision — and the sweet spot moves to much larger h. Use f=sin(x), exp(x), x^3 or sqrt(x) near 0, as in the book’s lab.

Type documents (one per line) and a query. The tool shows the vocabulary with document frequencies and idf = ln(N/df), the TF-IDF vectors as a table, and every document’s dot product, cosine and BM25 score, ranked by the one you choose. Raw counts turns off the idf weighting (the length-bias example). Stemming is a light suffix stripper (plurals, -ing, -ed), enough for “cats” to match “cat”. For BM25, move k₁ (term saturation) and b (length normalisation); the idf is the book’s ln(N/df) by default, or Lucene’s smoothed version.

Enter symbol probabilities (A:1/2,B:1/4,…) or paste text and count its letters. The tool gives the entropy, the Huffman code with each codeword’s length beside its surprise −log₂p, the average length against H, and a slider that builds the Huffman tree one merge at a time. Enter a wrong distribution q to build the code for q and pay for it with p: the cross-entropy H(p, q), D_KL in both directions, and the extra bits per symbol — exactly D_KL(p‖q) when the probabilities are powers of ½.

Scores (logits) and a temperature T: the softmax as bars, its entropy, how the probabilities move as T goes from 0.05 to 20, and what the naive formula does with scores like 1000 (it overflows; subtracting the maximum does not). Below, a tiny attention layer: tokens and the matrices Q, K and V, with the score matrix QKᵀ/√d_k, the row-wise softmax as a heat map, and the output rows as weighted averages of V. Untick √d_k to see the softmax saturate. Tick causal mask to see a decoder’s attention: every score above the diagonal becomes −∞, so a token gives zero weight to the tokens after it.

The sampling T buttons set the temperature to the usual sampling settings (greedy, 0.2, 0.7, 1, 1.5); the readout says what greedy decoding would pick and how often sampling picks it.

Heads sharing keys and values shows multi-head (MHA), grouped-query (GQA) and multi-query (MQA) attention: the query heads in coloured groups, each group reading one K/V head, and what that does to the KV cache — 2 × layers × KV heads × head dimension × bytes per token, for a context length you choose. Llama 3 70B (64 query heads, 8 KV heads, head dimension 128, 80 layers, FP16) keeps 327,680 bytes = 0.3125 MiB per token, 8× less than without sharing.

Related: Neural network designer (build and train whole networks), Matrix calculator (XᵀX θ = Xᵀy and other systems), Statistics (regression on your data, distributions, tests), Eigenvectors at work (PCA, SVD, LSI), and AI hardware calculator (the KV cache against GPU memory).

Companion to Mathematics, Reconnected: How computers compute derivatives (chapter 17), Vectors (chapter 26, TF-IDF and cosine similarity), Statistics and Inference (chapter 31, simple linear regression), and Horizons: Optimisation, Machine Learning, and Information (chapter 33): linear regression, logistic regression, backpropagation, search, information theory and softmax and attention.

Companion to Part IX of AI Acceleration: From Transistors to Trillion-Parameter Models: what self-attention actually computes (the attention equation, GQA and MQA) and decode (a softmax over the vocabulary, then sample the next token).

A link opens the tool with the given settings: /tools/ml-maths/?<parameters>#<tab>, where the tab is regression, logistic, autodiff, search, info or softmax. Share link (top right) builds one from whatever is on screen. Points are x,y pairs separated by semicolons; documents are separated by |. A ` typed straight into a link reaches the page as a space: in `f=` a space between two terms is turned back into `, so ?f=x^3+2*x works; elsewhere write %2B, or leave the + out.

Presets for example=:

From AI Acceleration (Part IX):

For example, https://geekychris.github.io/tools/ml-maths/?example=ch17-backprop#autodiff opens the book’s backpropagation example, and https://geekychris.github.io/tools/ml-maths/?f=x^3+2*x&at=x=2#autodiff builds the graph of x³ + 2x at x = 2.