PaperlayerView exact revision on arXiv

Source-linked AI summary

DiffIR: Efficient Diffusion Model for Image Restoration

Bin Xia, Yulun Zhang, Shiyin Wang, Yitong Wang, Xinglong Wu, Yapeng Tian, Wenming Yang, Luc Van Gool

arXiv:2303.09472v3cs.CV

TL;DR

Image restoration needs outputs consistent with low-quality references, whereas conventional diffusion applies many iterations to whole images or feature maps. DiffIR instead estimates a compact restoration prior through a two-stage CPEN, DIRformer, and diffusion design, achieving state-of-the-art performance with lower computational cost.

  • Problem

    Traditional diffusion models are inefficient for image restoration because they use many iterations on large models to estimate whole images or feature maps despite low-quality inputs already providing most image information.

  • Method

    DiffIR pretrains CPEN to extract a compact IPR that guides DIRformer, then trains diffusion to estimate the same prior from low-quality images with joint optimization of the denoising network and decoder.

  • Results

    DiffIR achieves state-of-the-art performance across image-restoration tasks while consuming much less computational resources than other diffusion-based methods.

  • Takeaways & Limitations

    Estimating a compact IPR rather than whole images enables fewer diffusion iterations and more efficient, stable restoration.

Abstract

from arXiv · show

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}.

1. Introduction

Image restoration must preserve consistency with a given low-quality image, making whole-image diffusion costly and potentially mismatched. DiffIR addresses this by estimating a compact restoration prior with diffusion and reports strong performance at lower computational cost.

  • IR reconstructs high-quality images from degraded low-quality inputs, but its ill-posed nature makes the task challenging.
  • Traditional diffusion models use about 50−1000 iterations on large denoising models, wasting computation and potentially generating details inconsistent with the input.
  • DiffIR uses CPEN to extract a compact IPR from ground-truth images, then trains diffusion to estimate that prior from low-quality images for DIRformer-guided restoration.
  • DiffIR uses diffusion to estimate a compact IPR rather than whole images or feature maps, reducing restoration cost and improving stability.
  • DiffIR jointly optimizes the denoising network and DIRformer, improving robustness to errors in the estimated restoration prior.
  • DiffIR achieves state-of-the-art restoration performance while consuming much less computational resources than other diffusion-based methods.

2. Related Work

Prior image-restoration research spans compact CNNs, architectural innovations, GANs, attention, and transformers, while diffusion models have recently extended to restoration tasks.

  • Early image-restoration methods used compact CNNs, followed by residual, GAN, attention, and knowledge-distillation designs.
  • Transformers model global interactions between image regions and have been adopted for image restoration alongside other vision tasks.
  • Diffusion models have achieved strong results in density estimation, sample quality, super-resolution, and image inpainting.

3. Preliminaries: Diffusion Models

Diffusion models transform data into Gaussian noise through a forward process and recover outputs through iterative denoising. Their denoising network is trained to estimate the noise added at sampled time steps.

  • The forward diffusion process transforms an input image x_0 into Gaussian noise x_T over T iterations.
  • In the proposed DiffIR overview, CPEN extracts an IPR during pretraining, while the diffusion model later estimates that prior without ground-truth inputs.
  • The noised image x_t uses β_t as a predefined scale factor, with N denoting the Gaussian distribution.
  • During inference, diffusion models sample Gaussian noise x_T and gradually denoise it until reaching a high-quality output x_0.
  • Training samples a time step t and Gaussian noise ε to create x_t, then optimizes ε_θ(x_t,t) to estimate that noise.

4. Methodology

DiffIR applies diffusion modeling to a compact IR prior rather than whole images or feature maps, combining CPEN, DIRformer, and a denoising network. Its two-stage training and few-step inference support efficient restoration and joint optimization.

  • Motivation: DiffIR estimates a compact IR prior representation instead of denoising whole images or feature maps, reducing model size and iteration requirements.The compact IPR guides image restoration while avoiding the inefficiency and artifact risk associated with applying image-synthesis diffusion directly to IR.
  • Architecture and training: DiffIR consists of CPEN, DIRformer, and a denoising network, and is trained through pretraining followed by diffusion-model training.The first stage extracts an IPR and trains DIRformer to use it; the second stage estimates the IPR from LQ images.
  • Pretraining: CPENS1 extracts an IPR from concatenated ground-truth and LQ images, which DIRformer uses as dynamic modulation parameters during restoration.The IPR is injected into both DGFN and DMTA, while DIRformer uses dynamic transformer blocks arranged in a Unet shape.
  • DIRformer: DIRformer aggregates global spatial information with transposed attention and local information with gated feed-forward processing.Its attention map has size R^Ĉ×Ĉ rather than R^(ĤŴ)×(ĤŴ), and DGFN combines channel and neighboring-pixel aggregation.
  • Diffusion inference: During inference, CPENS2 derives a conditional vector from LQ images, the denoising network estimates the IPR through T iterations, and DIRformer restores the image.Because the IPR is compact, DiffIR uses fewer iterations and a smaller model; all iterations can be jointly optimized with CPENS2 and DIRformer.

5. Experiments

DiffIR is evaluated on inpainting, 4× super-resolution, and single-image motion deblurring, achieving strong quantitative and qualitative results with reduced computation.

  • Inpainting: DiffIRS2 significantly outperforms LaMa on inpainting, improving FID by up to 0.2706 and 0.5583 with wide masks on Places and CelebA-HQ.It uses similar total parameters and Mult-Adds.
  • Inpainting: DiffIRS2 achieves better inpainting performance than RePaint while using 4.3% of its parameters and 0.1% of its computational resources.
  • Qualitative Results: Qualitative comparisons report more realistic details and better visual quality for DiffIRS2 in inpainting, super-resolution, and motion deblurring.
  • 4× Super-Resolution: DiffIRS2 achieves the best performance among the compared GAN-based super-resolution methods, surpassing BebyGAN by LPIPS margins up to 0.0151 and 0.0089.The margins are reported on DIV2K100 and Urban100, respectively.
  • 4× Super-Resolution: DiffIRS2 surpasses BebyGAN while consuming 63% computational resources and significantly outperforms LDM while consuming 2% computational resources.
  • Motion Deblurring: DiffIRS2 surpasses IPT and MIMI-Unet+ by 0.68 dB and 0.54 dB on GoPro, respectively, and surpasses Restormer by 0.28 dB on GoPro and 0.33 dB on HIDE.It consumes 78% computational resources in the Restormer comparison.

6. Ablation Study

Ablations show that diffusion-based IPR estimation, joint optimization, noise removal, the proposed diffusion loss, and few iterations each contribute to DiffIR’s performance or efficiency.

  • Diffusion Model: DiffIRS2-V3 has quite similar LPIPS to the ground-truth-input pretraining network DiffIRS1, indicating accurate IPR prediction by DM.
  • Diffusion Model: DiffIRS2-V3 significantly outperforms DiffIRS2-V1 without DM, showing that DM-learned IPR effectively guides DIRformer.
  • Training Scheme: Joint optimization enables CPENS2, DIRformer, and the denoising network to be optimized together because DiffIR uses few DM iterations.Traditional DM optimization cannot jointly optimize the later decoder when many iterations are required.
  • Variance Noise: Removing variance noise from the reverse process performs better than inserting it, indicating that noise cancellation preserves IPR estimation accuracy.
  • Loss Functions: Ldiff performs better than L2 and Lkl for learning accurate IPR estimates from LQ images on CelebA-HQ inpainting.
  • Iteration Count: Performance improves significantly through 3 iterations, stabilizes beyond 4, and converges faster than traditional DM methods requiring more than 200 iterations.The faster convergence is attributed to applying DM only to the compact one-dimensional IPR vector.

7. Conclusion

DiffIR adapts diffusion modeling to image restoration by estimating a compact restoration prior rather than whole images or feature maps. Its few iterations enable joint optimization and support strong restoration performance with lower computational cost.

  • Method: DiffIR combines CPEN, DIRformer, and a denoising network in a two-stage design that extracts and then estimates a compact IPR from LQ images.
  • Efficiency: Applying diffusion to the compact IPR lets DiffIR use fewer iterations than traditional DMs and reduce artifacts in restored images.
  • Optimization: Joint optimization of CPENS2, DIRformer, and the denoising network reduces the influence of IPR estimation errors.
  • Results: Extensive experiments show that DiffIR achieves general SOTA performance across image restoration tasks.

B. Evaluation on Real-world SR

DiffIRS2 is evaluated on real-world super-resolution using Real-ESRGAN settings and achieves strong quantitative and qualitative results with substantially lower computation.

  • Training uses Real-ESRGAN loss functions, including perceptual and adversarial losses alongside L1 loss, with Adam optimization and learning rate 2 × 10^-4.
  • DiffIRS2 outperforms KDSRS-GAN on LPIPS, DISTS, and PSNR while consuming fewer computational costs.
  • 63% Mult-Adds: DiffIRS2 surpasses Real-ESRGAN on LPIPS, DISTS, and PSNR while using less computation.
  • 2% Mult-Adds: DiffIRS2 achieves much better performance than LDM while consuming substantially less computation.
  • On smartphone-captured NTIRE2020 Track2 data, DiffIRS2 achieves the best reported qualitative performance.

D. More Training Details on Inpainting

The inpainting experiments follow LaMa-based mask generation and losses, with specified optimization, training duration, cropping, and evaluation settings.

  • Inpainting training uses LaMa to generate synthetic masks and incorporates multiple perceptual and adversarial losses with basic L1 loss.
  • DiffIR and discriminator learning rates are 0.0002 and 0.0001, respectively, using Adam optimization.
  • All inpainting models are trained for 1M iterations with batch size 30 and random 256 × 256 crops on Places and CelebA-HQ.
  • Testing uses fixed CelebA-HQ validation and testing subsets of 2k and 30k samples, respectively.

E. More Training Details on SR

The supplied material combines SR and motion-deblurring training settings with DiffIRS2 algorithm descriptions and qualitative comparisons across restoration tasks.

  • Super-resolution training: SR training uses a ×4 upsampling network appended to DIRformer and adopts ESRGAN losses combining L1, perceptual, and adversarial terms.
  • Super-resolution training: SR models are trained for 1M iterations with 256 × 256 crops from DIV2K and Flickr2K, using batch size 64 and Adam learning rates 0.0002 and 0.0001.
  • Motion deblurring training: Motion deblurring training uses only L1 loss, runs for 300K iterations, and applies cosine annealing from 2^-4 to 1^-6.
  • DiffIRS2 training: DiffIRS2 training initializes DIRformer from trained DiffIRS1, diffuses a compact prior, estimates it through CPENS2, and reconstructs HQ images with DIRformer.
  • DiffIRS2 inference: At inference, DiffIRS2 samples the terminal latent, applies CPENS2 and the reverse process, then outputs restored HQ images.
  • Qualitative comparisons: Qualitative comparisons report more realistic inpainting structures and details, best SR visual quality, and deblurring details close to corresponding HQ images.
Loading 2303.09472v3…