Module 6 · suggested week 6

The Mathematics Inside Modern AI

Decode the equations behind the models.

~3 focused hoursBuilds on modules 1, 4 (advisory)

0 of 9 sections read

What you will be able to do

Read the core equations of modern ML (vectors, dot products, matrices, gradients, loss, gradient descent, softmax and scaled dot-product attention) symbol by symbol, and know where each is an idealisation of real training and inference systems.

\theta

Theta

\eta

Eta

\nabla

Gradient operator (nabla, del)

\partial

Partial derivative

\mathcal{L}

Loss function (calligraphic L)

\mathbf{a}\cdot\mathbf{b}

Dot product (or multiplication)

x^{T}

Transpose

\|\mathbf{x}\|

Norm (length of a vector)

\sum

Summation (capital sigma)

e^{x}

Exponential

\log

Logarithm

\arg\min

Argmin (the input that minimises)

\frac{d}{dx}

Derivative

Vectors and embeddings: meaning as a list of numbers

A vector in ML is first of all a list of numbers: $\mathbf{x} = (x_1, x_2, \ldots, x_d)$ . Bold lower-case letters usually mean vectors; $d$ is the dimension (how many numbers). A token embedding is a vector of a few thousand numbers that the model learned to assign to that token.

The geometric view is what makes vectors useful for thought: each list is a point (or arrow) in $d$ -dimensional space, and directions in that space can carry meaning. Interpretability research finds that many concepts correspond, approximately, to directions.

Empirical regularity

In trained language models, many human-interpretable features can be found as approximately linear directions in activation space (for example via sparse autoencoders). This is well supported but not universal: some features appear nonlinear or entangled, and "how many features" depends on the method.

Source: Templeton et al. 2024, Scaling Monosemanticity

Dot products: multiply, then add

Recognition

A dot product is a Σ loop

The dot product $\mathbf{a}\cdot\mathbf{b}$ multiplies matching entries and adds them up. It is the summation loop from Module 1 with two arrays. Geometrically it measures how aligned two vectors are: big and positive when they point the same way, zero when perpendicular, negative when opposed.

\mathbf{a}\cdot\mathbf{b} = \sum_{i=1}^{d} a_i b_i = \|\mathbf{a}\|\,\|\mathbf{b}\|\cos\phi

let dot = 0;
for (let i = 0; i < a.length; i++) {
  dot += a[i] * b[i];
}

Vectors & dot products

Drag either arrow tip, or use the sliders.

\mathbf{a}\cdot\mathbf{b} = (3)(1.5) + (1)(3) = 7.5

a₁3a₂1b₁1.5b₂3

a · b7.50

‖a‖, ‖b‖3.16, 3.35

cos θ (cosine similarity)0.707

angle θ45° · roughly aligned

The dot product $\mathbf{a}\cdot\mathbf{b} = \sum_i a_i b_i$ is a Σ loop: multiply matching entries, then add. It equals $\|\mathbf{a}\|\|\mathbf{b}\|\cos\theta$ , so it is large when the vectors point the same way, zero when they are perpendicular, and negative when they oppose. Embeddings have hundreds or thousands of entries instead of two, but the loop is identical. That is why attention scores and “cosine similarity” are both dot products.

Loading interactive…

Drag the arrow tips. Watch the dot product go to zero when the vectors are perpendicular and negative when they oppose. Cosine similarity divides out the lengths so only direction matters, which is why retrieval systems use it to compare embeddings.

Three places you have already met a dot product without the name: a weighted average (weights · values), a neuron's pre-activation (weights · inputs, plus bias), and an attention score (query · key). The notation $\mathbf{q}^{\top}\mathbf{k}$ (a transposed column vector times a column vector) is the same dot product written in matrix form. See transpose and dot.

Matrices: many dot products at once

A matrix $W$ is a grid of numbers. Multiplying a matrix by a vector, $W\mathbf{x}$ , computes one dot product per row of $W$ : each row is a "detector" and the output lists how strongly $\mathbf{x}$ matches each detector. A neural-network layer is then:

A dense layerEstablished

\mathbf{h} = \sigma\!\left(W\mathbf{x} + \mathbf{b}\right)

Take input $\mathbf{x}$ , compute many weighted sums at once ( $W\mathbf{x}$ ), shift each by a bias $\mathbf{b}$ , then apply a nonlinear squashing function $\sigma$ elementwise. Here $\sigma$ is an activation function, not a standard deviation; Greek letters are reused.

import numpy as np
h = np.maximum(0, W @ x + b)   # ReLU as the nonlinearity

Does not say: Modern transformers use variants (GELU or SwiGLU activations, layer normalisation, residual connections) but the multiply-add-squash skeleton is the same.

Derivatives and gradients: which way is downhill?

A derivative $\frac{dL}{d\theta}$ is a rate: if you nudge $\theta$ a little, how much does $L$ change, per unit of nudge? A partial derivative $\frac{\partial L}{\partial \theta_j}$ asks the same thing about one knob $\theta_j$ while holding all the others fixed. The gradient $\nabla_\theta L$ collects every partial derivative into one vector, one entry per parameter:

The gradientEstablished

\nabla_\theta \mathcal{L} = \left(\frac{\partial \mathcal{L}}{\partial \theta_1}, \frac{\partial \mathcal{L}}{\partial \theta_2}, \ldots, \frac{\partial \mathcal{L}}{\partial \theta_n}\right)

A list of sensitivities: for every parameter, how much the loss would rise if that parameter increased slightly. It points in the direction of steepest increase, so its negative points steepest downhill.

Established

Backpropagation is the chain rule of calculus applied efficiently through the layers of a network to compute the gradient. It computes which direction to step; gradient descent is the rule for taking the step.

Loss functions: what "better" means

Before you can descend, you need a landscape. The loss $\mathcal{L}(\theta)$ is a single number that measures how wrong the model is with parameters $\theta$ . For language models the standard choice is cross-entropy: the average of $-\log p_\theta(\text{actual next token})$ over training text. It is small when the model assigns high probability to what actually came next. (Module 4 builds this from Shannon entropy; see the cross-entropy inspector.)

Next-token cross-entropyEstablished

\mathcal{L}(\theta) = -\frac{1}{N}\sum_{t=1}^{N} \log p_\theta(x_t \mid x_{<t})

Average, over $N$ positions in the text, of the "surprise" ( $-\log p$ ) the model felt at the token that actually occurred, given everything before it.

let loss = 0;
for (let t = 0; t < N; t++) {
  loss += -Math.log(model.prob(tokens[t], tokens.slice(0, t)));
}
loss /= N;

Does not say: This is the pre-training objective. Post-training (instruction tuning, RLHF, RLAIF) optimises different objectives on top of it.

Inspect symbol by symbol →

Gradient descent: the update rule, symbol by symbol

Gradient descentMathematical model

\theta_{t+1} = \theta_t - \eta\,\nabla_\theta \mathcal{L}(\theta_t)

New parameters = old parameters minus a small step (size $\eta$ ) in the direction the loss increases fastest. Subtracting means going downhill.

for (let t = 0; t < steps; t++) {
  const g = gradient(loss, theta);       // ∇θ L(θt)
  for (let j = 0; j < theta.length; j++) {
    theta[j] = theta[j] - eta * g[j];    // θ ← θ − η g
  }
}

Does not say: The textbook form. Real training uses minibatch estimates of the gradient and adaptive optimisers such as Adam, as discussed below.

Inspect symbol by symbol →

Every symbol

	Symbol	Read it as
$\theta$	All the model’s parameters (weights and biases) as one long vector	`theta` — a Float32Array with billions of entries
$t$ (subscript)	Which step of training we are on	Loop counter
$\eta$ (eta)	Learning rate: step size	A hyperparameter, e.g. `3e-4`
$\nabla_\theta$	"Gradient with respect to θ": all partial derivatives	`autograd` / `loss.backward()`
$\mathcal{L}$	The loss: one number scoring how wrong the model is	A function returning a scalar
$-$	Step against the gradient, i.e. downhill	The sign that makes it descent, not ascent

Gradient descent landscape

Learning rate

\eta

0.10Momentum0.00

step t0

θ(-2.60, 1.40)

L(θ)6.3870

‖∇L‖8.409

Each step applies $\theta_{t+1} = \theta_t - \eta\,\nabla_\theta L(\theta_t)$ : compute the slope at the current point and move a distance $\eta$ times that slope, downhill. In the narrow valley a large $\eta$ overshoots across the steep walls and zig-zags or diverges, while a small one crawls along the floor. Momentum (a running average of past steps) is one of the fixes used by real optimisers such as Adam. A real model has billions of θ dimensions rather than two, and the gradient is estimated from minibatches, so this is an idealisation.

Loading interactive…

Change the learning rate. Too small and progress crawls; well chosen and the point slides into the valley; too large and it overshoots, zig-zags across the valley, or diverges. Each step uses only the local slope.

Where the textbook equation stops describing real training

Minibatches: the gradient is estimated from a small random batch of examples, so each step is noisy (stochastic gradient descent).
Adaptive optimisers: frontier models typically use Adam-family optimisers that keep running averages of past gradients and rescale each parameter’s step. The update is not literally $\eta\nabla\mathcal{L}$ .
Schedules: $\eta$ changes over training (warm-up, then decay).
Precision and parallelism: mixed-precision arithmetic and computation sharded across thousands of chips change numerical details.
Multiple stages: pre-training on next-token loss is followed by fine-tuning and reinforcement learning from human or AI feedback (RLHF, RLAIF), with different objectives.
Proprietary details: exact recipes for frontier models are mostly unpublished. Treat any confident claim about "how GPT-x or Claude was trained" that goes beyond published papers as speculation.

Softmax: scores into probabilities

Softmax with temperatureEstablished

\operatorname{softmax}(z)_i = \frac{e^{z_i / T}}{\sum_{j} e^{z_j / T}}

Exponentiate each raw score (logit) $z_i$ so all are positive, then divide by the total so they sum to 1. Temperature $T$ sharpens ( $T<1$ ) or flattens ( $T>1$ ) the distribution.

const exps = logits.map(z => Math.exp(z / T));
const Z = exps.reduce((a, b) => a + b, 0);   // the Σ in the denominator
const probs = exps.map(e => e / Z);

Inspect symbol by symbol →

Token probabilities

The transformer was introduced in a paper titled “Attention Is All You▁

Need

95.9%z=9.1

Want

1.9%z=5.2

Get

0.9%z=4.4

Have

0.6%z=4.1

Know

0.3%z=3.2

Require

0.2%z=2.9

need

0.1%z=2.4

Deserve

0.1%z=1.6

Temperature

T

1.00

Entropy of next-token distribution0.33 bits

Loss if true token is “Need”0.042 nats

T → 0 behaves likeargmax (greedy)

A model outputs one score (logit) $z_i$ per vocabulary token. Softmax, $p_i = e^{z_i/T} / \sum_j e^{z_j/T}$ , turns the scores into probabilities that sum to 1. Temperature $T$ rescales the scores before the exponent: low T sharpens the distribution, high T flattens it. These eight logits were written by hand for teaching, and a real vocabulary has tens of thousands of tokens. The training loss on this step would be $-\ln p(\text{“ Need”})$ .

Loading interactive…

Adjust temperature and sample. Low temperature nearly always picks the top token; high temperature spreads probability onto unlikely tokens. The logits here are hand-set for illustration, not taken from a real model.

Go deeper: The same formula in physics

Replace $z_i/T$ with $-E_i/(k_B T)$ and softmax becomes the Boltzmann distribution of statistical mechanics: the probability of a state with energy $E_i$ at temperature $T$ , with the denominator called the partition function $Z$ . This is a genuine Level 2 connection: the same mathematical object. ML borrowed the word "temperature" from physics for that reason. It is not a claim that token sampling involves thermal energy. See the Boltzmann inspector.

Attention: the equation behind transformers

Scaled dot-product attention (Vaswani et al. 2017)Established

\operatorname{Attention}(Q,K,V) = \operatorname{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}\right) V

Each token asks a question (query), every token offers a label (key) and a payload (value). Score every query against every key with dot products, scale, turn each row of scores into weights with softmax, and output a weighted average of the values.

scores  = Q @ K.T / np.sqrt(d_k)       # every query · every key
weights = softmax(scores, axis=-1)     # each row sums to 1
out     = weights @ V                  # weighted average of values

Inspect symbol by symbol →

Unpacking attention

	Piece	Shape (one head, n tokens)
$Q, K, V$	$n \times d_k$ , $n \times d_k$ , $n \times d_v$	Three learned linear projections of the same token vectors
$QK^{\top}$	$n \times n$	All pairwise dot products: how relevant is token j to token i
$\div\sqrt{d_k}$	$n \times n$	Keeps scores from growing with dimension, so softmax doesn’t saturate
softmax (rows)	$n \times n$	Each row becomes a probability distribution: the attention pattern
$\cdot V$	$n \times d_v$	Each output is a weighted mix of value vectors

Scaled dot-product attention

Choose a query token:

Thecatsatonthematbecauseitwastired

Key dimension

d_k

(illustrative)64

“it” attends most to“cat” · 56%

Raw score (top)2.80

Row entropy (spread)2.07 bits

$\mathrm{softmax}\!\left(\tfrac{QK^{T}}{\sqrt{d_k}}\right)V$ . Each row of the grid is one query token. Its scores are dot products with every key, which softmax turns into weights summing to 1, and the output is that weighted mix of value vectors. These 4-dimensional vectors were written by hand so that “it” looks for an animate noun. Trained models learn their own features, and they are rarely this tidy. Turn off the $\sqrt{d_k}$ scaling and raise $d_k$ : the softmax saturates to one-hot, which is the problem the scaling was introduced to prevent.

Loading interactive…

A toy sentence with hand-chosen 2D query and key vectors. Select a token to see its row of scores become a softmax distribution. Real models learn Q, K and V projections in hundreds of dimensions, run many heads in parallel, and apply a causal mask so tokens cannot attend to the future.

Training versus inference

Training changes $\theta$ : forward pass, loss, backward pass, update, billions of times. Inference freezes $\theta$ and only runs forward passes: compute logits, softmax, sample a token, append it, repeat.

Production inference adds engineering that the equations omit. A KV cache stores past keys and values so each new token does not recompute attention over the whole prefix. Batching, quantisation, speculative decoding and sampling settings (temperature, top-p) all change cost or output without changing the learned parameters.

Established

At inference, a standard transformer’s parameters do not change. Apparent "learning" within a conversation (in-context learning) happens through activations computed from the prompt, not through weight updates.

To see these equations in their original habitat, open the guided reading of Attention Is All You Need.

Connections, with their levels

PhysicsLevel 2 · Shared math

Softmax and the Boltzmann distribution

Softmax with temperature is mathematically the Boltzmann distribution with logits playing the role of negative energies.

Why this level: Identical functional form and a shared vocabulary ("temperature", "partition function"). There is no shared physical mechanism: no heat bath in a GPU’s sampler.

PhysicsLevel 2 · Shared math

Gradient descent and energy minimisation

A ball in a viscous fluid rolling down a potential obeys an equation closely resembling continuous-time gradient descent (gradient flow).

Why this level: Gradient flow is the shared mathematics. Physical systems do not compute gradients; they respond to forces. Momentum and noise make the analogy richer, and also more approximate.

BiologyLevel 1 · Analogy

Artificial vs biological neurons

The weighted-sum-plus-threshold neuron was inspired by biological neurons, which integrate synaptic inputs and fire.

Why this level: Biological neurons spike in time, have dendritic computation, neuromodulation and plasticity rules unlike backpropagation. The historical inspiration is real; the equivalence is not.

Skeptic’s corner: Schematic equations vs. real systems

The equations in this module are correct as mathematics and accurate as descriptions of the core of modern ML. They are not complete descriptions of how any specific frontier model was trained or served, and much of that is unpublished.

Is the explanation using textbook SGD where the system actually uses Adam, schedules and minibatches, and does that difference matter for the claim?
Does a claim about "what the model attends to" rest on attention maps alone, or on causal interventions?
Is a claim about a proprietary model’s training sourced to a paper or system card, or inferred?
Do the matrix shapes in the explanation actually work out?

Watch & explore

Video10 min3Blue1Brown

Vectors (Essence of linear algebra, ch. 1)

Three views of a vector (arrow, list of numbers, abstract object) and why ML mostly uses the second while reasoning with the first.

Watch for: The "list of numbers" view. An embedding is exactly that, just with thousands of entries.