Math for AI from First PrinciplesPost 3 of 10 · 10 of 10 written

3. Probability

Math for AI · 3 of 106 min readOct 2026

A language model never outputs a word; it outputs a probability distribution over every possible next token. Probability is the mathematics of what a model actually produces.

Why it exists

The world is uncertain and data is noisy. The same input can have many valid outputs, and we rarely know anything for sure. Probability is the logic of uncertainty: it tells you how to represent, combine and update beliefs consistently.

Core ideas

1. Sample spaces, events and the three axioms

The sample space Ω is the set of all possible outcomes. An event is a subset of Ω, which is why Pillar 0’s set operations matter here. A probability P assigns each event a number following three rules.

  1. P(A) ≥ 0 for every event.
  2. P(Ω) = 1: something happens.
  3. For events that cannot happen together, probabilities add.

Everything else, including Bayes’ rule, follows from these three lines. There are two readings: frequency (long-run proportion) and belief (degree of confidence). Both use identical mathematics.

2. Random variables: neither random nor variables

A random variable is a function from outcomes to numbers: the result of a die roll, the next token’s index. The randomness lives in which outcome occurs; the function itself is fixed.

A discrete variable has a probability mass function: P(X = x) for each value. A continuous one has a density p(x), and probability is the area under it. A density can exceed 1; only its total area must equal 1.

3. Joint, marginal and conditional

The joint distribution P(A, B) covers both at once. Summing out one variable gives the marginal of the other. Conditioning restricts the universe to the cases where B happened, then renormalizes:

P(A∣B)=P(A,B)P(B)P(A \mid B) = \frac{P(A, B)}{P(B)}

Rearranged, this is the product rule: P(A, B) = P(A | B) P(B). Whenever you see a probability, train yourself to ask “given what?”

4. The chain rule of probability: language modeling in one line

Apply the product rule repeatedly and any joint distribution over a sequence factors into one conditional per step:

P(x1,…,xT)=∏t=1TP(xt∣x1,…,xt−1)P(x_1, \dots, x_T) = \prod_{t=1}^{T} P(x_t \mid x_1, \dots, x_{t-1})

This is not an approximation; it holds for every distribution. An autoregressive language model is simply a neural network that learns each factor: the next token given everything before it. In practice you sum log probabilities instead of multiplying, to avoid numbers too small for a computer.

5. Independence and the i.i.d. assumption

A and B are independent if P(A, B) = P(A) P(B): knowing one tells you nothing about the other. Conditional independence means they become independent once you know a third variable.

Almost all of machine learning assumes training examples are i.i.d.: independent and identically distributed, drawn from the same distribution as test data. When that fails, as with distribution shift, guarantees quietly disappear.

6. Bayes’ rule: updating beliefs with evidence

Bayes’ rule flips a conditional. It says your updated belief is proportional to how well a hypothesis explains the evidence, times how plausible it was beforehand:

P(H∣E)=P(E∣H) P(H)P(E)posterior∝likelihood×priorP(H \mid E) = \frac{P(E \mid H)\, P(H)}{P(E)} \quad\text{posterior} \propto \text{likelihood} \times \text{prior}

A worked example shows why intuition fails. A condition affects 1% of people, and a test catches 99% of cases with a 1% false-positive rate. A positive result means only a 50% chance of having it, because true and false positives are equally common.

The same trap hits AI systems. A detector that is “99% accurate” on a rare event, like AI-written text or fraud, can produce mostly false alarms. Always ask for the base rate.

7. Expectation, variance, covariance

The expectation 𝔼[X] is the probability-weighted average: the centre of mass of the distribution. Its most useful property is linearity: 𝔼[aX + bY] = a𝔼[X] + b𝔼[Y], which holds even when X and Y are dependent.

Variance measures spread: the expected squared distance from the mean. Scaling by a multiplies it by a². For independent variables, variances add.

Var(X)=E[(X−E[X])2]=E[X2]−E[X]2\mathrm{Var}(X) = \mathbb{E}\big[(X - \mathbb{E}[X])^2\big] = \mathbb{E}[X^2] - \mathbb{E}[X]^2

Covariance measures whether two variables move together; correlation is covariance scaled to between −1 and 1. For a vector of variables, the covariance matrix collects every pair. It is symmetric and positive semidefinite, linking straight back to Pillar 1.

8. The distributions you must know

Distribution Models Parameters Where in AI
Bernoulli one yes/no outcome p binary classifiers, dropout masks
Categorical one choice among K options p₁ … p_K every softmax output, every next-token prediction
Binomial number of successes in n tries n, p counting correct answers in an eval
Uniform every value equally likely a range random initialization, random seeds
Gaussian (normal) sums of many small effects mean μ, variance σ² noise, initialization, diffusion models, VAEs
Multivariate Gaussian correlated vectors mean vector, covariance matrix latent spaces, embedding statistics

The Gaussian appears so often for two reasons. The central limit theorem makes sums of many effects approximately Gaussian. And among all distributions with a given mean and variance, it assumes the least extra structure (maximum entropy, Pillar 5).

9. Law of large numbers and the central limit theorem

The law of large numbers says the average of many independent samples approaches the true expectation. The central limit theorem says how fast: the error of an average shrinks like σ/√n, and its distribution becomes Gaussian.

That square root is one of the most practical facts in AI. To halve the noise in an estimate, you need four times the samples. It governs minibatch gradient noise, and how many test questions an eval needs.

10. Sampling

Sampling turns a distribution into concrete outcomes. For a categorical distribution, you draw a uniform number and see which probability bucket it lands in. Every generated token is produced this way.

Temperature reshapes the distribution before sampling by dividing logits by T:

pi=ezi/T∑jezj/Tp_i = \frac{e^{z_i / T}}{\sum_j e^{z_j / T}}

T below 1 sharpens toward the top choice (T → 0 approaches greedy argmax); T above 1 flattens toward uniform. Top-k sampling keeps only the k likeliest tokens; top-p (nucleus) sampling keeps the smallest set whose probabilities reach p.

The reparameterization trick writes a Gaussian sample as z = μ + σ·ε with ε drawn from a standard normal. Randomness moves into ε, so gradients can flow through μ and σ. Variational autoencoders and diffusion models depend on it.

11. Markov chains, briefly

A Markov chain is a sequence where the next state depends only on the current one. Diffusion models are built on one: a chain that adds noise step by step, and a learned chain that removes it. Sampling methods such as MCMC also use them.

In AI

  • Language models learn P(next token | context); the chain rule turns that into a probability for whole texts.
  • Decoding is sampling: greedy, temperature, top-k and top-p are all ways of drawing from a categorical distribution.
  • Initialization is variance arithmetic. A neuron sums n inputs, so with independent zero-mean weights its output variance is n·Var(w)·Var(x). Choosing Var(w) ≈ 1/n (or 2/n for ReLU, He initialization) keeps signals from exploding or dying.
  • Calibration asks whether “80% confident” means right 80% of the time; many models are overconfident.
  • Generative models (diffusion, VAEs) are explicitly built from Gaussians, Markov chains and the reparameterization trick.

Questions to ask

  • What is the random variable, and what is its distribution?
  • Conditional on what? Is the conditioning stated or hidden?
  • Is independence assumed here, and is it actually true?
  • What is the base rate of the thing being detected?
  • Is this a probability or a density?
  • I know the expected value; what is the variance, and what does the tail look like?
  • Is my data i.i.d. from the distribution the model will face in production?

Fluency check

  • Solve the 1% / 99% / 1% test problem with Bayes’ rule, then with a table of 10,000 imaginary people.
  • Show that 𝔼[X + Y] = 𝔼[X] + 𝔼[Y] holds even when Y = X.
  • Derive that the average of n i.i.d. samples has variance σ²/n.
  • Given per-token probabilities for a 5-token sentence, compute its total log probability.
  • Implement temperature, top-k and top-p sampling over a logit vector in NumPy.
  • Simulate the central limit theorem: average n dice rolls many times and watch the histogram turn Gaussian.