Diffusion Model for Generative Image Denoising (Paper reading)

Diffusion Model for Generative Image Denoising

Yutong Xie, Peking University, arXiv, Cited:0, Code: 无, Paper.

1. 前言

在图像去噪的监督学习中，通常是收集或合成成对的干净图像和有噪声图像来训练去噪模型。使用L2范数损失或其他距离函数作为训练的目标函数。它通常会导致图像细节较少的过度平滑结果。在本文中，我们把去噪任务看作是一个以噪声图像为条件估计干净图像的后验分布的问题。应用扩散模型的思想实现了生成图像去噪。根据去噪任务中的噪声模型，我们重新定义了与原扩散过程不同的扩散过程。因此，后验分布的采样是从有噪声的图像开始的几十步的逆向过程。我们考虑了三种类型的噪声模型，高斯噪声，伽马噪声和泊松噪声。在理论的保证下，导出了统一的模型训练策略。通过对三种噪声模型的实验验证了该方法的有效性，取得了良好的效果。

2. 整体思想

本文观察到逆扩散过程与高斯噪声去除很像，如果扩散过程与噪声模型完全一致，就有可能从有噪声的图像开始反向过程，而不是完全的高斯噪声，从而大大减少迭代。为此，我们提出了一种新的任务去噪扩散模型。在我们的方法中，扩散过程是根据特定的噪声模型设计的，使它们是一致的。具体而言就是，推导出具体任务的后验概率$p_{\theta }(x_t|x_{t+1})$将采样过程替换。

3. 方法

3.1 高斯去噪

高斯噪声模型可以定义为:$x_N=x_0+\sigma z,z\sim N(0,I)$。选择一个单调递增的序列$0={\sigma }_0<{\sigma }_1<...<{\sigma }_{N-1}<{\sigma }_N=\sigma$，可以得到:
$$x_t=x_0+(\sum _{i=0}^{t-1}\sqrt{{\sigma }_{i+1}^2-{\sigma }_i^2}{z}_i)\text{\hspace{1em}(1)}$$
其中，$t=1,...,N$，该公式的推导可从Stochastic Image Denoising By Sampling from the Posterior Distribution获得。这就是前向扩散过程，那么如何定义逆扩散过程？与DDPM推导$p_{\theta }(x_t|x_{t+1})$过程一样，得到$x_t$的均值和方差分别为
$${\tilde{\mu }}_t=\frac{{\sigma }_t^2}{{\sigma }_{t+1}^2}{x}_{t+1}+\frac{{\sigma }_{t+1}^2-{\sigma }_t^2}{{\sigma }_{t+1}^2}{x}_0\text{\hspace{1em}(2)}$$
$${\tilde{\sigma }}_t=\frac{{\sigma }_t}{{\sigma }_{t+1}}\sqrt{{\sigma }_{t+1}^2-{\sigma }_t^2}\text{\hspace{1em}(3)}$$
而本文选择的是训练一个去噪器来获得(2)中的$x_0$，如下：

那么采样过程，是从噪声图像开始，一步一步去噪，如下：

3.2 伽马去噪

伽马噪声的去除过程与高斯噪声是相似的思想，推导过程建议看原文。

3.3 柏松去噪

4. 实验结果

4.1 数据集和实现细节

在两个基准数据集(Kodak数据集和CSet9数据集)中评估了所提出的方法。使用DIV2K和CBSD500数据集作为训练数据集。原始图像是彩色RGB自然图像，我们在训练和测试时将其转换为灰色图像。我们使用传统的监督学习，并将L2范数损失作为基线进行比较。所有方法都使用了相同的修改过的U-Net，该U-Net具有大约7000万个参数。训练时，我们将训练图像随机剪辑到分辨率为128 × 128的patch中。AdamW优化器用于训练网络。我们训练每个模型的批量大小为32个。为了减少内存，我们利用累积梯度和混合精确训练的技巧。学习率设置为1×10−4。所有的模型都是在PyTorch和NVidia V100中实现的。

4.2 实验结果

上图显示了不同噪声模型生成的样本。与监督学习相比，我们生成的结果在视觉上令人愉悦，包含更好的图像细节。生成样本的良好视觉质量验证了我们的方法是可行的生成降噪任务。我们还计算PSNR和SSIM来评估我们的方法，如表1所示。我们发现监督学习和生成样本之间在指标方面存在差距。这是可以理解的。由于噪声图像丢失了大量原始图像信息，很难生成与丢失图像完全相同的细节。

结论

之前的扩散模型用于low-level的工作，训练测试在特定类型的数据集，如人脸，猫狗。这与主流的low-level所用，如Set12，BSD68等数据集，不一致。因为生成模型需要训练在特定类别的数据集。本文，提出了将扩散模型在low-level的应用更加一般化。自认为的一些问题：
1）该类思想已经有了很多的工作，本文的想法并不是很新颖。此外，与类似“对于特定任务进行数学推导“的工作相比，本文采样过程并没有加入额外噪声，这样优点在于结果与原图整体保持一致，但是这已经不能算是生成模型了，顶多是多次迭代细化的去噪过程。
2）在结果图中可以看到，本文的PSNR指标并不高，可视化结果也并不算很出色。
3）基于数学推导的不好之处在于，本文要知道$\sigma$等噪声强度的先验，在理论上是可行的。用于真实世界的任务中，提出的模型就需要改进。