[论文解读] Refusion: Enabling Large-Size Realistic Image Restoration with Latent-Space Diffusion Models

Ziwei Luo, Fredrik Gustafsson+3 more|arXiv (Cornell University)|Apr 17, 2023

Advanced Image Processing Techniques被引用 6

一句话总结

Refusion 提出了一种基于 U-Net 的潜在扩散模型，通过在压缩的低分辨率潜在空间中执行去噪，同时通过解码器附加的潜在向量保留高分辨率细节，实现了高保真、大规模图像修复。该方法在真实世界图像修复任务中实现了最先进（SOTA）的感知性能，包括在 NTIRE 2023 阴影去除挑战赛中获得第二名，并且无需对抗性训练或数据集筛选即可成功处理 6000×4000 图像。

ABSTRACT

This work aims to improve the applicability of diffusion models in realistic image restoration. Specifically, we enhance the diffusion model in several aspects such as network architecture, noise level, denoising steps, training image size, and optimizer/scheduler. We show that tuning these hyperparameters allows us to achieve better performance on both distortion and perceptual scores. We also propose a U-Net based latent diffusion model which performs diffusion in a low-resolution latent space while preserving high-resolution information from the original input for the decoding process. Compared to the previous latent-diffusion model which trains a VAE-GAN to compress the image, our proposed U-Net compression strategy is significantly more stable and can recover highly accurate images without relying on adversarial optimization. Importantly, these modifications allow us to apply diffusion models to various image restoration tasks, including real-world shadow removal, HR non-homogeneous dehazing, stereo super-resolution, and bokeh effect transformation. By simply replacing the datasets and slightly changing the noise network, our model, named Refusion, is able to deal with large-size images (e.g., 6000 x 4000 x 3 in HR dehazing) and produces good results on all the above restoration problems. Our Refusion achieves the best perceptual performance in the NTIRE 2023 Image Shadow Removal Challenge and wins 2nd place overall.

研究动机与目标

• 提升扩散模型在具有复杂真实世界退化的大规模、真实图像修复任务中的适用性。
• 克服现有扩散模型的局限性，如计算成本高、对已知退化参数的依赖性以及潜在空间压缩中的不稳定性。
• 开发一种稳定、高效且可泛化的框架，适用于多种修复任务，且无需对抗性优化。
• 通过在低维潜在空间中操作，同时保留精细细节，实现高分辨率图像修复（例如 6000×4000）。

提出的方法

• 该方法采用均值回归的随机微分方程（SDE）作为反向去噪过程，受 IR-SDE 启发，隐式建模退化并实现任务泛化。
• 使用基于 U-Net 的压缩网络替代先前潜在扩散模型中使用的 VAE-GAN，实现更稳定、更精确的潜在表示学习，且无需对抗性训练。
• 模型采用无非线性激活函数的网络（NAFNet）作为噪声预测头，相比标准 U-Net 架构，提升了效率与性能。
• 潜在空间用于去噪，而高分辨率特征通过跳跃连接保留在解码器中，从而实现对精细细节的精确重建。
• 通过仅更换数据集和噪声网络，该方法可支持多种图像修复任务，确保灵活性与泛化能力。
• 潜在 Refusion 变体通过在潜在空间中执行扩散，将 FLOPs 减少约 15 倍，从而实现在 6000×4000 等大图像上的推理，而全模型 Refusion 和 IR-SDE 在此类图像上均无法运行。

实验结果

研究问题

• RQ1 基于 U-Net 的潜在扩散模型是否能在不依赖对抗性训练或 VAE-GAN 压缩的情况下实现稳定、高保真的图像修复？
• RQ2 将 U-Net 噪声网络替换为 NAFNet 对图像修复任务中的性能与效率有何影响？
• RQ3 在压缩潜在空间中进行潜在扩散是否能有效实现超大图像（如 6000×4000）的修复，同时保留高分辨率细节？
• RQ4 该方法在多种真实世界图像修复任务（如阴影去除、去雾、景深转换）中的泛化能力如何？
• RQ5 在大规模数据集上，该模型与现有基于扩散的方法相比，在感知质量与推理速度方面表现如何？

主要发现

• Refusion 在 NTIRE 2023 图像阴影去除挑战赛中获得第二名，LPIPS（0.149）和 MOS（8.94）两项感知评分最高，优于所有其他参赛团队。
• 在真实世界阴影去除任务中，该模型实现了卓越的感知性能，LPIPS 为 0.149，MOS 为 8.94，显著优于基线模型（如 DHAN）。
• 潜在 Refusion 变体相比完整 Refusion 模型将 FLOPs 减少约 15 倍，使 6000×4000 图像的推理成为可能，而完整 Refusion 和 IR-SDE 在此类图像上均无法运行。
• 在景深效果转换任务中，潜在 Refusion 的 LPIPS 分数优于 Restormer，且推理速度接近，生成结果更锐利。
• 基于 U-Net 的潜在压缩策略相比 VAE-GAN 方法显著更稳定，实现了无需对抗性优化的高精度图像恢复。
• 该方法成功实现了对 6000×4000 图像的非均匀去雾修复，而全分辨率扩散模型（如 IR-SDE 和 Refusion）由于显存与计算资源限制，无法在此类图像上运行。