图像复原论文解读：Image Restoration with Mean-Reverting Stochastic Differential Equations, ICML 2023

qq_20448873

已于 2023-05-01 18:11:04 修改

阅读量2.9k

点赞数 10

文章标签：机器学习人工智能深度学习

于 2023-05-01 18:08:49 首次发布

本文链接：https://blog.csdn.net/qq_20448873/article/details/130457619

版权

Paper： Image Restoration with Mean-Reverting Stochastic Differential Equations

Code： GitHub - Algolzw/image-restoration-sde: Image Restoration with Mean-Reverting Stochastic Differential Equations, ICML 2023.

Project： Image Restoration with Mean-Reverting Stochastic Differential Equations

Contribution

提出一种新的均值回归随机微分方程用于各种图像复原任务（IR-SDE），在训练中该方程有closed-form解可以算出正确的score function，并可以在测试中使用网络来估计它。
我们提出使用最大似然估计来计算最优的reverse-time路径，并以此作为loss function。与普通的score-matching loss比，新提出的loss训练更加稳定，且生成的图像也更加精确。
该方法不需要知道图像恢复任务的任何先验知识或参数设置，也不局限于线性退化或简单非线性退化，只需改变数据集即可用于任意任务，非常灵活。且该基于diffusion的方法在去雨任务中获得了SOTA效果。
对该方法进行简单扩展后（修改网络结果，调参数等）在NTIRE 2023 Shadow Removal Challenge中获得了第二名的好成绩（LPIPS和MOS获得第一）。

1. 动机：Motivation

随着扩散模型的火热，各种各样的任务都开始引入diffusion的概念，而同作为生成式任务的图像复原也是如此，引入了大量结合diffusion的工作，如DDRM ，DPS ，DDNM 等。然而，目前绝大多数基于diffusion的图像复原工作都需要显式地知道图像降质的具体过程，特别如下形式：

$y = Ax + n$

其中 $(x,y)$ 分别代表high-quality(HQ)和low-quality(LQ)图像，而 $(A, n)$ 代码某个(非)线性退化过程以及额外噪声。

这样的好处是可以对退化过程进行显式建模，并对A和n进行分解从而简化过程。但缺点也相当明显：

1）A 必须是已知的退化算子，并且通常需要可(SVD)分解。这些要求对于大多数传统图像复原任务都是不可能知道的，如去雨去雾去阴影去运动模糊等等。

2）每个任务的退化核A可能都不一样，需要针对不同任务设计不同求解方程，导致很难将其快速应用扩展到不同任务中。

3）很难处理非线性或不可分解的退化过程（有部分paper在试图解决这些问题但通常都使问题更复杂了）。

那么能不能设计一种比较直接的diffusion模型，可以像普通卷积网络一样，直接改一下数据集就可处理各种各样简单或复杂的图像复原任务呢？当然可以！如我们今年的ICML论文：Image Restoration with Mean-Reverting Stochastic Differential Equations，其核心思想是通过mean-reverting SDE建模任意的图像退化过程。

2. 随机微分方程介绍: An Introduction to Stochastic Differential Equations (SDE)

这个东西咋一听很高深的样子，但原理其实并不复杂。想象如何逐步的给一个信号添加高斯噪声，并使得整个加噪是连续的过程。一个简单的答案是布朗运动(Brownian motion/Wiener process): 一个连续的高斯过程，其在每一个连续瞬间都会给系统注入稳定，随机的噪声。而基本的SDE即可以写为如下形式：

$dx = f(x, t) dt + g(t)dw$

其中 $w$ 表示布朗运动(或者叫Wiener process)，而 $f$ 和 $g$ 分别为drift和dispersion方程。简单来说， $dw$ 是高斯噪声， $f$ 控制整个SDE的“流向”，而 $g$ 控制SDE的扩散程度或加噪速度。在score-based generative model中我们通常把上述方程叫做forward-SDE，对应于图像的退化过程，我们可以通过改变 $f$ 来控制SDE最后是变为纯噪声还是带噪图像等。而与之对应的则是reverse-time SDE，其也是score-based generative model的核心。具体来说，理论上任何SDE都会有一个对应的逆过程(从噪声变为干净图像)，如下方程：

$dx = \Bigl[ f(x, t) - g(t)^2\, \nabla_{x} \log p_t(x) \Bigr] d t + g(t) d \hat{w}$

其中 $\nabla_{x} \log p_t(x)$ 叫做score function，而 $\hat{w}$ 表示reverse-time Wiener process(即原本是加噪声现在变成减去噪声)。注意这里 $f$ 和 $g$ 都是已知的，未知项就只有 $\nabla_{x} \log p_t(x)$ ，我们可以通过多种方式来估计这个score，并模拟如上reverse-time SDE来恢复干净图像。具体推导Yang Song大佬在score-sde 这篇论文以及这个博客里都有详细说明，关于reverse-time SDE我个人也参考了这篇博客 .

3. 核心方法

3.1 均值回归SDE（Mean-Reverting SDE）

考虑如下特殊的一种SDE：

$d x = \theta_t \, (\mu - x) d t + \sigma_t d w$

与普通的SDE相比这里多了一个 $\mu$ 表示指定的值或者状态（通常叫做均值），所以叫做Mean-Reverting SDE。这样构造SDE的好处是在一定步数后，整个SDE一定会流向带有稳定高斯噪声的均值。这里的 $\mu$ 可以跟 $x$ 有关系也可以无任何关系，可以简单地看作从 $x(0)$ 到 $\mu$ 的插值并逐步添加噪声的过程。在图像恢复这个任务中，我们可以设置为退化后的LQ图像。只要有成对的Input和GT图像，我们不需要知道任何退化的先验知识或相关参数等，就可以通过这种特殊的SDE建模图像退化过程，并通过在LQ图像上添加固定噪声模拟reverse-time SDE来复原图像。其reverse-time SDE公式为：

$dx = \big[ \theta_t \, (\mu - x) - \sigma_t^2 \, \nabla_{x} \log p_t(x) \big] d t + \sigma_t d \hat{w}.$

而整个流程图如下所示：

注意与普通score估计的方式不同，这里我们提出的SDE是对score function 有确定解的。具体来说，令 $\sigma_t^2 \, / \, \theta_t = 2 \, \lambda^2$ ，此时可以证明forward-SDE的解为：

$x(t) = \mu + \bigl(x(s) - \mu \bigr) \, e^{-\bar{\theta}_{s:t}} + \int^t_s \sigma_z \, e^{-\bar{\theta}_{z:t}} d w(z)$

其中 $\bar{\theta}_{s:t} = \int^t_s \theta_z d z$ 是已知的（因为所有的 $\theta_t$ 都是预设好的，类似于DDPM里的 $\alpha_t$ 以及 $\beta_t$ ）。而其任意时刻的转移概率为正态分布 $p(x(t) | x(s)) = \mathcal{N}\bigl(x(t) | m_{s:t}(x(s)), v_{s:t}\bigr)$ ，且有

$m_{s:t}(x(s)) = \mu + (x(s) - \mu) \, x^{-\bar{\theta}_{s:t}}$

$v_{s:t} = \int^t_s \sigma_z^2 \, e^{-2\bar{\theta}_{z:t}} d z = \lambda^2 \, \Bigl(1 - e^{-2 \, \bar{\theta}_{s:t}}\Bigr).$

具体的推导请看paper的appendix部分。在训练中，给定LQ和GT图像对，根据以上解我们即可以算出正确的score为：

$\nabla_{x}\log p_t(x) = - \frac{x(t) - m_{t}(x)}{v_{t}}$

然后就可以通过网络来估计这个score了。与score-matching之类的方法比，这里估计的score更加准确且是针对图像复原任务本身的（包含了降质退化过程）。

3.2 最大似然学习：Maximum Likelihood Learning

这里的核心思想是在图像复原任务中直接学习噪声是相当困难的，因为图像退化本身就可能包含自然噪声以及各种复杂的变化过程。所以能不能通过直接学习图像逐步的降噪提质过程来稳定学习呢？或者说能不能找到一个最优路径，可以使得模型沿着最优的图像复原方向进行估计，从而得到更精确的HQ图像？

由于上面提到IR-SDE中状态转移概率（transition kernel）是可求解的，自然而然我们想到了贝叶斯法则：可以使用最大似然估计的方式来求解最优路径，即是求：

$x_{i-1}^{*} = \arg\min_{x_{i-1}} \Bigl[ -\log p \bigl(x_{i-1} \mid x_i, x_0 \bigr) \Bigr]$

其中表示最优的还原路径。而根据贝叶斯法则，我们可以计算出：

$x_{i-1}^{*} = \frac{1 - e^{-2 \, \bar{\theta}_{i-1}}}{1 - e^{-2 \, \bar{\theta}_i}} e^{-\theta_i} (x_i - \mu) + \frac{1 -e^{-2 \, \theta_i}}{1 - e^{-2 \, \bar{\theta}_i}} e^{-\bar{\theta}_{i-1}} (x_0 - \mu) + \mu.$

具体的推导请看paper的appendix部分。随后我们就可以使用网络估计噪声并学习最优的复原步骤了。

注意这里的方式其实也可以从另一个角度来推出DDPM的均值。感兴趣的同学可以自己推一下，我会把这个推导之后更新到camera ready以及新的arxiv版本的appendix中。

4. 实验结果

又到了喜闻乐见的实验贴表贴图部分了。但感觉这些还是直接看论文比较方便，这里就简单贴一些对比图好了。

编辑切换为居中

添加图片注释，不超过 140 字（可选）

编辑切换为居中

Real-world图像去雾

编辑切换为居中

Real-world图像去阴影

==============================================

======== 分割线 ============

最后我们的改进版本（Refusion）赢得了CVPR图像去阴影比赛（NTIRE 2023）的第二名（感知上第一名），感觉对于diffusion模型来说在这种real-world数据的比赛里能拿到这个成绩已经很不错了哈哈。同时在这篇论文里我们还提出了一种基于U-Net的latent-diffusion模型，也是非常的有意思^_^。下面给出paper和code的链接：

Paper：Refusion: Enabling Large-Size Realistic Image Restoration with Latent-Space Diffusion Models

Code：GitHub - Algolzw/image-restoration-sde: Image Restoration with Mean-Reverting Stochastic Differential Equations, ICML 2023.

qq_20448873

关注

10
点赞
踩
26

收藏

觉得还不错? 一键收藏
5
评论
图像复原论文解读：Image Restoration with Mean-Reverting Stochastic Differential Equations, ICML 2023

Image Restoration with Mean-Reverting Stochastic Differential Equations, ICML 2023. 基于diffusion的图像复原方法，只需切换数据集即可灵活应用到各种各样图图变换任务中。
复制链接

扫一扫

图像复原论文解读：Image Restoration with Mean-Reverting Stochastic Differential Equations, ICML 2023

Contribution

1. 动机：Motivation

2. 随机微分方程介绍: An Introduction to Stochastic Differential Equations (SDE)

3. 核心方法

3.1 均值回归SDE（Mean-Reverting SDE）

3.2 最大似然学习：Maximum Likelihood Learning

“相关推荐”对你有帮助么？