Probability Fundamentals

Learning Objectives

By the end of this section, you will:

1. Understand random variables as the mathematical objects that model uncertainty in images, noise, and neural network outputs
2. Master probability distributions including PDFs, joint distributions, and marginals that form the foundation of generative modeling
3. Compute expectations and variances that appear in every diffusion model equation from loss functions to sampling
4. Apply conditional probability to understand how diffusion models transform noisy observations into clean images
5. Use Bayes' theorem to see why diffusion models work: they learn to invert a known noising process

The Big Picture

Probability theory emerged from humanity's desire to reason about uncertainty. In the 17th century, mathematicians like Pascal and Fermat developed the foundations while analyzing games of chance. By the 18th century, Laplace and Bayes had extended these ideas to scientific inference: how do we update our beliefs about the world given new observations?

This question—how do we reason about what we cannot directly observe?—is precisely what diffusion models answer in the context of image generation. When we see noise, what was the original image? When we want to generate an image, how do we sample from the vast space of possibilities in a way that produces realistic results?

Diffusion models solve this problem through a beautiful probabilistic framework:

1. Forward process: Gradually corrupt images with noise according to a known distribution $q(x_t|x_{t-1})$
2. Reverse process: Learn to undo this corruption by approximating$p_\theta(x_{t-1}|x_t)$ with a neural network
3. Generation: Start from pure noise and iteratively denoise to create new images

Every step in this process relies on the probability concepts we'll develop in this section.

Random Variables

Definition and Intuition

A random variable is a function that assigns a numerical value to each outcome of a random experiment. We denote random variables with capital letters (X, Y, Z) and their realized values with lowercase letters (x, y, z).

Types of Random Variables

In diffusion models, we primarily work with continuous multivariate random variables. An image is a random variable $X \in \mathbb{R}^{H \times W \times C}$ where H, W, C are height, width, and channels respectively.

Implementation: Random Variables in PyTorch

X ~ N(μ=0, σ²=1)

samples = [x₁, x₂, ..., x₁₀₀₀₀]

E[X] ≈ (1/n)Σxᵢ

Var(X) ≈ (1/n)Σ(xᵢ - x̄)²

log p(0) for N(0,1) = log(1/√(2π)) ≈ -0.919

1import torch
2from torch.distributions import Normal
3
4# Define a random variable X ~ N(0, 1)
5X = Normal(loc=0.0, scale=1.0)
6
7# Sample realizations of X
8samples = X.sample((10000,))  # 10,000 samples
9
10# Compute statistics (estimates of theoretical values)
11mean_estimate = samples.mean()      # Estimates E[X] ≈ 0
12var_estimate = samples.var()        # Estimates Var(X) ≈ 1
13
14# Compute log probability at a specific point
15log_prob_at_zero = X.log_prob(torch.tensor(0.0))
16
17print(f"Sample mean: {mean_estimate:.4f}")
18print(f"Sample variance: {var_estimate:.4f}")
19print(f"Log probability at x=0: {log_prob_at_zero:.4f}")

Probability Distributions

Probability Density Functions (PDFs)

For continuous random variables, we describe probabilities using a probability density function (PDF), denoted $p(x)$ or $f(x)$.

Key properties of a PDF:

1. Non-negativity: $p(x) \geq 0$ for all x
2. Normalization: $\int_{-\infty}^{\infty} p(x) \, dx = 1$
3. Probability via integration: $P(a \leq X \leq b) = \int_a^b p(x) \, dx$

The Gaussian Distribution

The Gaussian (normal) distribution is central to diffusion models. Its PDF is:

$p(x) = \frac{1}{\sqrt{2\pi\sigma^2}} \exp\left(-\frac{(x - \mu)^2}{2\sigma^2}\right)$

where $\mu$ is the mean and $\sigma^2$ is the variance. We write $X \sim \mathcal{N}(\mu, \sigma^2)$.

Joint Distributions

When we have multiple random variables, we describe their joint distribution:

$p(x, y)$ = probability density of X = x AND Y = y

In diffusion models, we model the joint distribution of images at all timesteps:

$p(x_0, x_1, \ldots, x_T) = p(x_T) \prod_{t=1}^{T} p(x_{t-1}|x_t)$

Σ = [[1.0, 0.7], [0.7, 1.0]]

μ₁|₂ = μ₁ + (σ₁₂/σ₂²)(x₂ - μ₂)

1import torch
2from torch.distributions import MultivariateNormal
3
4# Define a 2D joint distribution
5mean = torch.tensor([0.0, 0.0])
6covariance = torch.tensor([
7    [1.0, 0.7],
8    [0.7, 1.0]
9])
10
11# Create joint distribution p(x1, x2)
12joint_dist = MultivariateNormal(mean, covariance)
13
14# Sample from the joint distribution
15joint_samples = joint_dist.sample((1000,))  # Shape: (1000, 2)
16
17# Marginal distribution p(x1)
18marginal_x1 = Normal(loc=mean[0], scale=torch.sqrt(covariance[0, 0]))
19
20# Conditional distribution p(x1 | x2 = observed_x2)
21observed_x2 = 1.5
22conditional_mean = mean[0] + covariance[0, 1] / covariance[1, 1] * (observed_x2 - mean[1])
23conditional_var = covariance[0, 0] - covariance[0, 1]**2 / covariance[1, 1]
24conditional_x1_given_x2 = Normal(loc=conditional_mean, scale=torch.sqrt(conditional_var))

Marginal Distributions

From a joint distribution, we obtain marginal distributions by integrating out (or summing over) other variables:

$p(x) = \int p(x, y) \, dy$

Geometrically, marginalization is like projecting a 2D distribution onto one axis—we "collapse" one dimension by summing all possibilities.

Expectation and Variance

Expected Value (Mean)

The expected value (or mean) of a random variable X is the probability-weighted average of all possible values:

$\mathbb{E}[X] = \int_{-\infty}^{\infty} x \cdot p(x) \, dx$

For a function g(X):

$\mathbb{E}[g(X)] = \int_{-\infty}^{\infty} g(x) \cdot p(x) \, dx$

Variance

Variance measures the spread of a random variable around its mean:

$\text{Var}(X) = \mathbb{E}[(X - \mathbb{E}[X])^2] = \mathbb{E}[X^2] - (\mathbb{E}[X])^2$

The standard deviation $\sigma = \sqrt{\text{Var}(X)}$ is in the same units as X.

Key Properties

Conditional Probability

Definition

The conditional probability of X given Y is:

$p(x|y) = \frac{p(x, y)}{p(y)}$

This reads as "the probability density of x given that we observed y." Conditioning "slices" through the joint distribution and renormalizes.

The Chain Rule of Probability

Joint distributions can be decomposed using the chain rule:

$p(x_1, x_2, \ldots, x_T) = p(x_1) \prod_{t=2}^{T} p(x_t|x_1, \ldots, x_{t-1})$

For Markov chains (which diffusion models are), each step only depends on the previous step:

$p(x_1, x_2, \ldots, x_T) = p(x_1) \prod_{t=2}^{T} p(x_t|x_{t-1})$

Bayes' Theorem

Bayes' theorem relates conditional probabilities in both directions:

$p(x|y) = \frac{p(y|x) \cdot p(x)}{p(y)}$

p(clean) = 0.8 means we believe 80% of images are clean

p(obs) = p(obs|clean)p(clean) + p(obs|noisy)p(noisy)

p(A|B) = p(B|A)p(A) / p(B)

1import numpy as np
2
3# Prior: P(image is clean)
4p_clean = 0.8
5p_noisy = 1 - p_clean  # = 0.2
6
7# Likelihood: P(high quality features | image type)
8p_high_quality_given_clean = 0.9
9p_high_quality_given_noisy = 0.3
10
11# Evidence: P(high quality) = sum over all hypotheses
12p_high_quality = (p_high_quality_given_clean * p_clean +
13                  p_high_quality_given_noisy * p_noisy)
14
15# Posterior: P(clean | high quality) via Bayes' theorem
16p_clean_given_high_quality = (p_high_quality_given_clean * p_clean) / p_high_quality
17
18print(f"Prior P(clean): {p_clean:.2f}")
19print(f"Evidence P(high quality): {p_high_quality:.2f}")
20print(f"Posterior P(clean | high quality): {p_clean_given_high_quality:.3f}")
21# Output: 0.923 - Observing high quality increases our belief!

Why Bayes' Theorem Matters for Diffusion

In diffusion models, we want to compute $q(x_{t-1}|x_t)$—given a noisy image, what was the slightly less noisy version? Direct computation requires knowing$q(x_{t-1})$, which is intractable.

However, if we also condition on the original image x₀, Bayes' theorem gives us:

$q(x_{t-1}|x_t, x_0) = \frac{q(x_t|x_{t-1}, x_0) \cdot q(x_{t-1}|x_0)}{q(x_t|x_0)}$

All terms on the right are Gaussians that we know! This closed-form posterior is the key insight that makes diffusion model training tractable—we'll derive it fully in Chapter 3.

Connection to Diffusion Models

Let's preview how every concept we've learned appears in diffusion models:

Summary

In this section, we built the probability foundations essential for understanding diffusion models:

1. Random Variables: Mathematical objects modeling uncertainty in images, noise, and model outputs. Continuous multivariate random variables represent entire images.
2. Probability Distributions: PDFs describe the likelihood of different values. Gaussians are central due to their mathematical properties and natural occurrence.
3. Joint and Marginal Distributions: Joint distributions describe multiple random variables together; marginals integrate out unwanted variables.
4. Expectation and Variance: Expected values appear in training objectives; variance controls noise schedules and model uncertainty.
5. Conditional Probability: The foundation of diffusion—forward process$q(x_t|x_{t-1})$ and reverse process $p_\theta(x_{t-1}|x_t)$.
6. Bayes' Theorem: Enables computing the tractable posterior$q(x_{t-1}|x_t, x_0)$ that makes training possible.

Exercises

Conceptual Questions

1. Why is $P(X = x) = 0$ for continuous random variables? What does this mean for image generation where we produce specific pixel values?
2. If X and Y are independent, what is $p(x|y)$? How does independence simplify joint distributions?
3. Derive the variance formula $\text{Var}(X) = \mathbb{E}[X^2] - (\mathbb{E}[X])^2$starting from the definition $\text{Var}(X) = \mathbb{E}[(X - \mathbb{E}[X])^2]$.
4. In diffusion models, why does conditioning on x₀ make the reverse distribution tractable? Hint: Think about what information x₀ provides.

Computational Exercises

1. Generate 10,000 samples from a standard normal distribution. Estimate the mean, variance, and the probability that $X \in [-1, 1]$. Compare with theoretical values.
2. Verify empirically that for $X \sim \mathcal{N}(\mu_1, \sigma_1^2)$ and$Y \sim \mathcal{N}(\mu_2, \sigma_2^2)$ independent, their sum follows$X + Y \sim \mathcal{N}(\mu_1 + \mu_2, \sigma_1^2 + \sigma_2^2)$.
3. Implement Bayes' theorem to update beliefs about whether an image contains a cat, given observations about edges and colors. Start with a uniform prior and observe how the posterior changes with evidence.
4. Create a simple 1D diffusion forward process: start with samples from a mixture of Gaussians and progressively add noise. Visualize how the distribution converges to a standard normal.

Challenge Problem

For a 2D Gaussian with mean $\mu = [0, 0]$ and covariance$\Sigma = \begin{bmatrix} 1 & \rho \\ \rho & 1 \end{bmatrix}$:

1. Derive the conditional distribution $p(x_1|x_2)$ analytically.
2. Show that as $\rho \to 1$, the conditional variance approaches 0 (perfect prediction).
3. Connect this to diffusion: How does the conditional variance change as the forward process progresses?