[DL]深度学习_针对图像恢复的高效扩散模型DiffIR

DiffIR: Efficient Diffusion Model for Image Restoration

Abstract

扩散模型(DM)通过将图像合成过程建模为去噪网络的顺序应用，实现了SOTA的性能。然而，与图像合成不同的是，图像恢复(IR)对生成符合ground-truth的结果有很强的约束。因此，对于IR，传统的DMs在大型模型上运行大量迭代来估计整个图像或特征映射是低效的。为了解决这个问题，我们提出了一种高效的图像恢复扩散模型 (DiffIR)，它由紧凑的图像恢复先验提取网络(Compact IR prior extraction network，CPEN)、动态图像恢复Transformer(DIRformer)和去噪网络组成。具体来说，DiffIR有两个训练阶段：预训练和训练扩散模型。

在预训练中，我们将真实图像输入到CPEN s1中，以捕获紧凑的图像恢复先验表示(IPR)来指导DIRformer。
在第二阶段，我们只使用低质量图像训练扩散模型直接估计与预训练的CPEN s1相同的紧凑图像恢复先验表示IRP。

我们观察到，由于IPR只是一个紧凑向量，DiffIR可以使用比传统扩散更少的迭代来获得准确的估计，并产生更稳定和真实的结果。由于迭代次数少，我们的DiffIR可以采用CPEN s2、DIRformer和去噪网络的联合优化，进一步降低了估计误差的影响。我们在几个图像恢复任务上进行了广泛的实验，并在消耗更少的计算成本的同时实现了SOTA性能。

1 Introduction

图像恢复(IR)由于其广泛的应用价值和不适定性，是一个长期存在的问题。图像恢复的目的是从其低质量(LQ)图像对应损坏的各种退化因素(例如，模糊，蒙版，降采样)来恢复高质量(HQ)的图像。目前，基于深度学习的人工智能方法已经取得了令人印象深刻的成功，因为它们可以从大规模数据集中学习强先验。

扩散模型(Diffusion Models, DMs)是由去噪自编码器层次构建的，在图像合成和图像恢复任务(如图像修复和超分辨率)中取得了令人印象深刻的结果。具体来说，训练扩散模型通过反转扩散过程来迭代地去噪图像。扩散模型已经表明，原则概率扩散建模可以实现从随机抽样的高斯噪声到复杂目标分布(如真实图像或潜在分布)的高质量映射，而不会像生成对抗网络GAN那样遭受模式崩溃和训练不稳定性。

作为一类基于似然的模型，DMs需要对大型去噪模型进行大量的迭代步骤(约50 ~ 1000步)来模拟数据的精确细节，这需要消耗大量的计算资源。与从头开始生成每个像素的图像合成任务不同，图像恢复任务只需要在给定的低质量图像上添加精确的细节。因此，如果扩散模型恢复图像采用图像合成的范式，不仅会浪费大量的计算资源，而且容易产生一些与给定低质量图像不匹配的细节。

图2 概述了DiffIR算法，该算法由diformer、CPEN和去噪网络组成。DiffIR有两个训练阶段:(a)第一阶段，CPEN s1以真实图像为输入，输出一个IPR Z来引导DIRformer恢复图像。(b)第二阶段，利用扩散模型强大的数据估计能力，对预训练CPEN s1提取的图像恢复先验表示进行估计。值得注意的是，我们没有将真实图像输入到CPEN S2和去噪网络中。在推理阶段，我们只使用扩散模型的反向过程。

本文旨在设计一个基于扩散模型的图像恢复网络，充分有效地利用扩散模型强大的分布映射能力来恢复图像。为此，我们提出DiffIR。由于Transformer可以模拟远距离的像素依赖关系，我们采用Transformer块作为DiffIR的基本单元。我们将Transformer块以Unet的形式堆叠形成动态Transformer(Dynamic IRformer，简称DIRformer)来提取和聚合多层次特征。我们分两个阶段训练DiffIR：

在第一阶段(图2 (a))，我们开发了一个紧凑的图像恢复先验提取网络(Compact IR prior extraction network，CPEN)，从真实图像中提取紧凑的图像恢复先验表示(Compact IR prior representation，IPR)，以指导DIRformer。此外，我们还开发了动态门控前馈网络(Dynamic gated feed-forward network，DGFN)和动态多头转置注意(Dynamic Multi-head transposed attention，DMTA)，以充分利用图像恢复先验表示IPR。值得注意的是，CPEN和DIRformer是一起优化的。
在第二阶段(图2 (b))，我们训练扩散模型直接从低质量图像中估计准确的图像恢复先验表示。由于图像恢复先验表示较紧凑，仅添加细节进行恢复，因此我们的扩散模型可以估算出相当准确的图像恢复先验表示，并经过多次迭代获得稳定的视觉结果。

图1 Multi - add是在256×256输入上测量的。我们的DiffIR在图像恢复任务上实现了SOTA性能。值得注意的是，LDM和RePaint是基于扩散模型的方法，DiffIR比RePaint效率高1000倍，同时获得更好的性能。

除了上述方案和架构的新颖性外，我们还展示了联合优化的有效性。在第二阶段，我们观察到估计的图像恢复先验表示IPR可能仍然存在较小的误差，这将影响到DIRformer的性能。但是，以前的扩散模型需要多次迭代，无法与解码器一起优化。由于我们的DiffIR只需要很少的迭代，我们可以运行所有的迭代，得到估计的图像恢复先验表示IPR，并与DIRformer联合优化。如图1所示，我们的DiffIR实现了SOTA性能，比其他基于扩散模型的方法(例如:、RePaint和LDM)。特别是，DiffIR比RePaint效率高1000倍。我们的主要贡献有三个方面：

我们提出DiffIR，一个强大的，简单的，有效的基于扩散模型的图像恢复基线。与图像合成不同，图像恢复输入图像的大部分像素是给定的。因此，我们利用扩散模型强大的映射能力来估计一个紧凑的图像恢复先验表示IPR来指导图像恢复，从而提高扩散模型在图像恢复中的恢复效率和稳定性。
为了充分利用图像恢复先验表示，我们提出了动态Transformer的动态多头转置注意DMTA和动态门控前馈网络DGFN。与以往单独优化去噪网络的潜在扩散模型不同，我们提出了去噪网络和解码器(即DIRformer)的联合优化，以进一步提高估计误差的鲁棒性。
大量的实验表明，与其他基于扩散模型的方法相比，所提出的DiffIR方法在图像恢复任务中可以达到SOTA性能，并且消耗的计算资源要少得多。

2 Related Work

图像恢复Image Restoration

作为先驱作品，SRCNN、DnCNN和ARCNN采用紧凑型CNN在图像恢复上取得了令人印象深刻的性能。之后，基于CNN的方法相对于传统的图像恢复方法变得更加流行。到目前为止，研究人员已经从不同的角度对CNN进行了研究，得到了更精细的网络架构设计和学习方案，如残差块、GAN、注意力机制、知识蒸馏等。

近年来，自然语言处理模型transformer在计算机视觉界获得了广泛的关注。与cnn相比，Transformer可以模拟不同区域之间的全局相互作用，并达到最先进的性能。目前，Transformer已被应用于许多视觉任务中，如图像识别、分割、目标检测和图像恢复。

扩散模型Diffusion Models

扩散模型(Diffusion Models, DMs)在密度估计和样本质量方面都取得了最先进的结果。扩散模型采用参数化马尔可夫链对似然函数的下变分界进行优化，使其生成的目标分布比GAN等生成模型更精确。近年来，扩散模型在图像恢复任务领域的影响力越来越大，如超分辨率和图像修复。SR3和SRdiff将扩撒模型引入图像超分辨率，取得了比基于GAN的SOTA的方法更好的性能。此外，Palette受到条件生成模型的启发，提出了图像恢复的条件扩散模型。LDM提出对潜空间进行扩散以提高恢复效率。此外，RePaint通过扩散模型中的重采样迭代设计了一种改进的去噪策略。然而，这些基于扩散模型的图像恢复方法直接使用扩散模型的范式进行图像合成。然而，图像恢复中的大部分像素是给定的，没有必要对整个图像或特征图执行扩散。我们的DiffI $\bar{\alpha }_{t}=\prod ^{t}_{i=0}\alpha _{i}$ R在紧凑的图像恢复先验表示IPR上执行扩散，这可以使扩散过程更加高效和稳定。

3 前言：扩散模型

在本文中，采用扩散模型生成准确的图像恢复先验表示。在训练阶段，扩散模型方法定义了一个扩散过程，通过T次迭代将输入图像 $x_{0}$ 转换为高斯噪声 $x_{T}\sim N(0,1)$ 。扩散过程的每次迭代可以描述如下：

式中， $x_{t}$ 为时间步长为 t 的噪声图像， $\beta _{t}$ 为预定义的比例因子，N 为高斯分布。(1)式可进一步简化为：

其中， $\alpha _{t}=1-\beta _{t}$ ， $\bar{\alpha }_{t}=\prod ^{t}_{i=0}\alpha _{i}$

在推理阶段(反向过程)，扩散模型方法对高斯随机噪声图 $x_{T}$ 进行采样，然后逐步对 $x_{T}$ 进行降噪，直到得到高质量的输出 $x_{0}$ ：

其中，均值 $\mu _{t}(x_{t},x_{0})=\frac{1}{\sqrt{\alpha _{t}}}(x_{t}-\epsilon \frac{1-\alpha _{t}}{\sqrt{1-\bar{\alpha }_{t}}})$ ，方差 $\sigma ^{2}_{t}=\frac{1-\bar{\alpha} _{t-1}}{1-\bar{\alpha }_{t}}\beta _{t}$ 。 $\epsilon$ 表示 $x_{t}$ 中的噪声，这是逆向过程中唯一的不确定变量。扩散模型采用去噪网络 $\epsilon _{\theta }(x_{t},t)$ 来估计 $\epsilon$ 。为了训练 $\epsilon _{\theta }(x_{t},t)$ ，给定一个干净的图像 $x_{0}$ ，扩散模型随机采样一个时间步长 t 和一个噪声 $\epsilon \sim N(0,I)$ ，根据公式(2)生成噪声图像 $x_{t}$ 。然后，扩散模型对 $\epsilon _{\theta }$ 的网络参数 $\theta$ 进行如下优化：

4 方法

传统的扩散模型需要大量的迭代、计算资源和模型参数来生成准确逼真的图像或潜在特征图。尽管扩散模型在从头生成图像(图像合成)方面取得了令人印象深刻的性能，但将图像合成的扩撒模型范例直接应用于图像恢复是浪费计算资源。由于图像恢复中的大部分像素和信息都是给定的，因此对整幅图像或特征图进行扩散不仅需要大量的迭代和计算，而且容易产生更多的伪影。
总体而言，扩散模型具有较强的数据估计能力，但将现有的扩散模式应用到图像恢复图像合成中效率不高。为了解决这一问题，我们提出了一种有效的图像恢复扩散模型(即DiffIR)，它采用扩散模型估计一个紧凑的图像恢复先验表示IPR来指导网络恢复图像。由于图像恢复先验表示相当轻量，因此与传统扩散模型相比，DiffIR的模型尺寸和迭代可以大大减少，从而产生更准确的估计。

在本节中，我们将介绍DiffIR。如图2所示，DiffIR主要由紧凑的图像恢复先验提取网络(CPEN)、动态图像恢复预处理网络(DIRformer)和去噪网络组成。我们分两个阶段训练DiffIR，包括DiffIR的预训练和扩散模型的训练。在接下来的章节中，我们首先介绍4.1节中的预训练DiffIR。然后，我们在第4.2节中提供了DiffIR的训练效率扩散模型的细节。

图2 概述了DiffIR算法，该算法由diformer、CPEN和去噪网络组成。DiffIR有两个训练阶段:(a)第一阶段，CPEN s1以真实图像为输入，输出一个IPR Z来引导DIRformer恢复图像。(b)第二阶段，利用扩散模型强大的数据估计能力，对预训练CPEN s1提取的图像恢复先验表示进行估计。值得注意的是，我们没有将真实图像输入到CPEN s2和去噪网络中。在推理阶段，我们只使用扩散模型的反向过程。

4.1 预训练DiffIR

在介绍预训练DiffIR之前，我们想在第一阶段介绍两个网络，包括紧凑的图像恢复先验提取网络(CPEN)和动态IRformer(DIRformer)。CPEN的结构如图2黄框所示，主要通过残差块和线性层进行叠加，提取紧凑的图像恢复先验表示(IPR)。之后，DIRformer可以使用提取的IPR来恢复低质量图像。变形Transformer的结构如图2粉色方框所示，由Unet形状的动态T $\bigodot$ ransformer块堆叠而成。动态Transformer模块由动态多头转置注意(DMTA，图2绿框)和动态门控前馈网络(DGFN，图2蓝框)组成，可将图像恢复先验表示IPR作为动态调制参数，将恢复细节添加到特征图中。

在预训练中(图2 (a))，我们将CPEN s1和DIRformer一起训练。具体来说，我们首先将Ground-truth和低质量图像连接在一起，并使用PixelUnshuffle操作对它们进行下采样，以获得CPEN s1的输入。然后，CPEN s1提取图像恢复先验表示IPR $Z\in \mathbb{R}^{{4C}'}$ 为：

然后将IPR Z作为动态调制参数送入DIRformer的DGFN和DMTA，指导恢复：

其中 $\bigodot$ 表示逐元素乘法，Norm表示层归一化， $W_{l}$ 表示线性层， $F$ 和 ${F}'\in \mathbb{R}^{\hat{H}\times \hat{W}\times \hat{C}}$ 分别是输入和输出特征映射， $W^{1}_{l}Z,W^{2}_{l}Z\in \mathbb{R}^{\hat{C}}$ 。

然后，再DMTA中聚合了全局信息。具体来说，将 ${F}'$ 投影为查询 $Q=W^{Q}_{d}W^{Q}_{c}{F}'$ ，键 $K=W^{K}_{d}W^{K}_{c}{F}'$ ，值 $V=W^{V}_{d}W^{V}_{c}{F}'$ ，其中 $W_{c}$ 为1 × 1的逐点卷积point-wise conv， $W_{d}$ 为3 × 3的深度可分离卷积depth-wise conv。之后，reshape查询 $\hat{Q}\in \mathbb{R}^{\hat{H}\hat{W}\times \hat{C}}$ ，键 $\hat{K}\in \mathbb{R}^{\hat{C}\times \hat{H}\hat{W}}$ 和值 $\hat{V}\in \mathbb{R}^{\hat{H}\hat{W}\times \hat{C}}$ 。之后，执行 $\hat{Q}$ 和 $\hat{K}$ 点积，生成大小为 $\mathbb{R}^{\hat{C}\times \hat{C}}$ 的转置注意力图A，比大小为 $\mathbb{R}^{\hat{H}\hat{W}\times \hat{H}\hat{W}}$ 的常规注意力图效率更高。DMTA的整体过程描述如下：

其中 $\gamma$ 是可以学习的缩放参数。与传统的多头自注意力一样，我们将通道划分为多头并计算注意力图。

接下来，在DGFN中，聚合局部特征。我们使用1 × 1卷积对不同通道的信息进行聚合，采用3 × 3深度可分离卷积depth-wise对空间相邻像素的信息进行聚合。此外，采用了门控机制来增强信息编码。DGFN整个过程定义为：

我们将CPEN s1和DIRformer一起训练，使DIRformer能够充分利用CPEN s1提取的图像恢复先验特征表示IPR进行恢复。训练损失定义如下:

其中， $I_{GT}$ 和 $\hat{I}_{HQ}$ 分别为Ground-truth和恢复好的高质量图像。 $\left \| \cdot \right \|_{1}$ 为L1范数。如果有些工作强调视觉质量，我们可以进一步增加感知损失和对抗性损失。在补充材料中提供了更多的细节。

4.2 图像恢复的扩散模型

在第二阶段(图2 (b))，我们利用扩散模型强大的数据估计能力来估计图像恢复先验表示IPR。具体来说，我们使用预训练的CPEN s1来捕获IPR $Z\in \mathbb{R}^{{4C}'}$ 。之后，我们将Z上的扩散过程应用到样本 $Z_{T}\in \mathbb{R}^{{4C}'}$ 上，可以描述为：

其中T是迭代的总次数， $\bar{\alpha }$ 和 $\alpha$ 在等式(1)和(2)中定义(即 $\bar{\alpha }_{T}=\prod ^{T}_{i=0}\alpha _{i }$ )。

在相反的过程中，由于图像恢复先验特征表示IPR是紧凑的，DiffIR s2可以使用比传统扩散模型更少的迭代和更小的模型尺寸来获得相当好的估计。由于传统的分布式去噪算法在迭代过程中计算成本巨大，只能随机采样一个时间步长t∈[1,T]，只在该时间步长对去噪网络进行优化(式(1)、(2)、(3)、(4))。由于去噪网络和解码器(即DIRformer)缺乏联合训练，意味着去噪网络引起的估计误差很小，使DIRformer无法发挥其潜力。相比之下，DiffIR从第 T 个时间步长(式(10))开始，运行所有去噪迭代(式(11))得到 $\hat{Z}$ ，并将其发送给DIRformer进行联合优化。

其中， $\epsilon$ 表示相同的噪声，我们使用CPEN s2和去噪网络来预测式(3)中的噪声。与式(3)中的传统扩散模型不同，我们的DiffIR s2删除了方差估计，并发现它有助于准确的图像恢复先验特征表示IPR估计和更好的性能。

在扩散模型的反向过程中，我们首先使用CPEN s2从低质量图像中得到一个条件向量 $D\in \mathbb{R}^{{4C}'}$ ：

其中，除了第一次卷积的输入维不同，CPEN s2与CPEN s1具有相同的结构。然后，我们使用去噪网络 $\epsilon _{\theta }$ 估计每个时间步 t 的噪声为 $\epsilon _{\theta }$ (Concat( $\hat{Z}_{t}$ , t, D))。将估计的噪声代入式(11)中，得到 $\hat{Z}_{t-1}$ ，开始下一次迭代。

然后，经过 T 次迭代，我们得到最终估计的图像恢复先验表示IPR $\hat{Z}\in \mathbb{R}^{{4C}'}$ 。使用 $L_{all}$ 对CPEN s2、去噪网络和DIRformer进行了联合训练：

我们可以在 $L_{all }$ 中进一步添加感知损失和对抗损失，以获得更好的视觉质量，如式(9)。

在推理阶段，我们只使用反向扩散过程(图2 (b)底部)。CPEN s2从低质量图像中提取条件向量D，随机抽取高斯噪声 $\hat{Z}_{T}$ 。去噪网络利用 $\hat{Z}_{T}$ 和D来估计 T 次迭代后的图像恢复先验特征表示IPR $\hat{Z}$ 。之后，DIRformer利用IPR恢复LQ映像。

5 实验

5.1 实验设置

我们将我们的方法分别应用于三种典型的图像恢复任务：

图像修复
图像超分辨率(SR)
单图像运动去模糊

我们的DiffIR采用4级编码器-解码器结构。从1级到4级，DMTA中的注意头为[1,2,4,8]，通道数为[48,96,192,384]。此外，在所有IR任务中，我们调整了DIRformer中动态Transformer块的数量，以比较DiffIR和SOTA方法在相似参数和计算成本方面的差异。具体来说，从第1级到第4级，我们将动态Transformer块的数量分别设置为[1,1,1,9]，[13,1,1,1]和[3,5,6,6]，分别用于图像修复，图像超分辨率和去模糊。此外，根据之前的工作，我们引入了用于图像修复和sr的对抗性损失和感知损失。CPEN的通道数 ${C}'$ 设置为64。

在训练扩散模型时，将总时间步长 T 设为4, 式(11)中的 $\beta _{t}$ ( $\alpha _{t}=1-\beta _{t}$ )从 $\beta _{1}=0.1$ 线性增加到 $\beta _{T}=0.99$ 。我们使用Adam优化器( $\beta _{1}=0.9$ ， $\beta _{2}=0.99$ )训练模型。更多的细节载于补充材料中。

5.2 图像修复评价

我们使用LaMa的相同设置对我们的DiffIR s2进行训练和验证。具体来说，我们分别在Places-Standard和CelebA-HQ数据集上训练了batch大小为30和patch大小为256的DiffIR。我们使用LPIPS和FID在验证数据集上比较了DiffIR s2与SOTA在图像恢复方法(ICT、LaMa和RePaint)中的差异。

表1 在基准数据集上图像恢复的定量比较(FID/LPIPS)。最佳和次佳成绩分别用粗体和下划线标出。下面用灰色标记的三种方法采用扩散模型。

定量结果如表1和图1(a)所示。我们可以看到我们的DiffIR s2明显优于其他方法。具体来说，我们的DiffIR s2超过竞争方法LaMa的FID差额高达0.2706和0.5583，在Places和CelebA-HQ上具有宽掩模，消耗相似的参数总数和Multi - Add。此外，与基于扩散模型的Re-Paint方法相比，我们的DiffIR s2在仅消耗4.3%的参数和0.1%的计算资源的情况下可以获得更好的性能。这表明DiffIR可以充分有效地利用扩散模型对图像恢复的数据估计能力。

图3 图像恢复方法的视觉比较。放大以获得更好的细节。

定性结果如图3所示。我们的DiffIR s2可以产生比其他竞争对手的图像恢复方法更真实合理的结构和细节。在补充材料中提供了更多的定性结果。

5.3 图像超分辨率评价

我们在图像超分辨率上训练和验证了DiffIR s2。具体来说，我们在4×super-resolution的DIV2K(800张图像)和Flickr2K(2650张图像)数据集上训练DiffIR s2。batchsize设置为64，低分辨率补丁大小设置为64×64。我们使用LPIPS和PSNR在5个基准(Set5，Set14， General100， Urban100和DIV2K100)上评估了我们的DiffIR s2和其他基于SOTA GAN的超分辨率方法。

表2 定量比较(LPIPS/DISTS)。用于基准数据集上的单幅图像超分辨率。最佳和次佳成绩分别用粗体和下划线标出。下面用灰色标记的两种方法采用扩散模型。

表2和图1(b)显示了DiffIR s2与基于SOTA GAN的超分辨率方法SFTGAN、SRGAN、ESRGAN、USRGAN、SPSR和BebyGAN的性能和Multi- Add比较。我们可以看到，DiffIR s2达到了最好的性能。与竞争对手的超分辨率方法BebyGAN相比，我们的DiffIR s2在DIV2K100和Urban100上的LPIPS差值高达0.0151和0.0089，而仅消耗63%的计算资源。此外，在消耗2%的计算资源的情况下，DiffIR s2不能显著优于基于扩散模型的方法LDM。

图4 4×image超分辨率方法的视觉对比。放大以获得更好的细节。

定性结果如图4所示。DiffIR s2实现了最好的视觉质量，包含了更多真实的细节。这些视觉对比与定量结果一致，显示了DiffIR的优越性。DiffIR可以有效地利用强大的DM来恢复图像。更多的视觉结果在补充材料中给出。

5.4 图像运动去模糊的评价

我们在GoPro数据集上训练DiffIR用于图像运动去模糊，并在两个经典基准(GoPro,HIDE)上评估DiffIR。我们将DiffIR s2与最先进的图像运动去模糊方法进行了比较，包括Restormer、MPRNet和IPT。

表3 基准数据集上单幅图像动态去模糊的定量比较。最佳和次优成绩分别用粗体和下划线标出。

定量结果(PSNR and SSIM)见表3，Multi - Add见图1(c)。我们可以看到，我们的DiffIR s2 优于其他运动去模糊方法。具体来说，DiffIR s2在GoPro上分别比 IPT 和 MIMI-Unet+ 高0.68 dB和0.54 dB。此外，DiffIR s2在GoPro和HIDE数据集上分别比 Restormer 高0.28 dB和0.33 dB，仅消耗78%的计算资源。这证明了DiffIR的有效性。

6 消融实验

图像恢复的有效扩散模型

表4 对CelebA-HQ用于图像修复的FID结果进行评价。性能和Multi-Adds是在尺寸为256×256的低质量图像上测量的。

在这一部分中，验证了DiffIR中扩散模型、训练方案、以及是否在扩散中插入方差噪声等成分的有效性(表4)。

DiffIR s2- v3实际为表1中采用的DiffIR s2，DiffIR s1为以真实图像为输入的第一阶段预训练网络。比较DiffIR s1和DiffIR s2-V3，我们可以看到DiffIR s2-V3的LPIPS与DiffIR s1非常相似，这意味着扩散模型具有强大的数据建模能力，可以准确预测图像恢复先验表示IPR。
为了进一步证明扩散模型的有效性，我们取消了在DiffIR s2- v3中使用扩散模型来获得DiffIR s2 V1。比较DiffIR s2-V1和DiffIR s2-V3，我们可以看到DiffIR s2-V3(使用扩散模型)明显优于DiffIR s2-V1。这意味着扩散模型学习到的图像恢复先验表示IPR可以有效地指导DIRformer恢复低质量图像。
为了探索更好的扩散模型训练方案，我们比较了两种训练方案：传统扩散模型优化和我们提出的联合优化。由于传统的扩散模型需要多次迭代来估计大图像或特征图，因此必须采用传统的扩散模型优化，通过随机采样一个时间步来优化去噪网络，而后者无法使用后面的解码器(即本文中的DIRformer)进行优化。由于DiffIR仅使用扩散模型来估计紧凑的一维向量图像恢复先验表示IPR，因此我们可以使用多次迭代来获得相当准确的结果。因此，我们可以采用联合优化的方式，运行去噪网络的所有迭代，得到IPR与DIRformer共同优化。比较DiffIR s2-V2和DiffIR s2-V3，DiffIR s2-V3显著优于DiffIR s2-V2，这证明了我们提出的联合优化训练扩散模型的有效性。这是因为扩散模型在IPR中较小的估计误差可能导致DIRformer的性能下降。联合训练扩散模型和DIRformer可以解决这个问题。
在传统的扩散模型方法中，它们会在反向扩散过程中插入方差噪声(式(3))，以生成更真实的图像。与传统的扩散模型预测图像或特征图不同，我们使用扩散模型来估计图像恢复先验表示IPR。在DiffIR s2-V4中，我们在反向扩散过程中插入噪声。我们可以看到，DiffIR s2-V3比DiffIR s2-V4实现了更好的性能。这意味着为了保证估计的图像恢复先验表示IPR的准确性，最好消除插入噪声。

扩散模型的损失函数

我们探索哪种损失函数最适合指导去噪网络和CPEN s2学习从低质量图像中估计准确的图像恢复先验表征IPR。这里，我们定义了三个损失函数。

定义 $L_{diff}$ 进行优化(式(13))
采用 $L_{2}$ (式(14)测量估计误差。
使用Kullback Leibler散度来度量分布相似性( $L_{kl}$ ，式(15))

其中， $\hat{Z}$ 和 $Z\in \mathbb{R}^{{4C}'}$ 分别为DiffIR s1和DiffIR s2提取的图像恢复先验特征表示。 $\hat{Z}_{norm}$ 和 $Z_{norm\in \mathbb{R}^{{4C}'}}$ 分别用 $\hat{Z}$ 和Z的softmax操作归一化。我们将这三个损失函数分别应用到DiffIR s2上，学习直接从低质量图像中估计准确的图像恢复先验表征IPR。然后，我们在CelebA-HQ上对他们进行了评估。结果如表5所示。我们可以看到 $L_{diff}$ 的性能优于 $L_{2}$ 和 $L_{kl}$ 。

表5 扩散模型损失函数比较(FID)

迭代次数的影响

图6 扩散模型中迭代次数的消融实验

在这一部分中，我们将探讨扩散模型中的迭代次数如何影响DiffIR s2的性能。我们在DiffIR s2中设置不同的迭代次数，并对式子(10)中的 $\beta _{t}(\alpha _{t}=1-\beta _{t})$ 进行调整，使Z在扩散过程(即扩散过程)后为高斯噪声 $Z_{T}\sim N(0,1)$ ，( $\bar{\alpha }\rightarrow 0$ )。结果如图6所示。当迭代次数增加到3次时，DiffIR s2的性能将显著提高。当迭代次数大于4次时，DiffIR s2基本保持稳定，即达到上界。此外，我们可以看到，我们的DiffIR s2比传统扩散模型(需要200多次迭代)具有更快的收敛速度。这是因为我们仅仅对图像恢复先验表征IPR(一个紧凑的一维向量)执行扩散。

7 结论

传统的扩散模型在图像合成方面取得了令人印象深刻的性能。与从头开始生成每个像素的图像合成不同，图像恢复给出了低质量图像作为参考。因此，将传统的决策模型直接应用于图像恢复是低效的。在本文中，我们提出了一种有效的图像恢复扩散模型(即DiffIR)，该模型由CPEN、DIRformer和去噪网络组成。具体来说，我们首先将真地图像输入到CPEN s1中，生成一个紧凑的图像恢复先验表征IPR来指导DIRformer。然后，我们训练扩散来估计CPEN s1提取的IPR。与传统的扩散模型相比，我们的DiffIR可以比传统的扩散模型使用更少的迭代来获得准确的估计，并减少恢复图像中的伪影。此外，由于迭代次数少，我们的DiffIR可以采用CPEN s2、DIRformer和去噪网络的联合优化，以减少估计误差的影响。大量实验表明，DiffIR可以达到一般的SOTA图像恢复性能。