DriftRec: Adapting diffusion models to blind image restoration tasks

Simon Welker, Universitat Hamburg, Germany, arXiv, Cited:0, Code: 无, Paper.

1. 前言

在这项工作中，我们利用扩散模型的高保真度生成能力来解决盲图像恢复任务，以高压缩级别的JPEG伪影去除为例。我们对扩散模型的前向随机微分方程（SDE）提出了一种优雅的修改，以使其适应恢复任务，并将我们的方法命名为DriftRec。将DriftRec与具有相同网络架构的L2回归baseline（JPEG重建的最新技术）进行比较，我们表明我们的方法可以避免两个基线生成模糊图像的倾向，并且显著更忠实地恢复干净图像的分布，同时只需要干净/损坏的图像对的数据集而不需要关于降质算子的知识。通过利用干净和损坏图像的分布比高斯先验更接近的思想，我们的方法只需要低水平的添加噪声，因此即使没有进一步的优化，也可以相对较少的采样步骤。

2. 整体思想

将扩散过程建模为特定的形式。作者把相同方法相同的论文，一篇应用在语音《Speech Enhancement with Score-Based Generative Models in the Complex STFT Domain》，这篇是图像。
在这里插入图片描述

3. 方法

相对于将清晰图像分布转变为纯噪声，本文采用将清晰图像分布转变为降质图像的分布，这么做有一下目的：

相对于纯噪声，使用损坏的图像（加上可处理的噪声）作为反向SDE的初始值，从而通过过程本身的公式来实现任务自适应，而不是仅提供损坏的图像作为条件信息。
由于添加的是高斯白噪声，因此在整个反向过程中，它充当所有可能的空间频率的连续源。然后，经过训练的分数模型对这些频率进行适当滤波，以在不损失高频细节的情况下从目标分布生成看似干净的图像。

即使图像受到其他噪声污染，其分布仍然适用于恢复任务，因为可以通过取一个受损图像并添加高斯噪声来从中抽样。现在提出以下线性前向SDEs模型family来实现本文的想法：
$\begin{equation} f(x_{t}, t) = \gamma F(t)(y-x_{t}), \quad g(t) = \nu (\frac{\sigma_{max}}{\sigma_{min}})^{2t} \end{equation}$
其中， $y$ 是对应 $x_{0}$ 的损坏图像。 $\gamma$ 是强度，控制着 $y$ 和 $x_{0}$ 之间的距离程度。 $F (t)$ 是控制曲线的形状。 $\nu$ 是正则化因子，确保 $\sigma_{T} \approx \sigma_{max}$ ， $\sigma_{T}$ 是（1）所描述的高斯过程的闭合形式方差。直觉上，提出的family结合了VE-SDE扩散的 $g$ 和一个添加的漂移项 $f$ ， $f$ 用于将 $x_{t}$ 拉向受损图像 $y$ 。有趣的是，逆向过程中，漂移项为负，这说明了逆向过程中 $x_{t}$ 不断远离受损图像 $y$ 。如下图所示：

在这里插入图片描述
如果没有分数项，反向过程的状态 $x_{t}$ 只会从 $y$ 开始偏离并发散。因此，分数模型可以被解释为一种适用于问题和数据集的学习控制策略，其必须在存在由于 $- f$ 而远离 $y$ 的这种排斥动作的情况下将 $x_{t}$ 转向合理的样本 $x_{0}$ 。

3.1 考虑两种SDEs

本文中，仅仅考虑了两种SDEs，既 $\ F(t)=t$ 。前一种称为Ornstein-Uhlenbeck Variance Exploding (OUVE) SDE，后一种称为 t-squared Variance Exploding (TSDVE) SDE。这里定义了两种方法的均值：
在这里插入图片描述
这两个表达式都描述了 $x_{0}$ 和 $y$ 之间的线性（逐像素）插值，插值参数由随时间t的指数（OUVE）或半高斯形状（TSDVE）衰减控制。对于两个SDE，对于所有有限 $t$ ，均值和 $Y$ 不等，这似乎是一个问题，因为其目的是使过程向 $y$ 移动。然而，设 $z ～ N （ 0 ； I ）$ ，只要求 $（ µ t + σ T z ）$ 和 $（ y + σ TZ ）$ 的分布接近，以便后者可以作为反向采样过程的合理初始值。我们可以通过增加刚度γ来控制这两个表达式的分布的匹配程度，代价是潜在地破坏反向过程的稳定，或者通过增加σmax来进一步平滑这两个分布的密度函数，代价是更多的反向迭代。