[NeurIPS 2021] 1편: Generative model - Diffusion model Review

‘NeurIPS 2021(Conference and Workshop on Neural Information Processing Systems 2021, 신경정보처리시스템학회)’은 1987년에 처음 시작된 학회로, AI와 머신러닝 분야의 가장 권위 있는 학회 중 하나입니다. 약 20%의 논문 승인율을 바탕으로 매우 높은 수준을 유지하고 있으며, NeurIPS에서 산출되는 연구 결과가 갖는 기술적 영향도가 높아 최신 연구 트렌드를 살펴볼 수 있습니다. LG AI연구원 Applied AI Research lab 연구원들이 지난 2021년 12월 온라인으로 개최되었던 ‘NeurIPS 2021’에서 소개된 의미 있는 논문을 소개합니다. - 1편: Generative model - Diffusion model Review - Applied AI Research lab 정혜민님 - 2편: Variational Deep Embedding 기반의 Deep Clustering Review - Applied AI Research lab 양진석님 - 3편: Neural Architecture Search Review - Applied AI Research lab 안원빈님

Introduction

LG AI연구원의 Applied AI Research Lab에서는 Generative model(생성 모델)에 대한 연구를 활발히 하고 있습니다. Generative model이란 Deep Learning(딥 러닝) 모델을 두 가지 큰 범주로 나누었을 때 Discriminative model(분별 모델)과 구분되는 개념입니다. 딥 러닝의 두 갈래에서도 LG AI연구원은 시계열 데이터에서의 Anomaly detection(이상 탐지), Representation learning(표현 학습), 음악 생성, 추천 등을 위해서 생성 모델 연구에 힘쓰고 있습니다.

이번 글에서는 다양한 종류의 Generative model이 갖는 특징들에 대해서 알아보고, Generative model 중 최근 활발히 연구되고 있는 Diffusion model(확산모델)에 대해 말씀드리려고 합니다. 아래 이어지는 내용은 Diffusion model의 대표적인 논문과 이를 기반으로2021 NeurIPS(신경정보처리시스템학회)에서 발표한 Variational diffusion model(변분 확산 모델)를 소개합니다.

Generative model

Generative model 이란?
Generative model은 주어진 데이터 x가 샘플링 된 분포(distribution)를 평가(estimate)하는 모델입니다. 딥러닝 모델이 distribution을 estimate 한다는 의미는, 모델의 출력값(output)이 distribution를 결정 짓는 변수(parameter)이거나, 주어진 데이터가 나올 확률, 즉 marginal likelihood인 것입니다. Generative model을 이용해서 주어진 test data의 likelihood 값을 계산할 수 있고, 새로운 data를 생성할 수 있습니다. 정리하면 학습한 data의 distribution을 따르는 새로운 data를 만들어내는 모델을 의미합니다.

Generative model의 성능은 주어진 데이터 x의 likelihood p(x) 값을 얼마나 잘 예측하는지, 생성된 데이터 x의 quality가 얼마나 우수한지로 판단할 수 있습니다. Generative model은 크게 6개로 분류되며 각 모델의 장단점은 다음과 같습니다.

Generative model 비교
Generative model들을 구분하면 크게 6개로 분류됩니다. 각각의 모델에 대해서 특징과 장단점은 아래와 같습니다.

Auto-regressive models(ARMs)
순서를 가지는 변수(variable)들의 조건부 확률(conditional probability)의 곱으로 데이터의 likelihood를 계산하는 모델입니다. 주어진 데이터의 likelihood를 계산할 수 있는 장점은 있지만 sampling이 느리고, 데이터에 내재된 잠재형상(latent feature)은 학습할 수 없는 단점이 있습니다.

Variational Auto encoders(VAEs)
잠재변수(Latent variable)기반의 generative model로 데이터 x와 latent variable z의 결합확률분포(joint distribution)를 구해서 x에 대해서 주변화(marginalize)하는 모델입니다. 빠른 학습과 sampling, 데이터의 latent feature 학습이 가능한 장점과, likelihood를 다루기 힘들고 사전확률(prior distribution)이 한정적이라는 단점이 있습니다.

Energy Based Models(EBMs)
에너지함수(Energy function)를 이용해서 distribution을 estimate합니다. 제약(Constraint)이 없어 간편하다는 장점이 있으나, likelihood와 sampling이 다루기 힘들다는 단점이 있습니다.

Generative Adversarial Networks(GANs)
Discriminator와 Generator를 서로 adversarial 방향으로 학습시켜서 데이터를 생성하는 모델입니다. Sample quality가 좋고, fast training과 fast sampling이 가능한 장점이, likelihood가 define되지 않고 unstable training이 단점이 있습니다.

Normalizing Flows
Simple한 base 분포 p(z)에서 복잡한 데이터 분포 p(x)로 가는 역사상(invertible mapping)함수를 이용해 distribution을 estimate하는 모델입니다. Exact likelihood 값을 계산할 수 있고, fast sampling이 가능하나, architecture에서 제한 조건이 부여된다는 단점이 있습니다.

Diffusion
데이터 x에서 점점 noise를 추가해서 noise data로 만들고, noise data에서 데이터 x로 돌아오는 과정을 학습해 distribution을 estimate하는 모델입니다. Likelihood는 다루기 힘들지만 좋은 likelihood 값을 가지고 특별한 제한 조건 없이 complex distribution을 잘 estimate하는 특징이 있습니다.

위의 다양한 Generative model 중 Diffusion model이 최근 활발하게 연구되고 있습니다. 2021 Neurips에서는 Autoregressive based Generative model 보다 더 좋은 likelihood 성능을 내고, GAN based model 보다 더 좋은 quality의 sample을 생성하는 Diffusion model들이 발표됐습니다.

Diffusion model이란?

Diffusion probabilistic models
Diffusion model을 제일 처음 제안한 논문은 ‘Deep Unsupervised Learning using Nonequilibrium Thermodynamics’(2015)^[1] 입니다. Diffusion model은 열역학에서 아이디어를 가져온 모델로 크게 2단계로 나뉩니다. 먼저 주어진 데이터 x에 점점 noise를 추가합니다. 이 과정을 Diffusion process라고 합니다. 그 다음 앞에서 정의한 Diffusion process를 reverse한 process를 계산합니다. 이 과정은 noise data에서 점점 noise를 제거하는 과정입니다. 딥러닝 모델을 이용해 이 reverse diffusion process를 학습하는 것으로 이 과정을 Reverse diffusion process라고 합니다.

1.Diffusion process

Diffusion process는 주어진 데이터 x에 time step마다 noise를 추가하는 것입니다. 0 step에서 noise를 추가하면서 T step으로 가고 각 step에서 noise가 추가된 것이 latent variable z_t입니다. 이런 Diffusion process를 수학적으로 표현하면 아래와 같은 Gaussian Distribution이 됩니다.

$q(z_t\overline{)z_{t-1}) = N(\sqrt{1-{\beta }_t}{z}_{t-1}},{\beta }_{t}I)$

위의 식의 의미는 이전 step의 latent variable $z_{t-1}$에서 $\sqrt{1-{\beta }_t}$만큼의 signal을 가져오고, ${\beta }_t$만큼 noise를 추가해서 $z_t$를 계산하는 것입니다. 여기서 ${\beta }_t$는 noise 양으로 직접 지정하는 값입니다. Time step이 커질수록 noise를 더 큰 값으로 정합니다. ${\beta }_0<{\beta }_1<\cdots <{\beta }_T$

T step까지 noise를 다 추가하면 최종적으로 z_T는 Isotropic Gaussian Distribution이 될 것입니다.

Diffusion process에는 Markov 성질을 가정합니다. 따라서 아래와 같이 Distribution의 곱으로 데이터 x가 주어졌을 때 latent variable들이 나올 분포를 표현할 수 있습니다.

$q\left(z_{1:t}\overline{)z_0}\right)=\prod _{t=1}^{T}q\left(z_t\overline{)z_{t-1}}\right)$

데이터 $\mathrm{x}\left(=z_0\right)$가 주어졌을 때의 latent variable $z_t$의 분포도 Gaussian distribution입니다.

$q\left(z_t\overline{)z_0}\right)=N\left(z_t;\sqrt{\overline{{\alpha }_t}}{z}_0,\left(1-{\overline{\alpha }}_t\right)I\right), {\alpha }_t=1-{\beta }_t {\overline{\alpha }}_t=\prod _{i=1}^T{\alpha }_i$

위의 식의 의미는 $\mathrm{x}\left(=z_0\right)$를 기준으로 두고 특정 noise를 추가한 것을 의미합니다.

위에서 정의된 Gaussian distribution들을 이용해서 아래와 같이 $z_t, \mathrm{x}\left(=z_0\right)$가 주어졌을 때 $z_{t-1}$가 나올 분포도 Gaussian distribution임을 도출할 수 있습니다.

$q\left(z_{t-1}\overline{)z_t,z_0}\right)=N\left(z_{t-1};{\tilde{\mu }}_t\left(z_t,z_0\right),{\tilde{\beta }}_{t}I\right)$
$\tilde{\mu }\left(z_t,z_0\right)=\frac{1}{\sqrt{{\alpha }_t}}\left(z_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\in }_t\right), {\tilde{\beta }}_t\left(z_t,z_0\right)=\frac{1-{\overline{\alpha }}_{t-1}}{1-{\alpha }_t}{\beta }_t$

${\in }_t~N\left(0,I\right)\mathrm{는} z_t~\mathrm{q}\left(z_t\overline{)z_0}\right)=N\left(z_t;\sqrt{{\overline{\alpha }}_t}{z}_0,\left(1-{\overline{\alpha }}_t\right)I\right)$를 reparameterization trick을 이용해서 계산할 때 사용되는 값입니다. 식의 의미는, 이전 time step의 latent variable이 나올 분포를 의미하며 나중에 Reverse diffusion process를 모델이 학습할 때 정답으로 사용되는 중요한 식입니다.

정리하면 Diffusion process는 지정한 noise ratio 값 ${\beta }_t$을 time step마다 추가하는 것이고 관련된 분포를 모두 가우스분포(Gaussian distribution)로 표현할 수 있습니다. 여기서는 따로 학습이 진행되지 않습니다.

2. Reverse Diffusion process

앞의 diffusion process 반대 방향의 process는 noise data에서 원래의 데이터 x로 돌아오는 과정입니다. 즉 generative model입니다. Diffusion process는 직접 Gaussian distribution으로 정의했지만 반대 방향인 Reverse Diffusion process를 표현하는 distribution들은 모르는 상태입니다. 모델을 통해 모르는 Reverse Diffusion process를 학습하는 것입니다.

Reverse diffusion process를 수식으로 표현하면 아래와 같습니다.

$p_{\theta }\left(z_{0:T}\right)=p\left(z_T\right)\prod _{t=1}^T{p}_{\theta }\left(z_{t-1}\overline{)z_t}\right)$

$p_{\theta }\left(z_{t-1}\overline{)z_t}\right)$는 반대 방향 분포를 모델이 estimate한 것입니다. 이 분포도 Gaussian distribution입니다. 모델이 Gaussian distribution의 parameter인 평균을 학습하는 것입니다. 수식은 아래와 같습니다.

$p_{\theta }\left(z_{t-1}\overline{)z_t}\right)=N\left(z_{t-1};{\mu }_{\theta }\left(z_t,t\right),{\Sigma }_{\theta }\left(z_t,t\right)\right)$

3. Loss
Reverse diffusion process, 즉 generative model을 학습하기 위한 loss식에 대해서 알아보겠습니다. Generative model의 목적은 Distribution estimate이고 이는 주어진 데이터 x의 likelihood를 잘 계산하는 것으로 판단할 수 있습니다. 따라서 목적식은 log likelihood of x를 잘 계산하는 것입니다. 이 목적식으로부터 loss 식을 도출할 수 있습니다. 자세한 도출 과정은 생략하겠습니다.


Loss 식을 Variational Lower bound(VLB)라고 하며 구성하는 각 term의 의미는 아래와 같습니다.

$D_{KL}\left(q\left(z_T\overline{)x}\right)\parallel p_{\theta }\left(z_T\right)\right)$ term은 두 Gaussian distribution을 가까이하는 것이고, $q\left(z_T\overline{)x}\right),p_{\theta }\left(z_T\right)$ 모두 Gaussian distribution이므로 closed form으로 계산 가능합니다.

$log\left(p_{\theta }\left(x\overline{)z_1}\right)\right)$ term은 latent variable $z_1$에서 원본 x를 reconstruct하는 것이고, 이 term도 직접 계산할 수 있습니다. 위의 두 term은 모델의 학습과는 무관하게 계산되는 term입니다.

$\sum _{t=1}^T{D}_{KL}\left(q\left(z_{t-1}\overline{)z_t,x}\right)\parallel p_{\theta }\left(z_{t-1}\overline{)z_t}\right)\right)$term이 Reverse diffusion process를 학습하게 하는 것이고, diffusion loss라고 합니다. $\left(\mathrm{x}=z_0\right)$

$q\left(z_{t-1}\overline{)z_t,z_0}\right)=N\left(z_{t-1};{\tilde{\mu }}_t\left(z_t,z_0\right),{\tilde{\beta }}_{t}I\right)$
${\tilde{\mu }}_t\left(z_t,z_0\right)=\frac{1}{\sqrt{{\alpha }_t}}\left(z_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\in }_t\right), {\tilde{\beta }}_t\left(z_t,z_0\right)=\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}{\beta }_t$

이 정답 분포이고 모델 분포 $p_{\theta }\left(z_{t-1}\right|z_t)=N(z_{t-1};{\mu }_{\theta }(z_t,t),{\Sigma }_{\theta }(z_t,t))$가 $q\left(z_{t-1}\overline{)z_t,z_0}\right)$ 분포에 가까워지도록 학습되는 것입니다. 두 분포 모두 Gaussian distribution 이므로 KL divergence를 아래처럼 계산합니다. 

$D_{KL}\left(q\left(z_{t-1}{\overline{)z}}_t,x\right)\parallel p_{\theta }\left(z_{t-1}\overline{)z_t}\right)\right)=E_{x,z}\left[\frac{1}{2{||{\Sigma }_{\theta }\left(z_t,t\right)||}_2^2}{||{\tilde{\mu }}_t\left(z_t,x\right)-{\mu }_{\theta }\left(z_t,t\right)||}^2\right]$

지금까지 Diffusion model의 개념과 loss식에 대해서 간략히 살펴보았습니다. 위 내용에 기반한 다양한 Diffusion model들이 제시되고 있고 2021 Neurips에서도 여러 편의 관련 논문들이 발표되었습니다. 이제 그 중 Variational diffusion model에 대해서 설명하겠습니다.

Variational diffusion model^[2]

Motivation
생성된 image quality 기준으로 좋은 성능을 내는 Diffusion model들이 많이 발표되었습니다. 하지만 likelihood 기준으로는 Diffusion model이 Autoregressive 기반 Generative model에 비해 낮은 결과를 보였습니다. 이 논문에서는 이런 한계점을 극복해서 likelihood 기준으로 SOTA 성능을 보여준 diffusion model을 제안합니다.


Contribution
이 논문의 contribution은 아래와 같습니다.

- Likelihood의 성능을 높이기 위해서 Fourier feature들을 사용할 것을 제안함
- Diffusion process에서 추가되는 noise schedule을 모델의 학습을 통해서 정함
- 간단한 형태의 Variational lower bound를 제안함
- 이론적으로 Diffusion process를 구성하는 steps들의 수가 많을수록 성능이 좋아지는 것을 증명함
- Time steps들이 무수히 많은 경우 (Continuous time steps)에서의 diffusion loss를 도출함
- Continuous time steps에서의 loss식과 기존의 발표된 diffusion model의 연결관계를 이론적으로 보임

Forward time Diffusion process


이 논문에서 제안하는 Diffusion process를 구성하는 첫번째 distribution은 아래와 같습니다.

$q\left(z_t\overline{)x}\right)=N\left({\alpha }_{t}x,{\sigma }_t^{2}I\right)$

위 식의 의미는 데이터 x가 주어졌을 때 latent variable $z_t$가 나올 distribution이 Gaussian 임을 뜻합니다. 원래의 데이터 x에서 ${\alpha }_t$만큼의 정보(signal)를 사용하고, ${\sigma }_t^2$의 noise를 추가해서 $z_t$가 계산되는 것입니다. ${\alpha }_t$를 signal, ${\sigma }_t^2$를 noise라고 합니다. ${\alpha }_t$과 ${\sigma }_t^2$의 관계는 아래와 같습니다.

${\alpha }_t=\sqrt{1-{\sigma }_t^2}$

논문에서는 새롭게 Signal to ratio at time t를 뜻하는 SNR(t)를 제안합니다.

$\mathrm{SNR}\left(\mathrm{t}\right)=\frac{{\alpha }_t^2}{{\sigma }_t^2}$

SNR(t) 값이 큰 것은 signal이 noise에 비해 많다는 것이고 SNR(t) 값이 작다는 것은 noise가 signal에 비해 많다는 것입니다. 따라서 SNR(t) 값은 t가 커질수록 즉 diffusion process가 진행될수록 점점 감소됩니다.


논문에서는 추가적으로 noise schedule을 미리 정하는 것이 아니라 모델로 학습할 것을 제안합니다.

${\sigma }_t^2=sigmoid\left({\gamma }_{\eta }\left(t\right)\right)$

${\gamma }_{\eta }\left(t\right)$는 monotonic neural network로 입력 값인 t가 커질수록 높은 output을 뱉는 모델입니다. Diffusion process는 데이터를 점점 noisy하게 해야 하므로 noise 값인 ${\sigma }_t^2$는 커져야 하니까 monotonic neural network로 정한 것입니다. 논문에서는 3개의 linear layer를 ${\gamma }_{\eta }\left(t\right)$로 사용했습니다.

따라서 ${\alpha }_t,{\sigma }_t^2,\mathrm{SNR}\left(\mathrm{t}\right)$를 아래와 같이 표현 가능합니다.

${\sigma }_t^2=sigmoid\left({\gamma }_{\eta }\left(t\right)\right)=\frac{\mathrm{SNR}\left(\mathrm{t}\right)}{1+\mathrm{SNR}\left(\mathrm{t}\right)}$
${\alpha }_t^2=sigmoid\left(-{\gamma }_{\eta }\left(t\right)\right)=\frac{1}{1+\mathrm{SNR}\left(\mathrm{t}\right)}$
$\mathrm{SNR}\left(\mathrm{t}\right)=\frac{sigmoid\left(-{\gamma }_{\eta }\left(t\right)\right)}{sigmoid\left({\gamma }_{\eta }\left(t\right)\right)}=exp\left(-{\gamma }_{\eta }\left(t\right)\right)$

(SNR(t)은 invertible monotonically decrease function)

Diffusion process를 구성하는 다른 distribution은 아래와 같습니다. (0<s<t<1)

$q\left(z_t\overline{)z_s}\right)=N\left({\alpha }_{t\overline{)s}}{z}_s,{\sigma }_{t\overline{)s}}^{2}I\right)$
${\alpha }_{t\overline{)s}}=\frac{{\alpha }_t}{{\alpha }_s}, {\sigma }_{t\overline{)s}}^2={\sigma }_t^2-{\alpha }_{t\overline{)s}}^2{\sigma }_s^2$

위 식의 의미는 이전 step의 latent variable이 주어졌을 때 다음 latent variable이 나올 분포가 Gaussian이고, $z_t$는 $z_s$에서 일정 정보를 취하고 noise가 추가됨을 의미합니다. ${\alpha }_{t\overline{)s}}, {\sigma }_{t\overline{)s}}^2$이 위와 같이 정해지는 과정도 도출 가능하지만 여기서는 생략했습니다.

위에서 정의된 Gaussian distribution들을 이용해서 아래와 같이 $z_t, x$가 주어졌을 때 과거 step인 $z_s$가 나올 분포도 아래와 같은 Gaussian distribution로 도출할 수 있습니다.

$q\left(z_s\overline{)z_t,x}\right)=N\left(z_s;{\mu }_Q\left(z_t,x;s,t\right),{\sigma }_Q^2\left(s,t\right)I\right)$
${\mu }_Q\left(z_t,x;s,t\right)=\frac{{\alpha }_{t\overline{)s}}{\sigma }_s^2}{{\sigma }_t^2}{z}_t+\frac{{\alpha }_s{\sigma }_{t\overline{)s}}^2}{{\sigma }_t^2}x, {\sigma }_Q^2\left(s,t\right)={\sigma }_{t\overline{)s}}^2\frac{{\sigma }_s^2}{{\sigma }_t^2}$

Reverse time Generative model


Diffusion process의 반대 방향 process에서는 T=1 → T=0 방향으로 latent variables들을 계산하면서 원본 데이터 x를 generate 합니다. 이 과정을 Hierarchical generative model로 볼 수 있습니다. T=1에서 T=0 사이에 T개의 discrete timesteps이 있는 경우인 Discrete time case 와 T=1에서 T=0 사이에 무수히 많은 timesteps T → ∞ 경우인 Continuous time case가 있습니다. 먼저 Discrete time case에 대해서 설명하겠습니다.

Discrete time case
0과 1사이를 uniformly discretize 해서 T개의 time step으로 구성한 것입니다.


이 경우 generative model for x를 아래의 식으로 쓸 수 있습니다.

$p\left(x\right)={\int }_{z}p\left(z_1\right)p\left(x\overline{)z_0}\right)\prod _{i=1}^{T}p\left(z_{s\left(i\right)}\overline{)z_{t\left(i\right)}}\right)$

위의 식은 noise인 $z_1$에서 시작해서, $z_{t\left(i\right)}$가 주어졌을 때 $z_{s\left(i\right)}$가 나오고, 마지막으로 $z_0$에서 원본 데이터 x로 가는 것을 표현한 것입니다. 구성하는 각 term의 의미는 아래와 같습니다.

$q\left(z_1\overline{)x}\right)=N\left(z_1;0,I\right)$이니까, $p\left(z_1\right)=N\left(z_1;0,I\right)$도 Gaussian으로 정의합니다. $q\left(x\overline{)z_0}\right)=\frac{q\left(z_0\overline{)x}\right)q\left(x\right)}{q\left(z_0\right)}$에서, SNR(0)에서는 noise 값이 작아서 $z_0$와 x가 비슷하므로 $q\left(z_0\overline{)x}\right)$ term의 값이 커서 $z\left(x\overline{)z_0}\right)$ 값이 $q\left(z_0\overline{)x}\right)$ 값에 dominate 됩니다. $q\left(z_0\overline{)x}\right)=N\left({\alpha }_{0}x,{\sigma }_0^{2}I\right)$이고, x에서 z₀를 계산할 때, elementwise scale/shift 하므로 다른 dimension에 영향을 안 받습니다. 따라서 $p\left(x\overline{)z_0}\right)=\prod _{i}p\left(x_i\overline{)z_{0,i}}\right)$로 정의합니다.

$p\left(z_{s\left(i\right)}\overline{)z_{t\left(i\right)}}\right)$는 반대 방향의 분포이고 모델이 학습을 통해 배워야 하는 것입니다. 즉 앞에서 정의한 $q\left(z_s\overline{)z_t,x}\right)=N\left(z_s;{\mu }_Q\left(z_t,x;s,t\right),{\sigma }_Q^2\left(s,t\right)I\right)$,  ${\mu }_Q\left(z_t,x;s,t\right)=\frac{{\alpha }_{t\overline{)s}}{\sigma }_s^2}{{\sigma }_t^2}{z}_t+\frac{{\alpha }_s{\sigma }_{t\overline{)s}}^2}{{\sigma }_t^2}x$, ${\sigma }_Q^2\left(s,t\right)={\sigma }_{t\overline{)s}}^2\frac{{\sigma }_s^2}{{\sigma }_t^2}$ 분포에 가까워지도록 학습되어야 합니다. $p\left(z_{s\left(i\right)}\overline{)z_{t\left(i\right)}}\right)$ 분포는 두 관점으로 도출 가능합니다. 첫번째 관점은 denoising model 관점입니다. (s(i)는 s, t(i)는 t로 표현)

$p\left(z_s\right|z_t)=N(z_s;{\mu }_{\theta }(z_t;s,t),{\sigma }_Q^2(s,t)I)$
${\mu }_{\theta }(z_t;s,t)=\frac{{\alpha }_{t|s}{\sigma }_s^2}{{\sigma }_t^2}{z}_t+\frac{{\alpha }_s{\sigma }_{t|s}^2}{{\sigma }_t^2}{\hat{\mathbf{x}}}_{\mathbf{\theta }}\mathbf{(}{\mathit{z}}_{\mathbf{t}}\mathbf{;}\mathit{t}\mathbf{)}, {\sigma }_Q^2(s,t)={\sigma }_{t|s}^2\frac{{\sigma }_s^2}{{\sigma }_t^2}$

${\hat{\mathrm{x}}}_{\mathrm{\theta }}(z_{\mathrm{t}};t)$는 denoising model로 $z_t$가 입력되면 원본 x를 predicts 합니다.

두번째 관점은 noise prediction model 관점입니다. $q(z_t│x)=N({\alpha }_{t}x,{\sigma }_t^{2}I)$에서 $z_t$를 reparameterization trick으로 계산하면, $z_t={\alpha }_{t}x+{\sigma }_t\odot ϵ$, $ϵ~N(0,I)$이고 이 식을 x기준으로 풀어 쓰면 $x=\frac{z_t-{\sigma }_t\odot \in }{{\alpha }_t}$이므로 ${\hat{x}}_{\theta }\left(z_t;t\right)$를 아래와 같이 표현할 수 있습니다.

${\hat{x}}_{\theta }\left(z_t;t\right)=\frac{z_t-{\sigma }_t\odot {\hat{\in }}_{\theta }(z_t;t)}{{\alpha }_t}$

${\hat{\in }}_{\theta }\left(z_t;t\right)$는 noise prediction model로 $z_t$가 입력되면 $\in$를 predicts 합니다. 이 관점에서 $p\left(z_s\overline{)z_t}\right)$는 아래와 같습니다.

$p\left(z_s\overline{)z_t}\right)=N\left(z_s;{\mu }_{\theta }\left(z_t;s,t\right),{\sigma }_Q^2\left(s,t\right)I\right)$
${\mu }_{\theta }\left(z_t;s,t\right)=\frac{1}{{\alpha }_{t\overline{)s}}}{z}_t-\frac{{\sigma }_{t\overline{)s}}^2}{{\alpha }_{t\overline{)s}}{\sigma }_t}{\hat{\in }}_{\theta }\left(z_t;t\right), {\sigma }_Q^2\left(s,t\right)={\sigma }_{t\overline{)s}}^2\frac{{\sigma }_s^2}{{\sigma }_t^2}$

직관적으로 이해하기 쉬운 것은 Denoising model 관점이지만 실험을 할 때는 Noise prediction 관점을 적용했습니다. 따라서 Generative model을 정리하면 아래와 같습니다.

$p\left(x\right)={\int }_{z}p\left(z_1\right)p\left(x\overline{)z_0}\right)\prod _{i=1}^{T}p\left(z_{s\left(i\right)}\overline{)z_{t\left(i\right)}}\right)$
$p\left(z_1\right)=N\left(z_1;0,I\right)$
$p\left(x\overline{)z_0}\right)=\prod _{i}p\left(x_i\overline{)z_{0,i}}\right)$
$p\left(z_s\overline{)z_t}\right)=N\left(z_s;{\mu }_{\theta }\left(z_t;s,t\right),{\sigma }_Q^2\left(s,t\right)I\right), {\mu }_{\theta }\left(z_t;s,t\right)=\frac{{\alpha }_{t\overline{)s}}{\sigma }_s^2}{{\sigma }_t^2}{z}_t+\frac{{\alpha }_s{\sigma }_{t\overline{)s}}^2}{{\sigma }_t^2}{\hat{x}}_{\theta }\left(z_t;t\right), {\sigma }_Q^2\left(s,t\right)={\sigma }_{t\overline{)s}}^2\frac{{\sigma }_s^2}{{\sigma }_t^2}$

또는

$p\left(z_s\overline{)z_t}\right)=N\left(z_s;{\mu }_{\theta }\left(z_t;s,t\right),{\sigma }_Q^2\left(s,t\right)I\right), {\mu }_{\theta }\left(z_t;s,t\right)=\frac{1}{{\alpha }_{t\overline{)s}}}{z}_t-\frac{{\sigma }_{t\overline{)s}}^2}{{\alpha }_{t\overline{)s}}{\sigma }_t}\hat{\in }\left(z_t;t\right), {\sigma }_Q^2\left(s,t\right)={\sigma }_{t\overline{)s}}^2\frac{{\sigma }_s^2}{{\sigma }_t^2}$

Loss 식
위의 Generative model을 학습하기 위한 Loss식은 아래와 같습니다.


Prior loss, Reconstruction loss는 Reparameterization trick을 이용해서 계산합니다. Diffusion loss 부분이 모델이 학습을 통해서 줄이는 term입니다. Discrete time case에서 Diffusion loss식은 아래와 같습니다.


즉 반대방향의 정답 분포와 모델 분포사이의 거리를 줄이는 것입니다. 위의 식을 앞에서 도입한 Denoising 관점에서 ($p\left(z_s\right|z_t)=N(z_s;{\mu }_{\theta }(z_t;s,t),{\sigma }_Q^2(s,t)I)$, ${\mu }_{\theta }(z_t;s,t)=\frac{{\alpha }_{t|s}{\sigma }_s^2}{{\sigma }_t^2}{z}_t+\frac{{\alpha }_s{\sigma }_{t|s}^2}{{\sigma }_t^2}{\hat{x}}_{\theta }(z_t;t)$, ${\sigma }_Q^2\left(s,t\right)={\sigma }_{t\overline{)s}}^2\frac{{\sigma }_s^2}{{\sigma }_t^2}$), 도출한 식은 아래와 같습니다.


위의 식의 의미는 다음과 같습니다. step 마다 모델에 latent variable $z_t$를 입력해서 나오는 결과 값과 실제 정답 x와의 차이를 계산하는데, 여기에 weight를 곱해서 이전 step에 비해 signal to ratio 값의 차이가 클수록 즉 더 noisy 한 data에 더 큰 penalty를 부여하는 것입니다.
Noise prediction 관점에서 ($p\left(z_s\right|z_t)=N(z_s;{\mu }_{\theta }(z_t;s,t),{\sigma }_Q^2(s,t)I)$, ${\mu }_{\theta }(z_t;s,t)=\frac{1}{{\alpha }_{t|s}}{z}_t-\frac{{\sigma }_{t|s}^2}{{\alpha }_{t|s}{\sigma }_t}{\hat{\in }}_{\theta }(z_t;t)$, ${\sigma }_Q^2(s,t)={\sigma }_{t|s}^2\frac{{\sigma }_s^2}{{\sigma }_t^2}$), ${\alpha }_t$, ${\sigma }_t^2$를 SNR(t)로 표현하고, $z_{t\left(i\right)}$를 reparameterization trick으로 계산한 것을 대입해서 정리하면 아래처럼 간단한 식으로 도출 가능합니다.



Continuous time case
논문에서는 수학적으로 Time step이 늘어날수록 Diffusion loss가 줄어드는 것을 증명했습니다. 이에 따라서 Time step이 무수히 많은 경우(T → ∞)에서의 Diffusion loss식도 도출했습니다. 이 경우에서는 t를 continuous 값으로 봅니다. 앞에서 구한 discrete case에서의 diffusion loss에 limit을 사용해서 아래의 식을 도출합니다.


Weight term 이 미분 값으로 바뀐 것을 확인할 수 있습니다. Denoising 관점에서 도출한 식은 아래와 같습니다. $\left(\mathrm{SNR}\left(\mathrm{t}\right)=exp\left(-{\gamma }_{\eta }\left(t\right)\right)\right)$


앞에서 Noise schedule을 학습한다고 했고 따라서 invertible monotonically decrease function인 SNR(t)을 학습합니다. $\left(\mathrm{SNR}\left(\mathrm{t}\right)=exp\left(-{\gamma }_{\eta }\left(t\right)\right)\right)$. 논문에서는 Continuous case에서 SNR(t)의 형태가diffusion loss식에 영향을 주지 않고 오직 SNR(0), SNR(1) 값만 diffusion loss에 영향을 주는 것을 증명을 통해 보였습니다. 따라서 Diffusion loss를 줄이기 위해서 SNR(t)의 parameter 인 $\eta$를 따로 학습할 필요는 없습니다. Noise prediction model의 parameter인 θ와 SNR(0), SNR(1) 값만 학습하면 됩니다. 위의 식을 Monte Carlo estimator로 계산해서 model의 parameter들을 학습합니다. MC estimator은 아래의 식과 같습니다.


논문에서는 Monte Carlo estimator의 분산을 줄이기 위해서 $\left(\mathrm{SNR}\left(\mathrm{t}\right)=exp\left(-{\gamma }_{\eta }\left(t\right)\right)\right)$의 parameter인 $\eta$를 학습할 것을 제안합니다.

Fourier Features
지금까지의 내용을 정리하면 Generative model은 아래와 같습니다.

$p\left(x\right)={\int }_{z}p\left(z_1\right)p\left(x\right|z_0\left)\prod _{i=1}^{T}p\right(z_{s\left(i\right)}\left|z_{t\left(i\right)}\right)$
$p\left(z_1\right)=N(z_1;0,I)$
$p\left(x\right|z_0)={\prod }_{i}p(x_i\left|z_{0,i}\right)$
$p\left(z_s\right|z_t)=N(z_s;{\mu }_{\theta }(z_t;s,t),{\sigma }_Q^2(s,t)I), {\mu }_{\theta }(z_t;s,t)=\frac{{\alpha }_{t|s}{\sigma }_s^2}{{\sigma }_t^2}{z}_t+\frac{{\alpha }_s{\sigma }_{t|s}^2}{{\sigma }_t^2}{\hat{x}}_{\theta }(z_t;t), {\sigma }_Q^2(s,t)={\sigma }_{t|s}^2\frac{{\sigma }_s^2}{{\sigma }_t^2}$

또는

$p\left(z_s\right|z_t)=N(z_s;{\mu }_{\theta }(z_t;s,t),{\sigma }_Q^2(s,t)I), {\mu }_{\theta }(z_t;s,t)=\frac{1}{{\alpha }_{t|s}}{z}_t-\frac{{\sigma }_{t|s}^2}{{\alpha }_{t|s}{\sigma }_t}{\hat{\in }}_{\theta }(z_t;t), {\sigma }_Q^2(s,t)={\sigma }_{t|s}^2\frac{{\sigma }_s^2}{{\sigma }_t^2}$

이를 학습하기 위한 loss식은 Diffusion loss는 discrete time case와 continuous time case가 있습니다. Discrete time case의 Diffusion loss는 다음과 같습니다.


Continuous time case의 Diffusion loss (MC estimator)는 다음과 같습니다.


즉, 각 time step 마다 latent variable $z_t$를 입력해서 $\in$를 predict 잘 predict하도록 모델을 학습합니다. 논문에서는 likelihood의 성능을 높이기 위해서 $z_t$의 high frequency furrier feature를 아래의 식으로 계산하고,


$z_t$에 concatenate하여 입력할 것을 제안했습니다. High frequency furrier feature 값은 input을 다양한 scale로 amplify한 것이고 이를 추가하면 fine scale detail들을 잘 capture하는 효과가 생깁니다.

Experiment
논문에서 제안한 모델인 VDM이 다른 Generative model에 비해 likelihood benchmark에서 SOTA 기록을 보였습니다.


생성된 이미지 quality는 아래와 같습니다.


결론

저는 주어진 데이터를 이용해 분포를 estimate해서 연관된 downstream task를 해결하는 것에 목표를 두고 이미지, sequential data, graph data의 Generative model에 대해서 연구하고 있습니다. 이번 블로그 글을 작성하면서 generative model 중 하나인 Diffusion model의 장점에 대해서 파악할 수 있었습니다. 그러나 Diffusion model의 likelihood 값도 결국 logp(x)의 lower bound 값이고 둘 사이에 gap이 여전히 존재합니다. 앞으로 이 gap을 줄이기 위한 방법을 고민할 필요가 있다고 생각합니다.