Path-Space ELBO in Diffusion Modeling
- Path-Space ELBO is a variational lower bound defined on continuous-time trajectories, extending traditional VAE bounds to stochastic processes.
- It leverages forward and reverse SDEs and employs the Hamilton–Jacobi–Bellman equation to integrate stochastic optimal control principles with generative modeling.
- This framework unifies diffusion-based generative models and deep Gaussian processes, improving score matching and posterior inference through innovative path-space regularizations.
The path-space Evidence Lower Bound (ELBO) is a variational lower bound on the marginal likelihood formulated not in terms of finite-dimensional random variables, but as a functional on measures over entire continuous-time trajectories—i.e., paths of a stochastic process. Path-space ELBOs arise naturally in the context of diffusion-based generative modeling, stochastic optimal control, and, more recently, in posterior inference for models such as deep Gaussian processes that utilize stochastic differential equations (SDEs) to define priors and generative mechanisms (Berner et al., 2022, Xu et al., 22 May 2026).
1. Forward and Reverse SDEs: Path-Space Laws
Diffusion-based generative models characteristically define a forward SDE (the inference SDE) on a time interval in via
where is the data-driven drift, is the volatility coefficient (typically independent of state), and is standard Brownian motion. The law of the path , denoted , lives on the function space .
The generative process is defined by reversing time in the SDE. The resulting process, , satisfies
0
with reversed drift
1
governed by the time-reversal of diffusion (Nelson–Haussmann–Föllmer theory). As the exact marginal score 2 is generally intractable, practical implementations substitute a (learned) approximation (Berner et al., 2022).
2. The Hamilton–Jacobi–Bellman Equation and Stochastic Control
The log-density of the time-reversed (generative) process admits a nonlinear PDE characterizing its evolution. Setting 3, one obtains 4, which evolves by the Hamilton–Jacobi–Bellman (HJB) equation:
5
where 6 and the terminal condition is 7, with 8 the prior density (Berner et al., 2022). The HJB equation directly links stochastic optimal control and diffusion-based generative modeling via the value functional over trajectories.
3. Verification Theorem: Derivation of the Path-Space ELBO
Given a family of admissible controls 9, define the controlled forward SDE:
0
Define the running cost as
1
and the terminal cost
2
The verification theorem shows that the value function at the process start is minimized over controls by the expected sum of these costs, with the minimizer given by a feedback control involving the log-density gradient. For any control 3,
4
The right-hand side is the path-space ELBO, representing a variational lower bound on the log marginal density at the endpoint (Berner et al., 2022).
4. Path-Space Kullback–Leibler Divergence and the ELBO Decomposition
Let 5 and 6 denote the path-space laws under a general and optimal control, respectively. Girsanov’s theorem gives
7
and taking expectations yields
8
Upon rearrangement, the path-space ELBO is
9
Maximizing the path-space ELBO is thus equivalent to minimizing the path-space KL divergence to the time-reversed (optimal) diffusion law (Berner et al., 2022).
5. Path-Space ELBOs in Diffusion-Based Generative Modeling and Deep Gaussian Processes
In diffusion probabilistic models, the path-space ELBO describes a fully continuous-time analogue of VAE bounds, with the bound defined on trajectory spaces. The running cost 0 penalizes deviation of a learned drift from the true marginal score (the pathwise optimal control), and the terminal cost employs the known generator prior density. Optimizing the ELBO by simulating sample paths and backpropagating through the running-cost integral recovers common losses (score-matching, denoising) used in diffusion models. The path-space ELBO also inherits the mode-seeking/mode-covering dichotomy from the directionality of the KL divergence. This framework unifies score-based diffusion models, VAEs, and Schrödinger-bridge methods under stochastic optimal control theory (Berner et al., 2022).
For deep Gaussian processes, recent formulations leverage Doob-bridged reference diffusions and Onsager–Machlup action regularization to derive strict path-space ELBOs and alternative objectives. For example, the "FFJORD log-det" and "OM-regularised CNF" ELBOs utilize probability-flow ODE variants and Onsager–Machlup actions as path priors. These formulations link the ELBO to negative log unnormalized path densities and, through the small-noise Freidlin–Wentzell LDP, to MAP path estimators under the formal tempered Doob-bridge posterior (Xu et al., 22 May 2026).
| Model Type | Path-Space ELBO | Notable Features |
|---|---|---|
| Diffusion-based generative models (Berner et al., 2022) | 1 | Unifies score models, VAEs, and Schrödinger bridges |
| Deep GP inference (Xu et al., 22 May 2026) | FFJORD/OM-regularised CNF | Deterministic samplers, Onsager–Machlup path regularization |
6. Assumptions, Implementation, and Empirical Perspectives
The path-space ELBO framework assumes Lipschitz continuity of drift and diffusion coefficients to ensure unique strong solutions and strictly positive densities in the corresponding Fokker–Planck PDEs. The volatility matrix 2 is often chosen independent of 3 (state) to simplify the Kolmogorov backward generator. The control function class is restricted to ensure Novikov’s condition and the applicability of the verification theorem. Boundary conditions require knowledge of the prior density at the trajectory endpoint (Berner et al., 2022).
In practical implementations, both stochastic and deterministic sampler paradigms are used. For instance, OM-Path (a deterministic sampler ODE) applies Onsager–Machlup regularization using bridge-marginal coefficients, and results empirically demonstrate domain-dependent advantages: on larger datasets with lower noise, pathwise regularized objectives outperform strict density-matching ELBOs, while stochastic formulations remain competitive in small-sample/high-noise settings (Xu et al., 22 May 2026).
7. Theoretical Insights and Unifying Principles
The path-space perspective on the ELBO reveals deep connections among generative modeling, stochastic control, and variational inference. Notably:
- The ELBO on path space is the negative of a KL divergence between trajectory measures.
- Formulations such as FFJORD log-det or OM-regularised ELBOs connect instantaneous change-of-variables with path-prior regularization.
- The Onsager–Machlup functional acts as the rate functional for the small-noise Freidlin–Wentzell LDP of the reference diffusion, making OM-regularized training amortized MAP estimation of the true path posterior rather than a variational density fit.
- The overall framework unifies Doob SDE anchoring, probability-flow ODEs, classical control theory, and Schrödinger bridge formulations.
A plausible implication is that path-space ELBOs provide a rigorous foundation to extend variational methods beyond densities over finite-dimensional variables, enabling principled learning and inference in diffusion-based and bridge-based generative models across domains (Berner et al., 2022, Xu et al., 22 May 2026).