An optimal control perspective on diffusion-based generative modeling

TL;DR

This paper links stochastic optimal control theory with diffusion-based generative models, deriving new theoretical insights and a novel sampling method that outperforms existing approaches in various examples.

Contribution

It introduces a control-theoretic framework for diffusion models, deriving a Hamilton-Jacobi-Bellman equation and a new diffusion-based sampling method for unnormalized densities.

Findings

Derived a Hamilton-Jacobi-Bellman equation for SDE log-densities

Showed the evidence lower bound follows from control verification theorem

Developed a time-reversed diffusion sampler (DIS) that outperforms other methods

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.