Stochastic Image Denoising By Sampling from the Posterior Distribution (Paper reading)

Stochastic Image Denoising By Sampling from the Posterior Distribution (Paper reading)

Bahjat Kawar, Haifa(Israel), ICCV Workshop2021, Cited:22, Code:无, Paper.

1. 前言

对于严重的噪音水平，最小均方误差 (MMSE) 解决方案可能会导致模糊输出图像。本文提出了一种新颖的随机去噪方法，该方法可产生可行且高感知质量的结果，同时保持较小的 MSE。本文的方法采用 Langevin 动力学，它依赖于重复应用任何给定的 MMSE 降噪器，通过从后验分布采样获得重建图像。

2. 整体思想

整体思想是条件扩散模型，也就是在原始扩散模型上加入了梯度修正，其他和扩散模型一致。

3. 方法

从概率分布中生成样本的方法之一是利用Langevin热力学的Markov Chain Monte Carlo方法：
$$x_{t+1}=x_t+\alpha {\nabla }_{x}logp(x_t)+\sqrt{2\alpha }{z}_t\text{\hspace{1em}(1)}$$
去噪任务可以表示为:$y=x+n$，其中$x\sim p(x)$，是干净图像，$n\sim N(0,{\sigma }_0^{2}I)$是已知噪声强度的加性高斯白噪声，去噪的目标是复原$x$。虽然p(x)是未知的，但是从Solving Linear Inverse Problems Using The Prior Implicit in a Denoiser这篇文章中可以知道，我们能够从一个MMSE的去噪器中获得$p(x)$。本文提出在给定噪声输入图像的情况下，通过从后验分布中采样来复原$x$，即$p(x|y)$。

本文基于退火Langevin动力学算法解决图像去噪问题。退火Langevin动力学算法通过分数函数${\nabla }_{x}logp(\tilde{x})$从$p(x)$中进行采样，对于不同大小的$\sigma$，其中$\tilde{x}=x+z,z\sim N(0,{\sigma }^{2}I)$。为了解决去噪问题，我们需要估计后验的分数函数${\nabla }_{\tilde{x}}logp(\tilde{x}|y)$。

3.1 后验分布的分数函数和去噪算法

我们考虑噪声图像中的噪声是逐步添加上去的，即先固定一组噪声强度序列${\left\{{\sigma }_i\right\}}_{i=0}^{L+1}$，如：${\sigma }_0>{\sigma }_1>...>{\sigma }_L>{\sigma }_{L+1}=0$，这里的${\sigma }_0$是$y$的噪声强度：
$$\begin{array}{rl}{\tilde{x}}_L& =x+z_L\\ {\tilde{x}}_{L-1}& ={\tilde{x}}_L+z_{L-1}\\ {\tilde{x}}_{L-2}& ={\tilde{x}}_{L-1}+z_{L-2}\\ ⋮& \\ {\tilde{x}}_1& ={\tilde{x}}_2+z_1\\ y={\tilde{x}}_0& ={\tilde{x}}_1+z_0\end{array}\text{\hspace{1em}(2)}$$
其中，$z_i\sim N(0,({\sigma }_i^2-{\sigma }_{i+1}^2)I)$，从上述过程可以得到：
$$y={\tilde{x}}_0=x+\sum _{i=0}^L{z}_i\text{\hspace{1em}(3)}$$
其中， ${\sum }_{i=0}^L{z}_i\sim N(0,{\sigma }_0^{2}I)$，注意到$y-{\tilde{x}}_i={\sum }_{j=0}^{i-1}{z}_j$，则${\sum }_{j=0}^{i-1}{z}_j\sim N(0,({\sigma }_0^2-{\sigma }_i^2)I)$，可以理解为在噪声图上减去当前$t$时刻状态后的剩余噪声量。后面对于任意$i$时刻的$x_i$用$\tilde{x}$来简化表示。接下来求${\nabla }_{\tilde{x}}logp(\tilde{x}|y)$：
$${\nabla }_{\tilde{x}}logp(\tilde{x}|y)={\nabla }_{\tilde{x}}logp(y|\tilde{x})p(\tilde{x})\frac{1}{p(y)}={\nabla }_{\tilde{x}}\left[logp(y|\tilde{x})+logp(\tilde{x})+log\frac{1}{p(y)}\right]$$
$y$是一个不依赖于$\tilde{x}$的固定观测，最后一项的梯度为0，则
$${\nabla }_{\tilde{x}}logp(\tilde{x}|y)={\nabla }_{\tilde{x}}logp(y|\tilde{x})+{\nabla }_{\tilde{x}}logp(\tilde{x})\text{\hspace{1em}(4)}$$
其中，第一项中，$p(y|\tilde{x})=g(y-\tilde{x})$，则：
$${\nabla }_{\tilde{x}}logp(y|\tilde{x})={\nabla }_{\tilde{x}}log\left[\frac{1}{\sqrt{2\pi ({\sigma }_0^2-{\sigma }_i^2)}}exp\left[-\frac{1}{2}\frac{||y-\tilde{x}|{|}^2}{({\sigma }_0^2-{\sigma }_i^2)}\right]\right]=\frac{y-\tilde{x}}{{\sigma }_0^2-{\sigma }_i^2}$$
因此：
$${\nabla }_{\tilde{x}}logp(\tilde{x}|y)={\nabla }_{\tilde{x}}logp(\tilde{x})+\frac{y-\tilde{x}}{{\sigma }_0^2-{\sigma }_i^2}\text{\hspace{1em}(5)}$$
这篇文章中Solving Linear Inverse Problems Using The Prior Implicit in a Denoiser，我们知道分数函数和MMSE去噪器的联系，既训练一个网络来估计分数函数可以解释为去噪器估计MMSE：
$${\nabla }_{\tilde{x}}logp(\tilde{x})=\frac{\hat{x}(\tilde{x})-\tilde{x}}{{\sigma }_i^2}$$
其中$\hat{x}(\tilde{x})=\mathbb{E}(x|\tilde{x})$是一个MMSE去噪器。随机去噪算法如下：

3.2 去噪实验

去噪网络选择 noise conditional score network version 2 (NCSNv2)。针对网络的训练，本文特意强调，是像图像合成(image synthesis)任务一样选择同一类别的数据集(CelebA, FFHQ, LSUN)对网络进行训练而不是自然图像集。超参数的设定保持和Improved techniques for training score-based generative models该文章一致。实验结果如下：

虽然说，相对于MMSE去噪器，生成去噪算法更真实，更具细节，但是相对于MMSE的去噪平均来说，生成算法具有不同的解，如左边的图3，会有机会不同的4张脸，这对于某些特殊的任务来说是不友好的。


从这里我们就能看到生成去噪器的优势了，图像更自然。

总结

这里假设了噪声强度${\sigma }_0$，也就是知道了噪声强度这个强先验，在模拟噪声数据集上的表现良好，但是该算法还无法处理真实噪声，或者盲噪声，这是扩散模型应用low-level的一个重要目标。