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

Area of Science:

• Artificial Intelligence
• Computer Vision
• Machine Learning

Background:

• Diffusion bridge models, utilizing Doob's $h$-transform, show promise in image translation and restoration.
• Existing methods often result in blurred details and lack theoretical grounding for these issues.

Purpose of the Study:

• Introduce UniDB, a unified and efficient diffusion bridge framework based on Stochastic Optimal Control (SOC).
• Address limitations of current diffusion bridges, focusing on detail preservation and theoretical completeness.

Main Methods:

• Reformulate diffusion bridges using SOC optimization, identifying Doob's $h$-transform as a special case.
• Incorporate a tunable terminal penalty coefficient for optimal balance between control costs and terminal penalties.
• Develop a training-free accelerated sampling algorithm with closed-form solutions for the reverse-time SDE.
• Employ a data prediction model and an SDE-Corrector mechanism for enhanced stability and perceptual quality in low-step regimes.

Main Results:

• UniDB achieves superior detail preservation and output quality in image restoration tasks.
• The framework demonstrates adaptability across diverse image restoration applications.
• The proposed methods significantly improve efficiency, bridging theoretical generality and practical application.

Conclusions:

• UniDB provides a theoretically robust and practically efficient solution for diffusion bridge models.
• The framework overcomes the blurring artifacts and enhances detail fidelity in image restoration.
• UniDB represents a significant advancement in diffusion bridge methodologies for generative tasks.