Conditioning diffusion models by explicit forward-backward bridging

Code

• github.com/zgbkdlm/fbs
  Officialjax★ 9

Abstract

Given an unconditional diffusion model targeting a joint model (x, y), using it to perform conditional simulation (x y) is still largely an open question and is typically achieved by learning conditional drifts to the denoising SDE after the fact. In this work, we express exact conditional simulation within the approximate diffusion model as an inference problem on an augmented space corresponding to a partial SDE bridge. This perspective allows us to implement efficient and principled particle Gibbs and pseudo-marginal samplers marginally targeting the conditional distribution (x y). Contrary to existing methodology, our methods do not introduce any additional approximation to the unconditional diffusion model aside from the Monte Carlo error. We showcase the benefits and drawbacks of our approach on a series of synthetic and real data examples.

Tasks

Reproductions