[paper review] Come-Closer-Diffuse-Faster: Accelerating Conditional Diffusion Models forInverse Problems through Stochastic Contraction

본 논문에서는 reverse process를 Gaussian noise에서 시작하지 않고 single forward diffusion에서 시작해 sampling steps를 상당히 줄일 수 있음을 보였다.

 

Introduction

Diffusion models는 forward diffusion process를 t=0에서 시작해 t=T까지 점진적으로 노이즈를 더하고 reverse diffusion process에서 trained score function을 활용해 t=T에서 t=0까지 점진적으로 denoising해 샘플을 생성한다.

Diffusion model들의 문제점은 sampling하기 너무 느리다는 점이다.

 

해당 문제를 해결하기 위한 Unconditional generative models에 대한 기존 연구들은 다음과 같다.

• deterministic sample paths를 규정 ex. DDIM, SDE
• score function 학습 이후 optimal steps 탐색
• Knowledge distillation을 이용해서 student network를 retraining하는 방식

본 연구에서는 reverse diffusion path의 contraction property를 이용해서 conditional diffusion model들을 accelerate하는 방법을 사용한다.

t_0 : reverse diffusion이 시작되는 time (< T)이라고 할때

즉 forward process를 t0까지만 진행하고 그 시점부터 reverse process를 시작하게 된다.

 

KEY : forward-diffused images는 초기 estimation errors를 증가시키지만 reverse conditional diffusion path는 stochastic difference equation(sde)의 contraction property 를 사용함으로써 지수적으로 빠르게 에러를 줄일 수 있게 된다.

 

Background

Score-based Diffusion models

Discrete Forms of SDEs

f_bar와 g_bar를 무엇으로 선택하느냐에 따라서 VP-SDE나 VE-SDE를 사용할 수 있다.

  (1) Variance Preserving SDE (VP-SDE)

     - DDPM의 continuous version

 

  (2) Variance Exploding SDE (VE-SDE)

     - SMLD(Score matching with Langevin dynamic)은 VE-SDE의 discrete version으로 볼 수 있음.

 

Main Contribution

1) The CCDF Algorithm

goal : Reverse diffusion 시작점 N’ = Nt_0 < N을 설정해 reverse diffusion step의 수를 줄이자.

1. 더 나은 초기치 x0에서 시작해 N’까지 forward diffusion 진행

    - a_{N'}, b_{N'} 은 diffusion model로 결정됨.

 

2. i=0까지 reverse conditonal diffusion 수행.

기본 reverse diffusion과 Data consistency를 부여하기 위한 operation을 번갈아 사용함.

이때 A가 linear하다고 가정.

2) Fast Convergence Principle of CCDF

Lemma 1. Forwad diffusion 이후 estimation error (GT clean 기반 - Initial Estimate 기반)

- Gaussian noise들의 independency의 결과

-

증명

Theorem 1. Contraction 

• SDE(Stochastic Difference Equation)의 stochastic contraction property를 활용함.
• \bar{\epsilon}_{0, r} : i=0까지 reverse 했을때 estimation error
• τ = Tr(A^TA) / n
• λ = contraction rate

  

  ⇒ conditional reverse diffusion 문제가 지수적으로 수렴할 수 있음을 보임.

Theorem 2. Shortcut path

 

• N’에서 시작함으로써 더 짧은 sampling path를 사용하는데, 이것을 이용해 우수한 결과를 내는 것이 tighter bound로 가능함을 보임.

-> 즉 굳이 N에서 sampling을 시작할 필요가 없고 N’에서 시작했을때 같은 point로 더 빨리 수렴할 수 있음을 보인 것이다.

만약 initialization을 잘해서 \tilde{x}0를 x_0와 비슷한 수준으로 estimate하면 \epsilon_0가 더 작아지고 즉 \bar{\epsilon}{0, r} 에 대한 bound가 더 tight해지기 때문에 더 빠르게 수렴이 가능하다. (N’이 더 작아진다 / reverse 시작지점을 더 빨리 끌어당길 수 있다)

 

Experiments

총 세가지 task에서 실험하였음

: super-resolution, inpainting, MRI reconstruction

• t_0를 줄여가면서 측정 (t_0 : reverse diffusion이 시작되는 time (< T))
• 데이터셋 = FFHQ 256x256, AFHQ 256x256

Super-resolution

- ε0 : x0와 x0 tilde의 차이가 작을수록(GT clean image, initial estimate) reverse diffusion의 수렴이 같은 지점에 더 안정적으로 이루어짐.

- Task가 어려울 수록 더 많은 reverse steps가 필요하다. 따라서 적절한 t_0를 선택하면 full step보다 좋은 성능이 가능하다.

- ILVR에서 iteration수를 줄일 수록 이미지의 퀄리티가 떨어지는 반면 제안된 방법론은 FID score를 iteration이 적을때에도 비슷한 수준으로 유지함.

 

DDIM과도 Incorporate해서 사용가능한데, 그 결과 5 steps만큼 적은 iterations로도 좋은 퀄리티를 보였다.(N’ = t0N = 5일때도 좋은 성능)

Limitations

Task마다/corruption마다 optimal한 t0값이 다르게 결정되는데, 어떤 명확한 rule이 없다는 한계점이 있다. (priori로 ε0를 알고있는게 아니기 때문에)

 

Conclusions

해당 논문에서는 stochastic contraction의 property를 통해서 conditonal diffusion model들을 accelerate할 수 있는 방법론을 제안했다.

NN을 통해서 x_0로 initialize하고 stochastic contraction theory를 통해서 shortcut path가 optimal하다는 것을 증명했다. 해당 방법론을 통해 향상된 안정성과 성능과 더불어 속도 향상이 가능했음을 여러 task들에 실험해서 확인했다.

 

Algorithms

 

원문 : https://arxiv.org/abs/2112.05146