UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control
Authors: Kaizhen Zhu, Mokai Pan, Yuexin Ma, Yanwei Fu, Jingyi Yu, Jingya Wang, Ye Shi
ICML 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | Extensive experiments across diverse image restoration tasks validate the superiority and adaptability of the proposed framework. |
| Researcher Affiliation | Academia | 1Shanghai Tech University 2Mo E Key Laboratory of Intelligent Perception and Human-Machine Collaboration 3Fudan University. Correspondence to: Ye Shi <EMAIL>. |
| Pseudocode | Yes | We provide pseudo-code Algorithm 1 and Algorithm 2 for the training and sampling process of Uni DB-GOU, respectively. |
| Open Source Code | Yes | Our code is available at https://github.com/ Uni DB-SOC/Uni DB/. |
| Open Datasets | Yes | Image 4 Super-Resolution Tasks. In super-resolution, we evaluated our models based on DIV2K dataset (Agustsson & Timofte, 2017), which contains 2K-resolution high-quality images. [...] Image Deraining Tasks. For image deraining tasks, we conducted the experiments based on Rain100H datasets (Yang et al., 2017). [...] Image Inpainting Tasks. In image inpainting tasks, we evaluated our methods on Celeb A-HQ 256 256 datasets (Karras, 2017). |
| Dataset Splits | Yes | As for the datasets of the three main experiments, we take 800 images for training and 100 for testing for the DIV2K dataset, 1800 images for training and 100 for testing for the Rain100H dataset, 27000 images for training and 3000 for testing for the Celeb A-HQ 256 256 dataset. |
| Hardware Specification | Yes | Our models are trained on a single NVIDIA H20 GPU with 96GB memory for about 2 days. |
| Software Dependencies | No | The paper mentions "Adam optimizer" but does not specify any key software components (like PyTorch or TensorFlow) along with their version numbers. |
| Experiment Setup | Yes | steady variance level λ2 = 302/2552, coefficient e θT = 0.005 instead of zero, sampling step number T = 100, 128 patch size with 8 batch size when training, Adam optimizer with β1 = 0.9 and β2 = 0.99 (Kingma, 2014), 1.2 million total training steps with 10 4 initial learning rate and decaying by half at 300, 500, 600, and 700 thousand iterations. With respect to the schedule of θt, we choose a flipped version of cosine noise schedule (Nichol & Dhariwal, 2021; Luo et al., 2023), θt = 1 cos( t/T +s cos( s 1+s π 2 )2 (85) where s = 0.008 to achieve a smooth noise schedule. |