Diffusion Models: Forward and Reverse Processes

A diffusion model defines two coupled stochastic processes over a sequence of latent variables $x_0, x_1, \dots, x_T$, where $x_0$ is a real data point (an image, say) and $x_T$ is approximately standard Gaussian noise. The forward process $q(x_{1:T} \mid x_0)$ is a fixed Markov chain that adds a small amount of Gaussian noise at each step according to a variance schedule $\beta_1, \dots, \beta_T \in (0,1)$:

This process has no learned parameters — it is a fixed corruption schedule, analogous to annealing a physical system toward thermal equilibrium. The reverse process is what we actually learn: a Markov chain $p_\theta(x_{0:T})$ that starts from $p(x_T) = \mathcal{N}(0, I)$ and applies learned Gaussian transitions

to walk back from noise to data. Because each forward step adds only a small perturbation (for large $T$ and small $\beta_t$), the reverse transition is also approximately Gaussian — this is the key mathematical fact that makes the reverse process learnable with a simple neural network rather than an intractable arbitrary distribution. The crucial computational trick, using the reparameterization trick, is that the marginal $q(x_t \mid x_0)$ has a closed form: letting $\alpha_t = 1-\beta_t$ and $\bar\alpha_t = \prod_{s=1}^t \alpha_s$,

This means we can sample a noised version of any data point at any timestep $t$ in a single step, with no need to iterate through $t-1$ intermediate steps. Training therefore never has to simulate the chain — it just picks a random $t$, generates $x_t$ directly, and asks a network to recover the noise $\epsilon$ that was added. This is what makes diffusion training embarrassingly parallel and efficient, distinguishing it sharply from the sequential nature of sampling.

Step 1 — Choose a noise schedule. Fix $T$ (commonly 1000) and a schedule $\{\beta_t\}_{t=1}^T$, e.g. linear from $10^{-4}$ to $0.02$, or a cosine schedule. Precompute $\alpha_t = 1-\beta_t$ and $\bar\alpha_t = \prod_{s\le t}\alpha_s$.

Step 2 — Training. For each minibatch, sample real data $x_0$, sample a random timestep $t \sim \text{Uniform}\{1,\dots,T\}$, sample noise $\epsilon \sim \mathcal N(0,I)$, form $x_t$ using the closed-form reparameterization above, and train a network $\epsilon_\theta(x_t, t)$ to predict $\epsilon$ via mean-squared error. The network is typically a U-Net with time-step embeddings (often sinusoidal, as in Transformers) injected at every resolution level.

Step 3 — Sampling. Start from $x_T \sim \mathcal N(0, I)$. For $t = T, \dots, 1$, use $\epsilon_\theta(x_t, t)$ to estimate the reverse mean $\mu_\theta(x_t, t)$, then sample $x_{t-1}$ from the Gaussian reverse transition (adding fresh noise except at the final step). After $T$ iterations, $x_0$ is a generated sample. ```

# ---------- Precompute schedule ---------- function PRECOMPUTE_SCHEDULE(T, beta_min, beta_max): beta[1..T] <- LINSPACE(beta_min, beta_max, T) alpha[1..T] <- 1 - beta alpha_bar[1..T] <- CUMULATIVE_PRODUCT(alpha) # alpha_bar[t] = prod_{s<=t} alpha[s] return beta, alpha, alpha_bar

# ---------- Training ---------- function TRAIN(data_distribution, T, beta, alpha, alpha_bar, epsilon_theta, num_steps): for step in 1..num_steps: x0 <- SAMPLE_BATCH(data_distribution) # shape [B, ...] t <- SAMPLE_UNIFORM_INT(1, T, size=B) # one timestep per example eps <- SAMPLE_GAUSSIAN(shape=x0.shape) # target noise sqrt_ab <- SQRT(alpha_bar[t]) # broadcast over batch sqrt_1mab <- SQRT(1 - alpha_bar[t]) x_t <- sqrt_ab x0 + sqrt_1mab eps # closed-form forward sample eps_pred <- epsilon_theta(x_t, t) # network's noise estimate loss <- MEAN_SQUARED_ERROR(eps_pred, eps) UPDATE_PARAMETERS(epsilon_theta, GRADIENT(loss)) return epsilon_theta

# ---------- Sampling (ancestral sampling, DDPM) ---------- function SAMPLE(T, beta, alpha, alpha_bar, epsilon_theta): x <- SAMPLE_GAUSSIAN(shape=DATA_SHAPE) # x_T ~ N(0, I) for t in T, T-1, ..., 1: eps_pred <- epsilon_theta(x, t) coef1 <- 1 / SQRT(alpha[t]) coef2 <- beta[t] / SQRT(1 - alpha_bar[t]) mu <- coef1 (x - coef2 eps_pred) # posterior mean estimate if t > 1: sigma <- SQRT(beta[t]) # or posterior variance, see below z <- SAMPLE_GAUSSIAN(shape=DATA_SHAPE) else: sigma <- 0 z <- 0 x <- mu + sigma z return x # approx sample from data distribution ``` The training loop's correctness hinges entirely on the closed-form line `x_t <- sqrt_abx0 + sqrt_1mabeps`: it is exactly* a sample from $q(x_t\mid x_0)$ by the reparameterization trick applied to a Gaussian, so no chain simulation, and no bias, is introduced. The sampling loop's mean formula `mu` comes from Bayes' rule applied to the forward process (derived in the next section) — it is the mean of $q(x_{t-1}\mid x_t, x_0)$ with $x_0$ replaced by its estimate implied by $\epsilon_\theta$.

Time complexity. Training one step costs $O(C)$ where $C$ is the cost of one forward/backward pass of the denoising network (typically a U-Net), independent of $T$ — this is because the closed-form reparameterization lets us jump directly to any $x_t$ without simulating intermediate steps. Sampling one image costs $O(T \cdot C)$, since the reverse chain must be simulated sequentially, one network evaluation per timestep; this is the well-known efficiency bottleneck of diffusion models relative to single-pass generators like GANs or VAEs. Accelerated samplers (e.g. DDIM, higher-order solvers) reduce the effective number of sampling steps to $O(T')$ with $T' \ll T$ (often 20–50), by treating the reverse process as a discretized ODE/SDE rather than a strict Markov chain, but this is a refinement built on the same forward/reverse formalism.

Space complexity. $O(P)$ for the network parameters $\theta$, plus $O(D)$ per sample for the data dimensionality $D$ (a single image tensor is stored and overwritten at each denoising step; no need to store the whole trajectory unless it is being visualized).

Correctness — the variational bound. The reverse process is trained to make $p_\theta(x_0)$ close to the true data distribution. Since $p_\theta(x_0) = \int p_\theta(x_{0:T})\, dx_{1:T}$ is intractable to compute directly, training instead maximizes a variational lower bound (an ELBO), exactly as in a hierarchical VAE:

This bound follows from Jensen's inequality applied to $\log p_\theta(x_0) = \log \mathbb E_{q(x_{1:T}|x_0)}\left[\frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)}\right]$, the same derivation pattern used for the standard ELBO in variational autoencoders. Because the forward posterior $q(x_{t-1}\mid x_t, x_0)$ is Gaussian (a consequence of $q$ being a linear-Gaussian chain, using Bayes' rule on jointly Gaussian variables), each KL term between two Gaussians has a closed form, reducible to a squared difference of means. Substituting the reparameterized forms of $x_t$ shows that this difference of means is proportional to $\|\epsilon - \epsilon_\theta(x_t,t)\|^2$ — precisely the training loss used in the pseudocode above. In other words, minimizing the simple noise-prediction MSE loss is provably minimizing (a reweighted version of) a true variational bound on data log-likelihood, not an ad hoc heuristic.

Honest tradeoffs. The forward process being fixed (not learned) sacrifices some flexibility for tractability and stable training — there is no adversarial game and no mode collapse in the GAN sense. The price is slow sampling ($T$ sequential network calls), high compute cost, and sensitivity to the noise schedule and number of steps: too few steps make the Gaussian reverse-step assumption inaccurate; too many steps waste compute. The variational bound is also a lower bound on likelihood, so a model can score well on the training objective while still producing samples with subtly wrong tails or texture — an accepted looseness in exchange for tractable, stable training.

Consider a toy 1-D "image" that is literally a single pixel value, $x_0 = 2.0$, so we can hand-trace the arithmetic. Use $T=4$ with a simple linear schedule for illustration.

Training-time forward sampling (closed form). Suppose training draws $t=3$ and samples noise $\epsilon = 1.0$ (a fixed value here for tracing, instead of random, to make the arithmetic checkable). Then

The network $\epsilon_\theta(x_3=2.124,\ t=3)$ is asked to output a prediction close to $1.0$; suppose it (imperfectly) outputs $\hat\epsilon = 0.8$. The squared-error loss contribution is $(1.0-0.8)^2 = 0.04$, and this drives a gradient update pushing $\epsilon_\theta$'s prediction upward, toward $1.0$, at input $(2.124, 3)$.

Sampling-time reverse walk. Now suppose after training, $\epsilon_\theta$ is reasonably accurate. Start from noise $x_4 = 1.5$ (drawn from $\mathcal N(0,1)$, here fixed for the trace). At $t=4$: suppose $\epsilon_\theta(1.5, 4) = 0.55$. Compute

Adding noise $\sigma_4 z = \sqrt{0.4}\times(-0.2) = -0.126$ gives $x_3 = 1.597 - 0.126 = 1.471$. Repeating this for $t=3,2,1$ (each time re-evaluating $\epsilon_\theta$ at the new $x_t$ and the corresponding $t$) walks the chain back toward values clustering near the original data mode $x_0=2.0$ — not exactly recovering $2.124$ or $2.0$ because sampling is stochastic and the network is imperfect, but statistically consistent with the data distribution the network was trained on. This mirrors the real algorithm: training touches one random $t$ per example using the closed form; sampling must traverse every $t$ from $T$ down to $1$ sequentially.

The mathematical seed of diffusion models is the observation, borrowed from non-equilibrium statistical physics, that a complex distribution can be reached by slowly perturbing a simple one and that the perturbation process is invertible if simulated finely enough. This connects diffusion models to two adjacent lineages: variational inference (making the reverse chain trainable via an ELBO) and score matching / Langevin dynamics (an alternative, closely related view developed independently and unified with diffusion models in later work, which is the bridge to Denoising and Score-Based Models).

Diffusion models now underpin most state-of-the-art image generation systems, including Stable Diffusion, Midjourney's backend, DALL·E 2/3, and Google's Imagen, all of which train a U-Net-style $\epsilon_\theta$ (often operating in a compressed latent space rather than pixel space, for efficiency) using exactly the forward/reverse machinery described here. The same framework generalizes beyond images: diffusion is used for audio and speech synthesis, video generation, 3D shape and molecule generation (important in drug discovery, where the data lives on non-Euclidean manifolds and the forward process must be adapted accordingly), and even as a component of reinforcement learning policies (diffusion policies for robotic control). The core forward-noising/reverse-denoising abstraction is also the conceptual bridge into score-based generative modeling, where the reverse process is reinterpreted as integrating an estimated score function — a connection formalized through stochastic differential equations and central to understanding modern accelerated samplers.

Ask the tutor: How does the score-based (score matching + Langevin dynamics) view of diffusion relate mathematically to the noise-prediction $\epsilon_\theta$ network described here?