编者整理中......

SDM: Spatial Diffusion Model for Large Hole Image Inpainting

原文：https://arxiv.org/abs/2212.02963

开源：https://github.com/fenglinglwb/SDM

生成对抗网络（GAN）在图像修复方面取得了巨大成功，但在处理大缺失区域方面仍然存在困难。相比之下，迭代算法，如自回归和去噪扩散模型，必须部署大量计算资源才能获得良好的效果。为了克服各自的局限性，我们提出了一种新的空间扩散模型（SDM），该模型使用几次迭代来逐步向整个图像提供信息像素，从而大大提高了推理效率。此外，由于所提出的解耦概率建模和空间扩散方案，我们的方法实现了高质量的大洞修复。在多个基准测试上，我们实现了最先进的性能。代码https://github.com/fenglinglwb/SDM

生成对抗网络（GAN）在图像修复方面取得了巨大成功，但在处理大缺失区域方面仍然存在困难。相比之下，迭代算法，如自回归和去噪扩散模型，必须部署大量计算资源才能获得良好的效果。为了克服各自的局限性，我们提出了一种新的空间扩散模型（SDM），该模型使用几次迭代来逐步向整个图像提供信息像素，从而大大提高了推理效率。此外，由于所提出的解耦概率建模和空间扩散方案，我们的方法实现了高质量的大孔修复。在多个基准测试上，我们实现了最先进的性能。

1.引言

图像修复是一项基本的计算机视觉任务，旨在用视觉上令人愉悦且语义上适当的内容填充图像中缺失的区域。它已广泛应用于图形和成像应用，如照片恢复[60，61]、图像编辑[3，21]、合成[32]、重新定位[8]和对象移除[10]。这项任务，尤其是填补大洞，比其他修复问题更为棘手，因此需要具有更强生成能力的模型。

在过去几年中，生成对抗网络（GAN）在图像修复方面取得了巨大进展[34，37，43，62，69，71]。通过通过最小-最大博弈对目标分布进行隐式建模，基于GAN的方法在视觉质量方面显著优于传统的基于示例的技术[9,10,16,55]。然而，这些方法的一次生成有时会导致不稳定的训练[14，28，51]，并使学习复杂分布变得困难，特别是在修复高分辨率图像中的大洞时。

相反，自回归模型[42，57，58]和去噪扩散模型[11，19，53]最近在内容生成中显示出显著的威力[44，49，52，73]。通过可追踪的概率建模，他们基于先验估计迭代地改进图像。在这方面，训练更加稳定，覆盖范围更广。众所周知，自回归模型逐像素处理图像，使得处理高分辨率数据变得困难。另一方面，去噪扩散模型通常需要数千次迭代才能达到良好的估计。因此，直接在图像修复中使用这些方法会产生各自的缺点——高质量大孔高分辨率图像修复的策略仍然不足。

为了完成修复图，在本文中，我们开发了一个新的空间扩散模型（SDM），该模型适合于大洞场景。SDM以迭代的方式工作，其中在每次迭代时并行预测所有像素，并且只保留合格的像素用于下一次迭代。它的作用是将可信像素逐渐扩散到未知位置。我们的核心设计在于一种简单但高效的解耦概率建模（见第3.1.1节），它具有GAN高效优化和概率模型可处理性的优点。详细地说，我们的模型同时预测修复结果（即平均项）和不确定性图（即方差项）。平均项通过隐式对抗性训练进行优化，以较少的迭代产生更准确的预测。相反，方差项是使用高斯正则化显式建模的。

这种解耦策略有几个优点。首先，由于对抗性优化，迭代步骤的数量显著减少。如图2所示，几次迭代可以产生有希望的结果。第二，高斯正则化产生的方差项自然作为不确定性度量，允许我们选择迭代细化的可靠估计（见第3.1.2节）。此外，显式分布建模使连续采样成为可能，从而提高预测的质量和多样性，如第4节所示。最后，不确定性度量有助于建立不确定性引导的注意力机制（见第3.2节），这鼓励网络利用更多信息像素进行有效推理。因此，我们的传感和诊断模块完成了大量缺失区域的照片逼真内容，如图1所示。

我们的贡献总结如下：

•我们开发了一种新的空间扩散模型（SDM），用于大孔图像修复。由于所提出的解耦概率建模和空间扩散方案，我们的模型实现了高效优化和高质量完成。

•我们的方法在Places[82]和CelebA HQ[23]基准数据集上都达到了尖端性能。值得注意的是，我们的SDM在很大程度上优于流行的去噪扩散模型，例如LDM[46]，在Places2[82]上产生1.1 FID改进，同时显著更轻（20%参数，10倍更快）。

2.相关工作

2.1传统方法

图像修复是一个经典的计算机视觉问题。早期的方法利用图像先验，例如自相似性和稀疏性。例如，基于扩散的方法[2,5]将附近未受损邻居的信息传递给孔洞。另一系列基于示例的方法[9，12，16，29，31，55]使用人类定义的距离度量来寻找高度相似的补丁以完成缺失区域。最具代表性的工作是PatchMatch[3]，它在多尺度图像空间中使用启发式搜索来大大加快修复速度。然而，由于缺乏上下文理解，它们不能保证视觉上有吸引力和语义上一致的结果。

2.2.基于深度学习的方法

使用大量训练数据显著提高高级理解能力，基于深度神经网络的方法[36，43，64，69，75]取得了成功。Pathak等人[43]将对抗性损失[13]引入到修复中，产生了视觉上真实的结果。沿着这条线的几种方法不断将性能推向新的高度。例如，为了获得局部细粒度细节和全局一致的结构，Iizuka等人[20]采用了两个鉴别器进行对抗性训练。此外，提出了部分[35]和门控[72]卷积层，以减少不规则掩模的伪影，例如颜色差异和模糊。此外，包括前景轮廓[68]、对象结构[40，45]和分割图[54]的中间线索用于多级生成。尽管小掩模的修补内容很好，但这些方法仍然不能保证大孔修补质量。

2.3.大孔图像修复

为了处理大面积的缺失区域，人们做出了大量努力来提高模型的性能。注意力技术[37，67，70，71]和Transformer架构[34，62，81]利用了上下文信息。当图像包含重复图案时，它们工作得很好。此外，Zhao等人[80]提出了一种新的架构，弥补了图像条件生成和无条件生成之间的差距，改进了自由形式的大规模图像完成。也有人试图研究进步一代。这条线是为了每次只选择高质量的像素并逐渐填充孔洞。我们注意到，这些方法严重依赖于专门设计的更新算法[15，33，41，78]，或消耗额外的模型容量来评估预测的准确性[77]，或在处理图像时需要更多的训练阶段[7]。

最近，受益于精确似然计算和迭代采样，自回归模型[62，66，74]和去噪扩散模型[1，38，46，48]在产生多样化和现实内容方面显示出巨大潜力。它们不可避免地需要数千步的高推理成本，并且需要大量的计算资源。在这项工作中，我们提出了一种新的框架，该框架基于固有的不确定性度量逐步扩散最可靠的预测，支持以显著更少的迭代完成高质量的图像。

3.我们的方法

我们的目标是使用照片逼真的材料来完成具有大量缺失区域的蒙版图像。在本节中，我们首先制定了空间扩散模型（SDM），并进行了全面分析。接下来是模型设计和损失函数的详细信息。

3.1空间扩散模型

尽管基于GAN的方法比传统方法取得了明显更好的结果，但它们仍然面临着处理大缺失区域的巨大困难。我们将其中一个原因归因于GAN的一次性性质，并提出了迭代修复。

在每次传递中，由于不可避免地会有一些好的预测，我们使用这些像素作为线索来帮助下一次生成时间。通过这种方式，我们的空间扩散模型逐渐向整个图像传播有价值的信息。在下文中，我们首先讨论单程建模，然后再讨论扩散过程。

3.1.1解耦概率建模

对于迭代修复，找到评估预测准确性的机制至关重要。一个直观的解决方案是引入一个可追踪的概率模型，以便可以分析计算不确定性信息。然而，这一要求通常导致假设近似目标分布是高斯分布，这对于解释真正复杂的分布来说过于简单。尽管一些迭代模型，如去噪扩散模型[19]，通过包含大量隐藏变量和优化变分证据下限，丰富了边际分布的表达，但这些方法通常会产生较高的推断成本。

为了解决这些关键问题，我们提出了一种适用于高效迭代修复的解耦概率模型。重要的见解是，我们以解耦的方式利用了基于隐式GAN的优化和显式高斯正则化的优点。因此，我们可以同时获得准确的预测和明确的不确定性度量。

如图3所示，给定时间t处的输入图像xt和大洞，我们的模型（参见第3.2节中的架构细节）预测了修补结果µt以及不确定性图σ2 t。我们使用对抗性损失（以及第3.3节中的其他损失）来监督图像预测µt，同时将（µt，σ2 t）共同处理为高斯分布的平均值和对角协方差。GAN的隐式优化使得尽可能接近真实分布成为可能，大大减少了迭代次数。它还为我们提供了均值项的显式不确定性度量，允许我们选择可靠的像素。高斯正则化主要应用于使用负对数似然（NLL）Lnll作为

其中D是数据维度，i是像素索引，θ表示模型参数，输入x和输出y缩放为[−1, 1]. δ+（y）和δ−（y） 定义为

具体而言，我们包括停止梯度运算（即sg[·]），它鼓励高斯约束仅优化方差项，并允许从更精确的隐式建模中估计平均项。

讨论

我们在扩散过程中使用估计的均值和方差进行采样，同时将确定性均值项作为最终迭代的输出。第4节中的实验证明了该设计的可行性。此外，概率建模使我们能够在扩散过程中应用连续采样，从而产生更高质量和更多样的估计。最后，我们发现不确定性度量也使我们能够在第3.2节中设计更有效的注意力机制。

3.1.2空间扩散方案

我们使用前馈网络，表示为fθ（·），从已知区域开始，将信息像素逐渐扩散到整个图像，如

其中t是时间步长，xt是掩码图像，mt是二进制掩码，ut是不确定性映射。输出包括更新的图像、遮罩和不确定性图。网络参数在所有迭代中共享。

我们在训练和测试中使用多次迭代来提高性能。具体而言，如图3和等式（4）所示，我们的方法在第t次迭代时运行如下。

1.预测。

给定蒙版图像xt−1，掩码mt−1和不确定度图ut−1，我们的方法估计所有像素的平均µt和方差σ2t。然后，通过转换方差图生成缩放为[0，1]的初步不确定性图u~t。

2.选择。

我们首先对未知像素的不确定性得分进行排序。根据预定义的掩模计划，我们计算将在该迭代中添加的像素数，并将具有最低不确定性的像素插入到已知类别中，将掩模更新为m。基于初步不确定性图u～t，通过将仍然缺失的位置标记为1，将初始已知像素标记为0，我们获得最终的不确定性图ut。

3.样品。

我们考虑两种情况。首先，对于初始已知位置m0，我们始终使用原始输入像素x0。其次，我们根据µt和σt对修复区域进行连续采样。结果公式如下

其中α是可调比率，z∼ N（0，I），F表示Hadamard乘积。注意，我们在最终迭代中不使用σtz项。

3.2.模型架构

我们使用深度U-Net[47]架构和StyleGAN[25，26]解码器，通过堆叠卷积到达大的感受野，以利用图像中的上下文信息[4，6，39，63]。此外，我们在不同分辨率下引入了多个注意力块，这是因为我们发现，在更高分辨率下，全局交互显著提高了更大、更多样的数据集上的重建质量[11，70，71]。

仅基于特征相似性，传统的注意力机制[59]为像素提供了交换信息的同等机会。然而，对于修复任务，丢失的像素将使用相同的指定值初始化，使它们彼此接近。因此，它通常无法有效利用来自可见区域的有用信息。更糟糕的是，有效像素被破坏，导致内容模糊和不清晰的瑕疵。

在这种情况下，如图4所示，我们考虑了像素的不确定性分数，以调整注意力中的聚合权重。它正确地解决了上述问题。注意力输出由以下公式计算

其中{q，k，v}是查询、键、值矩阵，dk表示缩放因子，F表示基于不确定性映射u预测偏置像素权重的可学习函数，还包括重塑操作。

3.3.损失函数

在每次迭代中，我们的模型输出均值和方差估计，如图3所示。使用对抗性损失[13]Ladv和感知损失[22，56]Lpcp优化均值项，其目的是生成自然外观的图像。损失描述如下。

对抗性损失。

我们将对抗性损失表述为

其中D是实现为[25]的鉴别器，x和x是真实图像和预测图像。

知觉损失。

我们采用了一种高接受场知觉损失[56]，其公式如下

其中φi是预训练的ResNet50[17]的层输出。

如第3.1.1节所述，我们应用负对数似然Lnll来约束不确定性建模的方差。因此，发电机的最终损失函数为

其中j是扩散迭代次数。我们根据经验将λ1=1、λ2=2和λ3设置为1×10−4.

4.实验

4.1.数据集和度量

我们在Places2[82]和CelebA HQ[23]上以512×512分辨率训练我们的模型，以充分评估所提出的方法。Places2是一个大型数据集，包含各种场景类别中的近800万张训练图像。此外，36500张图像构成了验证分割。在训练期间，图像进行随机翻转、裁剪和填充，而测试图像被集中裁剪为512×512大小。对于CelebA HQ，我们分别使用24183和2993张图像来训练和测试我们的模型。我们使用具有不同大小、形状和位置的矩形和笔刷笔划构建自由形式的遮罩，这遵循了广泛使用的设置[56，72，80]。此外，我们使用[34]提供的相同掩模来评估所有模型，以进行公平比较。

尽管在早期修复工作中采用了L1距离、PSNR和SSIM[65]，但在评估图像质量时，发现L1距离、PSTR和SSIM[30]与人类感知没有强烈关联[30，50]。在这项工作中，根据[34，80]，我们使用FID[18]、P-IDS[80]和U-IDS[79]作为更合适的度量，它们稳健地测量修复图像的感知保真度。

4.2实施细节

我们使用编码器-解码器架构。编码器由卷积块组成，而解码器采用StyleGAN2[26]。编码器的信道大小从64开始，在每次下采样后加倍，直到最大512。解码器具有对称配置。编码器和解码器采用32×32和16×16分辨率的注意力块。不确定性映射在第一次迭代时初始化为“1-掩码”。给定一个H×W输入，我们首先将特征大小降采样到H 32×W 32，然后返回到H×W。

我们使用8个NVIDIA A100 GPU在Places2和CelebA HQ上训练我们的模型，以获得2000万张图像。我们利用指数移动平均（EMA）、自适应鉴别器增强（ADA）和重量调制训练策略[24，34]。批量大小为32，学习率为1×10−3.我们使用Adam[27]优化器，β1=0，β2=0.99。我们根据实验结果在等式（5）中根据经验设置α=0.01。为了生成512×512的图像，我们迭代两次进行训练，四次进行测试。

之前的工作[34，80]在Place2上训练模型，并使用5000万张或更多的图像，这一事实证明了我们方法的优势。研究发现，额外的训练可以进一步改进我们的方法，但目前已有2000万张图像提供了尖端性能。

4.3.消融研究

在本节中，我们将对我们方法中提出的设计进行全面调查。为了快速评估，我们使用Places365 Standard（Places2[82]的子集）训练我们的模型，以获得分辨率为256×256的6M图像。我们从模型“A”开始，它采用了我们的完整设计，并在训练期间采用了三次迭代。

迭代数。

我们的核心思想是采用迭代优化来提高发电质量。我们调整迭代次数，并在训练和测试期间保持相同的设置。如表1所示，与模型“A”相比，具有一次和两次迭代（称为“B”和“C”）的模型产生0.59和0.19 FID。此外，如图2所示，采用更多迭代能够产生更美观的内容。第一个和第三个案例显示的伪影明显较少，第二个案例中的拱门在三次迭代后成功恢复。定量和定性结果都表明了迭代生成的重要性。

值得注意的是，我们可以使用与训练阶段不同的迭代次数来测试系统。如表2所示，使用更多的迭代会导致更高的FID性能，但代价是推断时间更长。因此，在推理速度和生成质量之间存在权衡。此外，当比较模型“A”和“B”时，很明显，在训练过程中引入更多迭代是有益的。但推理阶段的迭代次数更重要。

解耦概率建模。

为了在支持迭代修复的不确定性度量的同时提供准确的预测，我们提出了解耦概率建模。当对采样结果进行所有监督时，我们观察到训练减少了方差项（对于所有像素接近于0）。这是因为，与精确量化每个步骤的噪声水平的去噪扩散模型不同，我们基于GAN的方法不再为均值和方差项提供特定的优化目标。在这种情况下，对于微不足道的优化，方差项被低估了。这会降低拣货过程的效率。

如表1所示，与完整模型“A”相比，模型“D”获得的FID结果较差。此外，从图5中的视觉比较可以看出，模型“D”倾向于生成模糊的内容，而模型“A”产生更清晰的结构和细粒度的细节。所有这些观察结果证明了该设计的有效性。

连续取样。

我们的方法可以使用估计的方差来执行连续采样。表1表明，当不涉及连续取样（“E”型）时，FID降低了近0.1。此外，可以观察到，我们的完整模型导致了更具视觉一致性的内容。例如，根据图5中的可见像素，盒结构得到了很好的恢复。因此，连续采样为我们的结果带来了更高的保真度。如图6所示，我们的模型也支持多元化的一代，特别是在洞的中心。然而，当以低不确定性估计平均项或约束迭代次数时，结果的差异并不总是立即明显。

不确定性引导注意力。

为了充分利用遥远的背景，我们在框架中添加了注意力块。我们首先比较了32×32、16×16（模型“A”）和仅16×16的注意（模型“G”）。我们从表1中的定量比较中发现模型“G”中的FID下降0.28，证明了大孔图像修复中长距离相互作用的重要性。

此外，如第3.2节所述，传统的注意力机制可能导致颜色一致性和模糊性。为了支持这一说法，我们梳理了不确定性指导，并注意到表1中的性能略有下降。此外，我们在图5中进行了视觉比较。我们观察到模型“a”比模型“F”产生了更具视觉吸引力的窗口细节。

掩码计划。

如表4和图7所示，我们分析了各种掩码调度策略，发现统一策略实现了最佳FID。我们认为，这是因为输入图像的掩模比变化很大，并且统一的调度会为不同的迭代带来更稳定的训练。

4.4.与现有技术方法的比较

我们将所提出的空间扩散模型（SDM）与表3中基于GAN的模型[34、56、70、72、76、80、83]、自回归模型[62]和去噪扩散模型[46]进行了彻底比较。尽管最近两行甚至在商业用途上取得了显着进展，但大多数现有技术只能处理分辨率高达256×256的图像。我们使用512×512分辨率的可公开访问的模型，并在相同的掩模上测试它们，以进行公平的比较。

如表3所示，我们的SDM在Places2和CelebA HQ基准测试中，无论在大掩码设置还是小掩码设置下，都达到了最先进的性能。可以看出，我们的方法明显优于现有的基于GAN的模型。此外，即使只有强去噪扩散模型LDM[46]的20%的参数，我们的方法在所有度量方面都提供了优异的结果。例如，在Places2基准测试中，在大掩码设置下，我们的SDM在FID上实现了1.1的改进，在P-IDS和U-ID上实现了更大的增益。至于推理速度，我们的SDM花费了近250ms来获得512×512的图像，这比LDM快10倍(∼3秒）。值得注意的是，我们的模型使用的样本要少得多（我们的2000万张图像与CoModGAN的[80]5000万张图像相比）。我们观察到，更多的延长训练可以进一步提高成绩。所有这些结果都证明了我们方法的有效性。

我们还提供了图8中的视觉比较。在各种场景中，与现有方法相比，我们的方法生成的纹理更美观，伪影更少。例如，我们的方法可以更好地修复房间布局和建筑结构。补充材料中提供了更多示例。

5.结论

我们提出了一种新的用于大孔图像修复的空间扩散模型。利用所提出的解耦概率建模和空间扩散方案，我们的方法可以评估预测精度，并保留具有最低不确定性的像素作为后续处理的提示，从而产生高质量的完成。多个基准测试的最新结果证明了我们方法的有效性。

限制分析。

尽管我们的方法可以在缺失的区域生成不同的内容，但我们观察到，它倾向于在小细节而不是大结构中提供更多的变化。我们将进一步提高我们这一代人的多样性。此外，当只给出一些提示时，我们的方法有时难以理解对象。在补充材料中，我们展示了一些失败案例。

SDM:用于大孔图像修复的空间扩散模型（补充材料）

A、 体系结构详细信息

除了第3.2节和第4.2节中的描述外，我们在这里提供了更全面的架构细节说明。我们采用具有跳跃连接的U-Net架构，其中编码器将H×W输入的大小下采样到H×W 32，解码器将其上采样回H×W。在每个分辨率下，只有一个由两个3×3卷积层组成的残差块，除非另有说明。编码器和解码器都采用特征大小为H 16×W 16和H 32×W 32的注意力块，并且在这些尺度上还引入了早期卷积块。不同的注意力块使用图4中的自适应映射函数，每个映射函数由4个卷积层组成，内核大小为3×3。

输入由7个通道组成：3个通道用于彩色图像，1个通道用于初始掩模，1个用于更新掩模，一个用于不确定度图，一个通道用于时间步长。信道的数量最初被转换为64，然后在每次下采样后加倍，最多达到512，解码器采用对称设置。输出包含6个通道：3个用于平均项，3个用于对数方差项。

我们使用了一种权重调制技术，其中风格表示来自图像全局特征和随机潜码。对于全局特征，我们使用卷积层来进一步将特征大小从H 32×W 32降采样到H 256×W 256，并使用全局池化层来获得一维表示。通过8个完全连接的层从高斯噪声映射随机潜码。

B、 1024×1024分辨率的推广

为了评估模型的泛化能力，我们比较了我们的空间扩散模型（SDM）、MAT[34]和LaMa[56]，这些模型在1024×1024分辨率的512×512个地点2[82]上训练。如表B.1所示，尽管使用了较少的训练样本，但我们的SDM在所有指标上的表现明显优于MAT和LaMa。值得注意的是，我们的方法导致FID改进约1.9。我们不涉及去噪扩散模型（例如LDM）和其他基于GAN的模型（例如CoModGAN）进行比较，因为将它们放大到1024×1024分辨率是不切实际的。

C、 512×512 LPIPS结果

LPIPS[79]也是图像修复中广泛使用的感知度量。为了与现有技术方法进行全面比较，我们在表C.2中提供了LPIPS结果。我们认为，LPIPS可能不适合大孔图像修复，因为它是逐像素计算的。此措施仅供参考。

D、 256×256 CelebA HQ结果

我们还对256×256 CelebA HQ[23]数据集进行了定量比较。如表D.3所示，我们的方法在所有方法中取得了最佳性能。

E、 额外的复数生成

如图6和第5节所示，我们的方法也支持多元生成。从图E.1中的视觉示例中，我们看到差异主要在于细粒度细节。我们将致力于改善发电多样性。

F、 故障案例

如第5节所述，当提供有限的线索时，我们的模型有时无法恢复受损物体。我们在图F.2中显示了一些故障案例。例如，笔记本的缺失部分填充了背景，恢复的总线结构不完整。我们将原因之一归结为缺乏高级语义理解。我们将进一步提高模型的生成能力。

G、 其他定性比较

我们在图G.3、图G.4、图G.5、图G.6和图G.7中提供了512×512 Places2[82]和CelebA HQ[23]的更多视觉示例。与其他方法相比，我们的SDM生成了更逼真的照片和语义一致的内容。例如，我们的方法成功地恢复了人的腿、飞机结构以及更真实的室内和室外场景。