[Paper Review] Conditional Simulation Using Diffusion Schrödinger Bridges

Yuyang Shi, Valentin De Bortoli+2 more|arXiv (Cornell University)|Feb 27, 2022

Model Reduction and Neural Networks4 citations

TL;DR

This paper introduces Conditional Diffusion Schrödinger Bridges (CDSB), a novel framework that extends Schrödinger bridge theory to conditional generative modeling, enabling faster and more efficient sampling for inverse problems. By reformulating conditional score-based generative models as a Schrödinger bridge problem, CDSB achieves state-of-the-art performance in image super-resolution, Bayesian filtering, and refinement of pre-trained models with reduced inference time compared to standard denoising diffusion models.

ABSTRACT

Denoising diffusion models have recently emerged as a powerful class of generative models. They provide state-of-the-art results, not only for unconditional simulation, but also when used to solve conditional simulation problems arising in a wide range of inverse problems. A limitation of these models is that they are computationally intensive at generation time as they require simulating a diffusion process over a long time horizon. When performing unconditional simulation, a Schrödinger bridge formulation of generative modeling leads to a theoretically grounded algorithm shortening generation time which is complementary to other proposed acceleration techniques. We extend the Schrödinger bridge framework to conditional simulation. We demonstrate this novel methodology on various applications including image super-resolution, optimal filtering for state-space models and the refinement of pre-trained networks. Our code can be found at https://github.com/vdeborto/cdsb.

Motivation & Objective

• To address the computational inefficiency of score-based generative models (SGMs) in conditional simulation, particularly in inverse problems requiring long-time reverse diffusion processes.
• To extend the Schrödinger bridge framework—previously used for unconditional generation—to conditional settings, enabling faster and more efficient sampling.
• To develop a principled, iterative algorithm (CDSB) that approximates the conditional Schrödinger bridge solution using neural network-estimated scores.
• To demonstrate the method’s effectiveness across diverse applications, including image super-resolution, optimal filtering in state-space models, and refinement of pre-trained networks.
• To provide a theoretical and practical alternative to standard score-based conditional models that reduces generation time without sacrificing sample quality.

Proposed method

• Propose a conditional Schrödinger bridge (CSB) formulation that defines a finite-time stochastic process with data distribution at t=0 and reference distribution at t=T, conditioned on observed data Y=y^obs.
• Adapt the Diffusion SB algorithm from De Bortoli et al. (2021) to the conditional setting, using iterative optimization to approximate the conditional bridge process.
• Use neural networks to estimate the score functions of the conditional perturbed distributions, enabling end-to-end training and inference.
• Integrate the method into conditional score-based generative models (CSGMs), allowing sampling from the posterior p(x|y^obs) via time-reversed dynamics with learned scores.
• Apply dimension-wise scaling of the diffusion time step using the variance of the reference measure to improve stability and convergence.
• Use EnKF to estimate the posterior mean and variance for initializing the reference distribution in filtering tasks, improving robustness in chaotic systems.

Experimental results

Research questions

• RQ1 Can the Schrödinger bridge framework be successfully extended to conditional generative modeling to accelerate sampling in inverse problems?
• RQ2 How does the performance of the proposed CDSB method compare to standard conditional score-based models (CSGMs) in terms of generation speed and accuracy?
• RQ3 Can CDSB achieve state-of-the-art results in image super-resolution and Bayesian filtering for chaotic state-space models?
• RQ4 What is the impact of diffusion process length (N=20 vs. N=100) and ensemble size (M) on the stability and accuracy of CDSB in filtering tasks?
• RQ5 Does the use of a data-adaptive reference distribution (p_ref) improve performance in non-Gaussian or high-variance systems?

Key findings

• CDSB achieves lower RMSE in filtering mean estimation than CSGM and EnKF, with the lowest error of 0.331±0.006 at M=2000 and N=100, outperforming CSGM-C and EnKF.
• For filtering standard deviation estimation, CDSB-C achieves the lowest RMSE of 0.106±0.002 at M=2000 and N=100, significantly outperforming CSGM-C and EnKF.
• With N=100 and M≥500, CDSB and CDSB-C show stable performance without divergence, while CDSB with M=200 and long diffusion steps exhibits occasional large errors due to error accumulation in iterative regression.
• In image super-resolution, CDSB achieves state-of-the-art FID scores, demonstrating its effectiveness in high-dimensional conditional generation tasks.
• The use of a data-adaptive reference distribution (p_ref) tailored to the prior predictive distribution improves performance in chaotic systems like Lorenz-63, where a fixed p_ref would fail.
• CDSB enables faster inference than standard SGMs by avoiding long reverse-time diffusion processes, providing a complementary acceleration method to existing techniques like knowledge distillation or optimized solvers.