Maths II: probability, statistics & optimisation

AI systems make uncertain predictions from limited evidence.

Probability describes uncertainty. Statistics asks what we can learn from a sample. Optimisation changes model parameters. Information theory measures surprise and prediction quality. Numerical methods keep those calculations valid inside a computer.

Why this matters

An AI engineer must answer questions such as:

• How confident is this prediction?
• Did the new assistant actually improve resolution rate?
• Is an alert meaningful when the underlying event is rare?
• Why does training become unstable?
• Why does cross-entropy punish confident mistakes?

These are engineering questions with mathematical tools, not isolated textbook exercises.

The intuition

Part 1 — Probability distributions describe uncertainty

A probability distributionProbability distributionA rule that assigns probabilities to possible outcomes or ranges of values. A valid distribution gives no negative probability and assigns total probability one, allowing an AI engineer to describe uncertainty rather than only one prediction. assigns probability to possible outcomes.

For a support-ticket label:

Probabilities are non-negative and total 1.

A random variable maps an uncertain outcome to a value. Let X=1 when a ticket is resolved and X=0 otherwise. X follows a Bernoulli distribution with parameter p, the resolution probability.

Common distributions and their jobs:

• Bernoulli: one yes/no outcome.
• Binomial: count successes across a fixed number of independent Bernoulli trials with constant success probability.
• Categorical: one outcome from several categories.
• Normal: continuous values clustered symmetrically around a mean.
• Poisson: counts of events under a constant-rate model.
• Uniform: equal density across an interval or equal probability across listed outcomes.

A distribution is a model, not a fact. Real support messages may not be independent. Daily traffic rates may change. Choose a distribution because its assumptions are useful and checked, not because its curve is familiar.

Mean, expectation, and variance

The expected value is the probability-weighted long-run average:

E[X] = sum(value × probability)

For a reward:

E[X] = 0×0.5 + 10×0.3 + 20×0.2 = £7

The expectation need not be a possible outcome. You never receive exactly £7 in this example; it is the long-run average.

Variance measures average squared distance from the mean:

Var(X) = E[(X - E[X])²]

Squaring prevents positive and negative deviations from cancelling. Standard deviation is the square root of variance, returning to the original units.

For the reward above:

variance = (0-7)²×0.5 + (10-7)²×0.3 + (20-7)²×0.2
         = 24.5 + 2.7 + 33.8
         = 61

standard deviation ≈ 7.81

Expectation describes centre. Variance describes spread. Two systems can have the same average latency and very different tail behaviour.

Part 2 — Conditional probability changes the reference group

Conditional probabilityConditional probabilityThe probability of an event after restricting attention to cases where another event is known to have occurred. It is written P(A given B) and differs from P(B given A), a reversal that causes many diagnostic and classification mistakes. asks for the probability of A given that B occurred:

P(A | B) = P(A and B) / P(B)

The vertical bar means “given”.

Do not reverse the condition:

P(alert | urgent) is not P(urgent | alert)

The first asks what fraction of urgent messages trigger an alert. The second asks what fraction of alerts are truly urgent.

Bayes’ rule updates a base rate

Bayes’ ruleBayes' ruleA relationship that updates a prior probability using how likely the observed evidence is under competing possibilities. It combines a base rate with evidence to produce a posterior probability, making explicit why a sensitive test can still have many false alarms. connects the two directions:

P(urgent | alert)
= P(alert | urgent) × P(urgent) / P(alert)

The terms have useful names:

• Prior: P(urgent), before this alert.
• Likelihood: how compatible the alert is with each case.
• Posterior: P(urgent | alert), after observing the alert.

Counts make Bayes easier:

80 true alerts / (80 true alerts + 90 false alerts) ≈ 47%

The alert catches 80% of urgent messages, but an alert is not 80% likely to be urgent. The 10% base rate matters.

This same mistake appears in fraud detection, medical screening, content moderation, and anomaly alerts.

Part 3 — Maximum likelihood fits model parameters

Suppose a coin produced H, H, T, H, T. A Bernoulli model has unknown parameter p, the probability of heads.

The likelihood asks: for a proposed p, how likely were these observed data?

L(p) = p³ × (1-p)²

Maximum likelihood estimation, shortened to MLE, chooses the parameter that maximises likelihood.

For Bernoulli data, the MLE is the observed success fraction:

p_hat = 3/5 = 0.6

The hat means “estimate”.

Products of many small probabilities underflow, so software maximises log-likelihood:

log L(p) = 3 log(p) + 2 log(1-p)

The logarithm turns products into sums and preserves which positive likelihood is largest.

MLE finds the parameter that best explains observed data under the assumed model. It does not guarantee the model family is correct or that the sample represents future traffic.

Part 4 — Populations, samples, and sampling variation

A population is the full group you want to reason about: perhaps all future eligible support tickets.

A sample is the observed subset.

A parameter describes the population, such as true resolution rate. A statistic is calculated from a sample, such as the sample resolution rate.

Different random samples produce different statistics. That is sampling variation.

import random

random.seed(7)
true_rate = 0.60

for sample_size in (20, 200, 2_000):
    resolved = sum(
        random.random() < true_rate
        for _ in range(sample_size)
    )
    estimate = resolved / sample_size
    print(sample_size, round(estimate, 3))

A huge biased sample can be precisely wrong. Randomisation and representative inclusion matter as much as sample size.

Part 5 — Confidence intervals express estimation uncertainty

A confidence intervalConfidence intervalA range produced by a procedure that would contain the true population parameter at a stated long-run rate under repeated valid sampling. A 95 percent confidence interval is not a guarantee that this particular fixed interval has a 95 percent probability of containing the parameter. comes from a procedure with a long-run coverage rate.

For a sample proportion p_hat from n sufficiently independent observations, a rough standard error is:

SE ≈ square_root(p_hat × (1 - p_hat) / n)

A rough 95% interval is:

p_hat ± 1.96 × SE

Example: 540 resolutions from 1,000 tickets:

p_hat = 0.54
SE ≈ square_root(0.54×0.46/1000) ≈ 0.0158
interval ≈ 0.54 ± 0.031
         ≈ [0.509, 0.571]

Careful interpretation:

If we repeated the valid sampling procedure many times, about 95% of the intervals produced would contain the true parameter.

Once one interval is calculated, traditional frequentist language does not assign a 95% probability to the fixed unknown parameter being inside that fixed interval.

The rough normal interval can behave poorly for small samples or rates near 0 and 1. Real work may use Wilson intervals, bootstrap intervals, or a model suited to the design.

An interval reflects random sampling uncertainty under assumptions. It does not include unmeasured bias, broken randomisation, data leakage, or a changed product.

Part 6 — Hypothesis tests measure incompatibility with a null model

A hypothesis testHypothesis testA structured calculation that asks how incompatible observed data are with a stated null model. Its p-value is the probability, assuming the null model and test assumptions, of results at least as extreme as those observed; it is not the probability that the null hypothesis is true. starts with:

• Null hypothesis: a stated baseline model, often “no difference”.
• Alternative hypothesis: the relevant competing claim.
• Test statistic: a summary of the observed difference relative to expected variation.
• p-value: under the null and assumptions, the probability of a result at least as extreme as observed.

A p-value is not:

• the probability the null is true;
• the probability the result happened “by chance”;
• the size or importance of the effect;
• proof the experiment was unbiased.

“Statistically significant” at a threshold such as 0.05 means the result crossed a chosen decision rule. It does not mean practically important.

Report the estimated effect, uncertainty interval, sample size, design, and decision context—not only a p-value.

Part 7 — A/B tests require product and statistical discipline

An A/B test randomly assigns eligible experimental units to variants:

• A: current system;
• B: proposed system.

For an AI support assistant:

1. Define one primary outcome, such as resolution within 24 hours.
2. Define guardrails, such as escalation rate and safety complaints.
3. Choose the randomisation unit. User-level assignment may prevent one user seeing both experiences.
4. Estimate required sample size from minimum useful effect, baseline rate, power, and error threshold.
5. Run variants simultaneously.
6. Analyse according to the predeclared plan.
7. Consider practical value and harms, not only significance.

Common failures:

• Peeking: checking repeatedly and stopping the first time p < 0.05.
• Sample-ratio mismatch: assignment counts differ unexpectedly from the planned split.
• Interference: one user’s treatment affects another user’s outcome.
• Novelty effects: behaviour changes temporarily because the feature is new.
• Attrition: one variant causes missing outcomes.
• Metric movement without user value: a proxy improves while the real goal worsens.

Multiple comparisons

At a 5% false-positive threshold, testing many truly null metrics increases the chance that at least one crosses the threshold.

For 20 independent null tests:

P(at least one false positive) = 1 - 0.95²⁰ ≈ 64%

Independence is a simplification, but the danger remains.

Strategies include:

• choose one primary metric before running;
• label secondary analyses exploratory;
• control family-wise error, for example Bonferroni;
• control false discovery rate, for example Benjamini–Hochberg;
• use sequential-testing methods if continuous monitoring is required.

Do not choose a correction after seeing which one makes the desired result survive.

Expected output is close to:

simulated: 0.64
theory:    0.642

Part 8 — Gradient descent and SGD

The previous lesson introduced:

parameter = parameter - learning_rate × gradient

Batch gradient descent calculates a gradient using the complete training dataset before each update. That may be expensive.

Stochastic gradient descent, shortened to SGD, estimates the gradient from one example or, more commonly, a small mini-batch.

Mini-batch SGD:

1. Sample a batch.
2. Run predictions.
3. Calculate average batch loss.
4. Backpropagate a gradient estimate.
5. Update parameters.
6. Repeat.

The estimate is noisy. Noise makes each step less exact, but mini-batches are efficient on modern hardware and allow frequent updates. Shuffle data and ensure batches represent the intended training distribution.

An epoch is one pass through the training dataset. Batch size changes the number of updates per epoch and the gradient noise; it is not only a memory setting.

Part 9 — Momentum remembers recent direction

Plain SGD reacts to the current gradient. Momentum keeps an exponentially weighted moving average:

velocity = beta × velocity + (1 - beta) × gradient
parameter = parameter - learning_rate × velocity

beta near 1 remembers a longer history.

Momentum can:

• smooth alternating gradients;
• build speed along a consistent direction;
• pass through shallow local irregularities.

It can also overshoot. It adds state and another hyperparameter.

Part 10 — Adam and AdamW adapt step scales

Adam tracks:

• a first moment: moving average of gradients;
• a second raw moment: moving average of squared gradients.

Simplified:

m = beta1 × m + (1-beta1) × gradient
v = beta2 × v + (1-beta2) × gradient²

parameter -= learning_rate × corrected_m / (sqrt(corrected_v) + epsilon)

The square and square root operate entry by entry.

Why:

• m smooths direction like momentum.
• v records typical squared gradient size per coordinate.
• Dividing by sqrt(v) reduces steps where gradients have been persistently large.
• epsilon avoids division by zero.
• Bias correction compensates for moments starting at zero.

Adam often learns quickly with sparse or differently scaled gradients. It is not guaranteed to generalise better than SGD, and defaults are not universal.

Weight decay gently shrinks parameters:

parameter *= (1 - learning_rate × decay_rate)

AdamW decouples this shrinkage from Adam’s adaptive gradient transformation. With adaptive optimisers, adding an L2 penalty to the loss is not generally equivalent to separate weight decay.

Part 11 — Learning-rate schedules change training over time

A fixed learning rate treats every training stage alike. A learning-rate schedule changes it.

Common patterns:

• Warm-up: increase from a small rate during early unstable steps.
• Step decay: reduce at chosen milestones.
• Exponential decay: multiply by a fixed factor over time.
• Cosine decay: smoothly reduce following part of a cosine curve.
• Reduce on plateau: reduce when a validation metric stops improving.

Warm-up can protect large models from unstable early updates. Decay allows broad movement early and finer adjustment later.

The schedule must be defined in steps or epochs consistently with batch size, gradient accumulation, and distributed training. Changing batch size changes steps per epoch.

Part 12 — Convexity explains guarantees and their limits

A function is convex when the line between any two points on its graph lies on or above the graph. A bowl-shaped quadratic is convex.

For a differentiable convex function, any local minimum is global. Optimisation theory can provide strong guarantees.

Neural-network losses are generally non-convex because layers compose nonlinear transformations and parameters can interact symmetrically.

This does not make optimisation hopeless. High-dimensional networks often contain many useful low-loss solutions. But simple guarantees from convex optimisation do not transfer automatically.

Part 13 — Entropy measures expected surprise

Information theory uses:

surprise(outcome) = -log(probability)

A rare outcome is more surprising.

Entropy is expected surprise under a distribution:

H(P) = -sum P(x) × log P(x)

With natural logarithms, units are nats. With base-2 logarithms, units are bits.

For a fair coin:

H = -0.5 log₂(0.5) - 0.5 log₂(0.5) = 1 bit

For a coin that is always heads, entropy is 0: the outcome carries no surprise.

Entropy is uncertainty in the distribution, not model error by itself.

Part 14 — Cross-entropy scores predicted probabilities

Cross-entropyCross-entropyA loss that scores how much probability a model assigned to the observed class. For one correct class it is the negative logarithm of that class probability, so confident wrong predictions are penalised strongly and a probability of zero would create infinite loss. measures expected surprise when data follow distribution P but predictions use Q:

H(P, Q) = -sum P(x) × log Q(x)

For one correct class, the target is one-hot, so:

loss = -log(probability assigned to correct class)

Confident wrong predictions receive large loss.

Cross-entropy encourages probability on the observed class. It does not by itself guarantee calibrated probabilities, fairness, robustness, or product value.

Part 15 — KL divergence measures extra coding cost

Kullback–Leibler divergence, shortened to KL divergence, is:

KL(P || Q) = sum P(x) × log(P(x) / Q(x))

It measures the extra expected surprise from using Q when reality follows P.

Relationship:

cross_entropy(P, Q) = entropy(P) + KL(P || Q)

Properties:

• KL divergence is non-negative under valid conditions.
• It is zero when distributions match almost everywhere.
• It is not symmetric: KL(P || Q) usually differs from KL(Q || P).
• It can be infinite when Q assigns zero probability where P assigns positive probability.

Do not call KL a distance without qualification; it does not satisfy the symmetry requirement of a mathematical metric.

Part 16 — Perplexity summarises token prediction loss

PerplexityPerplexityThe exponential of average token cross-entropy. It can be interpreted as an effective branching factor under a fixed tokenisation and evaluation setup; lower is better on the same data and vocabulary, but values are not directly comparable across different tokenisers or datasets. is:

perplexity = exp(average token cross-entropy)

If average cross-entropy is log(4), perplexity is 4. Informally, the model is as uncertain as choosing among roughly four equally likely possibilities at each step.

Lower is better only under the same:

• evaluation text;
• tokenisation;
• token boundaries and masking;
• context handling;
• averaging convention.

A tokenizer that splits text differently changes the prediction task. Do not compare perplexities across incompatible setups as if they were one universal score.

Perplexity also does not directly measure factuality, usefulness, instruction following, or safety.

Part 17 — Numerical stability keeps formulas computable

Computers use finite-precision floating-point numbers. Very large magnitudes can overflow. Very small positive values can underflow to zero. Subtracting nearly equal values can lose significant digits.

Numerical stabilityNumerical stabilityThe property of an algorithm that avoids unnecessary loss of accuracy, overflow, underflow, or invalid values when represented with finite-precision numbers. Stable forms such as log-sum-exp compute mathematically equivalent results without creating enormous intermediate exponentials. means choosing an equivalent computation that behaves well in finite precision.

Stable softmax

Softmax turns logits into probabilities:

softmax_i = exp(logit_i) / sum_j exp(logit_j)

For logits [1000, 1001, 1002], direct exponentials overflow.

Subtract the maximum:

shifted = [-2, -1, 0]

Softmax does not change because the common factor cancels. Now every exponential is at most 1.

Log-sum-exp

The unsafe expression:

log(sum(exp(x_i)))

becomes:

m = max(x)
logsumexp(x) = m + log(sum(exp(x_i - m)))

Expected output:

[0.09, 0.2447, 0.6652]
sum: 1.0
log-sum-exp: 1002.4076

Other stability habits

• Use library log1p(x) for log(1+x) when x is tiny.
• Use expm1(x) for exp(x)-1 when x is tiny.
• Clip probabilities only with an explicit mathematical reason and documented epsilon.
• Prefer log-probabilities when multiplying many probabilities.
• Check for NaN and infinity during training.
• Use appropriate dtypes; lower precision saves memory but reduces range and resolution.
• Separate a formula’s mathematical meaning from the implementation used to compute it.

Worked example — did the AI assistant improve resolution?

Variant A resolves 500 of 1,000 tickets. Variant B resolves 540 of 1,000.

1. Estimate the effect

A rate = 0.500
B rate = 0.540
absolute difference = 0.040 = 4 percentage points
relative lift = 0.040 / 0.500 = 8%

Absolute and relative changes answer different questions. Report both with clear labels.

2. Estimate uncertainty

Approximate standard error for the difference of independent proportions:

SE_difference
≈ square_root(
    pA(1-pA)/nA + pB(1-pB)/nB
  )
≈ square_root(0.5×0.5/1000 + 0.54×0.46/1000)
≈ 0.0223

Approximate 95% interval:

0.040 ± 1.96×0.0223
≈ [-0.004, 0.084]

The interval includes zero under this rough method. The data are compatible with a small negative effect through an 8.4-point positive effect.

3. Make the product decision honestly

Do not translate that directly into “B does not work”. Ask:

• Was assignment random and stable?
• Was user-level clustering handled?
• Was this the declared primary metric?
• Were safety and escalation guardrails acceptable?
• Is the minimum useful effect inside the interval?
• Should the experiment continue under a valid preplanned design?

Statistical uncertainty is one input to a product decision.

In a real AI system

Common mistakes

Exercises

Knowledge check

Practice activity

from math import sqrt

resolved_a, total_a = 500, 1_000
resolved_b, total_b = 540, 1_000

# Calculate rates, difference, standard error, and interval.

from math import sqrt

resolved_a, total_a = 500, 1_000
resolved_b, total_b = 540, 1_000

rate_a = resolved_a / total_a
rate_b = resolved_b / total_b
difference = rate_b - rate_a
relative_lift = difference / rate_a

standard_error = sqrt(
    rate_a * (1 - rate_a) / total_a
    + rate_b * (1 - rate_b) / total_b
)
lower = difference - 1.96 * standard_error
upper = difference + 1.96 * standard_error

print(f"A: {rate_a:.1%}")
print(f"B: {rate_b:.1%}")
print(f"absolute: {difference:.1%}")
print(f"relative: {relative_lift:.1%}")
print(f"rough 95% interval: [{lower:.3f}, {upper:.3f}]")

A: 50.0%
B: 54.0%
absolute: 4.0%
relative: 8.0%
rough 95% interval: [-0.004, 0.084]

Key takeaways

Flashcards

One-page sketchnote summary

posterior ∝ likelihood × prior
parameter -= learning_rate * update
logsumexp(x) = max(x) + log(sum(exp(x-max(x))))

Resources