Probability and distributions

Problem

A spam filter cannot know whether the next email is spam. A language model cannot know which word comes next. A weather model cannot know tomorrow's temperature to the degree. What each of them can do is say how likely each outcome is, and that statement has to obey rules: the chances of all outcomes add up to one, and combining facts ("it is raining" and "you carried an umbrella") has to give consistent answers.

Probability is the language for those statements. Training a model mostly means making the probability it assigns to what actually happened as large as possible, so you need to be able to write that probability down, compute it stably, and draw samples from it.

Intuition

Picture probability as a fixed amount of sand, exactly one unit, spread over all the things that could happen. An event is a region; its probability is how much sand sits on it. A random variable is a rule that turns each outcome into a number (1 if spam, 0 if not; tomorrow's temperature), and its distribution says how the sand lands on those numbers.

When the possible values can be listed (spam or not, one of 50,000 tokens), the sand sits in piles and a probability mass function (PMF) gives each pile's height. When values are continuous (a temperature), any single exact value has zero sand on it; instead a probability density function (PDF) gives how thickly the sand is spread, and probability is the area under it over a range.

Discrete distributions put mass on listed values. The categorical one is what a classifier's softmax produces.

Mechanics

Random variables, PMFs and PDFs

A discrete random variable X has a PMF p(x) = P(X = x) with every p(x) ≥ 0 and Σ_x p(x) = 1. A continuous one has a PDF p(x) ≥ 0 with ∫p(x) dx = 1, and P(a ≤ X ≤ b) = ∫_a^b p(x) dx. A density can exceed 1 (a very narrow Gaussian does); only its area is bounded. The cumulative distribution function F(x) = P(X ≤ x) works for both kinds and climbs from 0 to 1.

The three distributions you meet first

Bernoulli(p):   P(X = 1) = p,   P(X = 0) = 1 − p

One yes/no outcome: spam or not, click or not. Binary classifiers output p through a sigmoid.

Categorical(π):   P(X = k) = π_k,   π_k = e^z_k / Σ_j e^z_j

One of K outcomes. Models emit unnormalized scores z (logits) and the softmax turns them into π. PyTorch's Categorical accepts either and normalizes in log space, subtracting the logsumexp of the logits, which never overflows even when the logits are in the thousands.

Normal(μ, σ):   p(x) = (1 / σ√(2π)) · exp(−(x − μ)² / 2σ²)

The bell curve, centred at μ with spread σ. Sums of many small independent effects tend toward it (the central limit theorem), which is why measurement noise and initial weights are usually Gaussian. Its log-density is a quadratic, log p(x) = −(x − μ)²/2σ² − log σ − ½ log 2π, so maximizing a Gaussian likelihood is the same as minimizing squared error.

The 68–95–99.7 rule: these areas are F(k) − F(−k) for k = 1, 2, 3, computed from the Gaussian CDF ½(1 + erf(x / √2)).

Joint, marginal and conditional

With two random variables, the joint distribution p(x, y) gives the probability of each pair. Summing out one variable gives the marginal of the other, and dividing a joint by a marginal gives a conditional:

p(x) = Σ_y p(x, y)     p(y | x) = p(x, y) / p(x)     p(x, y) = p(x) · p(y | x)

The last form, the product rule, chains: p(a, b, c) = p(a) p(b | a) p(c | a, b). A language model is exactly this chain over tokens, predicting each one conditioned on all the previous ones.

X and Y are independent when p(x, y) = p(x) p(y) for every pair, which is the same as saying p(y | x) = p(y): learning X tells you nothing about Y. Assuming independence is what lets a dataset's likelihood factor into a product over examples, and what lets a loss be an average over examples.

Sampling

To draw from a categorical distribution, line the probabilities up end to end on [0, 1), draw a uniform number u, and return the segment u lands in. That is inverse-CDF sampling: find the first k whose cumulative sum exceeds u. For a continuous distribution the same idea is x = F⁻¹(u), and for a Gaussian there is a shortcut: draw ε from Normal(0, 1) and return μ + σε.

Draws 0.1, 0.5, 0.7, 0.95 give cat, cat, dog, bird. Over many draws each label turns up in proportion to its segment's length.

In PyTorch

torch.distributions wraps each family in an object with the same interface: sample, log_prob, entropy, and for continuous families cdf, icdf and rsample. Normal.log_prob is the quadratic above, line for line; Normal.cdf is the erf formula. Categorical.__init__ normalizes probabilities by their sum or logits by their logsumexp, and Categorical.sample hands the probabilities to torch.multinomial. Bernoulli.log_prob is the negative of binary cross-entropy with logits, which is the first sign that the usual classification losses are log-probabilities in disguise.

Worked example

Implementation

The same numbers from torch.distributions:

import torch
from torch.distributions import Bernoulli, Categorical, Normal

coin = Bernoulli(probs=torch.tensor(0.3))
coin.log_prob(torch.tensor(1.0))     # tensor(-1.2040), log 0.3
coin.log_prob(torch.tensor(0.0))     # tensor(-0.3567), log 0.7

animal = Categorical(logits=torch.tensor([2.0, 1.0, 0.0]))
animal.probs                         # tensor([0.6652, 0.2447, 0.0900])
animal.log_prob(torch.tensor(1))     # tensor(-1.4076), cross-entropy 1.4076
animal.entropy()                     # tensor(0.8324)
animal.sample((5,))                  # five class indices, e.g. tensor([0, 0, 1, 0, 2])

z = Normal(loc=0.0, scale=1.0)
z.log_prob(torch.tensor(2.0))        # tensor(-2.9189)
z.cdf(torch.tensor(1.96))            # tensor(0.9750)

mu = torch.tensor(1.0, requires_grad=True)
x = Normal(mu, 2.0).rsample()        # mu + 2 * eps, differentiable in mu
x.backward()
mu.grad                              # tensor(1.)

And inverse-CDF sampling without any library, as in the diagram above:

import math

def softmax(logits):
    top = max(logits)
    exps = [math.exp(z - top) for z in logits]
    total = sum(exps)
    return [e / total for e in exps]

def sample(probs, u):
    total = 0.0
    for k, p in enumerate(probs):
        total += p
        if u < total:
            return k
    return len(probs) - 1

probs = softmax([2.0, 1.0, 0.0])
[sample(probs, u) for u in [0.10, 0.50, 0.70, 0.95]]   # [0, 0, 1, 2]

Subtracting the largest logit before exponentiating changes nothing mathematically (it cancels in the ratio) but keeps exp from overflowing; PyTorch does the same inside logsumexp.

Tradeoffs

• Work in log space. The probability of a 1,000-token sentence is a product of 1,000 numbers below 1, which underflows to 0 in floating point. Its log is a sum, which doesn't. That's why every API returns log_prob, and why you should pass logits rather than probs when you have them.
• The Gaussian is convenient, not always right. Its tails fall off so fast that a value 6σ out is treated as nearly impossible. Data with outliers or heavy tails (incomes, latencies) is better served by Laplace, Student-t or a mixture.
• Independence is an assumption. Treating examples as independent makes the loss an average and lets you shuffle and batch, but time series, duplicated data and leaked test examples break it, and the model's confidence ends up too high.
• Sampling is not differentiable by default. sample cuts the gradient. Continuous families offer rsample; discrete ones need score-function estimators or relaxations such as Gumbel-softmax, both noisier.

Check yourself

Connections

References

• Blitzstein and Hwang. Introduction to Probability, 2nd ed., chapters 2–5. CRC Press, 2019.
• Bishop. Pattern Recognition and Machine Learning, sections 1.2 and 2.1–2.3. Springer, 2006.
• Goodfellow, Bengio and Courville. Deep Learning, chapter 3 (probability and information theory). MIT Press, 2016.
• Kingma and Welling. Auto-Encoding Variational Bayes. ICLR 2014. arXiv:1312.6114 (the reparameterization trick).
• PyTorch v2.15.0: Normal, Categorical, Bernoulli.