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

TL;DR

UniDB introduces a unified, fast-sampling diffusion bridge framework based on Stochastic Optimal Control, improving image restoration quality by balancing control costs and terminal penalties, with a training-free accelerated algorithm.

Contribution

The paper reformulates diffusion bridges through SOC, unifies existing methods, and develops a training-free, efficient reverse-time SDE solution with enhanced detail preservation.

Findings

Outperforms existing diffusion bridge methods in image restoration tasks.

Achieves faster sampling with a closed-form solution, reducing computational costs.

Maintains high perceptual quality with a stable data prediction model and SDE-Corrector.

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.