Bi-Noising Diffusion: Towards Conditional Diffusion Models with Generative Restoration Priors (Paper reading)

Kangfu Mei, Johns Hopkins University, USA, arXiv, Cited:0, Code, Paper.

前言

条件扩散概率模型可以对自然图像的分布进行建模，并可以根据给定的条件生成不同的真实样本。然而，对于可观察到的颜色偏移和纹理，它们的结果往往是不现实的。我们认为，这个问题是由模型学习的概率分布和自然图像分布之间的差异造成的。在每个采样时间步长期间，微妙的条件逐渐扩大了发散。为了解决这个问题，我们引入了一种新的方法，该方法使用预训练的无条件扩散模型将预测样本带入训练数据流形。无条件模型充当正则化子，并减少条件模型在每个采样步骤引入的发散。我们进行了全面的实验，以证明我们的方法在超分辨率、彩色化、湍流去除和图像去噪任务上的有效性。通过我们的方法获得的改进表明，先验可以作为改进条件扩散模型的通用插件。

2. 整体思想

本文通过实验发现，条件扩散模型类似于SR3的方法，训练的去噪器会引入额外噪声。因为这个去噪器加入了条件，这个去噪器学习的不是自然图像的先验。因此本文，使用了两个扩散模型（有条件和无条件的），将无条件的丰富去噪先验融入到有条件中。具体方法是，在每一个中间的扩散时间步在有条件的扩散模型获得 $\hat x_{0}$ 的粗糙估计，然后利用前向扩散过程将估计采样到当前时间步，最后使用无条件的扩散模型逆扩散。

3. 方法

在本节中，我们讨论了所提出的将嵌入先验添加到扩散模型中的机制。为了一致性，我们将无条件扩散模型的中间输出表示为 $\epsilon_{\theta}(\cdot)$ ，在接下来的去噪扩散概率模型（DDPM）的讨论中用 $\theta$ 参数化。附加的、有条件的扩散模型用 $f_{\phi}(\cdot)$ 表示，条件（即退化图像）和自然图像对用 $y_{0}, x_{0}$ 表示；其中具有参数 $\phi$ 的条件扩散模型 $f_{\phi}(\cdot)$ 表示时间步长 $t$ 处concat条件 $y_{0}$ 对噪声图像 $x_{t}$ 进行去噪。

在本文的实验中，我们将最近的带有条件的学习优化扩散过程扩散模型表示为Learning to Refine Diffusion Process (LRDP)。LRDP对给定退化图像的干净图像的条件分布进行建模，用于恢复学习，因此它需要针对不同的任务或数据集进行单独的训练。此学习过程的目标如下：
$L_{vlb} = E_{t \sim [1, T], \epsilon \sim N(0,I)}[||\epsilon-\epsilon_{\theta}(y_{0}, x_{t}, t)||^{2}]$
这里的 $x_{t}$ 是前扩散过程获得的。LRDP需要针对不同的恢复任务从头开始重新训练，这进一步增加了训练成本。实验发现，与没有 $y_{0}$ 作为条件的协议相比，这样的协议降低了生成的视觉质量。对性能下降最直接的假设是，由于退化模型的不明确性，假设的后验 $p(y_{0}|x_{0})$ 与扩散过程 $p_{θ}(x_{0}|y_{0}) \propto p_{θ}(x_{0})p_{θ}(y_{0}|x_{0})$ 形成对比。因此，我们声称将扩散生成过程分解为不同的协议应该是处理恢复任务的一种更有前途的方法。

3.1 Conditioning on Diffusion Process (CDP)

这类方法通常优化采样过程的隐变量，如下(把 $y$ 换成 $x$ 阅读)，其中 $\sigma(\cdot)$ 是硬阈值转换函数，目的结合 $y_{0}$ 和 $x_{t-1}$ 用于精确的复原。比如ILVR。

在这里插入图片描述

3.2 Implicit Error-feedback Diffusion Priors

由于扩散模型遵循时间序列过程，每个步骤中的误差和视觉伪影传播并相加，因此严重降低了一些CDP结果的质量。然而，在无条件扩散模型中很少观察到这样的问题。我们认为，这种差异来自于条件反射，打破了每个采样时间步长噪声的固有概率分布，导致它们偏离自然图像的流形。因此，我们建议应用嵌入预训练的无条件模型中的生成先验来正则化在条件模型的每个时间步长预测的噪声。具体的方法是训练两个扩散模型，有条件的 $f_{\phi}(\cdot)$ 和无条件的 $\epsilon_{\theta}(\cdot)$ 。这里从采样的第1步举例，给定初始随机高斯分布 $x_{t}$ 估计 $x_{0}$ ，注意我们用有条件的去噪器估计：
$\tilde x_{0} = (x_{t} - \sqrt{1-\bar \alpha_{t}}f_{\phi}(y_{0}, x_{t},t)) / \sqrt{\bar \alpha_{t}}$
然后，我们将在有条件下估计的 $\tilde x_{0}$ 前向传播到当前步：
$q(x_{t}'|\tilde x_{0}):=x_{t}'=\sqrt{\bar \alpha_{t}}\tilde x_{0}+ \sqrt{1-\bar \alpha_{t}}\epsilon,\,\epsilon \sim N(0,I)$
注意这里的 $x_{t}'$ 和 $x_{t}$ 是不一样的， $x_{t}$ 是无条件扩散模型中的隐变量。这个过程可以看作加入了新的噪声，这个噪声是有指导性的，也就是条件。然后我们利用无条件扩散模型进行采样：
$x_{t-1}=\frac{1}{\sqrt{\alpha_{t}}}(x_{t}'-\frac{1-\alpha_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\epsilon_{\theta}(x_{t}',t))+\sigma_{t}z\tag{4}$

4. 实验

在FFHQ[30]数据集（70000张图像）上训练并在CelebA HQ数据集（前3000张图像）（分辨率为256×256像素）上评估的条件图像恢复。
条件图像恢复是在FFHQ数据集上训练的4×人脸超分辨率，并在CelebA HQ数据集（前3000张图像）上评估，分辨率为256×256像素。
在Rain800数据集和Jorder 200L数据集及其各自的训练集上进行的图像去雨。扩散模型的分辨率为256×256像素。

请注意，对于前两个任务，扩散模型是在FFHQ数据集上训练的，用于人脸生成。对于最后一个任务，在ImageNet数据集上训练扩散模型，用于自然图像生成。

在这里插入图片描述