UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control

• URL: http://arxiv.org/abs/2502.05749v2
• Date: Tue, 11 Feb 2025 08:33:03 GMT
• Title: 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,
• Abstract summary: We propose UniDB, a unified framework for diffusion bridges based on Optimal Control (SOC) UniDB formulates the problem through an SOC-based optimization and derives a closed-form solution for the optimal controller. We demonstrate that existing diffusion bridges employing Doob's $h$-transform constitute a special case of our framework.
• Score: 73.73585888013665
• License:
• Abstract: 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/.

Related papers

• Consistency Diffusion Bridge Models [25.213664260896103]
  Diffusion bridge models (DDBMs) build processes between fixed data endpoints based on a reference diffusion process. DDBMs' sampling process typically requires hundreds of network evaluations to achieve decent performance. We propose two paradigms: consistency bridge distillation and consistency bridge training, which is flexible to apply on DDBMs with broad design choices.
  arXiv  Detail & Related papers  (2024-10-30T02:04:23Z)
• Soft-constrained Schrodinger Bridge: a Stochastic Control Approach [4.922305511803267]
  Schr"odinger bridge can be viewed as a continuous-time control problem where the goal is to find an optimally controlled diffusion process. We propose to generalize this problem by allowing the terminal distribution to differ from the target but penalizing the Kullback-Leibler divergence between the two distributions. One application is the development of robust generative diffusion models.
  arXiv  Detail & Related papers  (2024-03-04T04:10:24Z)
• Double-Bounded Optimal Transport for Advanced Clustering and Classification [58.237576976486544]
  We propose Doubly Bounded Optimal Transport (DB-OT), which assumes that the target distribution is restricted within two boundaries instead of a fixed one. We show that our method can achieve good results with our improved inference scheme in the testing stage.
  arXiv  Detail & Related papers  (2024-01-21T07:43:01Z)
• Towards Unified and Effective Domain Generalization [42.579796989908914]
  We propose $textbfUniDG$, a novel and $textbfUnified framework for $textbfD$omain $textbfG$eneralization. Specifically, we encourage models to learn the distribution of test data in an unsupervised manner and impose a penalty regarding the updating step of model parameters.
  arXiv  Detail & Related papers  (2023-10-16T02:05:03Z)
• Generalized Schrödinger Bridge Matching [54.171931505066]
  Generalized Schr"odinger Bridge (GSB) problem setup is prevalent in many scientific areas both within and without machine learning. We propose Generalized Schr"odinger Bridge Matching (GSBM), a new matching algorithm inspired by recent advances. We show that such a generalization can be cast as solving conditional optimal control, for which variational approximations can be used.
  arXiv  Detail & Related papers  (2023-10-03T17:42:11Z)
• Direct Diffusion Bridge using Data Consistency for Inverse Problems [65.04689839117692]
  Diffusion model-based inverse problem solvers have shown impressive performance, but are limited in speed. Several recent works have tried to alleviate this problem by building a diffusion process, directly bridging the clean and the corrupted. We propose a modified inference procedure that imposes data consistency without the need for fine-tuning.
  arXiv  Detail & Related papers  (2023-05-31T12:51:10Z)
• I$^2$SB: Image-to-Image Schr\"odinger Bridge [87.43524087956457]
  Image-to-Image Schr"odinger Bridge (I$2$SB) is a new class of conditional diffusion models. I$2$SB directly learns the nonlinear diffusion processes between two given distributions. We show that I$2$SB surpasses standard conditional diffusion models with more interpretable generative processes.
  arXiv  Detail & Related papers  (2023-02-12T08:35:39Z)
• DSEE: Dually Sparsity-embedded Efficient Tuning of Pre-trained Language Models [152.29364079385635]
  As pre-trained models grow bigger, the fine-tuning process can be time-consuming and computationally expensive. We propose a framework for resource- and parameter-efficient fine-tuning by leveraging the sparsity prior in both weight updates and the final model weights. Our proposed framework, dubbed Dually Sparsity-Embedded Efficient Tuning (DSEE), aims to achieve two key objectives: (i) parameter efficient fine-tuning and (ii) resource-efficient inference.
  arXiv  Detail & Related papers  (2021-10-30T03:29:47Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.