Denoising Diffusion Probabilitic Models(DDPM)

김민솔·2024년 3월 6일
DDPMdiffusiongenerative

Diffusion

목록 보기
1/5

최근에 많은 generative 모델들이 고해상도의 이미지들을 생성해내고 있으며, 다양한 task에 응용되어오고 있습니다. 그 중에서도 Denoising diffusion 모델들은 generative 모델 중에서도 높은 성능을 보여주고 있습니다. Diffusion 모델이 지금과 같이 사용되기 전에 생성모델 분야를 주름 잡던 GAN 계열 모델들과의 비교에서도 압도하는 모습을 보여주고 있습니다. (https://arxiv.org/abs/2105.05233) 또한 mode coverage와 다양한 샘플 생성을 통해 복잡하고 다양한 데이터도 학습이 가능함을 보여줍니다. Denoising diffusion 모델을 공부하기 위해, 우선적으로 알아야 할 DDPM 논문을 리뷰하겠습니다.



Introduction

DDPM은 Forward diffusion process와 Reverse denoising process, 총 2개의 과정으로 구성되어 있습니다. Forward diffusion process에서는 각 timestep에서 input data에 noise를 더해갑니다. Reverse denoising process에서는 denoising을 통해 data를 생성하는 방법을 배우게 됩니다.

또한 모델 학습에 사용되는 Loss인 Variational upper bound에 대해서도 자세히 살펴보겠습니다.



Forward diffusion process

앞에서 짧게 언급했던 것처럼, input data에 각 timestep마다 가우시안 노이즈를 더하는 과정입니다. 이때 βt\beta_t는 noise schedule로 정의되고, 매우 작은 양의 값을 가집니다.

q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)q(x1:T∣x0)=∏t=1Tq(xt∣xt−1)q(\mathbf{x}_t|\mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_{t};\sqrt{1-\beta_t}\mathbf{x}_{t-1}, \beta_t\mathbf{I}) \\ q(\mathbf{x}_{1:T}|\mathbf{x}_{0}) = \prod^T_{t=1} q(\mathbf{x}_t|\mathbf{x}_{t-1})

forward 시에 한 번에 여러 번의 step을 적용하는 것이 가능합니다. 즉 여러 step에 걸쳐 이미지에 노이즈를 주지 않고, 한 번에 해당 과정을 수행할 수 있습니다. αt‾=∏s=1t(1−βs)\overline{\alpha_t} = \prod^t_{s=1} (1 - \beta_s)로 α‾\overline{\alpha}를 정의하면, q(xt∣x0)q(\mathbf{x}_t|\mathbf{x}_{0})를 다음과 같이 표현할 수 있고 이를 diffusion kernel로 정의합니다. noise schedule 역할을 하는 βt\beta_t를 통해 α‾\overline{\alpha}가 0이 되게 하여, q(xt∣x0)q(\mathbf{x}_t|\mathbf{x}_{0})가 std에 가까워지도록 학습합니다.

q(xt∣x0)=N(xt;αt‾x0,(1−αt‾)I)q(\mathbf{x}_t|\mathbf{x}_{0}) = \mathcal{N}(\mathbf{x}_{t};\sqrt{\overline{\alpha_t}}\mathbf{x}_{0}, (1-\overline{\alpha_t})\mathbf{I})

xt\mathbf{x}_t를 샘플링하려면, diffusion kernel 분포에 reparameterization trick을 적용하여 해결할 수 있습니다.

  • xt=αt‾x0+(1−αt‾)ϵ\mathbf{x}_t = \sqrt{\overline{\alpha_t}}\mathbf{x}_{0} + \sqrt{(1-\overline{\alpha_t})}\epsilon
  • ϵ∼N(0,I)\epsilon \sim \mathcal{N}(0, \mathbf{I})

이 샘플링이 의미하는 것은, xt\mathbf{x}_t를 샘플링할 때 q(xt)q(\mathbf{x}_t)에서 직접적으로 하는 것이 아닌 training data dist.에서 샘플링이 가능하다는 점입니다.



Forward process에서 일어나는 과정 파악하기

Forward 과정에서 이미지에 diffusion kernel을 적용하여 노이즈 x_t를 만드는 것을 다뤘습니다. 이번에는 Forward 과정이 Data dist.에 어떤 영향을 주는지 알아보겠습니다.

q(xt)=∫q(x0,xt)dx0=∫q(x0)q(xt∣x0)dx0q(\mathbf{x}_t) = \int q(\mathbf{x}_0, \mathbf{x}_t)d\mathbf{x}_0 = \int q(\mathbf{x}_0)q(\mathbf{x}_t|\mathbf{x}_0)d\mathbf{x}_0
  • q(xt)q(\mathbf{x}_t): Diffused data dist.
  • q(x0,xt)q(\mathbf{x}_0,\mathbf{x}_t): Joint dist.
  • q(x0)q(\mathbf{x}_0): Input data dist.
  • q(xt∣x0)q(\mathbf{x}_t|\mathbf{x}_0): diffusion kernel

Forward 과정을 거쳐서 생성된 diffused data dist. q(xt)q(\mathbf{x}_t)는 Input data dist. q(x0)q(\mathbf{x}_0)와 diffusion kernel q(xt∣x0)q(\mathbf{x}_t|\mathbf{x}_0) 곱의 적분으로 구해집니다.

또한 input data에 diffusion kernel을 적용하여 만들어지는 분포가 std의 형태를 띄게 되는데, 이는 Forward의 step이 분포를 더 부드럽게 만들어주는 효과를 적용함을 파악할 수 있습니다. 따라서 diffusion kernel을 Gaussian convolution로 바라보는 것도 가능합니다.



Reverse denoising process

diffusion 파라미터들이 q(xt)≃N(xt;0,I)q(\mathbf{x}_t) \simeq \mathcal{N}(\mathbf{x}_t;0, \mathbf{I})가 되게끔 학습되었습니다. 따라서 denoising을 통해 image generation을 학습하기 위해서 xt\mathbf{x}_t를 샘플링한 후에 반복적으로 xt−1\mathbf{x}_{t-1}를 샘플링해야 합니다. 해당 과정을 통해 diffusion step 이전의 분포를 학습하여 결과적으로 input data 분포를 학습하는 것입니다.

  • first sample xt∼N(xt;0,I)\mathbf{x}_t \sim \mathcal{N}(\mathbf{x}_t;0, \mathbf{I})
  • Iteratively sample xt−1∼q(xt−1∣xt)\mathbf{x}_{t-1} \sim q(\mathbf{x}_{t-1}|\mathbf{x}_t) / q(xt−1∣xt)q(\mathbf{x}_{t-1}|\mathbf{x}_t): Denoising dist.

denoising 분포를 구하기 위해 베이즈 이론을 적용하면 q(xt−1∣xt)q(\mathbf{x}_{t-1}|\mathbf{x}_t)와 q(xt−1)q(xt∣xt−1)q(\mathbf{x}_{t-1})q(\mathbf{x}_{t}|\mathbf{x}_{t-1})과 비례한다는 사실을 알 수 있습니다. 하지만 q(xt−1)q(\mathbf{x}_{t-1})가 적분 불가능하기 때문에 해당 방법으로는 denoising dist.를 구할 수 없게 됩니다.

따라서 q(xt−1∣xt)q(\mathbf{x}_{t-1}|\mathbf{x}_t)를 직접적으로 구하지 않고, density network를 통해 근사하는 방법을 이용합니다.

p(xT)=N(xT;0,I)pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),σt2I)→pθ(x0:T)=p(xT)∏pθ(xt−1∣xt)p(\mathbf{x}_T) = \mathcal{N}(\mathbf{x}_T;0, \mathbf{I}) \\ p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \mu_\theta(\mathbf{x}_t, t), \sigma^2_t\mathbf{I}) \\ \rightarrow p_\theta(\mathbf{x}_{0:T}) = p(\mathbf{x}_{T})\prod p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)
  • μθ(xt,t)\mu_\theta(\mathbf{x}_t, t): trainable network (U-net)


Variational upper bound

이번에는 ddpm의 목적 함수에 대해서 알아보겠습니다. ddpm의 목적 함수는 vae에서와 같이 ELBO를 도입하여 해결합니다.

Eq(x0)[−log⁡pθ(x0)]≤Eq(x0)q(x1:T∣x0)[−log⁡pθ(x0:T)q(x1:T∣x0)]E_{q(\mathbf{x}_0)}[-\log p_{\theta}(\mathbf{x}_0)] \leq E_{q(\mathbf{x}_0)q(\mathbf{x}_{1:T}|\mathbf{x}_0)} [-\log \frac {p_{\theta}(\mathbf{x}_{0:T})} {q(\mathbf{x}_{1:T}|\mathbf{x}_0)} ]
  • q(x0)q(\mathbf{x}_0): training data dist.
  • pθ(x0)p_{\theta}(\mathbf{x}_0): model data dist. (직접적으로 구할 수 없음)
  • q(x1:T∣x0)q(\mathbf{x}_{1:T}|\mathbf{x}_0): samples from latent space

해당 Loss function을 아래와 같이 총 3개의 term으로 표현할 수 있습니다. 전개 과정은 Appendix에 남겨두고, 다음으로는 각 term이 의미하는 것에 대해서 살펴보겠습니다.

(1) LT=DKL(q(xT∣x0)∣∣p(xT))L_T = D_{KL}(q(\mathbf{x}_T|\mathbf{x}_0) || p(\mathbf{x}_T))

  • q(xT∣x0)=N(xT;αT‾x0,(1−αT‾)I)q(\mathbf{x}_T|\mathbf{x}_{0}) = \mathcal{N}(\mathbf{x}_{T};\sqrt{\overline{\alpha_T}}\mathbf{x}_{0}, (1-\overline{\alpha_T})\mathbf{I})
  • p(xT)=N(xT;0,I)p(\mathbf{x}_T) = \mathcal{N}(\mathbf{x}_T;0, \mathbf{I})

ddpm을 구현할 때 βt\beta_t는 상수로 고정하여 사용합니다. 따라서 LTL_T term 내에서 근사 분포 q에 대해서 update할 파라미터가 존재하지 않으므로, 학습 시 무시됩니다.

(2) L0=−log⁡pθ(x0∣x1)L_{0}= - \log {p_\theta(\mathbf{x}_{0}|\mathbf{x}_1)}

이미지가 (0~255) 범위 -> [-1, 1]로 scaling되는 과정을 고려하여 L0L_0에 대한 object function을 별도로 구성하였습니다.

  • 마지막 time에 대한 reverse kernel: N(x0;μθ(x1,1),σ12I)\mathcal{N}(\mathbf{x}_0;\mu_\theta(\mathbf{x}_1,1),\sigma^2_1\mathbf{I})

(3) Lt−1=∑t>1DKL(q(xt−1∣xT,x0)∣∣pθ(xt−1∣xt))L_{t-1}= \sum_{t>1} D_{KL}(q(\mathbf{x}_{t-1}|\mathbf{x}_T, \mathbf{x}_0) || p_\theta (\mathbf{x}_{t-1}|\mathbf{x}_t))

  • q(xt−1∣xT,x0)=N(xt−1;μ~t(xt,x0),β~tI)q(\mathbf{x}_{t-1}|\mathbf{x}_T, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \tilde\mu_t(\mathbf{x}_t, \mathbf{x}_0), \tilde\beta_t\mathbf{I})
    • μ~t(xt,x0):=α‾t−1βt1−α‾tx0+αt(1−α‾t−1)1−α‾txt\tilde\mu_t(\mathbf{x}_t, \mathbf{x}_0) := \frac {\sqrt {\overline{\alpha}_{t-1}} \beta_t} {1 - \overline{\alpha}_{t}} \mathbf{x}_0 + \frac {\sqrt\alpha_{t}(1 - \overline{\alpha}_{t-1})} {1 - \overline{\alpha}_{t}} \mathbf{x}_t: weighted sum
    • β~t:=1−α‾t−11−α‾tβt\tilde\beta_t := \frac {1 - \overline{\alpha}_{t-1}} {1 - \overline{\alpha}_{t}} \beta_t
  • pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),σt2I)p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \mu_\theta(\mathbf{x}_t, t), \sigma^2_t\mathbf{I})

Loss 내에서 가장 중요한 denoising process term입니다. 1 < t ≤ T에 대해 두 정규분포인 q(xt−1∣xT,x0)q(\mathbf{x}_{t-1}|\mathbf{x}_T, \mathbf{x}_0)와 pθ(xt−1∣xt)p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) 사이의 KL divergence를 구해야 합니다. 이때 두 분포의 variance인 β~t\tilde\beta_t와 σt2\sigma^2_t를 같다고 두어도 비슷한 결과를 나타낸다고 합니다. 따라서 KL term을 두 분포의 평균 사이 l2 loss로 표현합니다. (수식 전개 과정은 Appendix에 남겨두었습니다.)

Lt−1=Eq[12σt2∣∣μ~t(xt,x0)−μθ(xt,t)∣∣2]+CL_{t-1} = \mathbb{E}_q[\frac 1 {2\sigma^2_t} ||\tilde\mu_t(\mathbf{x}_t, \mathbf{x}_0) - \mu_\theta(\mathbf{x}_t, t)||^2] + C

이때 reparameterization(xt=αt‾x0+(1−αt‾)ϵ\mathbf{x}_t = \sqrt{\overline{\alpha_t}}\mathbf{x}_{0} + \sqrt{(1-\overline{\alpha_t})}\epsilon)을 적용하면 μ~t(xt,x0)=11−βt(xT−βt1−αt‾ϵ)\tilde\mu_t(\mathbf{x}_t, \mathbf{x}_0) = \frac {1} {\sqrt{1- \beta_t}} (\mathbf{x}_T - \frac {\beta_t} {\sqrt{1 - \overline{\alpha_t}}}\epsilon)로 나타낼 수 있습니다. (수식 전개 과정은 Appendix에 남겨두었습니다.) 따라서 μθ(xt,t)\mu_\theta(\mathbf{x}_t, t)가 11−βt(xT−βt1−αt‾ϵ)\frac {1} {\sqrt{1- \beta_t}} (\mathbf{x}_T - \frac {\beta_t} {\sqrt{1 - \overline{\alpha_t}}}\epsilon)를 예측하도록 학습됩니다.

Lt−1=Ex0,ϵ[βt22σt2αt(1−αt‾)∣∣ϵ−ϵθ(αt‾x0+1−αt‾ϵ,t)∣∣2]L_{t-1} = \mathbb{E}_{\mathbf{x}_0, \epsilon}[\frac {\beta^2_t} {2\sigma^2_t\alpha_t(1 - \overline{\alpha_t})} ||\epsilon - \epsilon_\theta(\sqrt{\overline{\alpha_t}}\mathbf{x}_0 + \sqrt{1 - \overline{\alpha_t}}\epsilon, t)||^2]

각각의 식을 Loss에 대입한 식입니다. λt\lambda_t는 time-dependent한 성질을 띄며, t와 반비례합니다. 하지만 λt\lambda_t를 1로 두었을 때에도 같은 quality를 나타내기 때문에 Loss에선 생략가능합니다.

Lsimple(θ):=Et,x0,ϵ[∣∣ϵ−ϵθ(αt‾x0+1−αt‾ϵ,t)∣∣2]L_{simple}(\theta) := \mathbb{E}_{t, \mathbf{x}_0, \epsilon}[||\epsilon - \epsilon_\theta(\sqrt{\overline{\alpha_t}}\mathbf{x}_0 + \sqrt{1 - \overline{\alpha_t}}\epsilon, t)||^2]

최종적으로 구한 Loss 식은 다음과 같습니다! t=1일 때 L0L_0의 역할을 하며, t>1일 때 Lt−1L_{t-1}의 역할을 합니다.



Implementation

Algorithm

Training과 Sampling 알고리즘입니다. Training은 x0x_0, tt, ϵ\epsilon을 샘플링 한 후에, gradient step을 밟는 구조로 이루어집니다. Sampling은 xTx_T를 diffused dist.에서 샘플링 하여 denoising model을 통해 reparameterization 과정을 거쳐 x0x_0를 얻어내는 과정입니다.

Denoising model

Denoising model로는 U-net을 사용하였으며, Residual block과 self-attention layers를 추가하였습니다. U-net 모델을 통해 noise ϵθ(xt,t)\epsilon_\theta(\mathbf{x}_t, t)를 예측하였습니다. Time representation으로는 sinusoidal positional embedding 혹은 random Fourier features를 사용하였습니다.



Appendix

Appendix에서는 주로 ddpm의 주요한 식 전개 과정을 올리려고 합니다. 수식 전개에 대해 말로 자세히 설명하는 것보단, 제가 정리한 수식 전개 필기를 올리도록 하겠습니다!

ELBO 표현 유도 과정


VAE에서의 ELBO 유도 과정과 동일합니다. log likelihood를 최대화할 때 직접적으로 z에 대한 적분이 불가능하므로, ELBO를 통해 간접적으로 최대화하는 방법입니다.

Variational upper bound 전개 과정

Variational upper bound를 LTL_T, Lt−1L_{t-1}, L0L_0으로 나누는 과정입니다. 베이즈 이론을 통해 Lt−1L_{t-1} term이 유도되는 것이 중요한 점이었던 것 같습니다.

Gaussian KL divergence → l2 loss in Lt−1L_{t-1}

Gaussian KL의 증명과 이를 통해 Lt−1L_{t-1} term이 l2 loss로 표현되는 과정을 적어보았습니다.

μ~t(xt,x0)\tilde\mu_t(\mathbf{x}_t, \mathbf{x}_0) with reparameterization

posterior distribution q(xt−1∣xT,x0)q(\mathbf{x}_{t-1}|\mathbf{x}_T, \mathbf{x}_0)의 mean을 구하는 과정입니다. 이 과정에선 Gaussian function의 성질과 reparameterization trick 등이 사용되었습니다.

  • q(xt−1∣xT,x0)=N(xt−1;μ~t(xt,x0),β~tI)q(\mathbf{x}_{t-1}|\mathbf{x}_T, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \tilde\mu_t(\mathbf{x}_t, \mathbf{x}_0), \tilde\beta_t\mathbf{I})
    • μ~t(xt,x0):=α‾t−1βt1−α‾tx0+αt(1−α‾t−1)1−α‾txt\tilde\mu_t(\mathbf{x}_t, \mathbf{x}_0) := \frac {\sqrt {\overline{\alpha}_{t-1}} \beta_t} {1 - \overline{\alpha}_{t}} \mathbf{x}_0 + \frac {\sqrt\alpha_{t}(1 - \overline{\alpha}_{t-1})} {1 - \overline{\alpha}_{t}} \mathbf{x}_t: weighted sum
    • β~t:=1−α‾t−11−α‾tβt\tilde\beta_t := \frac {1 - \overline{\alpha}_{t-1}} {1 - \overline{\alpha}_{t}} \beta_t



Reference

[1] Jonathan Ho, Ajay Jain, Pieter Abbeel (2020), “Denoising Diffusion Probabilistic Models”, https://arxiv.org/pdf/2006.11239.pdf

[2] Arash Vahdat (in CVPR 2022), “Tutorial on Denoising Diffusion-based Generative Modeling: Foundations and Applications”, https://youtu.be/cS6JQpEY9cs?si=v0fybFzP4ajXCczB

[3] Lilian Weng, What are Diffusion Models?, https://lilianweng.github.io/posts/2021-07-11-diffusion-models/

[4] Wikipedia, Gaussian function, https://en.wikipedia.org/wiki/Gaussian_function

profile
김민솔
Interested in Vision, Generative, Neural Rendering
다음 포스트

Score-based Generative modeling with DEs

0개의 댓글