An optimal control perspective on diffusion-based generative modeling

Abstract: We establish a connection between stochastic optimal control and generative models based on stochastic differential equations (SDEs), such as recently developed diffusion probabilistic models. In particular, we derive a Hamilton-Jacobi-Bellman equation that governs the evolution of the log-densities of the underlying SDE marginals. This perspective allows to transfer methods from optimal control theory to generative modeling. First, we show that the evidence lower bound is a direct consequence of the well-known verification theorem from control theory. Further, we can formulate diffusion-based generative modeling as a minimization of the Kullback-Leibler divergence between suitable measures in path space. Finally, we develop a novel diffusion-based method for sampling from unnormalized densities -- a problem frequently occurring in statistics and computational sciences. We demonstrate that our time-reversed diffusion sampler (DIS) can outperform other diffusion-based sampling approaches on multiple numerical examples.

• The paper derives the evidence lower bound (ELBO) using a Hamilton-Jacobi-Bellman equation, linking control theory with diffusion-based generative modeling.
• It reformulates generative modeling as KL divergence minimization on path space, offering a clear theoretical framework and improved loss function insights.
• The novel time-reversed diffusion sampler (DIS) demonstrates superior performance in numerical experiments on challenging sampling problems.

Optimal Control and Diffusion-Based Generative Modeling

This paper establishes a connection between stochastic optimal control and diffusion probabilistic models, demonstrating that the evolution of the log-densities of SDE marginals is governed by a Hamilton-Jacobi-Bellman (HJB) equation. This perspective allows for the transfer of methods from optimal control theory to generative modeling, including deriving the evidence lower bound (ELBO) from the verification theorem and formulating diffusion-based generative modeling as a minimization of the Kullback-Leibler (KL) divergence between suitable measures in path space. The paper also introduces a novel time-reversed diffusion sampler (DIS) for sampling from unnormalized densities.

HJB Equation and ELBO Derivation

The paper leverages the connection between SDEs and partial differential equations (PDEs) in stochastic optimal control. By applying the Hopf-Cole transformation, the authors demonstrate that the time-reversed log-density of the diffusion process satisfies an HJB equation (Lemma 2.1). This connection enables the derivation of the ELBO, a fundamental objective in diffusion-based generative modeling, directly from control theory principles (Corollary 3.1). The ELBO is shown to be equivalent to the negative control costs, providing a lower bound on the negative log-likelihood of the generative model. The variational gap is explicitly stated in Proposition 3.1.

Path Space Perspective and KL Divergence

A key contribution of the paper is the interpretation of the variational gap in diffusion models as a KL divergence between measures on the path space of continuous-time stochastic processes (Section 3.3). This path space perspective offers an intuitive representation of the objective and the variational gap, allowing for the consideration of alternative divergences, such as the log-variance divergence, which can lead to improved loss functions and algorithms. The optimal path space measure is defined via a work functional and a Radon-Nikodym derivative (Proposition 3.1).

Figure 1: Illustration of importance sampling in path space. The left panel shows a histogram of the first component of the controlled process $X^u_T$ at terminal time $T$ compared to the corresponding marginal of the target density $\rho / \mathcal{Z}$.

Time-Reversed Diffusion Sampler (DIS)

The paper introduces a novel diffusion-based method, the time-reversed diffusion sampler (DIS), for sampling from unnormalized densities, a common problem in statistics and computational sciences (Section 4). DIS minimizes the reverse KL divergence between a controlled SDE and the reverse-time SDE, offering more flexibility in choosing the initial distribution and reference SDE compared to related approaches. The control objective is specified in Corollary 4.1, and the connection to Schr\"odinger half-bridges is discussed in Section 4.1.

Figure 2: Illustration of the DIS algorithm for the double well example, showing trajectories and histograms at initial and terminal times, along with a KDE density estimation of a 2d marginal.

Numerical Validation

The authors demonstrate the effectiveness of DIS through numerical experiments on several challenging problems (Section 5). DIS outperforms the Path Integral Sampler (PIS) on tasks such as estimating normalizing constants, expectations, and standard deviations for Gaussian mixture models, Funnel distributions, and double-well potentials. The numerical experiments highlight the practical relevance of the theoretical connections established in the paper.

Figure 3: Comparison of DIS and PIS in computing the log-normalizing constant $\log \mathcal{Z}$ using different numbers of Euler-Maruyama scheme steps.

Figure 4: Comparison of DIS and PIS in estimating expectations and average standard deviation, using different numbers of Euler-Maruyama scheme steps.

Relation to Denoising Score Matching

The paper also elucidates the connection between the ELBO and the denoising score matching objective commonly used for training continuous-time diffusion models (Section 3.4). By reparameterizing the generative model, the authors show that the ELBO can be expressed as a denoising score matching objective, providing a theoretical justification for this widely used training technique. The authors note that the setting in this section requires access to samples from the data distribution $\mathcal{D}$ in order to simulate the process $Y$.

Conclusion

This paper offers valuable insights into the theoretical foundations of diffusion models by connecting them to stochastic optimal control. The path space perspective and the novel DIS algorithm contribute to the advancement of generative modeling and sampling techniques. Numerical results demonstrate the practical benefits of the proposed approach, highlighting its potential for solving challenging problems in various scientific domains. Future research directions include exploring alternative divergences on path space, extending the framework to Schr\"odinger bridges, and developing efficient numerical methods for solving the HJB equation.