[Generative Model] Diffusion Model (1)

안녕하세요. 이번 포스팅은 생성 모델 중 하나로 diffusion model에 대해 포스팅하겠습니다. 
diffusion model은 최근 많은 이슈가 되었던 stable-Diffusion, DALL-E2, Imagen의 base가 되는 Generative Model입니다. 몇 년전만 하더라도 대표적인 Generative Model하면 GAN을 떠올렸지만 이제는 대부분 Diffusion model을 떠올리게 될 정도로 변화가 빠른 분야입니다.
Diffusion은 분자들의 확산을 뜻하는데 분자의 움직임은 Gaussian Distribution을 따릅니다. 즉, 평균과 표준 편차를 따르는 정규 분포로 움직임을 표현할 수 있습니다. 만약 매 스텝마다의 분자의 움직임을 알 수 있다면 이를 확산전의 모양으로 되돌리는 것도 가능할 것 입니다. 이 컨셉을 이용한 것이 Diffusion model입니다.
 

 

---

Diffusion Model?

overview

위의 그림은 diffusion model의 overview입니다.  입력이미지에 noise를 매 step마다 추가하는 부분이 Forward process이고, 이를 다시 매 단계마다  de-noise하는 부분이 Reverse Process로 입력 이미지와 유사한 이미지를 생성합니다.  각 단계의 $x_{t}$는 noise가 추가된 상태의 latent variable을 의미합니다. 유의할 점은 Forward은 고정된 Gaussian distribution의 noise가 추가가되고 reverse에서는 learnable parameter를 이용하여 denoise하게 되고 이 reverse과정을 통해 random noise로부터 image, text, graph등을 생성해 낼 수 있습니다.
 

---

Forward process

fix된 gaussian noise를 점진적으로 더해줍니다.

베타는 주입되는 gaussian noise의 크기를 의미하고 linear, sigmoid, quadratic 스케쥴링 등의 방식이 있습니다. 베타를 이용하여 스케쥴링하는 이유는 variance가 발산하는 것을 막기 위함입니다.

이를 직관적으로 해석해보면 현재 state의 pixel값을 정할때 $\beta$만큼의 스케일로 다른값을 선택하고 $\sqrt{1-\beta}$만큼 이전의 값을 선택한다고 상상할 수 있습니다.
즉 t(0~T)마다 점진적으로 가우시안 노이즈를 추가한다고 보면 되겠습니다. (베타는 일종의 regularization)

Reverse process

Diffusion model의 목표는 결국 $x_{t}$의 분포를 $x_{0}$인 초기 data의 분포로 복원하는 것 입니다.
하지만 실제로 t시점에서 이미지의 확률분포 q, $q(x_{t-1}|x_{t})$을 직접적으로 구할 수 없기때문에 $p_{\theta}(x_{t-1}|x_{t})$으로
approximation하여 구하게되고 $q(x_{t-1}|x_{t})$와 $p_{\theta}(x_{t-1}|x_{t})$이 같게 되도록 학습하게 됩니다.다시 정리하면 q에서 관측한 값으로 $q_{\theta}$의 likelihood를 구할 때 최대가 되는 분포를 찾는 문제이고 여기에 -를 붙여 Negative Likelihood를 최소화하도록 loss를 설계하게 됩니다. 이를 전개하여 정리하면 아래의 수식으로 나타낼 수 있습니다.
DPM(Diffusion Probabilistic Model) Diffusion 초기 논문에서 사용한 Loss fuction입니다.

 

 
 

 위의 Loss의 의미에 대해 살펴보면, Regularization term은 $x_{T}$에 대해 확률분포 q와 $p_{\theta}$의 KL Divergence를 최소화하는 것을 의미합니다. Denoising Process term은 given $x_{t}$일 때 $x_{t-1}$이 나올 확률 분포 q와  $p_{\theta}$의 KL Divergence를 최소화를 의미합니다.
DDPM에서는 위의 loss를 한번 더 정제하여 아래와 같이 loss를 설계합니다. (*상세 유도는 추후 포스팅하겠습니다.)

$\epsilon$은 실제의 noise(GT) $\epsilon_{\theta}$는 Diffusion model이 예측한 noise입니다. 즉 t단계에서의 실제 노이즈와 x, U-Net에 input(x0)에 noise를 적용한 값을 넣고 예측한 noise와의 차이를 sampling된 pixel마다 하여 mse합니다.
아래와 같이 실제적인 noise detector로는 U-Net을 사용하고 U-Net을 통해 estimation된 noise와 실제 노이즈를 위의 loss function을 이용하여 학습을 하게됩니다.

 
즉, 최초 state($x_{0}$)로 부터 step t만큼 이후에 noise detector(U-Net)가 GT noise($\epsilon$)을 따르게 하는 loss입니다. 
실제 코드화 하기위에 위의 t단계에서의 forward식을 t0을 기준으로 아래처럼 표현할 수 있습니다. (reparameter)

t에서의 실제 노이즈는 forward를 통해 알 수 있고, t단계의 노이즈와 베타값들을 알 수 있으면 t+1단계의 이미지 분포도 알 수 있습니다. 이를 통해 noise detector(U-Net)에 input으로 t+1의 이미지 분포를 인풋으로 하여 t단계에서의 noise detector를 추정하고 이를 다시 실제 noise와 mse하는 것을 loss function으로 하게됩니다.

 

두서없는 위의 내용을 다시 어떤 flow로 학습하게 되고 inference(sampling)하는지 아래의 sudo code를 통해 다

 

---

Application

diffusion model을 베이스로한 실제적인 application 및 변형들에 대해 살펴보겠습니다.

class guided Image Generation
사용자가 원하는 class 이미지만 생성하는 방법입니다. ImageNet으로 학습된 Diffusion Modeldms 1000개의 class를 생성할 것 입니다. <Diffusion Model Beat GANs on Image Synthesis>에서 제안한 방법은 아래와 같습니다.
우선 Noisy Image를 사용하여 Classifier를 학습시킵니다. 충분히 학습된 Classifier의 Gradient를 이용해서 Diffusion Model의 guidance가 되줍니다. 이 방법은 Classifier을 먼저 학습시켜야하기에 실제로 product등에 사용되기에는 쉽지 않습니다. 
<Classifier Free Guidance>방법은 별도로 Classifier를 학습해야하는 과정없이 한번에 학습할 수 있는 방법을 제안합니다. 간략하게 말하면 class에 관련된 condition(embedding-일종의 힌트)을 추가하여 학습한 noise와 condition없이 예측한 noise를 interpolation하여 최종 noise로 사용합니다.

Text Guided Image Generation

최근 매우 뜨거운 분야입니다. 대부분의 모델들이 아래와 같은 아키텍쳐를 따르고 있습니다. pretrained text encoder를 통해 caption에서 text embedding을 추출해주고 첫번째 Diffusion model에 text embedding과 noisy image를 더하여 image를 생성해내고 다시한번 Diffusion model을 통해 super-resolution을 수행하여 high resolution의 image를 생성해냅니다.
Google의 Imagen, OpenAI의 DALL E2, Stable Diffusion 등이 있습니다. DALL E2는 CLIP을 사용하여 Text-Image의 관계를 나타내는 embedding을 사용하였고,

Stable Diffusion은 Auto Encoder를 사용하여 이미지를 압축한뒤 Diffusion model을 학습시켜 computation cost를 많이 줄이고 sementic information을 보다 풍부하게 사용할 수 있게 되었습니다.

---

 
 
다음 포스팅에서는 loss function에 대한 수식과 code레벨에서의 구현에 대해 살펴보겠습니다.