95、Image Restoration with Mean-Reverting Stochastic Differential Equations

简介

主页：https://github.com/Algolzw/image-restoration-sde

扩散模型终于在去噪、超分辨率等应用了。

这是一种基于随机微分方程的通用图像恢复方法，关键结构包括均值还原SDE，该SDE将高质量图像转换为具有固定高斯噪声的平均状态的降级对应图像，通过模拟相应的逆时SDE，可以在不依赖任何特定任务的先验知识的情况下恢复低质量图像的原点，所提出的均值回归SDE具有封闭形式的解决方案，允许计算真实时间相关分数并使用神经网络学习它

贡献点

• 提出了一种通用的图像恢复方法，使用均值恢复SDE直接模拟图像退化过程。公式有一个封闭形式的解决方案，使能够计算地面真值与时间相关的分数函数，并训练神经网络来估计它。
• 提出了一个简单的替代损失函数来训练神经网络，基于最大化逆时间轨迹的可能性。与普通分数匹配目标相比，损失被证明可以稳定训练并持续提高图像恢复性能。
• 通过将其应用于六种不同的图像恢复任务来证明提出的方法的一般适用性:图像去除、去模糊、去噪、超分辨率、上漆和去雾。
• 在图像去噪、去模糊和去噪的定量比较中实现了极具竞争力的恢复性能，在两个去噪数据集上达到了新的水平。

实现流程

背景知识-SDE

前向过程

𝑓和𝑔分别为drift函数和dispersion函数，𝑤为标准Wiener过程，$x(0)\in R^d$为初始条件
逆向过程

$\hat{w}$为逆时Wiener过程，$p_t(x)$为x(t)在𝑡时刻的边际概率密度函数。${\nabla }_{x}log{p}_t(x)$

方法

使用均值回归SDE$dx={\theta }_t(\mu -x)dt+{\sigma }_t$进行图像恢复。SDE模型通过向低质量图像的带噪版本的扩散，来模拟从高质量图像的$x(0)$到低质量图像的$\mu +ϵ$的退化过程。通过模拟相应的逆时SDE，可以恢复高质量的图像。

核心思想是将均值恢复SDE与神经网络训练的最大似然目标相结合

分数函数定义为：

其中${\theta }_t$和${\sigma }_t$是时间相关的正参数，分别表征平均回归和随机波动的速度

为了进行图像退化，分别让$x(0)$和$\mu$为地面真实高质量(HQ)图像和其退化的低质量(LQ)对应图像。

为了使的SDE(3)具有封闭形式的解，设置${\sigma }_t^2/{\theta }_t=2{\lambda }^2$，其中${\lambda }^2$为平稳方差。

假设任意时刻t都满足 ${\sigma }_t^2/{\theta }_t=2{\lambda }^2$,给定任意时刻 x(s)，其中s<t，有

其中 ${\bar{\theta }}_{s:t}:={\int }_s^t{\theta }_{z}dz$是已知的，过度内核 $p(x(t)|x(s))=N(x(t)|m_{s:t}(x(s)),v_{s:t})$是高斯分布，均值 $m_{s:t}$ 方差 $v_{s:t}$表示为

任意时刻的边缘分布

当t->$\infty$，均值$m_t$收敛于低质量图像$\mu$，方差$v_t$收敛于平稳方差${\lambda }^2$

前向SDE(3)将高质量图像扩散为高斯噪声固定的低质量图像

反向SDE可以推导为：

唯一不知道的是分数函数 ${\nabla }_{x}log{p}_t(x)$

由于在训练期间可以获得地面真实高质量图像$x(0)$，因此可以训练神经网络来估计分数${\nabla }_{x}log{p}_t(x)$。具体来说，可以使用(6)来计算地面真值得分为

如果使 $x(t)=m_t(x)+\sqrt{v_t}{ϵ}_t$，其中 ${ϵ}_t$ 是标准高斯噪声 ${ϵ}_{t}N(0,I)$

使用一个条件时变神经网络，它将状态和时间作为输入和输出纯噪声

${\gamma }_1,{\gamma }_2,\cdots ,{\gamma }_T$是positive weights， { x i } i = 0 T \{x_i\}^T_{i=0} {xi​}i=0T​表示扩散过程的离散化

当应用于图像恢复中遇到的复杂退化时，训练往往变得不稳定,这源于试图学习给定时间的瞬时噪声。

在给定高质量图像𝑥0的情况下，试图找到最优轨迹$x_{1:T}$，最小化似然函数 $p(x_{1:T}|x_0)$

其中$p(x_T|x_0)=N(x_T:m_T(x_0),v_T)$是低质量图像

最小化负对数似然的最优反向状态

给定初始状态 $x_0$，任意时刻状态 $x_i$,i>0，从 $x_i->x_{i-1}$的IR-SDE表示为

优化噪声网络，使IR-SDE反向为最优轨迹，即

$(dx{)}_{\overline{ϵ}\Phi }$ 表示 公式7的反向SDE，分数由噪声模型${\overline{ϵ}}_{\Phi }$得到

其中期望${\int }_0^t{\sigma }_{s}d\hat{w}(s)$为零，那么只需要考虑 $(dx{)}_{{\widehat{ϵ}}_{\Phi }}$dirft 部分

实验