基于随机最优控制的统一快速采样扩散桥框架
计算机视觉与模式识别
2025-11-12 v2 人工智能
机器学习
摘要
扩散桥模型的最新进展利用 Doob 的 -变换 在分布间建立固定端点，在图像转换与复原任务中展现出可喜的结果。然而，这些方法常产生模糊或过度平滑的图像细节，且缺乏完善的理论基础来解释这些缺陷。为解决这些局限性，我们提出了 UniDB，一种基于随机最优控制 (Stochastic Optimal Control, SOC) 的统一快速采样扩散桥框架。我们通过基于 SOC 的优化重新表述该问题，证明了采用 Doob -变换的现有扩散桥构成一个特例，即当 SOC 代价函数中的终端惩罚系数趋于无穷大时出现的情况。通过引入可调的终端惩罚系数，UniDB 在控制成本与终端惩罚之间取得了最佳平衡，大幅提升了细节保留和输出质量。为避免 UniDB 中迭代 Euler 采样方法的高昂计算成本，我们通过推导 UniDB 逆时 SDE 的精确闭式解，设计了一种免训练的加速算法。此外，通过用更稳定的数据预测模型替代传统的噪声预测模型，并辅以在低步长机制下保持感知质量的 SDE-Corrector 机制，进一步减少了误差累积。在多种图像复原任务上的大量实验验证了所提框架的优越性和适应性，弥合了理论通用性与实际效率之间的差距。代码已在线开源：https://github.com/2769433owo/UniDB-plusplus。
引用
@article{arxiv.2505.21528,
title = {A Unified and Fast-Sampling Diffusion Bridge Framework via Stochastic Optimal Control},
author = {Mokai Pan and Kaizhen Zhu and Yuexin Ma and Yanwei Fu and Jingyi Yu and Jingya Wang and Ye Shi},
journal= {arXiv preprint arXiv:2505.21528},
year = {2025}
}