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
+
α
∇
x
l
o
g
p
(
x
t
)
+
2
α
z
t
(1)
x_{t+1}=x_{t}+\alpha \nabla_{x}logp(x_{t})+\sqrt{2\alpha}z_{t} \tag{1}
xt+1=xt+α∇xlogp(xt)+2αzt(1)
去噪任务可以表示为:
y
=
x
+
n
y=x+n
y=x+n,其中
x
∼
p
(
x
)
x\sim p(x)
x∼p(x),是干净图像,
n
∼
N
(
0
,
σ
0
2
I
)
n\sim N(0,\sigma_{0}^{2}I)
n∼N(0,σ02I)是已知噪声强度的加性高斯白噪声,去噪的目标是复原
x
x
x。虽然p(x)是未知的,但是从Solving Linear Inverse Problems Using The Prior Implicit in a Denoiser这篇文章中可以知道,我们能够从一个MMSE的去噪器中获得
p
(
x
)
p(x)
p(x)。本文提出在给定噪声输入图像的情况下,通过从后验分布中采样来复原
x
x
x,即
p
(
x
∣
y
)
p(x|y)
p(x∣y)。
本文基于退火Langevin动力学算法解决图像去噪问题。退火Langevin动力学算法通过分数函数 ∇ x l o g p ( x ~ ) \nabla_{x}logp(\tilde{x}) ∇xlogp(x~)从 p ( x ) p(x) p(x)中进行采样,对于不同大小的 σ \sigma σ,其中 x ~ = x + z , z ∼ N ( 0 , σ 2 I ) \tilde{x}=x+z,z\sim N(0,\sigma^{2}I) x~=x+z,z∼N(0,σ2I)。为了解决去噪问题,我们需要估计后验的分数函数 ∇ x ~ l o g p ( x ~ ∣ y ) \nabla_{\tilde{x}}logp(\tilde{x}|y) ∇x~logp(x~∣y)。
3.1 后验分布的分数函数和去噪算法
我们考虑噪声图像中的噪声是逐步添加上去的,即先固定一组噪声强度序列
{
σ
i
}
i
=
0
L
+
1
\left \{\sigma_{i} \right\}_{i=0}^{L+1}
{σi}i=0L+1,如:
σ
0
>
σ
1
>
.
.
.
>
σ
L
>
σ
L
+
1
=
0
\sigma_{0}>\sigma_{1}>...>\sigma_{L}>\sigma_{L+1}=0
σ0>σ1>...>σL>σL+1=0,这里的
σ
0
\sigma_{0}
σ0是
y
y
y的噪声强度:
x
~
L
=
x
+
z
L
x
~
L
−
1
=
x
~
L
+
z
L
−
1
x
~
L
−
2
=
x
~
L
−
1
+
z
L
−
2
⋮
x
~
1
=
x
~
2
+
z
1
y
=
x
~
0
=
x
~
1
+
z
0
(2)
\begin{aligned} \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}\\ \vdots\\ \tag{2} \tilde{x}_{1}&=\tilde{x}_{2}+z_{1}\\ y=\tilde{x}_{0}&=\tilde{x}_{1}+z_{0} \end{aligned}
x~Lx~L−1x~L−2⋮x~1y=x~0=x+zL=x~L+zL−1=x~L−1+zL−2=x~2+z1=x~1+z0(2)
其中,
z
i
∼
N
(
0
,
(
σ
i
2
−
σ
i
+
1
2
)
I
)
z_{i}\sim N(0, (\sigma_{i}^{2}-\sigma_{i+1}^{2})I)
zi∼N(0,(σi2−σi+12)I),从上述过程可以得到:
y
=
x
~
0
=
x
+
∑
i
=
0
L
z
i
(3)
y=\tilde{x}_{0}=x+\sum_{i=0}^{L}z_{i} \tag{3}
y=x~0=x+i=0∑Lzi(3)
其中,
∑
i
=
0
L
z
i
∼
N
(
0
,
σ
0
2
I
)
\sum_{i=0}^{L}z_{i}\sim N(0, \sigma_{0}^{2}I)
∑i=0Lzi∼N(0,σ02I),注意到
y
−
x
~
i
=
∑
j
=
0
i
−
1
z
j
y-\tilde{x}_{i}=\sum_{j=0}^{i-1}z_{j}
y−x~i=∑j=0i−1zj,则
∑
j
=
0
i
−
1
z
j
∼
N
(
0
,
(
σ
0
2
−
σ
i
2
)
I
)
\sum_{j=0}^{i-1}z_{j}\sim N(0,(\sigma_{0}^{2}-\sigma_{i}^{2})I)
∑j=0i−1zj∼N(0,(σ02−σi2)I),可以理解为在噪声图上减去当前
t
t
t时刻状态后的剩余噪声量。后面对于任意
i
i
i时刻的
x
i
x_{i}
xi用
x
~
\tilde{x}
x~来简化表示。接下来求
∇
x
~
l
o
g
p
(
x
~
∣
y
)
\nabla_{\tilde{x}}logp(\tilde{x}|y)
∇x~logp(x~∣y):
∇
x
~
l
o
g
p
(
x
~
∣
y
)
=
∇
x
~
l
o
g
p
(
y
∣
x
~
)
p
(
x
~
)
1
p
(
y
)
=
∇
x
~
[
l
o
g
p
(
y
∣
x
~
)
+
l
o
g
p
(
x
~
)
+
l
o
g
1
p
(
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]
∇x~logp(x~∣y)=∇x~logp(y∣x~)p(x~)p(y)1=∇x~[logp(y∣x~)+logp(x~)+logp(y)1]
y
y
y是一个不依赖于
x
~
\tilde{x}
x~的固定观测,最后一项的梯度为0,则
∇
x
~
l
o
g
p
(
x
~
∣
y
)
=
∇
x
~
l
o
g
p
(
y
∣
x
~
)
+
∇
x
~
l
o
g
p
(
x
~
)
(4)
\nabla_{\tilde{x}}logp(\tilde{x}|y)=\nabla_{\tilde{x}}logp(y|\tilde{x})+\nabla_{\tilde{x}}logp(\tilde{x}) \tag{4}
∇x~logp(x~∣y)=∇x~logp(y∣x~)+∇x~logp(x~)(4)
其中,第一项中,
p
(
y
∣
x
~
)
=
g
(
y
−
x
~
)
p(y|\tilde{x})=g(y-\tilde{x})
p(y∣x~)=g(y−x~),则:
∇
x
~
l
o
g
p
(
y
∣
x
~
)
=
∇
x
~
l
o
g
[
1
2
π
(
σ
0
2
−
σ
i
2
)
e
x
p
[
−
1
2
∣
∣
y
−
x
~
∣
∣
2
(
σ
0
2
−
σ
i
2
)
]
]
=
y
−
x
~
σ
0
2
−
σ
i
2
\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}}
∇x~logp(y∣x~)=∇x~log[2π(σ02−σi2)1exp[−21(σ02−σi2)∣∣y−x~∣∣2]]=σ02−σi2y−x~
因此:
∇
x
~
l
o
g
p
(
x
~
∣
y
)
=
∇
x
~
l
o
g
p
(
x
~
)
+
y
−
x
~
σ
0
2
−
σ
i
2
(5)
\nabla_{\tilde{x}}logp(\tilde{x}|y)=\nabla_{\tilde{x}}logp(\tilde{x})+\frac{y-\tilde{x}}{\sigma_{0}^{2}-\sigma_{i}^{2}} \tag{5}
∇x~logp(x~∣y)=∇x~logp(x~)+σ02−σi2y−x~(5)
这篇文章中Solving Linear Inverse Problems Using The Prior Implicit in a Denoiser,我们知道分数函数和MMSE去噪器的联系,既训练一个网络来估计分数函数可以解释为去噪器估计MMSE:
∇
x
~
l
o
g
p
(
x
~
)
=
x
^
(
x
~
)
−
x
~
σ
i
2
\nabla_{\tilde{x}}logp(\tilde{x})= \frac{\hat{x}(\tilde{x})-\tilde{x}}{\sigma_{i}^{2}}
∇x~logp(x~)=σi2x^(x~)−x~
其中
x
^
(
x
~
)
=
E
(
x
∣
x
~
)
\hat{x}(\tilde{x})=\mathbb{E}(x|\tilde{x})
x^(x~)=E(x∣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张脸,这对于某些特殊的任务来说是不友好的。
从这里我们就能看到生成去噪器的优势了,图像更自然。
总结
这里假设了噪声强度 σ 0 \sigma_{0} σ0,也就是知道了噪声强度这个强先验,在模拟噪声数据集上的表现良好,但是该算法还无法处理真实噪声,或者盲噪声,这是扩散模型应用low-level的一个重要目标。