[paper review] Diffusion Models for Image Restoration and Enhancement – A Comprehensive Survey

[2024/07/01] Image Restoration 연구 과정에서 연구 흐름을 recap하고자 해당 survey를 읽게 되었습니다. 아래 내용은 세부적인 내용은 포함하고있지 않으며 흐름을 파악하기 위한 용도로 작성되었습니다. (원문 :  https://arxiv.org/abs/2308.09388) 

 

Introduction

이전에 sptial/freqeuncy관점에서 hand-crafted 알고리즘들이 제안되었지만 여전히 texture generation에 한계가 존재했다.

이후 GAN이 등장하면서 MSE loss 등의 pixel-wise loss외에 adversarial loss를 함께 사용하게되었다.

GAN-based IR의 한계점

1. Mode corruption, Unstable optimization (minmax game이 object function으로 사용됨)
2. 생성된 이미지의 texture들이 fake로 보임.

새로운 breakthroughs로 DDPM, NCSN, SDE가 등장했다.

Diffusion model들은 GAN과 비교했을때 high-fidelity와 다양한 생성을 가능케함.

Texture recovery를 위한 Image restoration training 전략

1) supervised diffusion model- based IR

2) zero-shot diffusion model-based IR

• 사전학습된 diffusion model의 generative priors를 이용함

Background on Diffsion Model (DM)

1. NCSNs

= Annealed Langevin dynamics를 이용해 decreasing noise scale을 적용하여 data distribution을 모델링하는 방법론

2. DDPM

= Forward process에서 Gaussian noise를 점진적으로 더하고, reverse process에서 이를 제거해 이미지를 생성함. (Variational bound objective를 이용해 최적화)

이때 reverse process의 variance는 fixed values.

3. SDE

= continuous diffusion과 reverse를 stochastic differential equation으로 모델링하는 방법론. 본 논문에서 NCSNs와 DDPM은 결국 SDE의 two separate discretizations임을 증명함.

4. Improvements with Diffusion Model

• Optimization strategies : Forward/ backward processes에서 모두 variance와 noise schedule을 optimize하고자하는 시도들이 생겼고, 이에따라 DDIM(=IDDPM), DPM이 제안됨.
• Sampling Efficiency 향상을 위한 Works
  1. ODE와 관련된 Handcrafted sampling 전략들 
     • DDIM = Forward process에서 non-markovian chain → sampling over arbitrary steps
     • DPM-solver = black-box ODE solvers를 쓰는 대신 ODE solution의 linear한 부분을 계산해서 빠른 ODE solver를 구성함
  2. Diffusion process를 revise (ex. Early stopping을 도입해 non-Gaussian distribution에서 sampling)
  3. Knowledge Distillation
  4. Condition Strategy

• Model Architecture = U-Net은 resolution을 보존하면서도 cost가 덜들기 때문에 noise/score prediction시 많이 쓰임.      , 주로 CNN기반 U-Net과 Transformer 기반 모델을 사용.

Diffusion Model-based Image Restoration methods

1) supervised diffusion model- based IR

- Paired distorted / clean images of IR datasets을 이용함.

2) zero-shot diffusion model-based IR

- distorted images만 사용하고, 재학습 필요 없음.

- Pre-trained diffusion model로부터 structure와 texture priors를 추출해서 사용함.

→ 즉 여기서 어떻게 data structure를 보존하면서 perceptual priors를 추출할지가 중요하다.

 

3.1 SR3 – Representative Supervised DM for IR

Idea - Degraded image y를 generated image x_t와 concatenate하는 Condition Strategy를 사용함.

- Noise predictor로 U-Net을 사용함.

- Low-resolution image y를 upsamling한 다음 x_t와 concat해서 t-1에서의 noise를 예측함.

3.2 ILVR – Representative Zero-shot DM for IR

Idea - Pre-trained diffusion models의 intrinsic knowledge를 이용

- Denoised output x_t의 low-freqeuncy 요소를 reference image로 대체함

→ 생성한 이미지와 reference 이미지간의 semantic consistency 보장하게 해줌.

3.3 Supervised DM-based IR

SR3에서 모티브를 얻어 condition strategy enhancing에 집중하였음.

• Condition with Low-quality Reference Image = concat low-quality reference img + denoised result
• Condition with Pre-processed References = 학습된 모듈이나 pre-trained restoration nw를 사용해서 low-quality image들을 전처리
• Condition by Revising Diffusion Process
• Generate Residuals = Low-quality 이미지에 있는 structues/textures를 high-quality image를 만들때 같이 만들게 되다보니까 추가적인 resource cost가 발생한다는 점을 문제 삼아서, high-quality와 low-qulity 이미지들의 residuals를 generate하도록 하였음. 이때 Generation process를 image space 대신 residual space에서 진행함.
• Diffusion on Latent Space = 대표적으로 StableDiffusion이 있고, 이는 latent space에서 DM기반 generation을하는 모델. Pre-trained autoencoder으로 image-wise diffusion process를 perceptual space르 보내고, 해당 모델에 condition들을 부여하는 방식. ex. Refusion 

  

3.4 Zero-shot DM-based IR

Training/data-free Image Restoration.

1) Projection-based Methods

• Data consistency를 보장해줄 수 있도록 이미지 생성하는 각 스텝에서 structues/textures를 추출하는 것을 목표로함. ex. Image inpainting task에서 Mask region에서만 generating하도록 함.
• masking한 y에서 diffuse한 x^known와 denoised prediction에서 샘플링한 x^unknown

 

2) Decomposition-based methods

 - 대부분의 IR 문제들은 linear reverse problem으로 여겨짐(y로부터 original인 x를 estimate)

y = Hx + z

- x를 estimate할때 noise z때문에 바로 estimate하기 어려워서 SNIPS, DDRM에서는 Diffusion process를 spectral domain에서 SVD(Singular value decomposition)과 함께 적용함.

• DDRM = SVD decomposition 문제를 linear reverse problems의 variational objective로 확장해서, pre-trained DDPM, DDIM이 optimal solution이 될 수 있음을 보였다.

- Non-linear inverse problems

 : reverse에 noise가 없는 DDRM과 같은 특이 케이스, 그리고 JPEG artifact(왜곡) correction을 위한 pseudo-inverse를 위해 non-linear inverse problems에 집중하였음

ex. MRI reconstruction할때는 SVD decomposition이 적절하지 않았다. → unconditiondal DM + Matrix Decomposition 적용

Range-null space decomposition을 DDNM에서 사용

• range space = data consistency
• null space = improve reality

 

 

3) Posterior Estimation

= Posterior distribution인 p(x|y)를 estimate할때 bayes 정리 사용

이때 p(y|x_t)는 intractable하므로, 이 term을 잘 예측하는게 중요하다.

MCG, DPS = posterior p(y|x_t)를 p(y|x_hat)으로 approximate, data manifold관점에서 data consistency를 판단함.

 

 

 

Diffusion Models for blind/real-world Image Restoration

Image Restoration이 갖는 고질적인 문제점은 OOD(Out-of-distribution)에서 poorly perform한다는 점이다. 이는 real-worl IR에 존재하는 inherenct challenges가 원인인데

1. unknown degradation mode들은 잡아내기 어렵고
2. 실제 세계에서 distorted/ clean image쌍을 모으는건 현실적으로 힘들다.

⇒ solution = simulating real-world degradations, unsup learning

 

4가지 카테고리

1. Distortion Simulation
2. Kernel Estimation
3. Domain Translation - synthetic domain의 low-quality image를 real-world domain으로 옮김
4. Distortion-invariant Diffusion Model

 

Experimental Results

Challenges and Future Directions

1. Sampling Efficiency : SR3에서 224x224크기 이미지 복원에 50초 걸렸는데 DDRM은 DDIM의 Sampling strategy(deterministic) 방법을 써서 8초로 줄였음.
2. Model Compression
   1. model pruning 2) model quantization 3) knowledge distillation 4)low-rank decomposition