SDEdit GUIDED IMAGE SYNTHESIS AND EDITING WITH STOCHASTIC DIFFERENTIAL EQUATIONS

SDEdit: GUIDED IMAGE SYNTHESIS AND EDITING WITH STOCHASTIC DIFFERENTIAL EQUATIONS

TL; DR：

选择一个中间步 $t_0$​，将引导图进行加噪并重新执行生成。现在 diffusers 库中 SD Img2Img 的标准做法。

---

背景知识：基于SDE的图像生成模型

随机微分方程（Stochastic differential equations，SDEs）可以看作是常微分方程（ordinary differential equations，ODEs）通过向动力学过程中注入随机噪声得到的泛化版本。SDE 的解是一个关于时间的随机变量（即随机过程），我们将其定义为 $\mathbf{x}(t)\in {\mathbb{R}}^d$，其中 $t\in [0,1]$ 表示时间。在图像生成中，我们假设 $\mathbf{x}(0)\sim p_0=p_{\text{data}}$ 表示来自数据分布 $p_{\text{data}}$ 的的一个样本，我们可以通过一个高斯前向 SDE 来得到 $\mathbf{x}(t),t\in (0,1]$，即在给定 ${\mathbf{x}}_0$ 的条件下，${\mathbf{x}}_t$ 为一个高斯分布：
$$\mathbf{x}(t)=\alpha (t)\mathbf{x}(0)+\sigma (t)\mathbf{z},\mathbf{z}\sim \mathcal{N}(0,\mathbf{I})$$
其中 $\sigma (t):[0,1]\to [0,\infty ]$ 是一个标量函数，用来控制噪声 $\mathbf{z}$ 的幅度，$\alpha (t)$ 也是一个标量函数，用来控制数据 $\mathbf{x}(0)$ 的幅度。${\mathbf{x}}_t$ 的概率密度函数记作 $p_t$。

通常有两类图像生成 SDE，分别为 Variance Exploding SDE（VE SDE）和 Variance Preserving SDE（VP SDE）。

• VE SDE 对所有 $t\in [0,1]$ 有 $\alpha (t)=1$，而 $\sigma (1)$ 是一个极大的常数，从而使得当 $t=1$ 时，分布 $p_1\sim \mathcal{N}(0,{\sigma }^2(1)\mathbf{I})$；
• VP SDE 对所有 $t\in [0,1]$ 满足 ${\alpha }^2(t)+{\sigma }^2(t)=1$，而当 $t\to 1$ 时，$\alpha (t)\to 0$，从而分布 $p_1\sim \mathcal{N}(0,\mathbf{I})$

可以看到，这两种 SDE 都是在 $t$ 从 0 到 1 的过程中将数据分布 $p_0\sim p_{\text{data}}$ 逐渐转变为高斯分布 $p_1\sim \mathcal{N}$，区别是 VE SDE 方差非常大，故称为方差爆炸；VP SDE 方差收敛到 1，故称为方差保持。简单起见，以下正文我们基于 VE SDE 展开研究，原文附录中也介绍了 VP SDE 的过程。这两种 SDE 只是在形式上有所不同，但是本质上是一样的。

VE SDE 是如何进行图像生成的呢？基于以上定义，我们可以将图像生成的过程看作是逐渐从纯噪声 $\mathbf{x}(t)$ 中去除噪声，直到得到数据样本 ${\mathbf{x}}_0$ 的过程。这个过程可以通过解从时间 $t=1$ 到 $t=0$ 的反向 SDE 来实现。具体来说，VE SDE 的反向 SDE 为：
$$d\mathbf{x}(t)=[-\frac{d[{\sigma }^2(t)]}{dt}{\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})]dt+\sqrt{\frac{d[{\sigma }^2(t)]}{dt}}d\tilde{\mathbf{w}}$$
其中 $\tilde{\mathbf{w}}$ 表示时间从 1 到 0 的反向时的维纳过程。如前所述，在 VE SDE 中，我们将初始噪声设置为高斯分布 $\mathbf{x}(1)\sim p_1(\mathbf{x})=\mathcal{N}(0,{\sigma }^2(t)\mathbf{I})$。这个反向 SDE 中我们唯一未知的是噪声扰动分布 $p_t(\mathbf{x})$ 的得分函数 ${\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})$，我们通过训练一个神经网络 ${\mathbf{s}}_{\theta }(\mathbf{x}(t),t)$ 来预测它。如何训练这个神经网络呢，则是通过 denoising score matching 方法，其在时间 $t$ 的目标函数为：
$$L_t={\mathbb{E}}_{\mathbf{x}\sim p_{\text{data}},\mathbf{z}\sim \mathcal{N}(0,\mathbf{I})}[||{\sigma }_t{\mathbf{s}}_{\theta }(\mathbf{x}(t),t)-\mathbf{z}|{|}_2^2]$$
整体的目标函数则是每个 $L_t$ 的加权和。具体的加权函数在之前的工作中已有介绍。

神经网络 ${\mathbf{s}}_{\theta }(\mathbf{x}(t),t)$ 训练好之后，用来估计噪声扰动分布的得分 ${\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})$。有了这一项之后，整个反向 SDE 就确定下来，然后就可以通过 Euler-Maruyama 方法来求解这个 SDE。具体来说，从 $t+\Delta t$ 到 $t$ 的更新公式为：
$$\mathbf{x}(t)=\mathbf{x}(t+\Delta t)+({\sigma }^2(t)-{\sigma }^2(t+\Delta t)){\mathbf{s}}_{\theta }(\mathbf{x}(t),t)+\sqrt{{\sigma }^2(t)-{\sigma }^2(t+\Delta t)}\mathbf{z}$$
再具体实现时，我们可以对时间 $t$ 进行特定的离散化，然后初始化 $\mathbf{x}(1)\sim \mathcal{N}(0,{\sigma }^2(1)\mathbf{I})$ 并根据上式迭代，即可生成图像 $\mathbf{x}(0)$。

SDEdit

接下来，我们就介绍 SDEdit 如何基于无条件生成训练的 SDE 生成模型，来进行带引导的图像生成。

任务设定 用户需要提供一张像素空间的、全尺寸的图片 ${\mathbf{x}}^{(g)}$​，本文将其称为 引导图 (guide)。引导图中包含的细节层级可以有多种选择：1）Strock Painting to Image：粒度较粗的色块引导图；2）Stroke-based Editing 在真实图片上添加色块的引导图；3）Image Compisiting：在真实图片上添加图像块的引导图。下入展示了一些例子。

我们对生图结果有两方面的要求：

1. 真实性 Realism：生成结果确实一张真实图片（不能包含色块等）
2. 忠实性 Faithfulness：生成结果要与 ${\mathbf{x}}^{(g)}$ 相近（不能是毫无关系的图片）

考虑一张与引导图无关的随机图片，真实性很高却没有忠实性；而考虑引导图本身，忠实性很高却没有重视性。因此，真实性和忠实性这两点要求是相互影响的，我们需要在它们之间进行 trade-off。值得一提的是，不同于常规的 inverse problem，这里从真实图片到引导图的 inverse 映射是未知的，这里没有采用 inverse 技术和成对的数据集。

本文所提出的 SDEdit，基于一个事实：反向 SDE 不仅可以从 $t_0=1$ 开始，也可以从任意中间时间点 $t_0\in (0,1)$ 开始。我们需要选择一个合适的初始化点，从根据引导图生成出想要的真实且忠实的结果。具体来说，SDEdit 的过程可以表达为：对于任意给定的引导图 ${\mathbf{x}}^{(g)}$，采样（加噪声） ${\mathbf{x}}^{(g)}(t_0)\sim \mathcal{N}({\mathbf{x}}^{(g)},{\sigma }^2(t_0)\mathbf{I})$，然后迭代执行式 4，得到生成结果 $\mathbf{x}(0)$。我们用 $\text{SDEdit}({\mathbf{x}}^{(g)},t_0,\theta )$ 表示这一过程。

这里的关键就是如何选择初始时间点 $t_0$，$t_0$ 的选择事关真实性和忠实性二者之间的 trade-off。为了研究这一点，作者在 LSUN 数据集上，使用第一类引导图 Strock Painting to Image，对 $t_0\in [0,1]$ 生成结果的真实性和忠实性进行了量化实验。为了量化真实性，作者采用了 KID 指标，如果生成结果和真实图片的 KID 比较小，则可以认为其真实性较高；为了量化忠实性，作者选择计算生成结果和引导图之间的 L2 距离。从下图的结果来看，随着 $t_0$ 的增加，真实性上升而忠实性下降。

真实性和忠实性的权衡还有一个理解角度：如果引导图本身就就和所有的真实图像差距很大，那么为了生成一张真实性较高的图像，我们就必须接受一定程度地远离引导图（不忠实）。这可以表达为下述命题：

假设对于所有的 $\mathbf{x}\in \mathcal{X},t\in [0,1]$，都有 $||{\mathbf{s}}_{\theta }(\mathbf{x},t)|{|}_2^2\le C$，那么对于所有 $\delta \in (0,1)$，至少有 $1-\delta$ 的概率：
$$||{\mathbf{x}}^{(g)}-\text{SDEdit}({\mathbf{x}}^{(g)},t_0,\theta )|{|}_2^2\le {\sigma }^2(t_0)(C{\sigma }^2(t_0)+d+2\sqrt{-d\cdot \log \delta }-2\log \delta )$$
其中 $d$ 是引导图 ${\mathbf{x}}^{(g)}$ 的维度。

原文附录中提供了详细的证明。直观来讲，引导图和生成结果之间的差异可以看作由得分模型的输出和随机高斯噪声两部分组成，这两部分随着 $t_0$ 的增加都是增加的，从而引导图和生成结果之间的差异也会增加。上述命题指出，为了让生成结果具有较高的真实性，我们需要一个足够大的 $t_0$。反过来说，如果 $t_0$ 太大，生成结果对于引导图的忠实性会下降。极端情况下（$t_0=1$），SDEdit 相当于无引导条件的生成。

那么，究竟该如何选择合适的 $t_0$ 呢？作者首先指出，生成的质量与引导图的质量也有很大关系。对于合理的引导图，作者发现 $t_0\in [0.3,0.6]$ 就能工作得很好。然而，如果引导图比较 “异常”，比如是一张纯白图片，那么其与真实图片的数据分布差异很大，此时我们就必须选择一个较大的 $t_0$，牺牲掉一定的忠实性，来提高真实性。在实际应用中，如果是带交互式的应用，我们可以将 $t_0$ 初始化到 0.3 和 0.6 之间，然后询问用户需要提高忠实性还是提高真实性，根据用户的反馈进行二分搜索，找到最优的 $t_0$。如果是非交互式的引用，我们可以在随机选择的图像上执行类似的二分搜索以获得 $t_0$，然后随后为同一任务中的所有指导图固定相同的 $t_0$。尽管不同的指导图可能具有不同的最佳 $t_0$，但作者实验发现同一任务中共享的 $t_0$ 就能工作得很好。

VE SDE 具体的 SDEdit 算法如下。原文附录中还介绍了 VP SDE 的 SDEdit 算法，以及带掩码的 SDEdit。

总结

比较经典的 Img2Img 做法，关键是通过选择 $t_0$，来实现忠实性和真实性的权衡。