[论文解读] Diffusion models as plug-and-play priors

Alexandros Graikos, Nikolay Malkin+2 more|arXiv (Cornell University)|Jun 17, 2022

Bayesian Methods and Mixture Models被引用 42

一句话总结

本文展示了独立训练的去噪扩散概率模型（DDPMs）如何在可微约束下作为即插即用先验来进行推断，从而实现条件生成、图像分割以及组合问题的连续松弛。

ABSTRACT

We consider the problem of inferring high-dimensional data $\mathbf{x}$ in a model that consists of a prior $p(\mathbf{x})$ and an auxiliary differentiable constraint $c(\mathbf{x},\mathbf{y})$ on $x$ given some additional information $\mathbf{y}$. In this paper, the prior is an independently trained denoising diffusion generative model. The auxiliary constraint is expected to have a differentiable form, but can come from diverse sources. The possibility of such inference turns diffusion models into plug-and-play modules, thereby allowing a range of potential applications in adapting models to new domains and tasks, such as conditional generation or image segmentation. The structure of diffusion models allows us to perform approximate inference by iterating differentiation through the fixed denoising network enriched with different amounts of noise at each step. Considering many noised versions of $\mathbf{x}$ in evaluation of its fitness is a novel search mechanism that may lead to new algorithms for solving combinatorial optimization problems.

研究动机与目标

• 引入一个框架，其中一个预训练的 DDPM 作为模型中 p(x) 的先验，与可微约束 c(x,y) 共同作用。
• 证明推断可以通过对单一潜变量使用 DDPM 的去噪网络进行梯度优化来完成。
• 展示在条件图像生成、分割以及组合问题的连续松弛方面的应用。
• 强调对于这些即插即用用法，不需要对 DDPM 进行额外训练或微调。

提出的方法

• 通过自由能 F 将 p(x|y) 表示为一个带有归一化常数的形式，并用以潜变量 η 为中心的 delta 函数或高斯来近似 q(x)。
• 使用 DDPM 的前向加噪过程来计算期望并在 ε-空间推导重建的平方误差损失，形成一个实用的优化目标（Equation 12）。
• 对 η（或在潜变量-场景中对 y）进行基于梯度的优化，同时对时间步 t 进行退火以探索从粗到细的模态。
• 在未对约束模型进行联邦训练的情况下利用预训练的 DDPM，使得可以通过可微分约束 c(x,y) 直接进行就地条件化。
• 我们还讨论替代的后验近似和潜在空间表述以扩展适用性。

实验结果

研究问题

• RQ1 在存在可微约束且不重新训练的情况下，独立训练的 DDPM 是否可以用作即插即用的先验？
• RQ2 如何通过对 DDPM 的潜在变量或像素空间表示进行基于梯度的推断，在视觉任务中产生高保真且带有约束的样本？
• RQ3 在使用 DDPM 先验进行条件生成和分割时，哪些有效的条件化策略（如基于分类器、弱标签、颜色聚类等）？
• RQ4 DDPM 是否能够通过潜变量推断实现诸如旅行商问题之类的组合问题的连续松弛？
• RQ5 哪些实用的推断策略（如退火计划、初始化选择等）在跨领域中能产生鲁棒的结果？

主要发现

• 一个预训练的 DDPM 可以作为先验来推断满足可微约束的样本，从而在无需重新训练的情况下实现条件生成和分割。
• 对 DDPM 的潜在/去噪空间进行优化，并以约束项 log c(x,y) 引导，产生符合期望条件的现实样本（例如具备属性分类或弱标签）。
• 在 MNIST 和基于 CelebA 的实验中，该方法产生带有定向属性的条件数字和人脸，同时由于 DDPM 先验而保持图像真实度。
• 在语义分割方面，该方法在 EnviroAtlas 的多地理区域内实现了有竞争力的准确率和 IoU，展示了域迁移能力。
• 在连续松弛的 TSP 设置中，扩散模型推断潜在相邻结构，经过 2-opt 精炼后，得到的巡回路径距离最优解只有很小的百分比差，显示出在没有离散组合求解器的情况下的有竞争力的性能。
• 该方法完全在图像空间中工作，在某些编码方案下实现亚线性扩展，并利用视觉先验进行结构化推断。