[论文解读] DiffIR: Efficient Diffusion Model for Image Restoration

Bin Xia, Yulun Zhang+6 more|arXiv (Cornell University)|Mar 16, 2023

Image and Signal Denoising Methods被引用 16

一句话总结

DiffIR 使用紧凑的 IR 先验表示和两阶段训练方案来将扩散模型用于图像修复，在比以前基于 DM 的 IR 方法需要更少迭代和更低计算资源的情况下实现了业界最先进的结果。

ABSTRACT

Diffusion model (DM) has achieved SOTA performance by modeling the image synthesis process into a sequential application of a denoising network. However, different from image synthesis, image restoration (IR) has a strong constraint to generate results in accordance with ground-truth. Thus, for IR, traditional DMs running massive iterations on a large model to estimate whole images or feature maps is inefficient. To address this issue, we propose an efficient DM for IR (DiffIR), which consists of a compact IR prior extraction network (CPEN), dynamic IR transformer (DIRformer), and denoising network. Specifically, DiffIR has two training stages: pretraining and training DM. In pretraining, we input ground-truth images into CPEN$_{S1}$ to capture a compact IR prior representation (IPR) to guide DIRformer. In the second stage, we train the DM to directly estimate the same IRP as pretrained CPEN$_{S1}$ only using LQ images. We observe that since the IPR is only a compact vector, DiffIR can use fewer iterations than traditional DM to obtain accurate estimations and generate more stable and realistic results. Since the iterations are few, our DiffIR can adopt a joint optimization of CPEN$_{S2}$, DIRformer, and denoising network, which can further reduce the estimation error influence. We conduct extensive experiments on several IR tasks and achieve SOTA performance while consuming less computational costs. Code is available at \url{https://github.com/Zj-BinXia/DiffIR}.

研究动机与目标

• 在大多数输入像素已给定的前提下，推动高效地使用扩散模型进行图像修复。
• 开发紧凑的 IR 先验表示（IPR）来引导修复。
• 提出两阶段训练方案，以利用 CPEN 和扩散模型完成 IR。
• 实现 CPEN S2、DIRformer 和去噪网络的联合优化，以降低估计误差。

提出的方法

• 引入 CPEN，从真实图像中提取紧凑的 IR 先验表示。
• 提出 Dynamic IRformer（DIRformer），结合 DMTA 和 DGFN，以利用 IPR 进行修复。
• 在阶段1中，通过重建损失对 CPEN S1 与 DIRformer 进行联合优化进行训练。
• 在阶段2中，训练一个扩散模型来从劣质图像估计 IPR，使用紧凑潜在向量并进行联合优化。
• 在扩散框架内使用 CPEN S2 和去噪网络，以迭代方式细化 IPR 并修复图像。

Figure 1: The Mult-Adds are measured on 256 $\times$ 256 inputs. Our DiffIR achieves SOTA performance on IR tasks. Notably, LDM [ 50 ] and RePaint [ 40 ] are DM-based methods, and DiffIR is 1000 $\times$ more efficient than RePaint while achieving better performance.

实验结果

研究问题

• RQ1 扩散模型是否能在 IR 任务中有效地在紧凑的 IR 先验向量上运作，而不是在完整图像上？
• RQ2 两阶段训练（真实图像引导和低质量引导）是否会提升修复质量和稳定性？
• RQ3 CPEN S2、DIRformer 与去噪网络的联合优化是否能降低误差传播和伪影？
• RQ4 与基于扩散的最先进 IR 方法相比，DiffIR 在修补、超分辨和运动去模糊上的表现如何？

主要发现

• DiffIR 在多项 IR 任务上实现了最先进的性能，同时使用的迭代次数显著减少、计算资源更低。
• 在 CPEN 指导下，紧凑的 IPR 能实现轻量级 DIRformer 的有效修复。
• CPEN S2、DIRformer 与去噪网络的联合优化降低了估计误差对修复质量的影响。
• 在实验中，DiffIR 的效率显著高于 RePaint 和 LDM，在修补、超分辨和去模糊方面超过若干基于扩散的基线方法。
• 消融研究显示 DiffIR S2 设计、联合训练方案，以及在反向扩散中避免方差噪声对 IPR 估计的好处。

Figure 2: The overview of the proposed DiffIR, which consists of DIRformer, CPEN, and denoising network. DiffIR has two training stages: (a) In the first stage, CPEN S1 takes the ground-truth image as input and outputs an IPR $\mathbf{Z}$ to guide DIRformer to restore images. We optimize the CPEN S1