论文阅读笔记：Image Restoration with Mean-Reverting Stochastic Differential Equations

Paper
Repo
参考：作者在知乎上的解读

摘要

本文提出了一种用于通用图像复原的随机微分方程（SDE）方法：

• 关键的构造包含一个均值回复的SDE，它将高质量的图像转换为具有固定高斯噪声的均值状态的退化图像。然后，通过模拟相应的逆时SDE，我们能够在不依赖任何特定任务的先验知识的情况下，恢复低质量图像的原点。
• 关键的是，所提出的均值回复SDE具有闭式解，允许我们计算真实值的时间依赖分数，并使用神经网络对其进行学习。
• 此外，我们提出了一个最大似然目标来学习一个最优的反向轨迹，从而稳定训练并改善恢复结果。实验表明，我们提出的方法在图像去雨、去模糊和去噪上的定量比较中取得了极具竞争力的性能，在两个去雨数据集上设置了新的最先进水平。

最后，通过图像超分辨率、修复和去雾的定性结果进一步证明了我们方法的普遍适用性。

引言

扩散模型在各种图像生成任务中表现出令人印象深刻的性能，其基础是对一个扩散过程进行建模，然后学习其反向的过程。在常用的公式中，我们采用了由随机微分方程（SDE）定义的扩散模型。这就需要使用SDE将图像逐渐扩散到纯噪声分布，然后通过学习和模拟相应的反向时间SDE来生成样本。其实质是训练一个神经网络来估计噪声数据分布的得分函数。

图像复原是指从降质的低质量图像中恢复出高质量图像的总任务。常见的具体实例包括图像去雨、去模糊、去噪、超分辨率等。

最近，扩散模型已经被应用到不同的图像恢复任务中。这些方法都采用了标准的前向过程，将图像扩散到纯噪声中。因此，反向(生成)过程用高方差的采样噪声初始化，这可能导致高质量图像的真实恢复效果不佳。大量实验表明，扩散模型可以产生更好的感知分数，但在一些基于像素/结构的失真标准方面往往表现不尽人意。

为了解决这个问题，作者提出使用均值回归SDE来解决图像复原问题。如图1所示，这适应了前向过程，使其对图像退化本身进行建模，从高质量图像到低质量图像。通过模拟相应的逆时偏移，可以恢复出高质量的图像。重要的是，不需要任务特定的先验知识来建模测试时刻的图像退化，只需要一组图像对用于训练。

过程概述

主要贡献：

1. 提出了一种使用均值回归SDE的通用图像复原方法，该方法直接建模图像退化过程。该公式有一个封闭形式的解，使得我们能够计算真实时间依赖的得分函数，并训练一个神经网络来估计它。
2. 提出一个简单的替代损失函数来训练神经网络，基于最大化反向时间轨迹的可能性。与常见的分数匹配目标相比，该损失被证明能够稳定训练并一致地提高图像恢复性能。
3. 该方法不需要知道图像恢复任务的任何先验知识或参数设置，也不局限于线性退化或简单非线性退化，只需改变数据集即可用于任意任务，非常灵活（可应用于六种不同的图像复原任务：图像去雨、去模糊、去噪、超分辨率、修复和去雾，证明了我们提出的方法的普遍适用性）。
4. 该基于diffusion的方法在图像去雨、去模糊和去噪的定量比较中取得了极具竞争力的复原性能，在两个去雨数据集上开创了最新的研究水平。

方法

该方法的关键思想是将均值回归SDE与最大似然目标相结合用于神经网络训练。我们将其称为图像复原随机微分方程(IR-SDE)。我们首先描述均值回归SDE的正向和反向过程，并采用先前描述的、基于分数的训练方法来估计这个SDE。然后，我们将其与我们提出的基于最大似然目标的损失函数进行描述和对比。

用于图像退化的前向SDE

考虑一种特殊的SDE：

$$\mathrm{d}x=\theta _t\left( \mu -x \right) \mathrm{d}t+\sigma _t\mathrm{d}w \tag{1}$$

其中是状态均值，和是与时间相关的正参数，分别刻画了均值回复和随机波动的速度（这两个参数的选择会对恢复性能产生较大影响）。

一般情况下，和起始状态可以设置为任意一对不同的图像。然后，前向SDE将一幅图像传递给另一幅图像，作为一种噪声插值。为了进行图像退化，我们令和分别为真实的高质量(HQ)图像和其退化的低质量(LQ)图像（见图1）。值得注意的是，虽然依赖于（因为它们是同一对象或场景的配对HQ-LQ图像），但独立于布朗运动，因此SDE在Ito意义下仍然成立。

当时，均值收敛于低质量图像，方差收敛于平稳方差。也就是说，前向SDE将高质量图像扩散为固定高斯噪声的低质量图像。

用于图像修复的反向SDE

将SDE逆变换得到一个图像恢复SDE(IR-SDE)，即

$$\mathrm{d}x=\left[ \theta _t\left( \mu -x \right) -\sigma _{t}^{2}\nabla _x\log p_t\left( x \right) \right] \mathrm{d}t+\sigma _t\mathrm{d}\hat{w} \tag{2}$$

与普通score估计的方式不同，这里提出的SDE是对score function 有确定解的。具体来说，令，此时可以证明forward-SDE的解为：

$$x\left( t \right) =\mu +\left( x\left( s \right) -\mu \right) e^{-\bar{\theta}_{s:t}}+\int_s^t{\sigma _ze^{-\bar{\theta}_{s:t}}\mathrm{d}w\left( z \right)} \tag{3}$$

其中，其任意时刻的转移概率为正态分布

$$p\left( x\left( t \right) |x\left( s \right) \right) =\mathcal{N} \left( x\left( t \right) |m_{s:t}\left( x\left( s \right) \right) ,v_{s:t} \right) \tag{4}$$

其中均值和方差分别为

$$\begin{aligned} m_{s:t}\left( x\left( s \right) \right) &=\mu +\left( x\left( s \right) -\mu \right) e^{-\bar{\theta}_{s:t}}\\ v_{s:t}&=\int_s^t{\sigma _{z}^{2}e^{-2\bar{\theta}_{s:t}}\mathrm{d}z}=\lambda ^2\left( 1-e^{-2\bar{\theta}_{s:t}} \right)\\ \end{aligned} \tag{5}$$

在训练中，给定LQ和GT图像对，根据以上解可以算出正确的score为

$$\nabla _{\boldsymbol{x}}\log p_t\left( \boldsymbol{x} \right) =-\frac{\boldsymbol{x}\left( t \right) -m_t\left( \boldsymbol{x} \right)}{v_t} \tag{6}$$

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

基于最大似然的损失函数

当应用于图像复原中遇到的复杂退化时，训练往往变得不稳定。这个困难可能源于试图学习给定时刻的瞬时噪声。因此，我们基于在给定高质量图像的情况下，试图寻找最优轨迹的思想，提出了一个备选的最大似然目标。这个目标并不是为了学习一个更准确的得分函数而提出的。相反，它被用于稳定训练和恢复更准确的图像。

具体来说，我们希望最大化的可能性可以因式分解为

$$p\left( x_{1:T}|x_0 \right) =p\left( x_T|x_0 \right) \prod_{i=2}^T{p\left( x_{i-1}|x_i,x_0 \right)} \tag{7}$$

其中。反向过程则可以由贝叶斯法则推导出来：

$$p\left( x_{i-1}|x_i,x_0 \right) =\frac{p\left( x_i|x_{i-1},x_0 \right) p\left( x_{i-1}|x_0 \right)}{p\left( x_i|x_0 \right)} \tag{8}$$

由于所有的分布都是高斯分布，因此直接找到一个最优的反向状态使得负对数似然最小化来求解最优路径，即是求：

$$x_{i-1}^{*}=\mathrm{arg}\underset{x_{i-1}}{\min}\left[ -\log p\left( x_{i-1}|x_i,x_0 \right) \right] \tag{9}$$

其中，令表示从反转过来的理想状态。为了简化记号，我们令。

因此，给定一个初始状态，对于离散时间的任意状态，式的最优反向解为：

$$x_{i-1}^{*}=\frac{1-e^{-2\bar{\theta}_{i-1}}}{1-e^{-2\bar{\theta}_i}}e^{-\theta _{i}^{\prime}}\left( x_i-\mu \right) +\frac{1-e^{-2\theta _{i}^{\prime}}}{1-e^{-2\bar{\theta}_i}}e^{-\bar{\theta}_{i-1}}\left( x_0-\mu \right) +\mu \tag{10}$$

随后就可以使用网络估计噪声并学习最优的复原步骤了。

这篇论文的公式推导很多，详情可以看论文原文的Appendix部分。

实验

这一部分主要就是一些指标数据和可视化对比的呈现，具体的可以直接参阅论文。