PaperlayerView exact revision on arXiv

Source-linked AI summary

An optimal control perspective on diffusion-based generative modeling

Julius Berner, Lorenz Richter, Karen Ullrich

arXiv:2211.01364v3cs.LGmath.OCstat.ML

TL;DR

The paper addresses how SDE-based diffusion models can be understood and extended through stochastic optimal control. It derives an HJB-based framework, obtains ELBO and path-space interpretations, and develops a time-reversed diffusion sampler for unnormalized densities. The sampler shows significant numerical advantages over existing diffusion-based approaches across challenging, high-dimensional problems.

  • Problem

    The paper studies how to connect SDE-based diffusion models with stochastic optimal control and how to sample from unnormalized densities when only the target density is available.

  • Method

    The authors derive an HJB equation for time-reversed diffusion log-densities, interpret objectives using path-space KL divergences, and construct the time-reversed diffusion sampler DIS.

  • Results

    DIS provides significant numerical advantages over existing diffusion-based sampling approaches across a series of challenging, high-dimensional problems.

  • Takeaways & Limitations

    The framework transfers tools from optimal control to generative modeling and supports diffusion-based sampling of unnormalized densities.

  • Takeaways & Limitations

    The approach requires the target density up to a normalization constant, and training may amortize only when many samples are needed; control approximation can be sensitive to hyperparameters.

Abstract

from arXiv · show

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.

1 Introduction

The paper connects diffusion models with stochastic optimal control through an HJB equation, yielding theoretical interpretations and practical sampling methods. It develops path-space and ELBO perspectives and introduces diffusion-based sampling for unnormalized densities.

  • Theoretical connection: The authors connect SDE-based diffusion models to stochastic optimal control through an HJB equation governing the time-reversed log-density.This connection also links generative modeling to PDEs and reinforcement learning.
  • Theoretical connection: The HJB formulation derives the diffusion-model ELBO from control-theoretic principles and interprets it through KL divergences between path-space measures.The path-space perspective also motivates alternative divergences such as the log-variance divergence.
  • Sampling: The sampling strategy significantly outperforms related diffusion-based approaches across multiple numerical examples.The paper reports this advantage across a number of relevant examples rather than a single benchmark.
  • Sampling: The control perspective yields a novel diffusion-based method for sampling from unnormalized densities.The method transfers noise schedules, integrators, and other generative-modeling techniques to sampling.
  • Context: The work builds on prior SDE formulations of diffusion models and extends existing links among diffusion modeling, optimal control, and path-space methods.Related work includes diffusion-based sampling, Schrödinger bridges, and prior optimal-control perspectives.

2 SDE-based generative modeling as an optimal control problem

The paper formulates SDE-based generative modeling as a stochastic optimal-control problem by showing that reverse-time log-densities satisfy an HJB equation. This connection yields the ELBO, a path-space KL interpretation, and a control-based training perspective.

  • SDE-based generative modeling: The SDE-based modeling problem learns an initial condition and coefficient functions so the terminal distribution approximates the data distribution.Although the initial condition, drift, and diffusion can generally be learned, practical methods often learn only the drift.
  • PDE perspective: HJB equation for log-density: The reverse-time log-density satisfies an HJB equation, connecting diffusion generative models to stochastic optimal control.The paper obtains the HJB equation by applying the Hopf–Cole transformation to the time-reversed density equation.
  • Optimal control formulation: The verification theorem turns the HJB formulation into an optimal-control objective whose unique minimizer is u∗ := −σ⊤∇V.The controlled process uses terminal and running costs, and the stated control attains the unique minimum under the paper’s assumptions.
  • ELBO derivation: The resulting ELBO equals negative control costs and becomes a lower bound on the generative model’s negative log-likelihood.Equality holds for u∗ := σ⊤∇log⃗pX; maximizing the ELBO can be implemented by parameterizing u with a neural network.
  • Path space perspective: KL divergence in continuous time: The variational gap is represented as a KL divergence between measures on continuous trajectories, enabling an optimal path-space change of measure.The optimal path-space measure can also be interpreted through Doob’s h-transform or Bayes’ rule, with the controlled measure as posterior and the uncontrolled measure as prior.
  • Connection to diffusion-model training: The reparameterized control objective recovers the objective used to train continuous-time diffusion models and can reduce to denoising score matching.This reduction requires access to samples from the data distribution to simulate the inference process.

3 Sampling from unnormalized densities

The paper introduces the time-reversed diffusion sampler (DIS) for sampling from densities known only up to a normalizing constant. DIS minimizes a reverse-KL-based control objective and uses diffusion-model machinery while allowing flexible initial distributions and reference SDEs.

  • Problem setting: The target density is ρ/Z, where ρ is known but the normalizing constant Z is intractable and samples may be unavailable.This setting occurs in applications such as Bayesian statistics and computational physics.
  • DIS method: The time-reversed diffusion sampler (DIS) is a novel diffusion-based variational method for sampling from the target density.It is obtained by interchanging the roles of the generative and controlled processes in the optimal-control derivation.
  • DIS objective: DIS minimizes a reverse KL divergence objective involving the controlled path-space measure and the target distribution.The objective includes the reference density term and log Z, together with a path-space KL term.
  • Sampling guarantee: With X^u_0 approximately distributed as Y_T, the minimizing control u∗ := σ⊤∇log p_Y yields X^{u∗}_T ∼ D.A Gaussian initial distribution can be used when f and σ make p_YT approximately standard normal, at the cost of an irreducible prior loss.
  • Relation to Schrödinger bridges: Unlike the Schrödinger half-bridge, DIS requires an initial process distribution approximately matching Y_T rather than starting from an arbitrary fixed point.The two approaches also differ in their running and terminal costs and in how the reference function is chosen.
  • Relation to prior samplers: DIS can outperform previous diffusion-based state-of-the-art samplers, while transferring noise schedules, integrators, and related generative-model techniques to sampling.The paper also reports that its optimal control can be more numerically stable and may improve sample quality and normalizing-constant estimates.

4 Numerical examples

The numerical examples evaluate DIS against PIS for estimating normalizing constants, expectations, and standard deviations across multimodal and high-dimensional targets. DIS improves over PIS across the considered tasks and settings, while alternative divergences can further improve performance.

  • Experimental setup: DIS and PIS are evaluated on normalizing constants, expectations, and standard deviations using multimodal examples including GMM, Funnel, and double-well targets.The experiments use path-space importance sampling for log Z and Monte Carlo estimates for sample statistics.
  • Normalizing constants: DIS clearly outperforms PIS for log Z estimation on GMM and Funnel examples and offers a slight improvement on the double-well example.Figure 4 reports medians and interquartile ranges over 10 training seeds with 100, 200, 400, or 800 Euler-Maruyama steps.
  • Overall comparison: DIS improves over PIS for all considered tasks, including estimation of normalizing constants, expectations, and standard deviations.The comparison uses the same training techniques but different objectives for the two methods.
  • High-dimensional sampling: DIS converges to a better approximation and remains effective for multimodal, high-dimensional distributions, including d = 50 with 32 well-separated modes.This result holds even when the PIS reference process is already correct in d − w dimensions of the double-well example.
  • Practical trade-offs: DIS typically samples faster than classical MCMC and SMC methods, but its training cost may amortize only when many samples are needed.Approximating the optimal control can also be challenging and sensitive to hyperparameter choices.

5 Conclusion and outlook

The paper connects diffusion models with optimal control and transfers this perspective to ELBO derivation, path-space objectives, and sampling from unnormalized densities. The resulting framework yields numerical advantages over existing diffusion-based samplers, while its direct divergence objectives require access to the target density.

  • Contributions: The work connects diffusion models with optimal control and reinforcement learning, enabling transfer of tools between these fields.The connection is presented as a source of theoretical insights and practical algorithms.
  • Theoretical perspective: The framework derives the continuous-time ELBO from control arguments and interprets diffusion models through measures on path space.The path-space view also motivates improved loss functions.
  • Sampling: The framework extends diffusion modeling to sampling from unnormalized densities and shows numerical advantages over existing diffusion-based sampling approaches.The conclusion describes these advantages across challenging, high-dimensional problems.
  • Outlook: PDE techniques, tensor-based methods, SDE-based solvers, tensor trains, and policy iteration are identified as directions for approximating score functions and related equations.The paper also suggests Gaussian pretraining through linear-quadratic control and adding running costs to incorporate domain knowledge.

A.1 Setting

The setting defines densities and path-space laws for continuous stochastic processes and imposes regularity assumptions ensuring well-posed SDEs and smooth positive marginals. It also introduces time reversal and spatial differential-operator notation.

  • Densities and laws: For each time t, pY denotes the Lebesgue density of the random variable Yt, and PY denotes the law of Y on C([0, T], Rd).The path space is equipped with its Borel measure.
  • Assumptions: The SDE coefficients and initial conditions are assumed sufficiently regular for unique strong solutions and satisfaction of Novikov’s condition.The processes are also assumed to have smooth, strictly positive densities for t ∈ (0, T).
  • Time reversal: For f(x,t), the time-reversed function is defined by f(x,T − t).This notation is used for reversing time-dependent functions over the interval [0,T].
  • Differential operators: The notation ∇g, ∇2g, and div(g) denotes the spatial gradient, Hessian, and divergence, while matrix divergence is defined row-wise.The trace operator is also introduced for matrix-valued functions.

A.2 Identities for divergences

This section defines row-wise divergence for matrix-valued functions and states the resulting identities used later in the analysis.

  • Divergence definition: For a matrix-valued function A, divergence is defined row-wise using its rows and columns.The notation distinguishes Ai· as the i-th row and A·j as the j-th column.
  • Divergence identities: The defined divergence operators satisfy identities used in the paper’s subsequent calculations.The section introduces these identities after fixing the row and column notation.

A.3 Reverse-time SDEs

The section establishes that time-reversed Itô-process marginals can be represented by another Itô process, yielding a family of reverse-time SDEs whose marginals match the forward process in reverse order.

  • Reverse-time construction: The reverse-time construction derives a family of processes parameterized by a function λ.The formulation specializes to λ = 0 when div(D) = 0 because σ is state-independent.
  • Fokker–Planck relation: The reverse-time process has density ⃗p_Y determined by the corresponding Fokker–Planck equation.The proof uses the forward Fokker–Planck equation and divergence identities to establish the required form.
  • Reverse-time construction: For every λ ∈ C2([0, T], [0, 1]), the reverse-time SDE starts from Y_T and uses λ-dependent coefficients.Its initial law is ⃗Y_0 ∼ Y_T.
  • Marginal equivalence: For every t ∈ [0, T], the reverse-time process satisfies Y_{T−t} ∼ ⃗Y_t.Thus, the reverse-time SDE reproduces the forward process's marginals in reverse temporal order.
  • Marginal equivalence: The PDE in (41b) is a valid Fokker–Planck equation for the reverse-time SDE in (35), proving the reverse-time claim.This connects the coefficient construction to the claimed marginal representation.

A.4 Further details on the HJB equation

The section explains why directly solving the HJB equation is difficult: the optimal drift contains the unknown value function, and its terminal conditions typically depend on intractable inference marginals.

  • HJB formulation: Direct approximation of the HJB equation is complicated because the optimal drift contains the value function itself.Substituting the optimal drift into the HJB equation therefore produces an equation involving the unknown solution.
  • HJB formulation: Applying the Hopf–Cole transformation directly to p_Y yields the same PDE, but its signs do not match typical control-theoretic HJB equations.A time-reversed function ⃗V is introduced to obtain an HJB equation.
  • Boundary conditions: The terminal conditions in the resulting equations are typically unavailable because they correspond to intractable marginal densities of the inference process Y.The section notes a different situation for sampling from unnormalized densities.

A.5 Hopf–Cole transformation

The Hopf–Cole transformation links a linear Kolmogorov backward equation to the nonlinear HJB equation used in the paper's control formulation.

  • PDE connection: The Kolmogorov backward equation follows from the Fokker–Planck equation through divergence identities.The resulting expression includes diffusion, drift-gradient, and divergence terms.
  • Hopf–Cole transformation: The Hopf–Cole transformation converts a positive solution of a linear PDE into the nonlinear HJB formulation.The lemma specifies the transformation's relation between p, h, and the HJB variables.
  • PDE connection: When div(µ) = 0, the Kolmogorov backward equation can be viewed as the adjoint of the Fokker–Planck equation.This identifies a special structural relationship between the two PDEs.

A.6 Brief introduction to stochastic optimal control

The section introduces stochastic optimal control for SDEs, defines value functions through control costs, and shows that the HJB equation characterizes optimal Markov controls under stated regularity and existence assumptions.

  • Control formulation: Stochastic optimal control seeks strategies for noisy continuous-time dynamics represented by SDEs.The control process is progressively measurable and takes values in R^d.
  • Control formulation: For Markov controls, the control is a deterministic function of the current state and time, so its randomness comes only from the controlled SDE.The admissible controls are functions u ∈ U ⊂ C(R^d × [0, T], R^d).
  • Controlled dynamics: A common controlled SDE adds the control to the drift as eµ(x, s, u) := µ(x, s) + σ(s)u(x, s) while leaving the diffusion coefficient unchanged.The control acts as a steering force toward a target.
  • Cost and value functions: The objective is to minimize running and terminal costs over admissible controls, with the value function giving optimal costs conditioned on state and time.The cost-to-go and value function formalize the optimization problem.
  • HJB and verification: The dynamic programming principle yields an HJB PDE that determines optimality, reducing functional optimization to pointwise minimization over R^d.The verification theorem identifies the value function and optimal Markov control when the required assumptions hold.
  • HJB and verification: The verification result assumes an infimum-attaining measurable control and a strong solution for the corresponding controlled SDE.Without sufficient regularity, optimal-control theory uses weak-solution concepts such as viscosity solutions.
  • Quadratic running costs: With control-independent diffusion, additive drift control, and quadratic running costs, the general HJB equation reduces to a closed-form HJB equation.The minimizing control is u* = −σ^⊤∇V.

A.7 Verification theorem

The verification theorem shows that the HJB solution equals the minimum control cost, attained by a control defined from the gradient of the value function. Its proof follows from Itô’s lemma and martingale properties, while the variational gap quantifies deviation from the optimal control.

  • The value function V(Z^u_t,t) equals the minimum control objective almost surely, with unique optimizer u* = −σ^⊤∇V.The result identifies the optimal control directly through the gradient of the HJB solution.
  • The verification theorem is derived by applying Itô’s lemma to a solution of the HJB equation.The proof combines the HJB equation with an Itô expansion of the controlled process.
  • Under mild regularity assumptions, the stochastic integral has zero conditional expectation, completing the verification argument.This martingale property removes the stochastic-integral term after taking conditional expectations.
  • The term in the verification identity represents the variational gap between the current control objective and the optimal objective.In the cited setting, this gap receives a specific form that can be compared with earlier diffusion-model analyses.
  • The comparison with Huang et al. notes that their corresponding expression appears to omit a factor of 1/2.

A.8 Measures on path space

The path-space formulation represents diffusion models and sampling objectives through probability measures on continuous trajectories. It links optimal controls to KL divergences, supports alternative divergences, and motivates diffusion-based sampling methods with practical numerical trade-offs.

  • The path-space perspective formulates diffusion-based generative modeling through measures over continuous trajectories and enables an optimal-control interpretation.
  • The optimal controlled path-space measure equals the law induced by the control u* = σ^⊤∇log p̄_X.The identification follows from the path-space change-of-measure result.
  • The expected negative ELBO equals the denoising score-matching objective up to a control-independent constant.
  • The sampling formulation uses an optimization problem analogous to the generative-model control problem, with the roles of the generative and inference SDEs interchanged.The minimizing control is the scaled score function and guarantees the target terminal distribution when the initial condition matches the target marginal.
  • The DIS objective is valid for an arbitrary diffusion process, unlike the DDS formulation tied to particular drift and diffusion coefficients.
  • For Gaussian targets, DIS can have preferable numerical drift properties because PIS drift may become unbounded as t → 0 under an initial delta distribution.
  • DIS can improve over PIS across normalizing-constant, expectation, and standard-deviation estimation tasks, while PINNs can provide competitive results.The PINN comparison uses a 10-dimensional double-well example; PINNs avoid time discretization but require higher-order derivatives during training and sampling.
  • Importance-sampling estimators remain unbiased but can have variance that scales exponentially with dimension and with the control’s deviation from optimality.
Loading 2211.01364v3…