DiffIR: Efficient Diffusion Model for Image Restoration 利用扩散模型进行图像重建

一、主要贡献

•我们提出了DiffIR，一种强大、简单、高效的基于扩散模型的的图像修复方法。与图像生成不同的是，输入图像的大部分像素都是给定的。因此，我们利用DM强大的映射能力来估计一个紧凑的IPR（IR Prior Representation，图像修复的先验表示）来引导图像修复，从而提高DM在图像修复中的恢复效率和稳定性。

•我们建议为Dynamic IRformer提供DMTA(dynamic multi-head transposed attention，动态多头转置注意力机制)和DGFN(dynamic gated feed-forward network，动态门前馈网络)，以充分利用IPR。与以往单独优化去噪网络的LDM不同，我们提出联合优化去噪网络和解码器(即DIRformer)，进一步提高估计误差的鲁棒性。

•大量实验表明，与其他基于dm的方法相比，提出的DiffIR方法在实现图像修复任务SOTA性能的同时，消耗的计算资源少得多。

二、扩散模型

前向过程：

逆向过程：

目标函数：

三、方法

由于IR中的大部分像素和信息都是给定的，所以对整个图像或feature map进行dm不仅要花费大量的迭代和计算，而且容易产生更多的伪影。为了解决这一问题，我们提出了一种有效的图像重建算法，即DiffIR算法。该算法采用图像重建算法来估计一个紧凑的IPR，从而引导网络进行图像恢复。由于IPR非常小，DiffIR的模型大小和迭代可以大大减少，与传统的DM相比，产生更准确的估计。

注：（1）concat：连接两个数组

      （2）pixelunshuffle：一种下采样方法/pixelshuffle：一种上采样方法

      （3）GELU：激活函数

        (4) F、F_head：输入、输出的feature map

我们将DiffIR的训练分为两个阶段，分别是预训练DiffIR和训练扩散模型。

1.预训练DiffIR

CPEN（compact IR prior extraction network,紧凑图像修复先验提取网络）：将GT与低质量图像拼接在一起，经下采样后送入CPEN提取出特征Z即IPR。

DIRformer（dynamic IR transformer,图像修复动态transformer）：由DGFN与DMTA组成U-net结构，输入低质量图像，在z的指导下，输出重建的高质量图像。

DGFN(dynamic gated feed-forward network，动态门前馈网络)：作用是聚合局部特征。采用1×1 Conv对不同通道的信息进行聚合，采用3 × 3深度Conv对空间相邻像素的信息进行聚合。结构如下：

DMTA(dynamic multi-head transposed attention，动态多头转置注意力机制)：作用是聚合全局空间信息，计算attention map，结构如下图

联合训练CPEN和DIRformer，使DIRformer充分利用CPEN提取的图像修复表达IPR，损失函数为：

2.训练扩散模型

利用扩散模型来估计IPR

前向过程：固定CPEN的参数，提取图像修复先验表示Z，对其进行扩散，得到Zt

逆向过程：首先使用CPEN从LQ图像中得到条件向量D，指导去噪网络逐步去除Zt的噪声，经过T次迭代得到估计的Z。

由于IPR是紧凑的，DiffIR可以使用更少的迭代和更小的模型来获得比传统dm更好的估计。由于传统dm在迭代过程中具有巨大的计算成本，因此必须随机采样时间步长t∈[1,t]，并仅在该时间步长对去噪网络进行优化。去噪网络与解码器(即DIRformer)缺乏联合训练，意味着去噪网络造成的较小误差，也可能使DIRformer无法发挥其潜力。而DiffIR从第t个时间步长开始，经过所有去噪迭代，得到Zˆ，发送给DIRformer进行联合优化。

损失函数：

四、实验及结论

  在inpainting、超分辨率、去运动模糊方面实现了SOTA性能。