清华朱军团队提出DBIM:无需训练加速20倍的图像翻译扩散桥模型推理算法

清华朱军团队提出DBIM算法,无需额外训练即可加速扩散桥模型推理20倍,显著提升图像翻译和修复效率。

原文标题:ICLR 2025 | 无需训练加速20倍,清华朱军组提出用于图像翻译的扩散桥模型推理算法DBIM

原文作者:机器之心

冷月清谈:

清华大学朱军团队针对扩散桥模型(DDBM)在图像翻译和修复任务中计算效率低下的问题,提出了扩散桥隐式模型(DBIM)算法。DBIM通过引入非马尔科夫扩散桥,实现了随机性和确定性采样的灵活切换,并推导出一种新的常微分方程表达形式,以及高阶数值求解方法,有效提升了推理的精度与效率。此外,DBIM还提出了“启动噪声”机制,解决了确定性采样过程中的初始奇异性问题。实验结果表明,DBIM在图像翻译和修复任务中,仅需少量推理步骤即可达到甚至超越传统DDBM模型的生成质量,实现了显著的推理加速。

怜星夜思:

1、文章提到DBIM通过控制参数ρ在随机采样和确定性采样之间自由切换,那么在实际应用中,如何根据不同的图像翻译或修复任务来选择合适的ρ值?有没有一些通用的准则或经验可以参考?
2、DBIM中提到的“启动噪声”(booting noise)机制是为了解决确定性采样过程中的初始奇异性问题,那么除了这种方式,还有没有其他方法可以避免或缓解这个问题?
3、文章实验部分,DBIM在ImageNet图像修复任务中取得了很好的效果,那么这项技术在实际的图像处理领域有哪些潜在的应用场景?除了图像修复,还能应用到哪些其他领域?

原文内容


论文有两位共同一作。郑凯文为清华大学计算机系三年级博士生,何冠德为德州大学奥斯汀分校(UT Austin)一年级博士生。


扩散模型(Diffusion Models)近年来在生成任务上取得了突破性的进展,不仅在图像生成、视频合成、语音合成等领域都实现了卓越表现,推动了文本到图像、视频生成的技术革新。然而,标准扩散模型的设计通常只适用于从随机噪声生成数据的任务,对于图像翻译或图像修复这类明确给定输入和输出之间映射关系的任务并不适合。



为了解决这一问题,一种名为去噪扩散桥模型(Denoising Diffusion Bridge Models, DDBMs)的变种应运而生。DDBM 能够建模两个给定分布之间的桥接过程,从而很好地应用于图像翻译、图像修复等任务。然而,这类模型在数学形式上依赖复杂的常微分方程 / 随机微分方程,在生成高分辨率图像时通常需要数百步的迭代计算效率低下,严重限制了其在实际中的广泛应用。


相比于标准扩散模型,扩散桥模型的推理过程额外涉及初始条件相关的线性组合和起始点的奇异性,无法直接应用标准扩散模型的推理算法。为此,清华大学朱军团队提出了一种名为扩散桥隐式模型(DBIM)的算法,无需额外训练即可显著加速扩散桥模型的推理。



  • 论文标题:Diffusion Bridge Implicit Models

  • 论文链接:https://arxiv.org/abs/2405.15885

  • 代码仓库:https://github.com/thu-ml/DiffusionBridge


方法介绍


DBIM 的核心思想是对扩散桥模型进行推广,提出了一类非马尔科夫扩散桥(non-Markovian Diffusion Bridges)。这种新的桥接过程不仅与原来的 DDBM 拥有相同的边缘分布与训练目标,而且能够通过减少随机性,实现从随机到确定性的灵活可控的采样过程。


具体而言,DBIM 在模型推理过程中引入了一个方差控制参数 ρ,使得生成过程能够在随机采样与确定性采样之间自由切换。当完全采用确定性推理模式时,DBIM 能够直接以隐式的形式表示生成过程。这种思想是标准扩散模型的著名推理算法 DDIM 在扩散桥模型上的推广与拓展。


图片


更进一步,DBIM 算法可以导出扩散桥的一种全新的常微分方程(ODE)表达形式,相较于 DDBM 论文中的常微分方程形式更加简洁有效。


图片


在此基础上,作者首次提出了针对扩散桥模型的高阶数值求解方法,进一步提升了推理的精度与效率。



此外,为了避免确定性采样过程中出现的初始奇异性问题,作者提出了一种「启动噪声」(booting noise)机制,即仅在初始步骤中加入适当随机噪声,从而保证了模型的生成多样性,并同时保留了对数据的编码与语义插值能力。



实验结果


作者在经典的图像翻译和图像修复任务上进行了如下实验:


  • 在 Edges→Handbags(64×64)和 DIODE-Outdoor(256×256)图像翻译任务中,DBIM 仅需 20 步推理即可达到甚至超越 DDBM 模型 118 步推理的生成质量。当推理步数增至 100 步时,DBIM 进一步提升生成质量,在更高分辨率任务上全面领先。



  • 在更具挑战的 ImageNet 256×256 图像修复任务中,DBIM 仅需 20 步推理便显著超越了传统扩散桥模型 500 步推理的效果,实现了超过 25 倍的推理加速。在 100 步推理时,DBIM 进一步刷新了这一任务的 FID 记录。




通过参数 η 控制采样过程中的随机性大小,论文发现确定性采样模式在低步数时具备显著优势,而适当增加随机性能够在较高步数下进一步提升生成多样性与 FID 指标。这与标准扩散模型推理的性质相似。



此外,高阶采样器能够在不同采样步数下一致提升生成质量,增强图像细节。



论文所用训练、推理代码及模型文件均已开源。如果你对 DBIM 感兴趣,想要深入了解它的技术细节和实验结果,可访问论文原文和 GitHub 仓库。


© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:liyazhou@jiqizhixin.com


我觉着“启动噪声”有点像炼丹时候的“trick”hhh,感觉理论性不强,我提供一个思路,可以考虑从数学上对diffusion bridge的ODE做一些修改,比如添加一些正则化项,或者改变一下ODE的边界条件,这样可能可以从根本上解决奇异性问题。当然,这需要比较强的数学功底。

我觉得DBIM在安全领域也有应用前景。比如,可以用DBIM来修复被恶意篡改的图像或视频,恢复原始信息。 另外,可以用DBIM来提高图像识别的鲁棒性,例如,在图像受到噪声干扰的情况下,仍然可以准确识别出图像中的物体。

我个人的理解是,ρ值可以看作是一个trade-off的参数,在保真度和多样性之间做平衡。 如果希望结果更接近原始图像,那么就降低随机性,反之就增加随机性。

另外,我觉得可以考虑用机器学习的方法来自动调节ρ值,例如用强化学习来训练一个agent,根据当前图像和任务来动态调整ρ值,这样可能可以获得更好的效果。

DBIM在ImageNet图像修复上的出色表现让人眼前一亮!我觉得它在实际图像处理领域的应用潜力巨大:

* 老照片修复: 可以用于修复老照片中出现的划痕、污渍等问题,让老照片焕发新生。
* 视频修复: 修复视频中出现的噪点、马赛克等问题,提升视频的观看体验。
* 文物修复: 辅助文物修复工作者修复破损的文物,例如,修复壁画、瓷器等。
* 医学图像处理: 用于修复医学图像中的伪影,提高诊断的准确性。

除了图像修复,DBIM还可以应用于以下领域:

* 图像编辑: 可以用于图像的智能编辑,例如,去除图像中的某个物体,或者改变图像的风格。
* 图像生成: 可以用于生成高质量的图像,例如,根据文本描述生成图像。
* 视频生成: 可以用于生成流畅的视频,例如,根据文本描述生成视频。

我感觉DBIM在内容创作领域大有可为。比如现在很火的AI绘画,就可以用DBIM来提升生成图像的质量和细节。 另外,在游戏开发中,可以用DBIM来自动生成游戏场景,提高开发效率。

这个问题问得好!选择合适的ρ值确实是DBIM应用中的一个关键。从我的理解看,目前还没有完全通用的准则,但可以考虑以下几个角度:

* 任务的性质: 如果任务对图像细节要求非常高,或者需要保留更多的输入信息,可以尝试偏向确定性采样(ρ值较小),反之,如果更注重生成的多样性和创新性,可以适当增加随机性(ρ值较大)。
* 计算资源: 确定性采样通常需要更少的步数就能达到较好的效果,但在计算资源充足的情况下,增加一定随机性并增加步数,可能获得更好的生成质量和多样性。
* 实验调优: 最直接的方法还是通过实验来调整ρ值。可以先设置几个候选值(例如0.1, 0.5, 0.9),在验证集上进行测试,观察生成结果的质量和多样性,然后选择最优的ρ值。
* 参考论文实验: 论文中提到,确定性采样模式在低步数时具备显著优势,高步数下适当增加随机性可以提高生成多样性。可以以此为baseline进行调整与探索。

我觉得可以借鉴GAN中的一些思路,比如spectral normalization。通过限制模型参数的Lipschitz常数,可以提高模型的稳定性,从而减轻奇异性问题。当然,这需要对DBIM的模型结构进行一定的修改。

“启动噪声”确实是一个很巧妙的解决方案!关于其他避免奇异性的方法,我想到一些:

* 更精细的数值方法: 或许可以尝试使用更高阶、更稳定的数值求解器来处理常微分方程,从而减轻奇异性带来的影响。不过,这可能会增加计算复杂度。
* 修改模型结构: 从模型本身入手,例如,在模型中引入一些正则化项,或者改变模型的输出方式,使得模型在初始阶段更加稳定。
* 数据预处理: 对输入数据进行一些预处理,例如,进行归一化或者加入一些微小的扰动,使得输入数据更加平滑,从而降低奇异性出现的可能性。
* 混合方法: 结合多种方法,例如,同时使用“启动噪声”和更精细的数值方法,从而达到更好的效果。

我比较认同楼上的观点,ρ值的选择确实需要在保真度和多样性之间进行权衡。但是,我认为除了任务类型和计算资源外,还需要考虑数据集的特点。如果数据集本身就比较单一,那么可能需要增加ρ值来引入更多的随机性,避免生成结果过于相似。另外,可以考虑使用一些自适应的ρ值调整策略,例如根据生成过程中的loss变化来动态调整ρ值。