Image Inpainting via Tractable Steering of Diffusion Models https://arxiv.org/pdf/2401.03349
ability of TPMs to generate images that better match the semantics of the given pixels
摘要
扩散模型是目前用于生成逼真图像的最先进技术。然而，控制受限图像生成任务的采样过程，如修补，仍然具有挑战性，因为对这些约束进行精确条件的计算是棘手的。
尽管现有方法使用各种技术来近似受限后验，但本文提出利用可计算受限后验的可处理概率模型（TPMs）的能力，并利用这一信号来引导扩散模型的去噪过程。
具体来说，本文采用了一类称为概率电路（PCs）的表达式 TPM。在之前的进展基础上，我们进一步扩展PCs，使其能够引导扩散模型的图像生成过程。
实证结果表明，我们的方法可以持续改善三个自然图像数据集（即 CelebA-HQ、ImageNet 和 LSUN）中修复图像的整体质量和语义连贯性，仅增加了 10%的计算开销。TPM 带来的额外计算开销。此外，借助图像编码器和解码器，我们的方法可以轻松接受图像特定区域的语义约束，从而为更多受控图像生成任务打开潜力。
除了提出一种新的受限图像生成框架，本文还强调了更易处理的模型的好处，并激励了表达式 TPM 的发展。
1 引言
由于其表现力，扩散模型在生成逼真和高分辨率图像方面已经取得了最先进的成果（Ramesh等，2022年；Nichol＆Dhariwal，2021年；Rombach等，2022年）。然而，将无条件扩散模型引导到诸如图像修补等受限生成任务仍然具有挑战性，因为扩散模型在设计上不支持多种类型约束条件下后验样本分布的有效计算（Chung等，2022年）。这导致生成的样本未能与约束条件完全对齐。例如，在图像修补中，模型可能生成与给定像素语义不连贯的样本。之前的研究主要通过近似（受限）后验样本分布来解决这个问题。然而，由于扩散模型的不可计算性质，这样的方法引入了很高的偏差（Lugmayr等，2022年；Zhang等，2023a；Chung等，2022年），从而降低了使用高度表现力的扩散模型的好处。
观察到缺乏可计算性阻碍了我们充分利用扩散模型进行受限生成任务，我们研究了相反的问题：支持有效受限生成的模型的好处是什么？本文通过展示支持有效计算任意边际概率的一类表达力强的可计算概率模型（TPMs）——概率电路（PCs）（Choi等，2020年），积极证明了这一点，PCs可以有效地引导扩散模型的去噪过程，生成高质量的修补图像。我们将定义一类包括修补约束的约束，我们可以为这类约束提供以下保证。对于这一类约束中的任何约束c，给定噪声水平为t的样本xt，我们展示了一个在无噪声样本上训练的PC（即， 可以被用来有效地计算 ，这是扩散模型采样过程的关键步骤。然后可以使用这个由PC计算的分布有效地指导去噪过程，从而生成符合约束的逼真图像。图1说明了PCs在所提出的算法Tiramisu（通过引导扩散模型进行可计算图像修补）中的引导效果。具体地，我们绘制了扩散模型（Tiramisu的第一行）和PC（第三行）在去噪过程中的五个时间步的重建图像。与扩散模型基线相比，Tiramisu通过PC提供的指导生成了更具语义连贯性的图像。
总之，本文有三个主要贡献：
1. 一种新的受控图像生成方案。这是第一篇论文，演示了使用TPMs控制/约束自然和高分辨率图像生成过程的可能性。这不仅为受控图像生成开辟了新的途径，还突显了非标准学习架构（例如PCs）对现代图像生成任务的潜在影响。
2. 竞争性样本质量和运行时间。在三个具有挑战性的高分辨率自然图像数据集（即CelebA-HQ、ImageNet和LSUN）上的实证评估表明，所提出的方法Tiramisu在引入仅约10%的额外计算开销的情况下，持续提高了修补图像的整体质量，这是基于先前技术的PC模型进一步扩展的结果，并且对PC的训练和推理进行了改进的自定义GPU实现的联合努力。
3. 更复杂受控生成任务的潜力。在更一般的形式中，独立软证据约束包括超出图像修补范围的任务。作为一个说明性例子，我们演示了Tiramisu能够将一组参考图像的图像块的语义融合起来，并生成受语义约束条件限制的图像。这突显了Tiramisu在更具挑战性的受控图像生成任务上的潜力。
2 预赛
去噪扩散概率模型一个定义在变量X0上的扩散模型（Ho等人，2020年；Sohl-Dickstein等人，2015年）是形式为 的潜变量模型，其中 是潜变量，联合分布 被定义为一个马尔可夫链，称为逆向/去噪过程：
可计算概率模型可计算概率模型（TPMs）是一类生成模型，其设计支持对某些查询的高效且精确的计算（Poon＆Domingos，2011年；Kisa等，2014年；Choi等，2020年；Correia等，2023年；Sidheekh等，2023年；Rahman等，2014年；Kulesza等，2012年）。根据它们的结构，TPMs支持从边际/条件概率到对逻辑约束进行条件查询等各种查询（Vergari等，2021年；Bekker等，2015年）。由于其可计算性，TPMs能够支持各种下游应用，例如受限语言生成（Zhang等，2023b），知识图谱链接预测（Loconte等，2023年）和数据压缩（Liu等，2022a）。
3 引导扩散模型使用可计算概率模型
给定一个针对无条件生成训练的扩散模型，我们的目标是引导模型在不需要任务特定的微调的情况下生成给定不同条件/约束的样本。接下来，我们将重点放在图像修补任务上，以证明可计算概率模型可以引导扩散模型生成更一致且满足约束的样本。
其中 Z 是一个归一化常数， 是 xt 中的第 i 个变量，而 的因式分解遵循第 2 节中对扩散过程的定义。尽管右边看起来由于归一化常数而难以计算，但我们将在接下来的章节中展示存在一类表达力强的可计算概率模型可以高效且精确地计算它。
其中 是混合超参数。总之，作为利用扩散模型进行图像修补的关键步骤，我们从扩散模型和TPM计算 ,并且在去噪过程中使用它们的加权几何平均值。我们注意到胎压监测系统的使用独立于与扩散模型相关的设计选择，因此可以建立在任何先前的方法上。
4 使用概率电路的实际实现
前一节介绍了可计算概率模型如何帮助引导扩散模型的去噪过程，以生成高质量的修补图像。虽然有希望，但一个关键问题是 （方程4）能否被高效且精确地计算？我们肯定地回答了这个问题，通过展示一类被称为概率电路（PCs）的可计算概率模型能够回答这个查询，同时具有足够的表现力来建模自然图像。在接下来的内容中，我们首先提供概率电路的背景（第4.1节）。然后我们描述它们如何用来计算 （第4.2节）。
4.1 概率电路背景
概率电路（PCs）（Choi等，2020年）是一个广泛的可计算概率模型的统称，包括传统的模型，如隐马尔可夫模型（Rabiner＆Juang，1986年）和Chow-Liu树（Chow＆Liu，1968年），以及更近期的模型，包括和积网络（Poon＆Domingos，2011年），算术电路（Shen等，2016年）和割集网络（Rahman等，2014年）。我们定义概率电路的语法和语义如下。
利用概率电路回答查询相当于在其DAG上递归地以后序（即前馈）或前序（即反向）计算某些函数。例如，计算似然概率 p(x) 化简为在概率电路上进行前向传播：我们首先通过评估其密度/质量函数 f_n(x) 为每个输入节点 n 分配一个概率，然后对所有求和和乘积节点进行前馈传递（先处理子节点再处理父节点），根据方程（6）计算它们的输出概率。最后，根节点的输出值就是目标似然概率。在图2中，查询 的每个节点的输出概率都用蓝色标记在其顶部。
正如其定义所示，概率电路中可学习参数的集合包括（i）求和边的参数和（ii）输入节点/分布的参数。可以使用基于期望最大化的算法共同学习所有参数，该算法旨在最大化数据集 D 中所有样本的平均对数似然概率： 有关EM算法的详细信息请参见附录 C.1。
4.2 计算约束后验分布
前向传递 类似于第4.1节中介绍的似然查询算法，我们以后序遍历所有节点，并将每个节点n的输出存储在fwn中。对于求和和乘积节点，输出是按照方程（6）计算的；每个输入节点n的输出，该节点编码了 的分布，定义为 ，其中fn在方程（6）中定义。
后向传递 后向传递包括两个步骤：（i）以先序遍历（在子节点之前访问父节点）遍历所有节点来计算后向值bkn；（ii）使用所有输入节点的后向值计算目标概率。为了便于展示，我们假设PC在求和层和乘积层之间交替，并且任何输入节点的所有父节点都是乘积节点。首先，我们通过将根节点的 设置为1来计算后向值，然后递归地计算其他节点的后向值，如下所示：
5 迈向高分辨率图像修复
决定PC引导扩散模型有效性的另一个关键因素是PC p(X0)的表达能力，即它对目标图像分布的建模能力。最近的进展显著提高了PC的表达能力（刘等人，2022b；2023），在诸如CIFAR（Krizhevsky等人，2009）和下采样的ImageNet（Deng等人，2009）等数据集上取得了竞争性的可能性，这使我们能够直接将引导修复算法应用于它们。然而，直接对高分辨率（例如，256×256）图像数据进行建模仍有差距。尽管鉴于PC的快速发展，这种差距在不久的将来有可能被弥合，本文探索了一种替代方法，即使用（变分）自动编码器将高分辨率图像转换为较低维度的潜在空间。尽管通过这种方式我们失去了对每个像素的“完全可追溯性”，但正如我们将展示的，仍然可以实现相当程度的近似。关键的直觉是潜在空间简洁地捕获了图像的语义信息，因此可以有效地引导扩散模型生成语义上连贯的图像；诸如邻近像素颜色一致性等细节问题可以由基于神经网络的扩散模型适当处理。这一点在第6.1节中得到了实证证明。
本文的另一个主要贡献是进一步扩大基于刘等人（2022b；2023）的PC规模，以实现在由VQ-GAN生成的潜在图像空间上与GPT（Brown等人，2020）竞争的可能性。具体来说，对于256×256图像，潜在空间通常由16×16=256个分类变量组成，每个变量有2048-16384个类别。虽然变量的数量与之前PC学习方法考虑的数据集相似，但这些变量在语义上要复杂得多（例如，补丁语义与像素值）。我们在附录C.2中提供了完整的学习细节，包括模型结构和训练流程。
总的来说，类似于第3节和4.2节中介绍的像素空间引导修复算法，其潜在空间变体也计算 来引导扩散模型 ，使用方程（5），只是它是使用与VQ-GAN结合的潜在空间PC来近似的。
6 实验
在本节中，我们逐步分析并阐明我们的方法Tiramisu（通过引导扩散模型实现可行的图像修复）。具体来说，我们首先定性地研究TPM对去噪扩散过程的引导效应（第6.1节）。接下来，我们在三个高分辨率图像数据集上对Tiramisu进行实证评估，使用六个大洞掩模，这极大地挑战了它生成语义一致图像的能力（第6.2节）。最后，受到Tiramisu能够处理可以写成独立软证据的任意约束的事实启发（参见第4.2节），我们在称为图像语义融合的新的控制图像生成任务上测试它，其中我们得到了来自不同图像的部分（例如，补丁）并旨在将它们融合起来，既语义上连贯又具有高保真度（第6.3节）。
6.1 TPM提供的引导分析
由于我们很大程度上是受到TPM生成与给定像素语义更匹配的图像的能力的激励，因此自然要检查TPM生成的信号如何在去噪过程中引导扩散模型。回想第3节，每一步去噪t，重建分布pDM(˜x0|xt, x k 0)和pTPM(˜x0|xt, x k 0)分别使用扩散模型和TPM计算/估计。然后，这两个分布合并为p(˜x0|xt, x k 0)（方程（5）），并用于生成前一噪声水平（即xt−1）的图像。在所有实验中，我们采用CoPaint（张等人，2023a）生成pDM(˜x0|xt, x k 0)，这与TPM的设计选择无关。因此，定性比较Tiramisu和CoPaint的去噪过程使我们能够检查TPM提供的引导效应。
图1通过绘制上述分布（即 的期望值对应的图像，可视化了Tiramisu的去噪过程。为了最小化干扰，我们首先关注同一列中DM-和TPM生成的图像对。由于它们是从相同的输入图像xt生成的，比较图像对使我们能够检查两个分布中固有的归纳偏差。例如，在名人脸部图像中，我们观察到TPM生成的图像的面部特征轮廓更清晰。这在更大时间步的图像中更明显，因为TPM提供的引导在整个去噪过程中是累积的。
接下来，我们看第二行（即 的Tiramisu。尽管有些模糊，但在全球语义在去噪过程的早期阶段就出现了。例如，在右侧，我们可以在时间步217隐约看到两只鸵鸟。相比之下，CoPaint在t=217时刻的去噪图像并没有包含太多语义信息。基于这些模糊的内容，扩散模型可以进一步填充细粒度的细节。由于图像语义可以在几个去噪步骤中生成，我们只需要在早期时间步查询TPM，这也显著降低了Tiramisu的计算开销。有关定量分析，请参见第6.2节。因此，与基线相比，Tiramisu能够生成更高质量的修复图像。
6.2 与最新技术的比较
在本节中，我们在三个大规模高分辨率图像数据集上将Tiramisu与最先进的基于扩散的修复算法进行挑战：CelebA-HQ（刘等人，2015）、ImageNet（Deng等人，2009）和LSUN-Bedroom（Yu等人，2015）。为了进一步考验Tiramisu生成语义一致图像的能力，我们使用了七种类型的掩膜，这些掩膜仅揭示了原始图像的5-20%，因为修复算法很可能会忽略给定的视觉线索并生成语义不一致的图像。掩膜的详细信息可以在附录D中找到。
方法 我们考虑了以下六种基于扩散的修复算法：CoPaint（张等人，2023a）、RePaint（Lugmayr等人，2022）、DDNM（王等人，2022）、DDRM（Kawar等人，2022）、DPS（Chung等人，2022）和Resampling（Trippe等人，2022）。虽然不是全面的，但这组方法概括了图像修复的最新发展，并且可以被认为是最先进的。我们的方法Tiramisu基于CoPaint（即，使用CoPaint生成 。有关Tiramisu的详细信息，请参见附录（附录B和C.2），以及基线（附录D）。
定量和定性结果 表1显示了所有3×7=21个数据集-掩膜配置的平均LPIPS值（张等人，2018）。首先，在21个设置中的18个中，Tiramisu优于CoPaint，这表明TPM提供的引导一致性地提高了生成图像的质量。其次，与所有基线相比，Tiramisu在21个设置中的14个上实现了最佳的LPIPS值，这表明其优越性。这一结论得到了图3所示的样本修复图像的进一步支持，这表明Tiramisu生成了更多语义上一致的图像。有关更多样本，请参见附录E.2；有关用户研究，请参见附录E.1。
计算效率 如第6.1节所示，我们可以使用概率计算（PC）在早期阶段引导去噪步骤。虽然让PC参与更多的去噪步骤可能会带来更好的性能，但运行时间也会相应增加。为了更好地理解这种权衡，我们使用CelebA + Expand1掩膜作为例子来分析这种权衡。如图4所示，当我们在更多的去噪步骤中使用PC时，LPIPS得分首先下降然后上升，这表明在适度数量的步骤中加入PC可以获得最佳性能（在这种情况下大约是20%）。这种现象的一个解释是，在后期的去噪阶段，扩散模型主要关注细节的细化。然而，PC在早期的去噪阶段更擅长控制图像的全局语义。然后我们关注计算时间。当在20%的去噪步骤中使用PC时，由TPM引起的额外计算开销大约为10秒，这只占总计算时间的10%。
6.3 超越图像修复
前面的部分展示了在图像修复任务中使用TPM的有效性。一个自然的问题是，这个框架是否可以推广到其他受控/受限的图像生成任务？虽然我们没有确切的答案，但本节通过展示Tiramisu将语义信息从各种输入补丁/片段融合的能力，展示了将其扩展到更复杂任务的潜力。具体来说，考虑潜在空间软证据约束 （即方程（8））。对于各种最近的自动编码器模型，如VQ-GAN，大小为Hl × Wl的潜在变量是从大小为H × W的图像中编码出来的。直观地说， 每个潜在变量编码了一个大小为H/Hl × W/Wl的图像补丁的语义。因此，每个可以被视为对应图像补丁语义的约束。
我们引入了一个名为语义融合的受控图像生成任务，其中我们得到了多个参考图像，每个图像都配有一个掩膜。目标是生成图像，这些图像（i）在语义上与每个参考图像的未掩膜区域对齐，并且（ii）具有高质量和保真度。语义融合可以被视为更一般受控图像生成的初步任务，因为任何类型的视觉词汇信息（例如，语言条件）都可以转移到对 的约束上。
图5显示了Tiramisu在语义融合任务上的定性结果。对于每组参考图像，我们通过调整每个软证据函数 的温度，生成了五个语义一致性水平不同的样本。有关更多详细信息，请参阅附录F。
图5：CelebA-HQ在语义融合任务上的定性结果。在每个样本中，Tiramisu提供了两张参考图像及其掩膜。任务是生成图像，这些图像（i）在语义上与两张参考图像的未掩膜区域对齐，并且（ii）具有高保真度。对于每个输入，我们生成了五个语义一致性水平不同的样本。最左边的图像语义约束最少，几乎不匹配参考图像的语义模式。相比之下，最右边的图像严格匹配参考图像的语义。
7 相关工作与结论
现有的图像修复方法可以分为两类：监督方法和无监督方法。具体来说，监督方法要求模型在修复任务上进行显式训练，而无监督方法不需要特定任务的训练。监督方法广泛用于变分自编码器（Peng等人，2021；Zheng等人，2019；Guo等人，2019）、生成对抗网络（Iizuka等人，2017；Zhao等人，2020；Guo等人，2019）和Transformer（Yu等人，2021；Wan等人，2021）。监督修复算法的一个主要问题是它们高度偏向于训练期间观察到的掩膜类型，通常需要为每个修复任务单独微调（Xiang等人，2023）。然而，由于对于这些模型来说，近似给定已知像素的掩蔽像素的概率是非常困难的，我们很遗憾地受限于监督修复方法。
扩散模型的最新发展（Ho等人，2020；Song等人，2020）为无监督修复提供了可能性，因为这些模型提供了近似受限后验的潜在方法。具体来说，不同噪声水平上变量的存在使我们能够将已知像素的信息融入到扩散模型的去噪过程中（Song & Ermon，2019；Avrahami等人，2022；Kawar等人，2022）。此外，诸如在更高噪声水平上重新采样图像（Lugmayr等人，2022）和使用部分过滤来近似受限后验（Trippe等人，2022）的技术，极大地促进了生成高质量图像。
从监督到无监督修复方法的转变是一个明确的例子，展示了使用更可行模型的好处。基于这一观察，本文旨在进一步利用可行模型。具体来说，我们研究了“极端情况”，即使用TPMs（可精确计算受限后验的概率模型）。实证结果表明，TPMs可以有效提高修复图像的质量，且只需额外增加10%的计算开销。