A Unified Conditional Framework for Diffusion-based Image Restoration (Paper reading)

Yi Zhang, CUHK, CN, arXiv2023, Cited:0, Code, Paper

1. 前言

最近，扩散概率模型（Diffusion Probabilistic Models，DPMs）在图像生成任务中表现出了非凡的性能，能够生成高度逼真的图像。当将DPMs用于图像恢复任务时，关键的一点在于如何整合条件信息，以引导DPMs生成准确和自然的输出，这在现有的研究中往往被忽视。在本文中，我们提出了一个基于扩散模型的统一条件框架，用于图像恢复。我们利用一个轻量级的UNet来预测初始引导，并使用扩散模型学习引导的残差部分。通过精心设计扩散模型块的基本模块和集成模块，我们将引导和其他辅助条件信息整合到每个扩散模型块中，实现了空间自适应的生成条件。为了处理高分辨率图像，我们提出了一种简单而有效的逐步分块策略，可以生成任意分辨率的图像而不会产生网格伪影。我们在三个具有挑战性的任务上评估了我们的条件框架：极低光照去噪、去模糊和JPEG恢复，并展示了它在感知质量和泛化到恢复任务方面的显著改进。

2. 整体思想

不算新的思想，用扩散模型对初始估计进行增强。本文创新点应该在扩散模型的网络上，本文是自己设计的网络，其中用到了动态卷积。

3. 方法

我们的目标是为图像恢复任务设计一个统一的条件框架。该框架的输入条件信息由两个组成部分组成：降质图像和辅助标量信息。降质图像表示待恢复的图像，而辅助标量信息可以包括降质类型、强度或其他与每个恢复任务相关的细节。

为了增强条件信息的整合，我们首先采用一个轻量级的U-Net来预测初始输出，如图1（左侧）所示。这个初始输出捕捉到了最终恢复图像的低频和确定性方面，这些方面更容易恢复并包含了关键的结构信息。我们将这个初始输出作为扩散模型的空间引导。结合辅助标量信息（例如降级类型、扩散时间步长），我们将它们注入到扩散模型的每个块中，从而实现对扩散模型的更好控制和引导。这种注入不仅提供了全面的上下文，还增强了我们框架的灵活性。我们采用扩散模型来捕捉初始输出的残差分布。

基本模块：在我们的方法中，我们为图像恢复任务中使用的扩散模型设计了一个基本模块。我们的目标是通过利用现有的图像恢复主干网络使模块尽可能简单。我们尽量避免使用复杂的操作符，而是采用现有的图像恢复主干网络，使其尽可能简单。对于每个块，我们使用两个卷积层。在每个卷积层之前，我们引入了LayerNorm来稳定训练过程。我们使用Swish作为激活函数。我们应用了一条捷径来实现残差学习。为了能够注入条件信息，第二个卷积核的设计是基于条件动态变化的。

条件注入模块：为了更好地将条件信息整合到块中，我们提出了一个条件整合模块（Conditional Integration Module，CIM）。在CIM中，首先将引导信息进行缩放，以匹配块内特征图的分辨率。然后，这个缩放后的引导信息通过两个卷积层，并使用SimpleGate激活函数，有效地调整通道数并生成特征图$G$。

SimpleGate(x) = sigmoid(x) * x

同时，辅助标量信息通过两个线性层分支，并在其中使用Swish激活函数进行处理，生成特征图$S$。接下来，特征图$G$和$S$被传递给自适应核引导模块（Adaptive Kernel Guidance Module，AKGM），用于为基本模块中的第二个卷积层生成动态卷积核，如图1所示。AKGM的关键思想是自适应融合卷积核基，使每个空间位置能够根据融合的多源条件信息处理特征图。

如图2（左侧）所示，每个AKGM都有$N$个可学习的卷积核基，表示为$W_b\in R^{C\times C\times k\times k}$，其中$C$表示通道数，$k$表示卷积核大小。这些卷积核基经过训练，用于处理不同的情况和场景。特征图$G\in R^{H\times W\times N}$和$S\in R^{1\times 1\times N}$通过逐点乘积融合生成多源融合权重$M\in R^{H\times W\times N}$。这里，$H$和$W$分别表示特征图的高度和宽度，$N$表示卷积核基的数量。对于特定位置$(i，j)$的融合卷积核$F(i，j)$，通过线性融合该位置上的多源融合权重得到。具体而言，可以表示为：
$$F_{i,j}={\sum }_{N-1}^{b=0}{M}_{i,j}[b]W_b$$

4. 实验