The Math Behind AI, Learned From What You Already Build

Here is one line from jepa.py, the model behind my JEPA code writer:

return z + self.mlp(torch.cat([z, self.act(a)], -1))

It’s the predictor. Given a summary of a program’s current state, it guesses the summary of the state after one more step. It looks like plumbing. But almost every piece of math that matters in AI is hiding inside it.

• z is a vector: a list of 128 numbers that describes the current state.
• self.act(a) looks up another vector, 64 learned numbers that describe the action a.
• torch.cat([...], -1) glues the two lists end to end into one list of 192 numbers. That’s concatenation.
• self.mlp(...) is a small neural network. Inside it, the 192 numbers are multiplied by a grid of learned numbers, a matrix, passed through a simple bend (a non-linear function), and multiplied by another matrix, so that 128 numbers come out. That’s matrix multiplication, the most common operation in all of AI.
• z + ... adds those 128 numbers to the old z. The network only has to learn the change, not the whole new state.

And the line doesn’t explain where those learned numbers come from. During training, the model measures how far its guess was from the truth (a loss), works out which way to nudge every number to shrink that loss (a gradient, computed with calculus), and takes a small step that way (optimization). Underneath sit probability and information theory, which tell you what “how wrong” should even mean.

So one line contains linear algebra, calculus, optimization and, one level down, probability and information theory. You already use this math every time you call a model. This post is about learning to read it.

It’s written for developers and builders who use AI and want the math without the fear. Every symbol gets explained once. Every example is short Python with numpy only, and I ran each one; the outputs shown are what it printed.

The map

There are six areas. Learn them in this order, because each one leans on the one before.

1. Vectors and linear algebra

The idea. A vector is a list of numbers. AI uses vectors to describe things: a word, an image, the state of a program. The number of entries is the vector’s dimension. My JEPA’s z has dimension 128.

Three operations do most of the work.

The dot product multiplies two vectors entry by entry and adds the results:

a · b = a1*b1 + a2*b2 + a3*b3 + ...

It’s large when the two vectors point the same way. Cosine similarity is the dot product divided by both lengths, so only direction counts. It runs from -1 (opposite) through 0 (unrelated) to 1 (same direction):

cos(a, b) = (a · b) / (|a| * |b|)

Here |a| is the vector’s length: the square root of the sum of its squared entries.

Matrix multiplication is many dot products at once. A matrix W with 3 rows and 2 columns takes a vector of 3 numbers and returns 2: each output is the dot product of the input with one column. That’s what a neural network layer is, plus a bend.

import numpy as np
cat    = np.array([0.9, 0.8, 0.1])
kitten = np.array([0.8, 0.9, 0.2])
mat    = np.array([0.1, 0.2, 0.9])
def cosine(a, b):
    return (a @ b) / (np.linalg.norm(a) * np.linalg.norm(b))
print(round(cat @ kitten, 2))        # dot product
print(round(cosine(cat, kitten), 3)) # similar
print(round(cosine(cat, mat), 3))    # not similar
W = np.array([[1.0, 0.0],
              [0.0, 1.0],
              [1.0, -1.0]])          # a tiny layer: 3 numbers in, 2 out
print(cat @ W)

Output:

1.46
0.99
0.303
[1.  0.7]

The @ sign is numpy’s operator for both dot products and matrix multiplication.

The intuition to keep: a matrix is a function that moves vectors. It stretches, rotates and squashes space. A neural network is a stack of these moves with bends in between, and training picks the moves.

Where it shows up. Embeddings are vectors, and semantic search ranks documents by cosine similarity. Attention scores every pair of tokens with dot products. The inside of an LLM is mostly matrix multiplication, which is why GPUs, chips built to multiply matrices fast, run AI. In RIAI, the JEPA planner ranks candidate programs by the distance between a predicted vector and the goal’s vector, and checking progress meant asking whether the predicted vector was closest to the true next state among 100 candidates. It was, 59% of the time at the first checkpoint, against 1% for guessing.

2. Calculus for learning

The idea. You only need one piece of calculus to understand training: the derivative. It tells you how much an output changes when you nudge an input a tiny bit. If the loss goes up when you nudge a weight up, the derivative is positive, so you should move the weight down.

A model has many weights. The list of derivatives, one for each weight, is the gradient. It points in the direction that makes the loss grow fastest, so training steps the opposite way.

The other piece is the chain rule. When one function feeds another, the derivatives multiply. If loss = (w*x - y)^2, then:

d loss / d w = 2 * (w*x - y) * x

The 2 * (w*x - y) is the derivative of the outer square; the x is the derivative of the inner w*x. A neural network is a long chain of functions, and backpropagation is just the chain rule applied from the loss backwards, layer by layer, reusing work as it goes.

You can always check a derivative with arithmetic: nudge the input up and down a tiny amount and see how much the output moved. That’s a numeric gradient check.

import numpy as np
x, y = 2.0, 7.0                     # one example: input 2 should give 7
def loss(w):
    return (w * x - y) ** 2         # squared error
def grad(w):                        # chain rule by hand
    return 2 * (w * x - y) * x
w, h = 1.0, 1e-5
numeric = (loss(w + h) - loss(w - h)) / (2 * h)
print(grad(w), round(numeric, 4))

Output:

-20.0 -20.0

The two agree. The negative sign says: increase w and the loss falls. That’s true, since w = 3.5 is the answer.

The intuition to keep: the gradient is a blame report. For every weight, it says how much that weight contributed to the error, and in which direction.

Where it shows up. Every training run of every neural network. PyTorch’s loss.backward() is backpropagation, done for you. In jepa.py the gradient flows through the predictor and the encoder, but not through the target encoder: “no gradient” there means the chain rule is deliberately cut at that point, so the targets can’t be bent to make the loss easy.

3. Probability

The idea. A probability distribution lists every possible outcome with a number between 0 and 1, and the numbers add up to 1. A language model’s job, at every step, is to output a distribution over its whole vocabulary: how likely each token is to come next.

The network itself outputs raw scores, called logits, which can be any size and don’t add to 1. Softmax turns them into a distribution:

softmax(z)_i = exp(z_i) / (exp(z_1) + exp(z_2) + ...)

exp is the exponential function, e (about 2.718) raised to that power. It makes every score positive and makes big scores much bigger.

Temperature divides the logits before softmax. Below 1, the top choice gets more dominant. Above 1, the distribution flattens. Then the model samples: it picks a token at random, weighted by those probabilities.

import numpy as np
logits = np.array([2.0, 1.0, 0.1])  # scores for "mat", "rug", "moon"
def softmax(z, temperature=1.0):
    z = z / temperature
    e = np.exp(z - z.max())         # subtract max for stability
    return e / e.sum()
for t in [0.5, 1.0, 2.0]:
    print(t, softmax(logits, t).round(3))
rng = np.random.default_rng(0)
print(rng.choice(["mat", "rug", "moon"], size=10, p=softmax(logits)))

Output:

0.5 [0.864 0.117 0.019]
1.0 [0.659 0.242 0.099]
2.0 [0.502 0.304 0.194]
['mat' 'mat' 'mat' 'mat' 'rug' 'moon' 'mat' 'rug' 'mat' 'moon']

Subtracting the maximum doesn’t change the result, but it stops exp from overflowing on large scores. Every real implementation does it.

The intuition to keep: a model doesn’t output an answer, it outputs odds. The text you see is one draw from those odds. Ask again and you can get a different draw, which is why the same prompt gives different replies.

Where it shows up. The temperature setting in every LLM API. Attention uses the same softmax to turn its dot-product scores into weights that add up to 1. Classifiers, rerankers and “confidence scores” are all distributions underneath.

4. Statistics

The idea. Probability says what a known process will do. Statistics works backwards: from data you observed, what can you say about the process?

Two summaries come first. The mean is the average. The variance is the average squared distance from the mean, a measure of spread; its square root is the standard deviation.

Then the idea that matters most for builders: train/test split. You fit a model on one set of examples and measure it on a different set it has never seen. A model that does far better on the training set than on the test set has overfit: it memorised the examples instead of learning the pattern.

The last idea is that a score from a small test is noisy. Suppose a system truly solves 70% of tasks of some kind. Give it 12 tasks. It won’t always score 70%. Each task is like a weighted coin flip, and the count of successes follows a binomial distribution. You can see how wide that is by simulating it:

import numpy as np
rng = np.random.default_rng(0)
skill = 0.7                          # true chance of solving any one task
for n in [12, 167]:
    scores = rng.binomial(n, skill, size=100_000) / n
    low, high = np.percentile(scores, [2.5, 97.5])
    print(n, "tasks:", round(low, 2), "to", round(high, 2))

Output:

12 tasks: 0.42 to 0.92
167 tasks: 0.63 to 0.77

With 12 tasks, the same 70% system scores anywhere from 5 of 12 to 11 of 12 in about 95% of runs. With 167 tasks, the range tightens to 63% to 77%.

The intuition to keep: the error bar shrinks with the square root of the number of tests. To halve the noise, you need four times as many tasks.

Where it shows up. Every benchmark and every “our new prompt is better” claim. In RIAI, the JEPA was tested on 12 exercises (10 of 12 in closed loop) and on 167 held-out random programs (91% in closed loop). The simulation above is why the 167-program number is the steadier one, and why a change from 9 of 12 to 10 of 12 on its own proves little. The VICReg loss in jepa.py is also statistics: it measures the variance of each of the 128 dimensions across a batch, as we’ll see below.

5. Optimization

The idea. Gradient descent is the loop that trains almost every model:

w = w - learning_rate * gradient

Compute the gradient, take a small step against it, repeat. The learning rate is the step size. Too small and training crawls. Too big and every step overshoots the bottom, so the loss bounces or blows up.

Real training adds two refinements, in words:

• Momentum keeps a running average of recent gradients and steps along that, like a ball that keeps rolling. It smooths out noisy steps and rolls through small bumps.
• Adam (Kingma and Ba, 2014) adds momentum and also gives each weight its own step size, scaled down for weights whose gradients are large and jumpy. It’s the default optimizer for most deep learning today.

Both lean on one tool: the exponential moving average (EMA). It’s a running average that forgets slowly:

average = 0.99 * average + 0.01 * new_value

Here’s gradient descent finding the line y = 2x + 1 from four points, then an EMA slowly following the learned weight:

import numpy as np
x = np.array([0., 1., 2., 3.])
y = 2 * x + 1                        # the rule to discover
w, b, lr = 0.0, 0.0, 0.05
for step in range(500):
    err = (w * x + b) - y
    w -= lr * 2 * (err * x).mean()   # gradient of mean squared error
    b -= lr * 2 * err.mean()
print(round(w, 3), round(b, 3))
target = 0.0                         # EMA, like a target encoder
for step in range(100):
    target = 0.99 * target + 0.01 * w
print(round(target, 3))

Output:

2.0 1.0
1.268

Gradient descent recovered the slope 2 and the intercept 1. The EMA, after 100 updates at rate 0.99, has only covered about 63% of the way from 0 to 2. That slowness is the point.

The intuition to keep: training is walking downhill in fog. You can’t see the bottom, only the slope under your feet. The learning rate is your stride.

Where it shows up. Every training loop. In jepa.py, the target encoder is an EMA of the encoder: after every training step its weights move 1% of the way toward the encoder’s (rate 0.99). Because it trails behind, the targets change slowly and can’t chase the predictor, which is one of the two guards against the model collapsing into a single vector.

6. Information theory

The idea. Information theory, started by Claude Shannon in 1948, measures surprise.

The surprise of an outcome with probability p is -log(p). A sure thing (p = 1) has zero surprise; a 1-in-1,000 event has a lot. With base-2 logs, surprise is measured in bits.

Entropy is the average surprise of a distribution: H(p) = -(p1*log p1 + p2*log p2 + ...). A fair coin has 1 bit. A coin that lands heads 90% of the time has less, because it’s more predictable.

Cross-entropy is the surprise when reality follows one distribution but you predicted with another. For a language model with one correct next token, it’s simply -log(probability the model gave to the right token). That is the loss almost every LLM is trained on.

KL divergence is the gap between cross-entropy and entropy: the extra surprise you pay for predicting with the wrong distribution. It’s zero only when the two match:

KL(p || q) = p1*log(p1/q1) + p2*log(p2/q2) + ...

import numpy as np
def entropy(p):
    return -(p * np.log2(p)).sum()
coin  = np.array([0.5, 0.5])
loaded = np.array([0.9, 0.1])
print(round(entropy(coin), 3), round(entropy(loaded), 3))
p = np.array([0.7, 0.2, 0.1])       # what the model predicts
print(round(-np.log(p[0]), 3))       # cross-entropy if token 0 was right
print(round(-np.log(p[2]), 3))       # if token 2 was right
q = np.array([1/3, 1/3, 1/3])
print(round((p * np.log(p / q)).sum(), 3))  # KL(p || q)

Output:

1.0 0.469
0.357
2.303
0.297

The fair coin has 1 bit of entropy, the loaded coin 0.469. If the model gave the right token 70%, the loss is 0.357. If it gave the right token only 10%, the loss is 2.303: being confidently wrong costs a lot. (The last three numbers use natural logs, as training code does, so they’re in units called nats, not bits.)

The intuition to keep: a good predictor is a good compressor. If you can guess the next symbol well, you need fewer bits to write it down. Lowering an LLM’s cross-entropy is, quite literally, teaching it to compress text. A 2023 DeepMind paper, “Language Modeling Is Compression”, built on exactly this link.

Where it shows up. The training loss of every LLM. Perplexity, a common LLM metric, is cross-entropy turned back into a number of choices: exp(loss). KL divergence appears in fine-tuning with reinforcement learning, where a penalty keeps the tuned model from drifting too far from the original, and in distillation, where a small model learns to match a big one’s distribution.

How much math you actually need

It depends on what you build.

For builders, the two most useful pieces are cosine similarity and small-sample statistics. The first explains how retrieval works and fails. The second stops you from shipping a “better” prompt because it won 7 of 10 tests.

What to skip at first: eigenvalues beyond the intuition, formal proofs, integrals (you’ll meet them in probability, but sums carry the same ideas), and any course that spends its first month on notation. Learn each idea by computing it in code, then read the formula. The formula will look like a summary of something you already did.

A 12-week learning path

Two weeks per area, a few hours a week. Each block ends with something small you build yourself, in plain Python and numpy, no frameworks.

The projects chain together. The bigram model from weeks 5-6 gets its gradients from your autograd in weeks 3-4, its training from weeks 9-10, and its score from weeks 11-12. By the end, you’ve built a whole (tiny) language model training stack and know what each line does.

Free resources

All of these are free online. I checked each link and title.

• 3Blue1Brown, Essence of linear algebra (YouTube series). The best visual explanation of what a matrix does to space. Weeks 1-2.
• 3Blue1Brown, Essence of calculus (YouTube series). Derivatives and the chain rule as pictures. Weeks 3-4.
• Andrej Karpathy, “The spelled-out intro to neural networks and backpropagation: building micrograd” (YouTube). Builds backprop from scratch in about two and a half hours. The model for the weeks 3-4 project.
• Marc Peter Deisenroth, A. Aldo Faisal and Cheng Soon Ong, Mathematics for Machine Learning (Cambridge University Press, 2020). The authors keep the PDF free on the book’s site. Linear algebra, calculus, probability and optimization, written for exactly this purpose. Use it as the reference beside the videos.
• Tengyu Ma and Andrew Ng, CS229 Lecture Notes (Stanford). The main notes for Stanford’s machine learning course, free as a PDF. Good for seeing the math applied to real algorithms once weeks 1-10 are done.
• David J. C. MacKay, Information Theory, Inference, and Learning Algorithms (Cambridge University Press, 2003). Free to download from the author’s site. Deep and readable; the first chapters are enough for weeks 11-12.

What it means for RIAI

RIAI is my lab project on AI that learns new skills by writing and checking code. My research map listed the experiments I want to run next. Each one leans on specific math from above, so these are my own learning goals, not results.

Multi-step rollout losses (optimization and calculus). jepa.py trains on single steps but plans up to 4 steps ahead, which is the likely cause of the 1 of 12 open-loop score. A rollout loss feeds the predictor’s own output back in for several steps and adds up the error at each one. To get that right I need to understand how gradients flow back through a chain of repeated calls to the same function, which is the chain rule over time, and why long chains can make gradients shrink or blow up.

The VICReg terms, concretely (statistics and linear algebra). Take a batch of encoded vectors, each with 128 dimensions.

In the VICReg paper’s form:

• The variance term looks at each dimension separately, across the batch, and computes its standard deviation. If it’s below 1, the loss adds 1 - std for that dimension; above 1, nothing. Weight 1.0 in my code. It stops any dimension from going flat.
• The covariance term computes the 128 by 128 covariance matrix of the batch: how each pair of dimensions moves together. It adds up the squares of everything off the diagonal, then scales the sum by the dimension count. Weight 0.04. It pushes dimensions to carry different information, so the 128 numbers aren’t 128 copies of one.

Together they make collapse, every input getting the same vector, expensive. Understanding them properly means understanding covariance matrices, which is weeks 1-2 and 7-8 together.

Cosine versus distance for planning (linear algebra). The planner ranks candidates by distance to the goal’s vector. Plain distance cares about both direction and length; cosine cares only about direction. Which one tracks “this program is closer to correct” better is a question I can test, and I need the geometry to read the answer.

Calibration for an ambiguity flag (probability and statistics). The live test showed the writer confidently returning a program that fit the examples but wasn’t the one meant. A flag that says “another program fits these examples too” needs a score whose confidence means something. A model is calibrated when things it rates 80% likely happen about 80% of the time. Guo and others showed in 2017 that modern neural networks are often overconfident, and that a single learned temperature on the softmax, the same temperature from section 3, fixes much of it. Measuring calibration on a few hundred ambiguous cases is a statistics exercise straight out of weeks 7-8.

Key ideas

• One line of model code holds most of AI’s math: vectors, matrix multiplication, a learned function, and the gradient that trains it.
• Linear algebra is how AI represents things; cosine similarity and matrix multiply are the two operations to know cold.
• Calculus gives the gradient, a blame report for every weight; backprop is the chain rule run backwards.
• A model outputs odds, not answers. Softmax makes them, temperature reshapes them, sampling picks one.
• Small tests are noisy. On 12 tasks, a 70% system scores anywhere from 5 to 11 most of the time.
• Cross-entropy is the LLM loss, and a better predictor is a better compressor.
• Learn by building: twelve weeks, six small projects, one tiny training stack you fully understand.

Sources: 3Blue1Brown, Essence of linear algebra and Essence of calculus; Andrej Karpathy, “The spelled-out intro to neural networks and backpropagation: building micrograd”; Deisenroth, Faisal and Ong, Mathematics for Machine Learning (Cambridge University Press, 2020); Ma and Ng, CS229 Lecture Notes (Stanford); MacKay, Information Theory, Inference, and Learning Algorithms (Cambridge University Press, 2003); Shannon, “A Mathematical Theory of Communication” (Bell System Technical Journal, 1948); Kingma and Ba, “Adam: A Method for Stochastic Optimization” (2014); Bardes, Ponce and LeCun, “VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning” (ICLR 2022); Delétang et al., “Language Modeling Is Compression” (2023); Guo, Pleiss, Sun and Weinberger, “On Calibration of Modern Neural Networks” (ICML 2017). Links checked November 2026. All code outputs were produced by running the snippets with numpy.

Related: How AI Turns Words Into Numbers: Embeddings Explained, How AI Knows What “It” Means: Attention Explained, and What’s a Token, and Why AI Has Limits.