The Math Behind AI

A neural network can only work with numbers. A pixel, a word and a sound all get turned into lists of numbers called vectors. This one has two entries.

Papers write $v \in \mathbb{R}^2$, which means $v$ is a list of two real numbers. Its entries are $v_1 = 3$ and $v_2 = 2$.

The same two numbers are also directions: 3 steps to the right, then 2 up. An arrow from the origin to where you end up shows the vector as a point in space, with a direction.

Every vector is both a list of numbers to compute with and an arrow to think with.

The arrow has a length, written $\lVert v \rVert$. The two entries are the sides of a right triangle, so Pythagoras gives it.

With more entries the rule stays the same: square each entry, add them up, take the square root. This length is called the norm.

Multiplying a vector by a number scales every entry. The arrow keeps its direction and changes length. A negative number also flips it to point the other way.

Adding two vectors adds their entries. In the picture, put $w$ at the tip of $v$. The sum is the arrow from the origin to the end of $w$.

Scaling and adding are all you need for most of what follows. A weighted sum such as $0.7\,a + 0.3\,b$ is the two operations together.

Language models give every token a vector called its embedding. Directions in that space can carry meaning. In a famous result from early word embeddings, the step from “man” to “woman” is roughly the same arrow as the step from “king” to “queen”.

Real embeddings have hundreds or thousands of entries, not two. Nobody can draw them, but length, scaling and addition work exactly the same way.

Take two vectors with the same number of entries. Multiply the entries that line up, then add the products. One number comes out.

Papers also write it as $a^{\top} b$ or $\langle a, b \rangle$. In NumPy it is a @ b.

That number also has a meaning in the picture. It is the product of the two lengths and the cosine of the angle between the arrows.

The lengths set the scale. The cosine measures agreement: 1 when the arrows point the same way, 0 at a right angle, −1 when they point in opposite directions.

From the tip of $b$, drop a line onto the line through $a$, meeting it at a right angle. The stretch from the origin to where it lands, shown in amber, is the projection of $b$ onto $a$. Think of it as the shadow $b$ casts on $a$. It is 1.58 long.

The dot product equals the length of that shadow times the length of $a$.

The dot product measures how far one vector goes in the direction of the other.

Turn $b$ until it makes a right angle with $a$. The shadow disappears and the dot product is 0.

Vectors at right angles are called orthogonal. Neither has any part that points along the other.

Turn further and the shadow falls behind the origin. The dot product becomes negative: -6.12 at 150°.

Now turn $b$ yourself. The value follows $\cos\theta$: largest when the arrows line up, zero at 90°, most negative when they point opposite ways.

The arithmetic does not care how many entries there are. In 4,096 dimensions the dot product still means agreement in direction, even though nobody can draw the angle.

Search engines, recommendation systems and retrieval for chatbots all rank items by the dot product between their embeddings and a query vector $q$.

Here the rose arrow $q$ is a query about puppies. The raw dot product ranks “dog” first, with 2.83 against 1.19, only because its arrow is longer. In real embeddings, length often reflects how common a word is, not what it means.

The fix is to divide every vector by its length first. All arrows then end on the unit circle, and the dot product is exactly the cosine of the angle between them.

Now “puppy” wins with 1.00 against 0.93. This is cosine similarity.

Stack all item embeddings as the rows of a matrix $E$. One product, E @ q, then scores every item at once, and sorting the scores gives the best matches.

Attention in a Transformer uses the same step. Every token has a query, and it scores every token, itself included, by a dot product. The Transformer essay builds that up from here.

An artificial neuron holds a weight vector $w$ and a number $b$ called the bias. Given an input $x$, it computes one score.

Here $w = [1, 1]$ and there are two kinds of points. We want the score to tell them apart.

Ignore the bias for now. Which points have $w \cdot x = 0$? Exactly the ones at a right angle to $w$. Together they form a line through the origin, perpendicular to $w$.

In three dimensions this is a plane. In more, it is called a hyperplane, and the idea is the same.

On the side $w$ points to, the score is positive. On the other side it is negative. The marked point scores 4.5.

The sign says which side of the line a point is on, and the size grows with its distance from the line. This line is in the wrong place, though: most rose points are on the positive side too.

The bias moves the line without turning it. With $b = -3$, the boundary moves 2.12 along $w$, and the sign of the score now separates the two classes.

Training a neuron means finding the $w$ and $b$ that put the line in the right place.

A layer has many neurons, each with its own weights and bias, and so each with its own line. Stack the weight vectors as the rows of a matrix $W$, and the whole layer becomes one operation.

A layer is many dot products computed at once. That is what a matrix does, which is the next chapter.

A matrix is a grid of numbers. Multiplying it by a vector gives a new vector. Read row by row, each output entry is the dot product of one row with the input, exactly like a layer of neurons.

The same product can also be read column by column. $Ax$ takes 1 of the first column and 2 of the second, and adds them.

In the picture, $x$ is 1 step along $e_1 = [1, 0]$ followed by 2 steps along $e_2 = [0, 1]$.

Now apply $A$ to every point of the plane at once. The grid tilts and stretches, but its lines stay straight, parallel and evenly spaced, and the origin does not move. That is what linear means.

The columns of a matrix are where $e_1$ and $e_2$ land.

Every other point follows the same recipe as before: 1 step along the new $e_1$, 2 along the new $e_2$.

The unit square becomes a parallelogram. Its area is the determinant, and every area in the plane is scaled by the same factor.

A determinant of 0 means the plane is squashed onto a line or a point, and some information is lost for good. A negative determinant means the plane is flipped over.

Follow $A$ with a quarter turn $R$. The two steps together are a single matrix, the product $RA$.

Multiplying matrices composes transformations, and $(RA)x = R(Ax)$. The order matters: $AR$ turns first and gives a different result.

Matrices do not have to be square. A matrix with 2 rows and 3 columns has shape (2, 3) and lives in $\mathbb{R}^{2 \times 3}$.

Two matrices can be multiplied when the inner sizes match. Those sizes are summed over and disappear.

Each entry of the result is a row of $A$ dotted with a column of $B$. Entry $(i, j)$ uses row $i$ and column $j$.

Every row meets every column. An $m \times n$ by $n \times p$ product is $m \cdot p$ dot products, each of length $n$.

All of those dot products are independent of each other, so they can run in parallel. GPUs are built to do exactly that.

This is how a layer runs on a batch. Each row of $X$ is an input, each column of $W$is one neuron's weights, and one product scores every input against every neuron.

Papers often write a single input as a column and use $Wx$, with one row of $W$ per neuron, as in the previous chapter. Code puts inputs in rows and writes X @ W, with one column per neuron. It is the same computation with everything transposed, because $(Wx)^{\top} = x^{\top} W^{\top}$.

The transpose $A^{\top}$ swaps rows and columns: $(A^{\top})_{ij} = A_{ji}$. PyTorch's nn.Linear stores its weight with shape (out, in), so it computes x @ W.T + b.

When shapes do not line up, a missing transpose is the first thing to check.

Here is the plane again, with a circle drawn on it. We will push it through a tiny network with two layers and watch what happens to the shape.

The first layer, $W_1 x + b_1$, turns and shrinks the plane, and the bias shifts it. The circle shrinks and moves, and the grid lines stay straight.

Apply the second layer straight away and the grid is still straight. Two linear layers multiply into one:

A hundred linear layers in a row can do no more than a single one. Depth alone adds nothing.

Put a ReLU between them. ReLU keeps positive numbers and replaces negative ones with zero, entry by entry.

Every point with a negative coordinate is pushed onto an axis. The plane folds, and everything outside the top-right quarter is pressed flat against its edges.

The second layer then moves the folded sheet. The result is no longer one straight grid. It is made of flat pieces joined at creases, and each ReLU neuron adds a crease.

This two-layer network is an MLP. Every Transformer block contains one, with thousands of dimensions instead of two and often a smoother cousin of ReLU such as GELU. With enough creases it can bend space into almost any shape.

Matrices move space. Nonlinearities bend it. Training chooses the matrices.

Training means turning a model's weights until a loss, a number that measures its mistakes, is as small as possible. Start with a single weight $x$ and a loss $f(x)$.

We are at $x = 2.5$, where the loss is 1.05. Should $x$ go up or down? From the picture it is obvious. A computer has no picture. It can only evaluate $f$ at points.

So it tries a nearby point. Move $x$ by $h = 1.5$ and compare the two values of $f$. Rise over run gives the slope of the line through both points.

A smaller step, $h = 0.5$, gives a slope of 1.17. The line turns and hugs the curve more closely around $x$.

Let $h$ shrink towards zero. The line ends up touching the curve at a single point. This is the tangent, and its slope is the derivative.

Nobody takes these limits by hand. A few rules, such as $x^n \to n x^{n-1}$, $e^x \to e^x$ and $\ln x \to 1/x$, together with the chain rule from the next chapter, cover nearly everything in deep learning.

A positive slope means the loss goes up to the right, so we step left, against the slope. This is gradient descent.

The step size $\eta$ is the learning rate. With $\eta = 0.75$, each step here halves the distance to the minimum: $2.50,\ 1.75,\ 1.38,\ 1.19,\ 1.09$.

Make the learning rate too large, $\eta = 3.3$, and every step overshoots to the other side and lands further out than it started: $2.50,\ -0.80,\ 3.16,\ -1.59,\ 4.11$. The loss gets worse and worse.

The derivative gives the direction. The learning rate decides how far to trust it.

With two weights the loss becomes a landscape over the plane. Here it is drawn as a contour map: each ellipse joins points with the same loss, like height lines on a hiking map. The minimum sits in the middle.

Freeze $w_2$ and move only along $w_1$. That is an ordinary one-dimensional slope, called the partial derivative $\partial L / \partial w_1$. Then do the same in the other direction.

The curly $\partial$ means: change one input and hold all the others still.

Put the partial derivatives in a vector and you get the gradient, $\nabla L = [-2.5,\ 3.2]$. It points in the direction where the loss rises fastest, which is always at right angles to the contour line through the point.

The gradient is the slope in many dimensions: one number per weight, pointing uphill.

Gradient descent takes a small step against the gradient, computes a new gradient, and repeats. It is the same update as before, now with vectors.

With $\eta = 0.2$ the path curves down into the valley and settles at the minimum.

With $\eta = 0.45$ the path zigzags. The valley is steep across $w_2$ and shallow along $w_1$, so a step that suits one direction overshoots in the other.

Real loss landscapes are full of valleys like this, and better optimizers reduce the zigzag. Momentum keeps a running average of past gradients, so the back-and-forth cancels out. Adam does the same and also gives each weight its own step size.

In PyTorch you never write the gradient yourself. loss.backward() fills w.grad with $\nabla L$, for millions of weights at once. It can do that because of the chain rule.

A neural network is a chain of functions: each layer feeds the next. Take a short chain, $u = x^2$ followed by $y = 3u + 1$.

At $x = 1.5$ we get $u = 2.25$ and $y = 7.75$. How strongly does $y$ react when $x$ changes?

Nudge $x$ by a small amount. Near $x = 1.5$, the square changes about $2x = 3$ times as fast as $x$, so the nudge is stretched 3 times on its way to $u$.

Then $y = 3u + 1$ stretches it 3 times more. Every function in the chain stretches the nudge by its own derivative.

So $y$ moves about 3 × 3 = 9 times as much as $x$. That is the chain rule: the rate of a chain is the product of the rates of its links.

The notation makes it look as though the $du$'s cancel like in a fraction. They are not really fractions, but it is a good way to remember the rule.

Here is one neuron with a sigmoid, scored by a squared error against a target $y$, drawn as a computation graph. Each circle is one simple operation.

The forward pass computes left to right. The values are written in blue above the edges, ending in a loss of 0.072.

Now ask how the loss depends on every value, starting from the end. The last operation is a square, whose derivative is simple.

It is negative: raising $a$ would lower the loss.

Next comes the sigmoid. Its local derivative is $\sigma'(z) = \sigma(z)(1 - \sigma(z)) = 0.20$. By the chain rule, multiply it with what arrived from the right.

An addition passes the gradient on unchanged to both of its inputs, because nudging either input nudges the sum by the same amount. So $\partial L / \partial b = -0.11$.

For a multiplication $p = wx$, the local derivative with respect to each input is the other input. So the gradient for $w$ is scaled by $x$, and the gradient for $x$ by $w$.

Each operation only needs its own local derivative. Multiply them together on the way back.

That is backpropagation. During the forward pass, autograd records the graph. loss.backward() then walks it in reverse and fills in the gradient of every weight. One backward pass costs only a small multiple of a forward pass, typically about two, however many weights there are.

Deep networks multiply many local derivatives together, and that can go wrong. The sigmoid's derivative is never more than 0.25, so a long chain of sigmoids drives gradients towards zero. This vanishing gradient problem is one reason modern networks use ReLU and residual connections.

Here are seven measurements. Networks constantly summarize and rescale numbers like these: their inputs, their activations inside every layer, and their weights when they are first set up.

The mean is the balance point. If every value were a weight on a ruler, the ruler would balance there. To compute it, add the values and divide by how many there are.

Papers write the mean as an expectation, $\mathbb{E}[x]$. Read $\mathbb{E}$as “the average of”.

How far are the values from the mean? Draw each distance as a square. The variance is the average area of those squares.

Squaring makes every distance positive, and it makes values far from the mean count much more than values close to it.

Variance is measured in squared units. Its square root, the standard deviation $\sigma = 2.45$, is the side of the average square, back in the original units. It is the typical distance from the mean.

Many quantities bunch up around their mean in a bell shape. This is the normal distribution $\mathcal{N}(\mu, \sigma^2)$, and the mean and variance are all it needs.

It is everywhere in AI. Weights often start out as random draws from it, and variational autoencoders and diffusion models generate new samples by starting from it.

Subtract the mean and divide by the standard deviation. The data now has mean 0 and standard deviation 1, and its shape is unchanged.

LayerNorm does exactly this to every token vector inside a Transformer, and then applies a learned scale and shift. It keeps numbers in a steady range as they pass through dozens of layers.

A language model ends with one score per word in its vocabulary. These raw scores are called logits. Here are four candidates for the next word in “the cat sat on the”.

Logits can be any number, negative ones included. They are not probabilities yet.

First, exponentiate. $e^z$ is always positive and keeps the order of the scores. It also stretches the gaps: a logit that is 1 higher becomes $e \approx 2.7$ times larger.

Then divide by the total. The results are positive and sum to 1, so they form a probability distribution. “mat” gets 0.68.

Dividing the logits by a temperature $T$ first changes the shape. A low temperature sharpens the distribution towards the top word. A high temperature flattens it towards uniform.

This is the temperature setting of a chatbot. Drag it and watch the probabilities move.

The right next word is “mat”, and the model gave it a probability of 0.68. The loss is the negative log of that probability.

The curve shows this loss for every possible probability. It is 0 when the model gives the right answer probability 1, and grows as that probability shrinks.

A confident mistake is expensive. If the model gives “mat” only 0.05, the loss jumps to 3.00. The curve rises steeply near zero, so wrong answers given with confidence are punished hard.

A confident correct answer costs almost nothing: 0.03. Averaged over every token in the training data, this is the cross-entropy loss.

Read $p_\theta(y \mid x)$as “the probability the model with weights $\theta$ gives to $y$ when it sees $x$”. Minimizing this loss is the same as maximizing the likelihood of the data. Language model papers often report $e^{L}$ instead, called perplexity.

With softmax and cross-entropy combined, the gradient of one example's loss with respect to the logits is remarkably simple. With $y$ the one-hot vector of the right answer:

The right word's logit gets a gradient of -0.32, so a descent step raises it. Every other logit gets its own probability as gradient, so it is lowered in proportion to how likely the model thought it was.

Prediction minus truth. That is the error signal backpropagation carries into the network.