Chapter 2: Mathematical Foundations

Week 2

Master the essential mathematical toolkit underpinning diffusion models: Gaussian distributions, Markov chains, KL divergence, variational inference, stochastic differential equations, and Langevin dynamics. These foundations are critical for understanding both the theory and implementation of modern diffusion-based generative models.

Topics: 3
Videos: 3
Min read: 19

Chapter Overview

Diffusion models rest on a surprisingly elegant mathematical framework. At their core, they define a forward process that gradually destroys data by adding Gaussian noise, and a learned reverse process that reconstructs data from pure noise. Understanding why this works---and how to train it efficiently---requires fluency in several branches of mathematics.

We begin with Gaussian distributions, the building block of the noise process. Every step of diffusion adds Gaussian noise, and the remarkable closure properties of Gaussians make the entire forward process analytically tractable. Next, Markov chains formalize the sequential structure: each noising step depends only on the previous state, enabling efficient factorization of joint distributions.

KL divergence measures how one probability distribution differs from another, and it is the natural loss function for training generative models. Variational inference provides the framework for optimizing intractable likelihoods by maximizing a tractable lower bound (the ELBO), which is exactly how diffusion models are trained.

Finally, the continuous-time perspective connects diffusion models to stochastic differential equations and Langevin dynamics. Score-based generative models use the score function (gradient of the log-density) to reverse the noising process, unifying denoising diffusion with classical MCMC sampling.

This chapter covers:

Gaussian Distributions: Properties that make the forward process analytically tractable
Markov Chains: The sequential structure underlying discrete diffusion steps
KL Divergence: The divergence measure used to train the reverse process
Variational Inference: Deriving the ELBO training objective
Stochastic Differential Equations: Continuous-time formulation of diffusion
Langevin Dynamics: Score-based sampling and its connection to diffusion models

Chapter Roadmap

Click any topic to jump in

Explainer video

Gaussian Distributions

Properties of Gaussians and reparameterization — the building block of diffusion.

Univariate and Multivariate GaussianClosure Under Linear TransformationsThe Reparameterization Trick

Process and objective

Sequential structure and optimization target

Markov Chains

Transition kernels and stationary distributions — the sequential structure of diffusion.

Transition ProbabilitiesStationary DistributionErgodicity and Mixing Time

KL Divergence

Measuring distribution distance — the objective connecting forward and reverse processes.

Definition and AsymmetryForward vs Reverse KLKL Between Gaussians

Topics

3 topics in this chapter

Before we can add noise to an image, we need to say precisely what noise means --- and in diffusion models the answer is always the same distribution: the Gaussian. This topic opens the mathematical-foundations chapter because everything downstream depends on it. Chapter 1 argued that diffusion models generate by reversing a gradual noising process; that process adds Gaussian noise at every step, and the reverse network predicts Gaussian parameters. Why does one distribution get to play both roles? Because Gaussians have algebraic properties no other distribution shares: they stay Gaussian under addition, under linear maps, under conditioning, and under marginalization. Those closure properties are what make the entire forward process solvable in closed form --- you can jump straight to any noise level without simulating the steps in between, a fact we will exploit constantly. This topic builds that toolkit: the univariate and multivariate density, closure under affine transformations, and the reparameterization trick that makes sampling differentiable. Master these and the DDPM equations in later chapters read as routine algebra rather than magic.

Definition

A Gaussian (normal) distribution $\mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma})$ is the distribution over $\mathbb{R}^d$ specified entirely by a mean vector $\boldsymbol{\mu}$ and covariance matrix $\boldsymbol{\Sigma}$ . It is closed under affine transformation, marginalization, and conditioning --- properties that make the diffusion forward process analytically tractable at every timestep.

In this topic

1Univariate and Multivariate Gaussian
2Closure Under Linear Transformations
3The Reparameterization Trick

1 of 3

Univariate and Multivariate Gaussian

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

The univariate Gaussian is parameterized by a mean $\mu$ and variance $\sigma^2$ ; the density peaks at $\mu$ and its width scales with $\sigma$ . The multivariate generalization replaces the scalar mean with a vector $\boldsymbol{\mu} \in \mathbb{R}^d$ and the variance with a covariance matrix $\boldsymbol{\Sigma} \in \mathbb{R}^{d \times d}$ , whose off-diagonal entries encode correlations between dimensions. In diffusion models the covariance is almost always isotropic, $\boldsymbol{\Sigma} = \sigma^2 \mathbf{I}$ , so noise is added independently and equally to every dimension. This is deliberate: correlated noise would bias the forward process toward certain directions in data space, breaking the symmetry that later lets a single scalar schedule describe corruption.

Mathematical Intuition

The Gaussian is the maximum entropy distribution for a given mean and variance — it makes the fewest assumptions beyond those two moments. In diffusion, isotropic noise $\mathcal{N}(0, \sigma^2 I)$ treats all dimensions equally, which is essential: the forward process should not favor any direction in data space. The normalization constant $1/\sqrt{(2\pi)^d |\Sigma|}$ ensures the density integrates to 1 over $\mathbb{R}^d$ .

Example:

Write the PDF of a 2D isotropic Gaussian with mean $\boldsymbol{\mu} = [1, 2]^T$ and variance $\sigma^2 = 0.5$ .

2 of 3

Closure Under Linear Transformations

$\text{If } \mathbf{x} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma}), \text{ then } \mathbf{A}\mathbf{x} + \mathbf{b} \sim \mathcal{N}(\mathbf{A}\boldsymbol{\mu} + \mathbf{b},\; \mathbf{A}\boldsymbol{\Sigma}\mathbf{A}^T)$

Any affine transformation of a Gaussian random variable is again Gaussian: if $\mathbf{x} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma})$ then $\mathbf{A}\mathbf{x} + \mathbf{b}$ has mean $\mathbf{A}\boldsymbol{\mu} + \mathbf{b}$ and covariance $\mathbf{A}\boldsymbol{\Sigma}\mathbf{A}^T$ . This is the property that makes diffusion tractable. The forward step $\mathbf{x}_t = \sqrt{\alpha_t}\,\mathbf{x}_{t-1} + \sqrt{1 - \alpha_t}\,\boldsymbol{\epsilon}$ is an affine combination of a Gaussian input and independent Gaussian noise, so $\mathbf{x}_t$ is Gaussian too. Applying closure repeatedly collapses the whole chain, letting us write $q(\mathbf{x}_t \mid \mathbf{x}_0)$ in closed form for any $t$ --- no need to iterate through intermediate steps, which is what makes training on random timesteps possible.

Mathematical Intuition

Closure means the family of Gaussians is preserved under affine maps: if $x \sim \mathcal{N}(\mu, \Sigma)$ , then $Ax + b \sim \mathcal{N}(A\mu + b, A\Sigma A^T)$ . In DDPM, $x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon$ is affine in both $x_0$ and $\epsilon$ . Since both are Gaussian, $x_t$ is Gaussian with analytically computable moments — no simulation needed to characterize the distribution at any timestep.

Example:

If $\mathbf{x} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ and $\mathbf{y} = 0.8\mathbf{x} + 0.6\boldsymbol{\epsilon}$ where $\boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ is independent, what is the distribution of $\mathbf{y}$ ?

3 of 3

The Reparameterization Trick

$\mathbf{x} = \boldsymbol{\mu} + \boldsymbol{\sigma} \odot \boldsymbol{\epsilon}, \quad \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$

Sampling from $\mathcal{N}(\boldsymbol{\mu}, \text{diag}(\boldsymbol{\sigma}^2))$ is not differentiable with respect to $\boldsymbol{\mu}$ and $\boldsymbol{\sigma}$ : the randomness sits between the parameters and the output, blocking gradient flow. The reparameterization trick rewrites the sample as $\mathbf{x} = \boldsymbol{\mu} + \boldsymbol{\sigma} \odot \boldsymbol{\epsilon}$ with fixed noise $\boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ , moving the stochasticity to an input that carries no parameters. Now the sample is a deterministic, differentiable function of $\boldsymbol{\mu}$ and $\boldsymbol{\sigma}$ , so gradients pass straight through it. This is what makes end-to-end training of both VAEs and diffusion models possible; without it the sampling step would break backpropagation entirely.

Mathematical Intuition

Sampling $x \sim \mathcal{N}(\mu,\sigma^2)$ is a stochastic operation with no gradient through $\mu$ or $\sigma$ . The trick $x = \mu + \sigma\epsilon$ decomposes sampling into a deterministic differentiable path (multiply and add) plus fixed randomness ( $\epsilon$ ). The pathwise gradient $\partial x/\partial \mu = 1$ and $\partial x/\partial \sigma = \epsilon$ enable backpropagation through the sampling step — critical for training both VAEs and diffusion models.

Example:

Show how to compute gradients of $\mathbb{E}_{x \sim \mathcal{N}(\mu, \sigma^2)}[x^2]$ with respect to $\mu$ using the reparameterization trick.

Theory Exercise

Problem:

In the DDPM forward process, $q(\mathbf{x}_t | \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t}\,\mathbf{x}_0,\; (1 - \bar{\alpha}_t)\mathbf{I})$ where $\bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s$ . Derive this closed-form expression starting from the single-step transition $q(\mathbf{x}_t | \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{\alpha_t}\,\mathbf{x}_{t-1},\; (1 - \alpha_t)\mathbf{I})$ .

Hints:

Write x_t in terms of x_{t-1} and noise using the reparameterization trick
Recursively substitute x_{t-1} in terms of x_{t-2}, and so on down to x_0
Use the fact that a sum of independent Gaussians is Gaussian, with variances adding
Check that the coefficients squared sum to 1 at each step (variance preservation)

Papers

Ho et al. (2020) — Denoising Diffusion Probabilistic Models

Blogs

What are Diffusion Models? (Lilian Weng)Normal Distribution (Wikipedia)

Definition

A Markov chain is a sequence of random variables in which each state depends only on its immediate predecessor: $q(\mathbf{x}_t \mid \mathbf{x}_{t-1}, \ldots, \mathbf{x}_0) = q(\mathbf{x}_t \mid \mathbf{x}_{t-1})$ . This property factors the joint distribution into a product of transition kernels, and under mild conditions the chain converges to a unique stationary distribution.

In this topic

1Transition Probabilities
2Stationary Distribution
3Ergodicity and Mixing Time

1 of 3

Transition Probabilities

$q(\mathbf{x}_{1:T} | \mathbf{x}_0) = \prod_{t=1}^{T} q(\mathbf{x}_t | \mathbf{x}_{t-1})$

A Markov chain is fully described by its transition kernel $q(\mathbf{x}_t \mid \mathbf{x}_{t-1})$ , the probability of moving to state $\mathbf{x}_t$ from the current state. The Markov property asserts $q(\mathbf{x}_t \mid \mathbf{x}_{t-1}, \ldots, \mathbf{x}_0) = q(\mathbf{x}_t \mid \mathbf{x}_{t-1})$ : the future is conditionally independent of the past given the present. Multiplying the kernels along the chain gives the factorized joint $q(\mathbf{x}_{1:T} \mid \mathbf{x}_0) = \prod_t q(\mathbf{x}_t \mid \mathbf{x}_{t-1})$ . In DDPM each transition is a small Gaussian perturbation, so the entire forward trajectory is a product of Gaussian conditionals --- the representation that made closure so powerful in the previous topic.

Mathematical Intuition

The Markov property $P(x_t|x_{t-1},\ldots,x_0) = P(x_t|x_{t-1})$ means the future is conditionally independent of the past given the present. For the joint distribution, this yields $P(x_{0:T}) = P(x_0)\prod_{t=1}^T P(x_t|x_{t-1})$ , factoring an exponentially complex joint into $T$ simple transitions. In DDPM, each transition is a known Gaussian, making the entire forward process analytically tractable.

Example:

Write the joint distribution $q(\mathbf{x}_0, \mathbf{x}_1, \mathbf{x}_2)$ for a 2-step Markov chain and explain how the Markov property simplifies it.

2 of 3

Stationary Distribution

$\boldsymbol{\pi} = \boldsymbol{\pi} \mathbf{P} \quad \text{where } P_{ij} = q(x_j | x_i)$

A stationary (or invariant) distribution $\boldsymbol{\pi}$ is one that a single step of the chain leaves unchanged: $\boldsymbol{\pi} = \boldsymbol{\pi}\mathbf{P}$ , so $\boldsymbol{\pi}$ is a left eigenvector of the transition matrix with eigenvalue 1. For the DDPM forward process the stationary distribution is the standard normal $\mathcal{N}(\mathbf{0}, \mathbf{I})$ : after enough variance-preserving noising steps, any starting data converges to it regardless of where it began. This is exactly why the reverse process can begin from $\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ --- the forward chain has forgotten the input entirely, so pure noise is a faithful sample of the chain's endpoint.

Mathematical Intuition

A distribution $\pi$ is stationary if $\pi P = \pi$ — it is a left eigenvector of the transition matrix with eigenvalue 1. For the DDPM forward process, the stationary distribution is $\mathcal{N}(0,I)$ : applying the noising step $x_t = \sqrt{\alpha_t}x_{t-1} + \sqrt{1-\alpha_t}\epsilon$ to standard normal noise returns standard normal noise. The spectral gap $1 - \lambda_2$ of the transition operator controls how fast convergence occurs.

Example:

Consider a 2-state Markov chain with transition matrix $\mathbf{P} = \begin{pmatrix} 0.7 & 0.3 \\ 0.4 & 0.6 \end{pmatrix}$ . Find the stationary distribution.

3 of 3

Ergodicity and Mixing Time

$\|\mathbf{P}^t(\mathbf{x}_0, \cdot) - \boldsymbol{\pi}(\cdot)\|_{TV} \leq C \cdot e^{-\lambda t}$

An ergodic Markov chain converges to its stationary distribution from any starting state. The mixing time $t_{\text{mix}}$ is the number of steps needed for the chain to be within a small total variation distance of stationary. In diffusion models, the number of forward steps $T$ must exceed the mixing time to ensure $q(\mathbf{x}_T) \approx \mathcal{N}(\mathbf{0}, \mathbf{I})$ . Typical DDPM uses $T = 1000$ , which is far beyond the mixing time of the Gaussian noising chain.

Mathematical Intuition

The total variation distance $\|P^t(x_0, \cdot) - \pi\|_{TV} \leq C e^{-\lambda t}$ decays exponentially with rate $\lambda$ determined by the spectral gap. Mixing time $t_{\text{mix}} = O(1/\lambda)$ tells us how many steps are needed for the chain to forget its initial state. DDPM uses $T=1000$ steps with $\bar{\alpha}_T \approx 10^{-4}$ , ensuring $q(x_T) \approx \mathcal{N}(0,I)$ with negligible total variation error.

Example:

If a noise schedule has $\bar{\alpha}_{1000} = 0.0001$ , how close is $q(\mathbf{x}_{1000} | \mathbf{x}_0)$ to pure noise?

Theory Exercise

Problem:

Show that the reverse of a Markov chain is also a Markov chain. Specifically, given the forward chain $q(\mathbf{x}_{1:T} | \mathbf{x}_0) = \prod_{t=1}^{T} q(\mathbf{x}_t | \mathbf{x}_{t-1})$ , derive the form of the reverse transition $q(\mathbf{x}_{t-1} | \mathbf{x}_t, \mathbf{x}_0)$ using Bayes' theorem.

Hints:

Apply Bayes' theorem: q(x_{t-1} | x_t, x_0) = q(x_t | x_{t-1}, x_0) * q(x_{t-1} | x_0) / q(x_t | x_0)
By the Markov property, q(x_t | x_{t-1}, x_0) = q(x_t | x_{t-1})
All three distributions in Bayes' theorem are Gaussian---use the Gaussian product formula
The result is a Gaussian whose mean is a linear combination of x_t and x_0

Papers

Sohl-Dickstein et al. (2015) — Deep Unsupervised Learning using Nonequilibrium Thermodynamics

Blogs

Markov Chain (Wikipedia)What are Diffusion Models? (Lilian Weng)

Definition

The Kullback-Leibler divergence $D_{KL}(p \| q) = \mathbb{E}_p[\log(p/q)]$ measures the expected extra information needed to encode samples from $p$ using a code built for $q$ . It is non-negative, zero only when $p = q$ , and asymmetric --- $D_{KL}(p \| q) \neq D_{KL}(q \| p)$ in general. It is the objective used to train the diffusion reverse process.

In this topic

1Definition and Asymmetry
2Forward vs Reverse KL
3KL Between Gaussians

1 of 3

Definition and Asymmetry

$D_{KL}(p \| q) = \int p(\mathbf{x}) \log \frac{p(\mathbf{x})}{q(\mathbf{x})} \, d\mathbf{x} = \mathbb{E}_{p}\left[\log \frac{p(\mathbf{x})}{q(\mathbf{x})}\right]$

KL divergence measures the expected extra information needed to encode samples drawn from $p$ using a code optimized for $q$ . Here $p$ is the reference distribution the expectation is taken under, and the integrand $\log(p/q)$ is the pointwise log-likelihood ratio. It is non-negative by Gibbs' inequality and zero only when $p = q$ almost everywhere. Crucially it is asymmetric: $D_{KL}(p \| q) \neq D_{KL}(q \| p)$ in general, because the expectation weights the ratio by $p$ , not $q$ . It also blows up to infinity wherever $q = 0$ but $p > 0$ --- the code cannot describe an event the model thinks impossible --- which drives the mode behavior explored next.

Mathematical Intuition

KL divergence $D_{\text{KL}}(p\|q) = \mathbb{E}_p[\log(p/q)]$ measures the expected log-likelihood ratio under $p$ . It equals the extra bits needed to code $p$ -samples using a code optimized for $q$ . The asymmetry arises from the weighting: $D_{\text{KL}}(p\|q)$ is infinite if $q(x)=0$ where $p(x)>0$ (missing support), but finite if $p(x)=0$ where $q(x)>0$ (wasted support). This asymmetry determines whether training produces mode-covering or mode-seeking behavior.

Example:

Compute $D_{KL}(p\|q)$ and $D_{KL}(q\|p)$ where $p = \text{Bernoulli}(0.7)$ and $q = \text{Bernoulli}(0.5)$ .

2 of 3

Forward vs Reverse KL

$\text{Forward: } D_{KL}(p_{\text{data}} \| p_\theta) \quad \text{Reverse: } D_{KL}(p_\theta \| p_{\text{data}})$

The two orderings of KL induce opposite behaviors when a model $p_\theta$ cannot perfectly fit the data. Forward KL, $D_{KL}(p_{\text{data}} \| p_\theta)$ , takes the expectation under the data, so it penalizes the model for putting low probability anywhere the data has mass --- it forces mode-covering, spreading density to include every mode even if that covers empty regions between them. Reverse KL, $D_{KL}(p_\theta \| p_{\text{data}})$ , weights by the model, so it punishes mass where data is sparse --- encouraging mode-seeking collapse onto one mode. Diffusion models effectively minimize a forward-KL objective through the ELBO, a central reason they cover modes and generate diverse samples rather than collapsing like some GANs.

Mathematical Intuition

Forward KL $D_{\text{KL}}(p_{\text{data}}\|p_\theta)$ penalizes the model for having zero density where data exists, enforcing mode coverage at the cost of spreading mass into empty regions. Reverse KL $D_{\text{KL}}(p_\theta\|p_{\text{data}})$ penalizes the model for placing mass where no data exists, producing sharp but mode-collapsed fits. The ELBO in diffusion models optimizes a sum of forward KL terms at each timestep, which is why they achieve excellent mode coverage.

Example:

A bimodal target distribution has two peaks at $x = -3$ and $x = 3$ . A unimodal Gaussian model must choose its mean. What happens under forward vs reverse KL?

3 of 3

KL Between Gaussians

$D_{KL}(\mathcal{N}_1 \| \mathcal{N}_0) = \frac{1}{2}\left[\log\frac{|\boldsymbol{\Sigma}_0|}{|\boldsymbol{\Sigma}_1|} - d + \text{tr}(\boldsymbol{\Sigma}_0^{-1}\boldsymbol{\Sigma}_1) + (\boldsymbol{\mu}_0 - \boldsymbol{\mu}_1)^T \boldsymbol{\Sigma}_0^{-1}(\boldsymbol{\mu}_0 - \boldsymbol{\mu}_1)\right]$

When both distributions are Gaussian, KL divergence has a closed form, which is why it is so useful in diffusion training. For isotropic Gaussians $\mathcal{N}(\boldsymbol{\mu}_1, \sigma_1^2\mathbf{I})$ and $\mathcal{N}(\boldsymbol{\mu}_0, \sigma_0^2\mathbf{I})$ in $d$ dimensions the general matrix expression collapses to $\frac{d}{2}\left[\log\frac{\sigma_0^2}{\sigma_1^2} + \frac{\sigma_1^2}{\sigma_0^2} - 1\right] + \frac{\|\boldsymbol{\mu}_1 - \boldsymbol{\mu}_0\|^2}{2\sigma_0^2}$ . The first bracket measures variance mismatch and the second measures mean mismatch. In DDPM the variances are fixed by the noise schedule, so only the mean term survives --- and that term is a scaled squared error, which is exactly how the KL objective reduces to predicting the noise $\boldsymbol{\epsilon}$ .

Mathematical Intuition

For $\mathcal{N}(\mu_1, \sigma_1^2 I)$ vs $\mathcal{N}(\mu_0, \sigma_0^2 I)$ in $d$ dimensions: $D_{\text{KL}} = \frac{d}{2}[\log(\sigma_0^2/\sigma_1^2) + \sigma_1^2/\sigma_0^2 - 1] + \|\mu_1 - \mu_0\|^2/(2\sigma_0^2)$ . The first term measures variance mismatch; the second measures mean mismatch. In DDPM training, the variance is fixed by the schedule, so the KL reduces to a weighted MSE on means — which further reduces to predicting the noise $\epsilon$ .

Example:

Compute $D_{KL}(\mathcal{N}(1, 0.5) \| \mathcal{N}(0, 1))$ in 1D.

Theory Exercise

Problem:

Derive the relationship $D_{KL}(p \| q) = H(p, q) - H(p)$ where $H(p,q)$ is the cross-entropy and $H(p)$ is the entropy. Then explain why minimizing cross-entropy is equivalent to minimizing KL divergence when $p$ is fixed.

Hints:

Write out the definition of cross-entropy: H(p, q) = -E_p[log q(x)]
Write out the definition of entropy: H(p) = -E_p[log p(x)]
Subtract and compare with the KL definition
When p is fixed (the data distribution), H(p) is a constant

Papers

Kullback & Leibler (1951) — On Information and Sufficiency

Blogs

Kullback-Leibler Divergence (Wikipedia)What are Diffusion Models? (Lilian Weng)

Chapter 2: Mathematical Foundations

Chapter Overview

Chapter Roadmap

Gaussian Distributions

Markov Chains

KL Divergence

3 topics in this chapter

Gaussian Distributions

Definition

In this topic

Univariate and Multivariate Gaussian

Closure Under Linear Transformations

The Reparameterization Trick

Theory Exercise

Problem:

Hints:

Related Problems on PixelBank

Markov Chains

Definition

In this topic

Transition Probabilities

Stationary Distribution

Ergodicity and Mixing Time

Theory Exercise

Problem:

Hints:

Related Problems on PixelBank

KL Divergence

Definition

In this topic

Definition and Asymmetry

Forward vs Reverse KL

KL Between Gaussians

Theory Exercise

Problem:

Hints:

Related Problems on PixelBank