Research map / Image generation / Fast diffusion sampling
UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes UniDB, a stochastic-optimal-control framework that generalizes Doob's h-transform diffusion bridges (DDBMs, GOUB) and adds a tunable terminal penalty gamma to improve detail preservation in image restoration.
Problem
Diffusion bridges built on Doob's h-transform force exact endpoint matching, which produces blurred or oversmoothed details in image translation and restoration, and lack a theoretical account of why.
Approach
Formulate the bridge as a linear-quadratic stochastic optimal control problem with a terminal penalty gamma; derive a closed-form optimal controller and show Doob's h-transform is the gamma-to-infinity limit, so a finite gamma trades small endpoint mismatch for smoother trajectories.
Key results
- On DIV2K 4x super-resolution, UniDB (SDE) with gamma=1e7 reaches LPIPS 0.179 and FID 16.21, versus GOUB (SDE) at 0.220 and 20.85, though its PSNR (25.46) is lower than GOUB (SDE) at 26.89.
- On Rain100H deraining, UniDB (SDE) improves LPIPS to 0.045 and FID to 17.65 over GOUB (SDE) at 0.046 and 18.14.
- On CelebA-HQ inpainting, UniDB (SDE) reports FID 4.08 and LPIPS 0.036 versus GOUB (SDE) at 4.30 and 0.037.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace GOUB (ODE) with UniDB (ODE) | GOU diffusion bridge · DIV2K 4× super-resolution | PSNR: 28.50 -> 28.64 | Helped |
| replace GOUB (ODE) with UniDB (ODE) | GOU diffusion bridge · Rain100H deraining | PSNR: 34.56 -> 34.68 | Helped |
| replace GOUB (ODE) with UniDB (ODE) | GOU diffusion bridge · CelebA-HQ inpainting | PSNR: 31.39 -> 31.67 | Helped |
| replace GOUB (SDE) with UniDB (SDE) | GOU diffusion bridge · DIV2K 4× super-resolution | LPIPS: 0.220 -> 0.179 | Helped |
| replace GOUB (SDE) with UniDB (SDE) | GOU diffusion bridge · DIV2K 4× super-resolution | PSNR: 26.89 -> 25.46 | Hurt |
| replace GOUB (SDE) with UniDB (SDE) | GOU diffusion bridge · Rain100H deraining | FID: 18.14 -> 17.65 | Helped |
| replace GOUB (SDE) with UniDB (SDE) | GOU diffusion bridge · CelebA-HQ inpainting | FID: 4.30 -> 4.08 | Helped |
| tune γ from ∞ to 1×10^7 (ablation) | UniDB-GOU · DIV2K 4× super-resolution | LPIPS: 0.220 -> 0.179 | Helped |
| tune γ from ∞ to 1×10^7 (ablation) | UniDB-GOU · CelebA-HQ inpainting | FID: 4.30 -> 4.08 | Helped |
| tune γ from ∞ to 1×10^8 (ablation) | UniDB-GOU · Rain100H deraining | FID: 18.14 -> 17.65 | Helped |
| tune γ from ∞ to 1×10^8 (ablation) | UniDB-GOU · DIV2K 4× super-resolution | LPIPS: 0.220 -> 0.289 | Hurt |
| tune γ from ∞ to 1×10^8 (ablation) | UniDB-GOU · CelebA-HQ inpainting | FID: 4.30 -> 4.64 | Hurt |
Surprising findings
- The gamma-infinity Doob h-transform setting (GOUB) gives the highest PSNR/SSIM on super-resolution (26.89 / 0.7478), while the finite-gamma UniDB improves only perceptual metrics (LPIPS, FID) there.
Limitations
- Sampling remains computationally intensive, especially for high-resolution images, as the authors acknowledge.
- Best gamma varies by task and must be tuned per task, with no principled selection rule beyond the sweep.
Experimental setup
Datasets: DIV2K CelebA-HQ Rain100H
Benchmarks: DIV2K Rain100H CelebA-HQ
Metrics: PSNR SSIM LPIPS FID
Baselines: Bicubic DDRM IR-SDE GOUB (SDE) GOUB (ODE) MAXIM MHNet PromptIR
Builds on: DDBMs GOUB Doob's h-transform Diffusion Schrödinger Bridge stochastic interpolants
Closest papers by meaning
- Exploring the Design Space of Diffusion Bridge Models (NeurIPS 2025 · similarity 0.93)
- Stochastic Optimal Control for Diffusion Bridges in Function Spaces (NeurIPS 2024 · similarity 0.90)
- Image Restoration Through Generalized Ornstein-Uhlenbeck Bridge (ICML 2024 · similarity 0.89)
- Diffusion Bridge Implicit Models (ICLR 2025 · similarity 0.89)
- Taming Diffusion Prior for Image Super-Resolution with Domain Shift SDEs (NeurIPS 2024 · similarity 0.89)
- Solving Diffusion ODEs with Optimal Boundary Conditions for Better Image Super-Resolution (ICLR 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.