Conditioning diffusion models by explicit forward-backward bridging

TL;DR

This work tackles the problem of exact conditional sampling $\pi(x|y)$ when a diffusion model has been trained for the joint distribution $\pi(x,y)$. It introduces forward-backward bridging (FBB), an MCMC framework that casts conditioning as inference on an augmented space and alternates forward noising with backward bridging via a Gibbs sampler, allowing asymptotically exact conditioning without modifying the unconditional model beyond Monte Carlo error. The authors develop Gibbs-CSMC and PMCMC variants, with a separable-dynamics specialization enabling memory-efficient, in-place operation, and a pseudo-marginal implementation for unbiased marginal likelihood estimation. Empirical results across synthetic GP conditioning, non-separable Schrödinger-bridge settings, and inpainting/super-resolution tasks show improved sample quality and reduced bias compared with standard particle methods, highlighting the method’s versatility and practical impact for conditional diffusion-based inference. The approach provides a principled, scalable path to exact conditioning in diffusion models and related bridge-based generative frameworks, with potential extensions to broader inverse problems and nonlinear likelihoods.

Abstract

Given an unconditional diffusion model targeting a joint model $π(x, y)$, using it to perform conditional simulation $π(x \mid 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 \emph{exact} conditional simulation within the \emph{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 \mid 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.

Figures (13)

• Figure 1: Autocorrelations (of the worst dimension) in Section \ref{['sec:experiment-gp-regression']}. For each of the 100 dimensions the autocorrelation is averaged over four chains. We see that Gibbs-CSMC outperforms PMCMC, with higher particle counts reducing autocorrelation.
• Figure 2: Conditional sampling errors on a Gaussian Schrödinger bridge: PF (ideal) uses the exact posterior (generally unavailable), while PF (approximate) uses a standard Normal. Gibbs sampling significantly improves quality, even over PF (ideal), underscoring PF's inaccuracies.
• Figure 3: MNIST super-resolution (x4) on a non-separable Schrödinger bridge noising process. Each method shows two samples (more are in Figure \ref{['fig:app-sb-imgs']}). We find that PF is significantly affected by its initialisation of $X_0$, while Gibbs-CSMC is largely unaffected.
• Figure 4: Examples of inpainting (first and third panels) and super-resolution (second and fourth panels) on MNIST and CelebA-HQ. In each panel, the first to the last rows show the results of PF, Gibbs-CSMC, PMCMC-0.005, TPF, and CSGM, respectively. We see that the samples generated by Gibbs-CSMC and PMCMC have overall better quality.
• Figure 5: Contour plot of two-dimensional marginal densities for the true and approximate distributions in Section \ref{['sec:experiment-gp-regression']}. The contour level lines are consistent in all the figures. The first and second rows show the results when using 10 and 100 particles, respectively. The figure shows that the Gibbs-CSMC method is visibly the best, and that it works well, even when using a small number of particles.

Theorems & Definitions (6)

• Remark 2.1
• Remark 2.2
• Remark 2.3
• Remark 3.1
• Remark 3.2
• Remark 3.3