UniDB++: Fast Sampling of Unified Diffusion Bridge

2025-05-23

AI Summary

While the UniDB framework achieves high-fidelity image generation via stochastic optimal control (SOC), its reliance on Euler-based iterative sampling incurs high computational cost and slow inference. Existing acceleration methods fail because they cannot accommodate UniDB’s missing terminal mean constraint and SOC-specific penalty coefficients. Method: We propose a retraining-free, efficient sampling algorithm: (i) deriving the closed-form analytical solution of UniDB’s reverse stochastic differential equation (SDE); (ii) replacing noise prediction with data prediction; and (iii) introducing an SDE-Corrector mechanism to preserve perceptual quality within only 5–10 steps. Results: Experiments demonstrate that our method matches Euler sampling in generation quality while reducing sampling steps by 20×, achieving state-of-the-art performance in image inpainting and significantly lowering inference latency.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Computer Vision: Diffusion Models for Vision
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Diffusion Bridges enable transitions between arbitrary distributions, with the Unified Diffusion Bridge (UniDB) framework achieving high-fidelity image generation via a Stochastic Optimal Control (SOC) formulation. However, UniDB's reliance on iterative Euler sampling methods results in slow, computationally expensive inference, while existing acceleration techniques for diffusion or diffusion bridge models fail to address its unique challenges: missing terminal mean constraints and SOC-specific penalty coefficients in its SDEs. We present UniDB++, a training-free sampling algorithm that significantly improves upon these limitations. The method's key advancement comes from deriving exact closed-form solutions for UniDB's reverse-time SDEs, effectively reducing the error accumulation inherent in Euler approximations and enabling high-quality generation with up to 20$ imes$ fewer sampling steps. This method 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 (5-10 steps). Additionally, we demonstrate that UniDB++ aligns with existing diffusion bridge acceleration methods by evaluating their update rules, and UniDB++ can recover DBIMs as special cases under some theoretical conditions. Experiments demonstrate UniDB++'s state-of-the-art performance in image restoration tasks, outperforming Euler-based methods in fidelity and speed while reducing inference time significantly. This work bridges the gap between theoretical generality and practical efficiency in SOC-driven diffusion bridge models. Our code is available at https://github.com/2769433owo/UniDB-plusplus.

Problem

• Improves slow, expensive inference in UniDB framework
• Addresses missing terminal constraints in SDEs
• Enables high-quality generation with fewer steps

Innovation

• Exact closed-form solutions for reverse-time SDEs
• Stable data prediction replaces noise prediction
• SDE-Corrector maintains quality in low-step regimes