# UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control

arXiv: https://arxiv.org/abs/2502.05749 · HTML: https://mlsanity.com/paper/2502.05749/

## Key facts

- Cited by 20 papers in the map (data through 14 Oct 2026).
- Builds on 47 of its 63 references that are in the map.
- First posted on arXiv 9 Feb 2025.

## Builds on (most cited first, 30 of 47)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [A Style-Based Generator Architecture for Generative Adversarial Networks](https://mlsanity.com/paper/1812.04948/) (2018)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://mlsanity.com/paper/1710.10196/) (2017)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Denoising Diffusion Restoration Models](https://mlsanity.com/paper/2201.11793/) (2022)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Multi-Stage Progressive Image Restoration](https://mlsanity.com/paper/2102.02808/) (2021)
- [MaskGAN: Towards Diverse and Interactive Facial Image Manipulation](https://mlsanity.com/paper/1907.11922/) (2019)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Deep Joint Rain Detection and Removal from a Single Image](https://mlsanity.com/paper/1609.07769/) (2016)
- [MAXIM: Multi-Axis MLP for Image Processing](https://mlsanity.com/paper/2201.02973/) (2022)
- [Progressive Image Deraining Networks: A Better and Simpler Baseline](https://mlsanity.com/paper/1901.09221/) (2019)
- [DiffIR: Efficient Diffusion Model for Image Restoration](https://mlsanity.com/paper/2303.09472/) (2023)
- [I2SB: Image-to-Image Schr"odinger Bridge](https://mlsanity.com/paper/2302.05872/) (2023)
- [Image Restoration with Mean-Reverting Stochastic Differential Equations](https://mlsanity.com/paper/2301.11699/) (2023)
- [Denoising Diffusion Bridge Models](https://mlsanity.com/paper/2309.16948/) (2023)
- [An optimal control perspective on diffusion-based generative modeling](https://mlsanity.com/paper/2211.01364/) (2022)
- [Score-Based Diffusion Models as Principled Priors for Inverse Imaging](https://mlsanity.com/paper/2304.11751/) (2023)
- [PromptIR: Prompting for All-in-One Blind Image Restoration](https://mlsanity.com/paper/2306.13090/) (2023)
- [Policy Representation via Diffusion Probability Model for Reinforcement Learning](https://mlsanity.com/paper/2305.13122/) (2023)
- [Diffusion-based Reinforcement Learning via Q-weighted Variational Policy Optimization](https://mlsanity.com/paper/2405.16173/) (2024)

## Built on by (most cited first, 20 of 20)

- [Towards Immersive Human-X Interaction: A Real-Time Framework for Physically Plausible Motion Synthesis](https://mlsanity.com/paper/2508.02106/) (2025)
- [System-Embedded Diffusion Bridge Models](https://mlsanity.com/paper/2506.23726/) (2025)
- [Training-Free Reward-Guided Image Editing via Trajectory Optimal Control](https://mlsanity.com/paper/2509.25845/) (2025)
- [Diffusion Bridge or Flow Matching? A Unifying Framework and Comparative Analysis](https://mlsanity.com/paper/2509.24531/) (2025)
- [Sample from What You See: Visuomotor Policy Learning via Diffusion Bridge with Observation-Embedded Stochastic Differential Equation](https://mlsanity.com/paper/2512.07212/) (2025)
- [Efficient Image Restoration with State-Dependent Forward Diffusion](https://mlsanity.com/paper/2505.16733/) (2025)
- [A Unified and Fast-Sampling Diffusion Bridge Framework via Stochastic Optimal Control](https://mlsanity.com/paper/2505.21528/) (2025)
- [Inference-Time Attribute Distribution Alignment for Unconditional Diffusion](https://mlsanity.com/paper/2605.07456/) (2026)
- [Where Should Action Generation Begin? A Learnable Source Prior for Generative Robot Policies](https://mlsanity.com/paper/2606.17408/) (2026)
- [Energy-oriented Diffusion Bridge for Image Restoration with Foundational Diffusion Models](https://mlsanity.com/paper/2604.10983/) (2026)
- [Distributional Reinforcement Learning with Diffusion Bridge Critics](https://mlsanity.com/paper/2602.05783/) (2026)
- [Optimally Bridging Semantics and Data: Generative Semantic Communication via Schr"odinger Bridge](https://mlsanity.com/paper/2604.17802/) (2026)
- [Unifying Deep Stochastic Processes for Image Enhancement](https://mlsanity.com/paper/2605.01568/) (2026)
- [Mind the Gap: Quantifying the Domain Gap in Cross-Sensor Diffusion Super-Resolution](https://mlsanity.com/paper/2606.28039/) (2026)
- [Stochastic Optimal Control Sampling for Diffusion Inverse Problems](https://mlsanity.com/paper/2606.28785/) (2026)
- [PRISM: Principled Reference Identification for Schrodinger Bridge Model](https://mlsanity.com/paper/2608.06893/) (2026)
- [SDDBMs: Soft Denoising Diffusion Bridge Models](https://mlsanity.com/paper/2608.08594/) (2026)
- [Residual Diffusion Implicit Models](https://mlsanity.com/paper/2609.33020/) (2026)
- [Graph Anomaly Detection as Finite-Horizon Control: Training-Free Scoring via Empirical Bayes](https://mlsanity.com/paper/2609.38424/) (2026)
- [FASTER: Fast Adjoint Stochastic Transport for Endpoint Refinement in Reward-Guided Image Editing](https://mlsanity.com/paper/2610.04538/) (2026)

## Abstract (arXiv)

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 frequently 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 framework for diffusion bridges based on Stochastic Optimal Control (SOC). UniDB formulates the problem through an SOC-based optimization and derives a closed-form solution for the optimal controller, thereby unifying and generalizing existing diffusion bridge models. We demonstrate that existing diffusion bridges employing Doob's h-transform constitute a special case of our framework, 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. Notably, UniDB seamlessly integrates with existing diffusion bridge models, requiring only minimal code modifications. Extensive experiments across diverse image restoration tasks validate the superiority and adaptability of the proposed framework. Our code is available at https://github.com/UniDB-SOC/UniDB/.