arXiv:2502.05749cs.CVcs.AI2025-02ICML被引 20

统一扩散桥框架,用随机最优控制提升图像修复细节质量

UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control

  • 基于随机最优控制构建统一框架,推导出最优控制器闭式解
  • 通过可调终端惩罚系数平衡控制成本与输出精度,减少模糊失真
  • 兼容现有模型,仅需少量代码修改,适用于多种图像修复任务

近期扩散桥模型利用杜布的h-变换在分布间建立固定端点,已在图像翻译和修复任务中展现良好效果。然而,这些方法常导致图像细节模糊或过度平滑,且缺乏解释其局限性的完整理论基础。为此,本文提出基于随机最优控制(SOC)的统一扩散桥框架UniDB。UniDB将问题建模为SOC优化问题,并推导出最优控制器的闭式解,从而统一并推广了现有扩散桥模型。我们证明,采用杜布h-变换的现有方法是本框架在终端惩罚系数趋于无穷时的特例。通过引入可调节的终端惩罚系数,UniDB实现了控制代价与终端惩罚间的最优平衡,显著改善了细节保留与输出质量。值得注意的是,UniDB可无缝集成至现有扩散桥模型,仅需极少代码修改。在多种图像修复任务上的大量实验验证了该框架的优越性与适应性。代码已开源:https://github.com/UniDB-SOC/UniDB/

原文摘要 · Abstract (English)

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/.

扩散模型图像修复随机控制统一框架

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。