1 概述

1.1 题目

  2022CVPR：用于图像修复的去噪扩散概率模型 (RePaint: Inpainting using Denoising Diffusion Probabilistic Models)

1.2 代码

  Torch：https://github.com/andreas128/RePaint

1.3 摘要

  自由形式修复 (Free-form inpainting) 是在任意二进制蒙版指定的区域中，向图像添加新内容的任务。大多数现有方法：
  1）针对特定的蒙版分布进行训练，这限制了它们在未知蒙版上的泛化能力；
  2）像素级和感知损失的训练通常导致缺失区域的简单纹理扩张，以致不能引导语义上有意义的生成。
  本文提出一种基于去噪扩散概率模型 (DDPM) 的修复方法，以适用于多种类型的蒙版：
  1）使用预训练的无条件DDPM作为生成先验 (Generative prior)；
  2）为了调节生成过程，仅采样给定图像的无蒙版区域，以调控逆扩散迭代；
  3）不修改或者不调节原始的DDPM网络，因此对于任意的修复形式可以产生高质量和多样性的输出图像。

1.4 Bib

@inproceedings{Lugmayr:2022:1146111471,
author		=	{Andreas Lugmayr and Martin Danelljan and Andres Romero and Fisher Yu and Radu Timofte and Luc Van-Gool},
title		=	{Repaint: {I}npainting using denoising diffusion probabilistic models},
booktitle	=	{{IEEE/CVF} Conference on Computer Vision and Pattern Recognition},
pages		=	{11461--11471},
year		=	{2022}
}
1
2
3
4
5
6
7

2 去噪扩散修复模型

  DDPM可以习得训练集中图像的分布：
  1）推理阶段：随机采样噪声向量$x_T$并逐渐去噪，直到输出高质量图像$x_0$；
  2）训练阶段定义扩散过程，在$T$个时间步内将图像$x_0$转换为高斯白噪声$x_T\sim \mathcal{N}(0,1)$。每一步的前进方向给定如下：
$$q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_t\mathbf{I})\text{\hspace{1em}(1)}$$  $x_t$包含两个部分，在时间步$t$关于变量${\beta }_t$的独立同分布高斯噪声和通过缩放的先前样本$x_{t-1}$。
  逆过程使用神经网络建模，该网络描述了带有参数${\mu }_{\theta }(x_t,t)$和${\Sigma }_{\theta }(x_t,t)$的高斯分布：
$$p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))\text{\hspace{1em}(2)}$$  模型(2)的学习目标通过变分下界 (Variational lower bound) 得出：
E [ − log ⁡ p θ ( x 0 ) ] ≤ E q [ − log ⁡ p θ ( x 0 : T ) q ( x 1 : T ∣ x 0 ) ] = E q [ − log ⁡ p ( x T ) − ∑ t ≥ 1 log ⁡ p θ ( x t − 1 ∣ x t ) q ( x t ∣ x t − 1 ) ] = L (3) \tag{3}

​E[−logpθ​(x0​)]≤Eq​[−logq(x1:T​∣x0​)pθ​(x0:T​)​]=Eq​[−logp(xT​)−t≥1∑​logq(xt​∣xt−1​)pθ​(xt−1​∣xt​)​]=L​(3)  通过Ho等人的扩展，该损失可以分解为：
E q [ D K L ( q ( x T ∣ x 0 ) ∥ p ( x T ) ) ⏟ L T + ∑ t > 1 D K L ( q ( x t − 1 ∣ x t , x 0 ) ∥ p θ ( x t − 1 ∣ x t ) ) ⏟ L t − 1 − log ⁡ p θ ( x 0 ∣ x 1 ) ⏟ L 0 ] (4) \tag{4}

Eq[DKL(q(xT∣x0)‖p(xT))⏟LT+∑t>1DKL(q(xt−1∣xt,x0)‖pθ(xt−1∣xt))⏟Lt−1−log⁡pθ(x0∣x1)⏟L0]" role="presentation" style="position: relative;">Eq[DKL(q(xT∣x0)∥p(xT))LT+∑t>1DKL(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))Lt−1−logpθ(x0∣x1)L0]$$\begin{array}{rl}& {\mathbb{E}}_q[\underset{L_T}{\underbrace{D_{\mathrm{K}\mathrm{L}}\left(q\left({\mathbf{x}}_T\mid {\mathbf{x}}_0\right)\Vert p\left({\mathbf{x}}_T\right)\right)}}\\ & +\sum _{t>1}\underset{L_{t-1}}{\underbrace{D_{\mathrm{K}\mathrm{L}}\left(q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)\Vert p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)\right)}}\underset{L_0}{\underbrace{-\log p_{\theta }\left({\mathbf{x}}_0\mid {\mathbf{x}}_1\right)}}]\end{array}$$

​Eq​[LT​ DKL​(q(xT​∣x0​)∥p(xT​))​​+t>1∑​Lt−1​ DKL​(q(xt−1​∣xt​,x0​)∥pθ​(xt−1​∣xt​))​​L0​ −logpθ​(x0​∣x1​)​​]​(4)  $L_{t-1}$项训练网络(2)以执行一次逆扩散步骤。此外，由于$q(x_{t-1}|x_t,x0)$是高斯的，因此它保证了目标表达式的封闭。模型参数化的最佳方式是预测添加到当前图像$x_t$的累计噪声${ϵ}_0$。因此，均值${\mu }_{\theta }$预测如下：
$${\mu }_{\theta }\left(x_t,t\right)=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }\left(x_t,t\right)\right)\text{\hspace{1em}(5)}$$  简化的优化目标如下：
$$L_{simple}=E_{t,x_0,ϵ}[\Vert ϵ-{ϵ}_{\theta }(x_t,t){\Vert }^2]\text{\hspace{1em}(6)}$$  如Nicho 和Dhariwal所述，学习逆过程(2)中的方差${\Sigma }_{\theta }(x_t,t)$有助于将采样步骤的次数降低一个数量级，因此他们添加了变分下界损失。后继将通过因子4 (Factor four.) 来降低推理的时间开销。
  为了训练DDPM，需要样本$x_t$和一个能够转换$x_0$到$x_t$的噪声。通过使用在每个步骤(1)中添加噪声的独立性质，总损失变量计算为：
$${\overline{\alpha }}_t=\prod _{s=1}^t(1-{\beta }_s)$$  则公式(1)重写为：
$$q(x_t|x_0)=\mathcal{N}(x_t;{\overline{\alpha }}_t{x}_0,(1-{\overline{\alpha }}_t)\mathbf{I})\text{\hspace{1em}(7)}$$  这允许有效采样成对训练数据以训练逆转换。

3 方法

  首先说明如何调整无条件DDPM逆扩散以进行图像修复，然后介绍一种改进的逆过程方法。

3.1 调整已知区域

  修复的目的是预测图像以蒙版作为条件的缺失像素，接下来考虑训练一个无条件DDPM。令$x$表示原始真实图像、$m\odot x$表示未知像素，以及$(1-m)\odot x$表示已知像素。
  由于公式(2)中$x_t$到$x_{t-1}$的每个逆步骤都仅依赖于$x_t$，因此需要在保持对应分布的正确性质下更改已知区域$(1-m)\odot x_t$。由于前进方向由添加了高斯噪声的马尔可夫链，即公式(1)定义，可以通过公式(7)采样任意时间点的中间图像$x_t$，同理可以采样已知区域$m\odot x_t$。因此，基于公式(2)的未知区域采样和公式(7)的已知区域采样，可以实现每个逆步骤的以下表达：
x t − 1 known  ∼ N ( α ˉ t x 0 , ( 1 − α ˉ t ) I ) x t − 1 unknown  ∼ N ( μ θ ( x t , t ) , Σ θ ( x t , t ) ) x t − 1 = m ⊙ x t − 1 known  + ( 1 − m ) ⊙ x t − 1 unknown 

xt−1known ​xt−1unknown ​xt−1​​∼N(αˉt​ ​x0​,(1−αˉt​)I)∼N(μθ​(xt​,t),Σθ​(xt​,t))=m⊙xt−1known ​+(1−m)⊙xt−1unknown ​​  $x_{t-1}^{\text{known}}$是在给定图像$m\odot x_0$上使用已知像素的采样，而$x_{t-1}^{\text{unkonwn}}$是在先前迭代$x_t$上基于模型的采样。结合蒙版后，将生成新样本$x_{t-1}$，如图2。

图2：RePaint更高标准去噪过程以调整给定图像的内容。在每一步中，将从输入采样已知区域，以及采样DDPM输出的已修复部分

3.2 重采样

  直接应用章节3.1描述的方法时，只是将内容的类型与已知区域匹配。例如图3中，当$n=1$时，修复区域的质地与狗狗的头发匹配，这在语义上是不正确的。因此，DDPM虽然利用了已知区域的上下文，但它并不能很好地与图像的其余部分协调一致。接下来讨论这种现象的可能原因。

  图3：应用$n$采样步的性能。$n=1$作为DDPM基准线，$n=2$是一次采样步，随着采样步的增加生成的图像更加协调。性能大约在10次采样步后满足需求

  图2分析DDPM如何调节已知区域。公式(8)说明模型使用$x_t$来预测$x_{t-1}$，其由DDPM的输出和已知区域的采样组成。然而，基于公式(7)的已知像素采样并未考虑图像的生成部分，这将导致不和谐。尽管模型于每一步中尽量生成和谐的图像，但是这似乎并不奏效因为下一步开始时依然没有解决上述问题。此外，在每个逆步骤中，${\beta }_t$ 的方差限制会减小图像的最大变化。最终，在下一步去噪开始前，模型需要更多的时间使用$x_{t-1}^{\text{unknown}}$来和谐$x_{t-1}^{\text{known}}$。
  训练后的DDPM用以生成服从数据分布的图像，其自然也需要生成一致的结构。 在重采样方法中，DDPM这个属性被用于协调模型的输入。因此，通过公式(1)来采样$x_t\sim \mathcal{N}(\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_t\mathbf{I})$以将输入$x_{t-1}$扩散回$x_t$。尽管该操作缩放输出以及添加噪声，一些合并在生成区域$x_{t-1}^{\text{unknown}}$的信息依然保留其中，这使得新生成的$x_t^{\text{unknown}}$能够更和谐于$x_t^{\text{known}}$且包含一些其本身的调节信息。
  由于此操作只能协调一个步骤，因此可能无法在整个去噪过程中合并语义信息。为了克服这个问题，将此操作的时间范围表示为跳跃长度，$j=1$对应于前一种情况。类似于扩散速度的标准变化，也就是减速，重采样也增加了逆扩散的运行时间。减速通过减少每个去噪步骤中增加的方差来应用更小但更多的重采样步骤。这是一种本质不同的方法，因为减缓扩散仍然存在不协调图像的问题。RePaint的伪代码如算法1。