基于随机最优控制的统一扩散桥架构 UniDB

最近的扩散桥模型研究利用 Doob 的 $h$-变换在分布之间建立固定端点，展示了在图像翻译和恢复任务中的前景成果。然而，这些方法常常产生模糊或过度平滑的图像细节，并且缺乏充分的理论基础来解释这些短comings。为此，我们提出了 UniDB，一个基于随机最优控制 (SOC) 的扩散桥统一框架。UniDB 通过 SOC 基于的优化问题进行建模，并推导出最优控制器的闭式解，从而统一和概括现有的扩散桥模型。我们展示，采用 Doob 的 $h$-变换的现有扩散桥是本框架的一个特例，当 SOC 成本函数中的终端惩罚系数趋于无穷时即成立。通过纳入可调节的终端惩罚系数，UniDB 在控制成本与终端惩罚之间实现了最优平衡，显著改善了细节保留和输出质量。值得注意的是，UniDB 可无缝集成到现有的扩散桥模型中，只需极少的代码修改。广泛的实验在多样化的图像恢复任务中验证了该框架的优越性和适应性。我们的代码可在 https://github.com/UniDB-SOC/UniDB/ 上获取。