Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A Unified and Fast-Sampling Diffusion Bridge Framework via Stochastic Optimal Control

Mokai Pan, Kaizhen Zhu, Yuexin Ma, Yanwei Fu, Jingyi Yu, Jingya Wang, Ye Shi

TL;DR

UniDB reframes diffusion-bridge models as stochastic optimal control with a tunable terminal penalty $\gamma$, unifying prior Doob-based approaches and enabling a principled trade-off between control cost and endpoint accuracy. It introduces UniDB-GOU as a concrete instantiation and UniDB++ to achieve fast, training-free reverse-time sampling via exact reverse-SDE solutions plus a data-prediction model and an SDE-Corrector. The framework yields state-of-the-art performance on image restoration tasks while achieving 5x–20x inference speedups, bridging theory and practice for diffusion bridges. This work broadens the applicability of diffusion bridges to real-time applications and diverse inverse problems, including potential medical-imaging use cases.

Abstract

Recent advances in diffusion bridge models leverage Doob's $h$-transform to establish fixed endpoints between distributions, demonstrating promising results in image translation and restoration tasks. However, these approaches often produce blurred or excessively smoothed image details and lack a comprehensive theoretical foundation to explain these shortcomings. To address these limitations, we propose UniDB, a unified and fast-sampling framework for diffusion bridges based on Stochastic Optimal Control (SOC). We reformulate the problem through an SOC-based optimization, proving that existing diffusion bridges employing Doob's $h$-transform constitute a special case, emerging when the terminal penalty coefficient in the SOC cost function tends to infinity. By incorporating a tunable terminal penalty coefficient, UniDB achieves an optimal balance between control costs and terminal penalties, substantially improving detail preservation and output quality. To avoid computationally expensive costs of iterative Euler sampling methods in UniDB, we design a training-free accelerated algorithm by deriving exact closed-form solutions for UniDB's reverse-time SDE. It is further complemented by replacing conventional noise prediction with a more stable data prediction model, along with an SDE-Corrector mechanism that maintains perceptual quality for low-step regimes, effectively reducing error accumulation. Extensive experiments across diverse image restoration tasks validate the superiority and adaptability of the proposed framework, bridging the gap between theoretical generality and practical efficiency. Our code is available online https://github.com/2769433owo/UniDB-plusplus.

A Unified and Fast-Sampling Diffusion Bridge Framework via Stochastic Optimal Control

TL;DR

UniDB reframes diffusion-bridge models as stochastic optimal control with a tunable terminal penalty , unifying prior Doob-based approaches and enabling a principled trade-off between control cost and endpoint accuracy. It introduces UniDB-GOU as a concrete instantiation and UniDB++ to achieve fast, training-free reverse-time sampling via exact reverse-SDE solutions plus a data-prediction model and an SDE-Corrector. The framework yields state-of-the-art performance on image restoration tasks while achieving 5x–20x inference speedups, bridging theory and practice for diffusion bridges. This work broadens the applicability of diffusion bridges to real-time applications and diverse inverse problems, including potential medical-imaging use cases.

Abstract

Recent advances in diffusion bridge models leverage Doob's -transform to establish fixed endpoints between distributions, demonstrating promising results in image translation and restoration tasks. However, these approaches often produce blurred or excessively smoothed image details and lack a comprehensive theoretical foundation to explain these shortcomings. To address these limitations, we propose UniDB, a unified and fast-sampling framework for diffusion bridges based on Stochastic Optimal Control (SOC). We reformulate the problem through an SOC-based optimization, proving that existing diffusion bridges employing Doob's -transform constitute a special case, emerging when the terminal penalty coefficient in the SOC cost function tends to infinity. By incorporating a tunable terminal penalty coefficient, UniDB achieves an optimal balance between control costs and terminal penalties, substantially improving detail preservation and output quality. To avoid computationally expensive costs of iterative Euler sampling methods in UniDB, we design a training-free accelerated algorithm by deriving exact closed-form solutions for UniDB's reverse-time SDE. It is further complemented by replacing conventional noise prediction with a more stable data prediction model, along with an SDE-Corrector mechanism that maintains perceptual quality for low-step regimes, effectively reducing error accumulation. Extensive experiments across diverse image restoration tasks validate the superiority and adaptability of the proposed framework, bridging the gap between theoretical generality and practical efficiency. Our code is available online https://github.com/2769433owo/UniDB-plusplus.

Paper Structure

This paper contains 28 sections, 4 theorems, 61 equations, 4 figures, 3 tables, 5 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Diffusion Bridges
  4. Diffusion with Stochastic Optimal Control
  5. Diffusion Solvers for Fast Sampling
  6. Preliminaries
  7. Denoising Diffusion Bridge Models
  8. Generalized Ornstein-Uhlenbeck Bridge
  9. Stochastic Optimal Control
  10. UniDB Framework
  11. Constructing Diffusion Bridges via SOC Problem
  12. Diffusion Bridges via Doob's $h$-transform is A Special Case
  13. An Example: UniDB-GOU
  14. UniDB++ for Fast Sampling
  15. Solutions to the Reverse-Time SDE with Data Prediction
  16. ...and 13 more sections

Key Result

Theorem 1

Consider the SOC problem SOC_problem_generalized_ode, denote $d_{t, \gamma} = \gamma^{-1} + e^{2\bar{f}_{T}} \bar{g}^2_{t:T}$, $\bar{f}_{s:t} = \int_{s}^{t} f_z dz$, $\bar{h}_{s:t} = \int_{s}^{t} e^{-\bar{f}_{z}} h_z dz$ and $\bar{g}^2_{s:t} = \int_{s}^{t} e^{-2\bar{f}_{z}}g^2_z dz$, denote $\bar{f} and the transition of $\mathbf{x}_t$ from $x_0$ and $x_T$ is

Figures (4)

  • Figure 1: We formulate the forward process of diffusion bridges as a stochastic optimal control problem and employs the penalty coefficient $\gamma$, balancing realistic SDE trajectories and target endpoint matching, to produce images with more realistic details. We also find that Doob’s $h$-transform is a special case of UniDB when $\gamma \rightarrow \infty$. To overcome the low-efficiency in sampling process, we develop a fast training-free sampling method UniDB++ by deriving the exact solution of the reverse SDE with data prediction model, which directly estimates the fixed and smooth target $\textbf{x}_0$, ensuring stable predictions of the solvers even under few sampling-step regimes.
  • Figure 2: Qualitative comparison of visual results among Ground-Truth (GT) images, GOUB, UniDB (ours) and UniDB++ (ours) on Image 4$\times$Super-Resolution, Image Deraining, Image RainDrop Restoration, and Image Inpainting tasks with zoomed-in image local regions.
  • Figure 3: Ablation Study on SDE-Corrector. Qualitative comparison between UniDB++ without SDE-Corrector and UniDB++ with SDE-Corrector (both in first-order) (simply denoted as UniDB++ & Corrector in the figures).
  • Figure 4: Qualitative comparison of visual results between GOUB, UniDB (ours), UniDB++ (ours) and UniDB++ with SDE-Corrector (simply denoted as UniDB++ & Corrector in the figures) (ours) on Image Super-Resolution and Raindrop Image Restoration tasks.

Theorems & Definitions (8)

  • Theorem 1
  • Theorem 2
  • Proposition 1
  • Theorem 3
  • proof
  • proof
  • proof
  • proof