Deep Unsupervised Learning using Nonequilibrium Thermodynamics 논문 리뷰

Diffusion model의 시초가 되는 'Deep Unsupervised Learning using Nonequilibrium Thermodynamics'에 대해 리뷰해보고자 한다. 관련 연구에 대한 설명은 생략하고 논문의 주된 논의 내용 위주로 소개하고자 한다.

 

Diffusion Probabilistic Models

Diffusion Probabilistic Models의 장점은 이 논문에서 4가지로 정리해 소개하고 있다.

(1) 모델 구조의 유연함

(2) 명확한 샘플링이 가능

(3) 다른 분포와의 곱을 하는 등 연산이 간편함

(4) 모델의 log likelihood나 각각의 state에서의 확률을 구하기 쉬움

 

Diffusion Probabilistic Models은 연속직인 Markov chain을 이용해 간단하고 알고 있는 분포를 우리가 모르는 타겟 분포로 변환하는 모델로 diffusion process를 이용한다.  각각의 chain에 대한 확률을 analytic하게 계산할 수 있으므로, full chain 역시 계산 가능하다.

 

Diffusion Probabilistic Model을 도식화하면 아래와 같다.

Forward Trajectory(data distribution to Gaussian Noise)

아래와 같이 정의하자.

여기서 $q(x^{t}|x^{t-1})$은 Gaussian diffusion에서 Gaussian diffusion으로 가거나 binomial에서 binomial로 가는 분포를 나타낼 수 있다. 아래 표를 보면 forward diffusion kernel과 reverse diffusion kernel, training target 등이 분포에 따라 달라짐을 알 수 있다.

Reverse Trajectory(Gaussian Noise to data distribution)

forward trajectory에서 우리는 $q$가 무엇인지 알고 있지만, reverse trajectory에서의 $p$는 analytic하게 계산할 수 없음을 생각하자. 다음은 $p$에 대한 식이다.

다만 Feller에 의하면 상당히 작은 $\beta$에 대해 reverse of the diffusion process 역시 forward process와 같은 형태를 가진다는 사실을 알 수 있다. 앞에서 우리가 $q(x^{t}|x^{t-1})$을 Gaussian(binomial) distribution으로 정의했기 때문에 만약 ${\beta}_t$가 작다면 $q(x^{t-1}|x^{t})$ 역시 Gaussian(binomial) distribution임을 알 수 있다.

 

여기서 $p(x^{t-1}|x^{t})$가 Gaussian distribution을 갖는다면  mean과 variance가 학습되어야 하며, 이를 각각 $f_{\mu}(x^{(t)}, t)$와 $f_{\Sigma} (x^{(t)}, t) $로 놓자. 이 논문에서는 MLP를 사용해 학습하였다. 

 

Model Probability

모델이 부여한 데이터의 확률은 다음과 같이 나타낼 수 있다.

위 적분은 intractable하다. 아래와 같이 식을 변형해보자.

만역 $\beta$가 상당히 작다면 forward distribution과 reverse distribution이 동일할 것이므로 위 식은 하나의 샘플을 통해 $q(x^{(1 ... T)}|x^{(0)})$에 의해 계산될 수 있다.

 

Training

log likelihood를 최대화 하기 위해 (=cross entropy를 최소화 하기 위해) $L$을 다음과 같이 정의해보자.

Jensen's inequality에 의해 다음과 같이 식이 변형된다.

더보기에 있는 증명에 의해 다음과 같은 식을 얻을 수 있다.

위 식은 analytic하게 계산할 수 있으며, 만약 forward와 reverse trajectories가 동일하다면 $L=K$가 성립한다. training과에서 reverse transition을 찾는 것은 다음과 같이 각각의 스텝에서 $K$를 최대화하는 파라미터를 찾으면 된다.

위와 같이 식이 바뀌면 지금까지의 확률 분포 추정 작업이 특정 함수의 회귀 문제로 단순화 되었음을 의미한다. 

 

Setting the diffusion rate ${\beta}_t$

${\beta}_t$는 확산 과정에서 각 시간에 따른 노이즈의 크기를 결정하기 때문에 ${\beta}_t$의 값에 따라 모델의 성능도 달라진다. Gaussian noise의 경우 논문에서는 gradient ascent를 통해 스케쥴링한다. 다만 첫번째 노이즈인 ${\beta}_1$은 작은 상수로 고정하고, 그 다음 매개변수인 ${\beta}_{2...T}$부터 최적화하는 frozen noise 기법을 사용한다.

 

Multiplying Distributions, and Computiong

signal denoising이나 missing values의 inference를 위해 이 논문에서는 또 다른 distribution $r(x)$를 도입하여 새로운 distribution $\tilde{p(x^{(0)})} \propto p(x^{(0)})r(x^{(0)})$를 도입한다. 

 

일반적으로 distribution을 곱하는 것은 다른 모델에서는 까다로운 일이지만, diffusion probablistic modle에서는 second distribution을 일종의 small perturbation으로 간주할 수 있어 다루기 용이하다.

여기서 $\tilde{Z_t}$는 normalizing constant이다.

 

$\tilde{p}$를 이용해 앞에서 진행한 논의들을 수정해보자.

왼쪽 식이 오른쪽 식과 같이 되기 위해서는 다음을 만족해야 한다.

(20)번 식은 아래 식이 만족할 때 만족한다.

다음 식을 만족하도록 $\tilde{p}$는 많지만, normalized되지 않았을 수 있으므로 아래와 같이 골라준다.

(22)

각각의 Gaussian noise는 분산이 매우 작기 때문에 상당히 뾰족한 모양을 가지고 있게 된다. 이는 $\frac{r(x^{(t)})}{r(x^{(t+1)})}$이 $p(x^{(t)}|x^{(t+1)})$에게 있어 작은 perturbation임을 알 수 있다. 이는 Gaussian noise의 평균에만 영향을 줄 뿐, variance에는 영향을 주지 않으며, normalization constant로 바꾸지 않는다. (증명은 더보기에서 확인)

Applying $r(x^{(t)})$

각각 $r(x^{(t)})$를 지정했을 때에 대한 설명이 주어져 있다. 우리는 $r(x^{(t)})$가 시간에 따라 빠르게 바뀌는 것을 원하지 않기 때문에 constant로 정의하거나  ${r(x^{(0)})}^{\frac{T-t}{T}$ 등으로 정의하기도 한다. 

 

Entropy of Reverse Process

forward process는 잘 알려져 있기 때문에 우리는 conditional entropy의 상한과 하한을 다음과 같이 구할 수 있다.

여기서의 entropy는 analytic하게 계산될 수 있으며, 상한과 하한 모두 $q(x^{(1...T)}|x^{(0)})$에만 의존한다.