Technical Summary: A Mathematical Introduction to Diffusion Models

Problem Statement

The paper addresses the fundamental challenge of sampling from complex, high-dimensional probability distributions $\pi$ (the target data distribution) when direct sampling is infeasible. In many modern applications, $\pi$ is known only up to a normalizing constant or only through empirical data samples, rather than an explicit tractable density formula. The core difficulty lies in constructing a random process that converges to $\pi$ efficiently, particularly when the target is multimodal, singular, or defined on discrete spaces.

The paper frames this problem through the lens of indirect sampling: instead of sampling the target directly, one constructs a stochastic process that is easy to simulate and whose distribution gradually drifts toward the target. The work focuses on two primary mechanisms for this drift:

1. Langevin Dynamics: Continuous-time Markov processes driven by the gradient of the log-density (score) and Brownian motion.
2. Diffusion Models: A framework that adds noise to data to form a path of smoother laws, then learns to reverse this process using score estimation.

Methodology

The notes proceed in five distinct "movements," building a rigorous mathematical bridge from classical sampling theory to modern generative modeling.

1. Sampling Language and Langevin Dynamics

The foundation is laid by defining the sampling problem via discrepancies (Total Variation, Wasserstein-2, and KL divergence). The authors introduce Overdamped Langevin Diffusion:
$$dX_t = -\nabla U(X_t)dt + \sqrt{2}dB_t$$
where $\pi(x) \propto e^{-U(x)}$.

• Convergence Mechanism: The paper proves that the KL divergence $D_{KL}(q_t \| \pi)$ acts as a Lyapunov function. Its time derivative is the negative of the relative Fisher information ($-FI$). Under a Log-Sobolev Inequality (LSI), this yields exponential convergence rates.
• Discretization (ULA): The Unadjusted Langevin Algorithm (ULA) is introduced as the Euler–Maruyama discretization. The paper analyzes its discretization bias, showing that while the chain mixes, its stationary distribution $\pi_h$ differs from $\pi$ by an error of order $O(h)$.
• Correction (MALA): The Metropolis-adjusted Langevin algorithm (MALA) is presented as a method to remove this bias via an accept-reject step, restoring exact stationarity at the cost of requiring density evaluations.

2. Continuous-Time Score-Based Diffusion Models

The framework shifts to scenarios where the target density is unknown, but data is available.

• Forward Noising: Data $X_0$ is corrupted by a continuous-time Gaussian process (e.g., Variance-Exploding or Variance-Preserving SDEs), creating a path of marginals $p_t$.
• Tweedie's Identity: A central identity links the score $\nabla \log p_t(x)$ to the posterior mean (denoiser) of the clean data:
  $$\nabla \log p_t(x) = \frac{1}{\sigma_t^2} \mathbb{E}[a_t X_0 - X_t | X_t = x]$$
  This justifies training neural networks to predict the noise (or clean data) as a proxy for learning the score.
• Reverse Dynamics: Two equivalent reverse-time processes are derived:
  1. Reverse SDE: A stochastic process with a drift term containing the score, $dY_s = [-f_{T-s} + g_{T-s}^2 \nabla \log p_{T-s}]ds + g_{T-s} d\bar{B}_s$.
  2. Probability Flow ODE: A deterministic ODE, $\frac{dX_t}{dt} = f_t(X_t) - \frac{1}{2}g_t^2 \nabla \log p_t(X_t)$, which preserves the same marginals as the SDE but allows for exact likelihood computation via the continuity equation.

3. Stochastic Localization and Polchinski Flow

The paper provides alternative theoretical perspectives on the same Gaussian channel:

• Stochastic Localization: Views the process as a Bayesian inference problem where the posterior mean of the hidden signal evolves as a martingale. This perspective yields powerful covariance identities (e.g., the posterior covariance dissipates over time), which are crucial for later error bounds.
• Polchinski Flow: Interprets the smoothing of the density as a renormalization group flow. The effective potential $U_t = -\log p_t$ satisfies a nonlinear PDE (the Polchinski equation), connecting diffusion models to statistical physics and coarse-graining.

4. Discretization and Error Analysis

The continuous dynamics are discretized into implementable samplers (DDPM). The paper provides a rigorous error decomposition for the KL divergence between the true data distribution (or a smoothed version) and the sampler output:
$$D_{KL}(p_1 \| \hat{p}_1) \leq \underbrace{\epsilon^2_{init}}_{\text{Initialization}} + \underbrace{\epsilon^2_{disc}}_{\text{Discretization}} + \underbrace{c \sum \eta_k \epsilon^2_{k,score}}_{\text{Score Estimation}}$$

• Initialization Error: Arises from starting the reverse process at a high-noise level rather than the exact terminal distribution.
• Discretization Error: Analyzed via Euler–Maruyama and refined using Hessian control. Instead of assuming a global Lipschitz constant for the score, the paper uses posterior covariance budgets (derived from stochastic localization) to bound the error, yielding nearly linear dependence on dimension $d$.
• Score Estimation Error: The error from using a learned score $s_t$ instead of the true score $s^\star_t$.
• High-Accuracy Sampling: The paper introduces First-Order Rejection Sampling (FORS) as a score-only correction mechanism (analogous to MALA) that achieves high accuracy with polylogarithmic dependence on the target error $\epsilon$, avoiding the polynomial dependence of standard Euler discretization.

5. Discrete Diffusion Models

The framework is extended to finite state spaces (e.g., text, molecules) where gradients and Brownian motion are undefined.

• Mechanism: Forward corruption is modeled as a Continuous-Time Markov Chain (CTMC) (e.g., masked diffusion).
• Reverse Dynamics: The reverse process is also a CTMC, with rates determined by ratio scores $s^\star_t(x, y) = p_t(y)/p_t(x)$ rather than gradients.
• Training: The objective is score-entropy, a Bregman divergence derived from the path-space relative entropy, which can be optimized using supervised learning on clean-token prediction.
• Sampling: Algorithms like $\tau$-leaping (approximate) and Uniformization (exact) are analyzed. Uniformization is shown to eliminate discretization error, leaving only initialization and score-estimation errors.

6. Inference-Time Control

The final section addresses modifying the sampler to satisfy conditions or maximize rewards without retraining.

• Reward Tilting: The target distribution is reweighted as $p_\beta(x) \propto p_0(x)e^{\beta r(x)}$.
• Guidance: The paper derives that the score of the tilted distribution is the base score plus a correction term $\nabla \log h^\beta_t$, where $h^\beta_t$ is a posterior expectation of the reward. This unifies Classifier Guidance and Classifier-Free Guidance under a single mathematical framework.
• Path-Space Control: The problem is framed as a Doob transform of the base path law, linking inference-time control to stochastic control theory.

Key Results and Contributions

1. Unified Error Analysis: The paper provides a comprehensive, proof-oriented error decomposition for diffusion models that separates initialization, discretization, and score estimation errors. It demonstrates that under weak assumptions (finite second moments), the discretization error scales nearly linearly with dimension $d$ via covariance budgets, rather than requiring strong convexity or global Lipschitz conditions.
2. High-Accuracy Guarantees: It establishes that First-Order Rejection Sampling (FORS) and Uniformization (for discrete models) can achieve high-accuracy sampling with complexity depending polylogarithmically on the target error $\epsilon$, significantly improving upon standard Euler-based bounds.
3. Theoretical Unification: The work connects diffusion models to Stochastic Localization (martingale theory), Polchinski Flow (renormalization group), and Optimal Transport (Wasserstein gradient flows), providing a deeper geometric and analytic understanding of why these models work.
4. Discrete State Spaces: It rigorously extends diffusion theory to finite state spaces using CTMCs and ratio scores, offering a complete error analysis for masked diffusion and discrete samplers.
5. Inference-Time Control: It formalizes guidance and reward tilting as changes of measure (Doob transforms), clarifying the mathematical relationship between the base model, the reward function, and the resulting guided dynamics.

Significance and Claims

The paper claims to provide a proof-oriented introduction suitable for beginning graduate students with a background in probability but no prior exposure to stochastic differential equations (SDEs) or diffusion models. Its significance lies in:

• Rigorous Foundations: Moving beyond heuristic explanations to provide full proofs for core definitions (Fokker-Planck, Tweedie's identity, reverse SDE) and representative estimates for convergence.
• Bridging Gaps: Connecting classical MCMC theory (Langevin, MALA) with modern generative modeling (DDPM, Score-Based), showing that diffusion models are essentially a specific, learned implementation of reverse-time sampling dynamics.
• Clarifying Assumptions: Explicitly stating the assumptions required for convergence (e.g., Log-Sobolev inequalities, smoothness) and distinguishing between "research-level theorems" (stated with roadmaps) and "representative estimates" (proved under simplifying assumptions).
• Self-Contained Toolkit: The notes serve as a self-contained reference for the analytic tools (Itô calculus, Girsanov theorem, Gaussian identities) necessary to understand and analyze diffusion models, avoiding the need to consult disparate literature for basic machinery.

The authors explicitly state that the notes do not survey the entire landscape of generative models (excluding flow matching, consistency models, etc., in detail) but focus on the specific mathematical mechanisms of diffusion and Langevin dynamics, treating learned score error as an input to sampler guarantees rather than a statistical learning problem.