An optimal control perspective on diffusion-based generative modeling
This work provides a theoretical framework for generative modeling, potentially benefiting researchers in machine learning and statistics, though it appears incremental in linking existing control theory to diffusion models.
The paper tackles the problem of connecting stochastic optimal control to diffusion-based generative models, deriving a Hamilton-Jacobi-Bellman equation for SDE marginals and developing a novel diffusion sampler for unnormalized densities, demonstrating that their method outperforms other approaches on multiple numerical examples.
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.