[论文解读] Learning A Coarse-to-Fine Diffusion Transformer for Image Restoration

Liyan Wang, Qinyu Yang+4 more|arXiv (Cornell University)|Aug 17, 2023

MRI in cancer diagnosis被引用 5

一句话总结

本文提出一种用于图像修复的级联式扩散Transformer（C2F-DFT），引入一种新颖的训练方案，通过将采样输出与真实图像进行对比优化，而非仅约束噪声，从而提升修复质量。该方法通过利用时间步长条件化的扩散自注意力和前馈网络，在去雨、去模糊和去噪任务中均实现了最先进性能。

ABSTRACT

Recent years have witnessed the remarkable performance of diffusion models in various vision tasks. However, for image restoration that aims to recover clear images with sharper details from given degraded observations, diffusion-based methods may fail to recover promising results due to inaccurate noise estimation. Moreover, simple constraining noises cannot effectively learn complex degradation information, which subsequently hinders the model capacity. To solve the above problems, we propose a coarse-to-fine diffusion Transformer (C2F-DFT) for image restoration. Specifically, our C2F-DFT contains diffusion self-attention (DFSA) and diffusion feed-forward network (DFN) within a new coarse-to-fine training scheme. The DFSA and DFN respectively capture the long-range diffusion dependencies and learn hierarchy diffusion representation to facilitate better restoration. In the coarse training stage, our C2F-DFT estimates noises and then generates the final clean image by a sampling algorithm. To further improve the restoration quality, we propose a simple yet effective fine training scheme. It first exploits the coarse-trained diffusion model with fixed steps to generate restoration results, which then would be constrained with corresponding ground-truth ones to optimize the models to remedy the unsatisfactory results affected by inaccurate noise estimation. Extensive experiments show that C2F-DFT significantly outperforms diffusion-based restoration method IR-SDE and achieves competitive performance compared with Transformer-based state-of-the-art methods on $3$ tasks, including image deraining, image deblurring, and real image denoising. Code is available at https://github.com/wlydlut/C2F-DFT.

研究动机与目标

• 为解决基于扩散的图像修复模型在采样过程中因噪声估计不准确而导致修复质量不佳的问题。
• 通过一种两阶段训练方案，超越简单的噪声约束，学习复杂退化模式，从而提升模型容量。
• 通过嵌入扩散机制的Transformer架构，增强图像修复中的长距离依赖建模与层次化特征学习。
• 开发一种微调策略，利用采样得到的修复结果而非噪声来优化模型，从而纠正因噪声估计不准确导致的误差。
• 在多个图像修复基准上实现具有竞争力的性能，包括去雨、去模糊和真实世界去噪。

提出的方法

• 该模型采用时间步长嵌入条件化的扩散Transformer架构，包含扩散自注意力（DFSA）和扩散前馈网络（DFN），以建模扩散过程中的时序动态。
• 提出一种级联式训练方案，包含（a）粗粒度训练，约束噪声预测；（b）细粒度训练，将采样过程生成的模型输出与真实图像进行对比约束。
• 在细粒度训练阶段，模型使用粗粒度训练模型生成的采样修复结果进行优化，而非使用噪声，以纠正因噪声估计不准确导致的误差。
• 微调过程采用固定采样步数（4步），以确保一致性与效率，在感知质量与计算成本之间实现平衡。
• 在粗粒度训练中采用补丁循环训练策略，以提升泛化能力与效率，其性能优于具有相似训练时间的固定补丁基线模型。
• 时间步长条件化被嵌入DFSA与DFN中，使模型能够根据不同扩散步骤自适应调整行为，显著提升修复性能。

实验结果

研究问题

• RQ1 是否可通过一种级联式训练方案，将优化目标从噪声转向采样输出，从而提升扩散模型中的图像修复质量？
• RQ2 扩散Transformer中时间步长嵌入如何影响模型捕捉长距离依赖与修复保真度的能力？
• RQ3 在微调阶段对生成输出进行约束，是否能带来优于标准噪声约束训练的性能提升？
• RQ4 在微调流程中，为平衡修复质量与推理速度，最优采样步数是多少？
• RQ5 与固定补丁训练相比，补丁循环训练策略在提升模型泛化能力与效率方面的有效性如何？

主要发现

• 所提出的C2F-DFT在Rain100H去雨基准上达到31.38 dB的PSNR与0.921的SSIM，优于基线粗粒度训练（30.93 dB PSNR）与IR-SDE。
• 在Test2800去模糊数据集上，模型达到34.03 dB的PSNR与0.944的SSIM，相比仅粗粒度训练提升0.29 dB。
• 微调阶段显著提升了视觉质量，相比粗粒度训练，有效减少了伪影与噪声，如图9与图10所示。
• 时间步长嵌入使PSNR提升4.08 dB（从26.85提升至30.93 dB），证明其在性能提升中的关键作用。
• 在微调中使用4步采样实现了感知质量（LPIPS）与推理时间之间的良好平衡，优于步数更少或更多的设置。
• 补丁循环训练策略在保持相近训练时间的前提下，性能优于固定补丁训练，证实其有效性。