Understanding Diffusion Models: A Unified Perspective (Part 1: 생성 모델의 개념과 ELBO)

해당 글은 Calvin Luo가 2022년 arXiv에 게재한Understanding diffusion models라는 논문이다. Diffusion 모델이 초기에는 서로 다른 notation으로 정리되어 있어 개별 논문마다 이해가 조금 어려운데, 이 논문은 그런 수식들을 통합하여 작성해서 개인적으로는 diffusion 모델의 전반을 이해하기 참 좋은 논문이였다.

 

내용이 너무 길어 쪼개서 상세히 기록을 남기고자 한다. 먼저 part1, part2는 diffusion 모델의 기본이 되는 VAE와 VAE를 이해하기 위한 ELBO에 관한 내용이다. 이것이 어떻게 전개되는지 알 수 있다.

 

1. Introduction: Generative Models

어떤 분포로부터 샘플링한 $x$가 주어졌을 때, generative model의 목표는 그것의 실제 분포 $p(x)$를 모델링 하는 것이다. 그 모델을 학습하면, 모델로부터 새로운 샘플을 생성할 수 있다. 기존에 잘 알려진 모델들은 다음과 같다.

• GANs : 복잡한 분포의 샘플링 절차를 모델링 하는데, adversarial한 방식으로 학습된다.
• Likelihood-based : 관찰된 데이터 샘플에 높은 likelihood를 부여하는 모델을 학습한다. 자기회귀모델, normalizing flows, VAEs 등이 있다.
• Energy-based modeling : 임의의 정규화된 에너지 함수에 의해 학습되는 분포이다.
• Score-based generative models : Energy-based 모델의 score를 인공신경망을 통해 학습한다.

여기서는 diffusion model을 살펴보는데, diffusion model은 likelihood-based와 score-based 해석을 모두 가지고 있다는 것을 보여준다.

 

2. Background: ELBO, VAE, and Hierarchical VAE

다양한 형태의 데이터에서 우리가 관찰하는 데이터는 연관된 보이지 않는 잠재 변수에 의해 표현되거나 생성된다. 이를 임의의 변수 $z$로 표현한다. 일반적으로 실제 세계에서 우리가 마주하는 객체들은 더 높은 수준의 표현의 함수로 생성될 수 있다. 예를 들어 플라톤의 동굴 비유를 들자면, 평생 동굴안에 갇혀있는 사람들은 그들 앞 벽에 투사된 2차원의 그림자만 볼 수 있다. 이 그림자들은 불 앞을 지나는 보이지 않는 3차원 객체에 의해 생성된다. 동굴안의 사람들이 보는 2차원 그림자는 뒤쪽의 불, 불의 강도, 그리고 그림자를 생성하는 객체와 불 사이의 거리, 그 객체를 들고 있는 사람들의 움직임 등 아주 많은 요소에 의해 2차원 그림자가 결정된다. 즉, 그들이 관찰하는 모든 것은 실제로 그들이 결코 볼 수 없는 더 높은 차원의 추상적 개념에 의해 결정된다.

 

2.1. Evidence Lower Bound(ELBO)

• ELBO가 필요한 이유

잠재 변수와 관찰된 데이터를 joint distribution  $p(x, z)$ 로 모델링 할 수 있다. 즉, 위의 예에서 2차원 그림자 $x$는 수 많은 여러 가지 요소들인 $z$ 에 의해 결정되므로, 이를 joint distribution으로 모델링 하는 것이다. Likelihood-based 관점에서 모든 관측치 $x$ 는 likelihood $p(x)$ 를 최대화하는 모델을 학습한다. 그런데 우리는 2차원 그림자가 나타날 확률을 알고 싶은 것이다. 즉, 여기서 순수하게 우리가 관찰한 데이터 $p(x)$ 의 likelihood를 찾을 수 있는 두 가지 방법이 있다.(즉, $p(x, z)$ 에서 $p(x)$ 를 찾는 방법) 

 

1. 모든 잠재 변수에 대해서 적분하여 구할 수 있다.
   
   $p(x) = \displaystyle\int p(x,z)dz$ (1)
   
   
2. Chain rule을 통해 구할 수 있다.
   
   $p(z|x) = \cfrac{p(z\cap x)}{p(x)}$     (2) → 베이즈 정리
   
   ▼ 전개
   $p(x) = \cfrac{p(z\cap x)}{p(z|x)}\approx\cfrac{p(x,z)}{p(z|x)}$
   $p(x, z) \approx p(x\cap z)$
   $p(x) = \cfrac{p(x,z)}{p(z|x)}$

(1)처럼 $z$의 모든 공간을 적분 하거나, (2)처럼 $p(z|x)$에 접근하는 것은 매우 어렵기 때문에 직접적으로 likelihood $p(x)$를 계산하는 것은 매우 어렵다.

따라서 Evidence of Lower BOund(ELBO)를 통해서 추론할 수 있다. 

 

• ELBO equation
  

(1)과 (2)의 식을 사용하여 ELBO를 도출할 수 있다. ELBO의 의미는 관찰된 데이터의 log likelihood의 하한(Lower Bound)을 나타낸다. 그러면 우리는 ELBO를 최대화하는 것이 결국 잠재 변수 모델을 최적화하는 것임을 알 수 있다. 최적의 경우에는 ELBO가 evidence, 즉 관찰된 데이터의 log likelihood와 정확히 동일해진다. ELBO의 식은 아래와 같이 나타낼 수 있다.

$\mathbb{E}{q\phi(z|x)}\Bigg[\text{log}\cfrac{p(x, z)}{q_\phi(z|x)}\Bigg]$ (3)

→ 일반적으로 variational inference에서 $q_\phi(z|x)$라는 표기를 사용, $\phi$는 파라미터로 $q_\phi(z|x)$는 $p(z|x)$를 근사. 즉 $q_\phi(z|x)$가 파라미터$(\phi)$에 의해 $p(z|x)$에 가까워지도록 하는 표현

→ $q_{\phi}(z|x)$에 따라 분포되는 $f(z)$가 다수 생성될 것 이고, 이 기대값을 구하면 $p(z|x)$에 근사하는 과정

 

Evidence와의 관계를 관계를 명확히 하기 위해 수학적으로 다음과 같이 쓸 수 있다.

$\text{log}\;p(x)\geq\mathbb{E}{q\phi(z|x)}\Bigg[\text{log}\cfrac{p(x, z)}{q_\phi(z|x)}\Bigg]$ (4)

→ (2)식 양변에 log를 취하고, 우변의 ELBO가 추정치 이므로 위와 같이 표현할 수 있음. 즉, $\text{log}\;p(x)$가 ELBO 이상 되어야 한다.

여기서 $q_\phi(z|x)$는 최적화 하고자 하는 파라미터$(\phi)$에 의해 표현되는 variational 분포이다. 즉 이것은 실제 posterior인 $p(z|x)$를 찾는 과정이라고 할 수 있다. $\phi$를 조정하여 ELBO를 최대화하면, Lower Bound가 높아지고 이로 인해 실제 데이터 분포를 모델링하고 그로부터 샘플링을 할 수 있게 된다. 그렇게 생성 모델을 학습하게 된다. 다음은 ELBO가 왜 최대화 되어야 하는 목표인지 알아본다.

 

• ELBO 최대화

ELBO를 유도하기 위해 식 (1)을 사용한다.

$\text{log}\;p(x) = \text{log}\displaystyle\int p(x,z)\,dz$                             (5)

$\quad\quad\quad= \text{log}\displaystyle\int\cfrac{p(x,z)q_\phi(z|x)}{q_\phi(z|x)}\,dz$                 (6)

$\quad\quad\quad= \text{log}\,\mathbb{E}_{q_\phi(z|x)}\Bigg[\cfrac{p(x,z)}{q_\phi(z|x)}\Bigg]$                      (7)

$\quad\quad\quad\geq\mathbb{E}_{q_\phi(z|x)}\Bigg[\text{log}\cfrac{p(x,z)}{q_\phi(z|x)}\Bigg]$                       (8)

 

(5) ELBO의 식에서 $z$에 대해서 적분을 수행한다.

(6) 분자와 분모에 $q_\phi(z|x)$를 곱한다.

(7) $E[X] = \int xf(x)\,dx$ 이므로, 기대값의 정의에 따라 정리할 수 있다.

(8) Jensen’s Inequality에 의해서 성립한다. (log 함수는 오목 함수 이기 때문에, $f(E[X])\geq E[f(X)]$

이 증명은 ELBO가 실제로 evidence의 하한인 이유에 대한 직관을 제공하지 않으며, Jensen’s 부등식은 속임수일 뿐이다. 또한 ELBO를 왜 최대화 하는지에 대한 설명이 충분하지 않아, 이번에는 식 (2)를 통해 유도를 수행해본다.

 

$\text{log}\;p(x) = \text{log}\;p(x)\displaystyle\int q_\phi(z|x)\,dz$                                                      (9)

$\quad\quad\quad = \displaystyle\int q_\phi(z|x)(\text{log}\;p(x))\,dz$                                                     (10)

$\quad\quad\quad= \mathbb{E}_{q_\phi(z|x)}[\text{log}\;p(x)]$                                                                 (11)

$\quad\quad\quad= \mathbb{E}_{q_\phi(z|x)}\Bigg[\text{log}\;\cfrac{p(x,z)}{p(z|x)}\Bigg]$                                                         (12)

$\quad\quad\quad= \mathbb{E}{q\phi(z|x)}\Bigg[\text{log}\;\cfrac{p(x,z)q_\phi(z|x)}{p(z|x)q_\phi(z|x)}\Bigg]$                                        (13)

$\quad\quad\quad= \mathbb{E}{q\phi(z|x)}\Bigg[\text{log}\;\cfrac{p(x,z)}{q_\phi(z|x)}\Bigg] + \mathbb{E}{q\phi (z|x)}\Bigg[\text{log}\;\cfrac{q_\phi(z|x)}{p(z|x)}\Bigg]$    (14)

$\quad\quad\quad= \mathbb{E}{q\phi(z|x)}\Bigg[\text{log}\;\cfrac{p(x,z)}{q_\phi(z|x)}\Bigg] + D_{KL}(q_\phi(z|x)\; ||\;p(z|x))$      (15)

$\quad\quad\quad\geq \mathbb{E}{q\phi(z|x)}\Bigg[\text{log}\;\cfrac{p(x,z)}{q_\phi(z|x)}\Bigg]$                                                     (16)

 

(9) $\displaystyle\int q_\phi(z|x)\,dz$ 는 확률밀도함수인 $q_\phi(z|x)$를 $z$공간에 대해 적분한 것이므로 1과 같다. 따라서 $\text{log}\;p(x) = \text{log}\;p(x) \times1$ 과 같다.

(10), (11) $\text{log}\;p(x)$를 적분기호 안으로 가져올 수 있고, 그렇게 되면 (7)과 동일하게 기대값의 정의에 따라 표현할 수 있다.

(12) 식 (2)에서 증명한 베이즈 정리 $p(x) = \cfrac{p(x,z)}{p(z|x)}$를 적용한다.

(13) 분자와 분모에 $\cfrac{q_\phi(z|x)}{q_\phi(z|x)}$를 곱한다.

(14) 로그의 성질에 따라 기대값을 분리할 수 있다.

(15) KL-Divergence의 정의 $KL(p||q) = \sum_i\,p_i\,\text{log}(\cfrac{p_i}{q_i})$ 에 따라 두 번째 항을 KL-Divergence로 표현 $D_{KL}(q_\phi (z|x)||p(z|x))$ 할 수 있다(여기서 확률분포의 기대값은 확률의 가중합이므로 $\sum$으로 이해할 수 있다). 결과적으로 KL-Divergence는 $q_\phi(z|x)$와 $p(z|x)$의 거리를 측정한다.

(16) KL-Divergence 정의에 의해 KL-Divergence는 항상 0보다 크므로, (15) $\geq$ (16)이 성립한다.

결과적으로 evidence ($p(x)$의 loglikelihhod) 가 ELBO와 추정된 posterior $q_\phi(z|x)$와 실제 posterior $p(z|x)$ 사이의 KL-Divergence의 합으로 표현된다.

→ Evedence = ELBO + posterior간 KLD

ELBO가 최대화한다는 것은 결국 $q_\phi(z|x)$가 $p(x, z)$와의 비율을 최대로 만든다는 것이고, 이 의미는 $q_\phi(z|x)$가 실제 분포를 잘 성명한다는 뜻이다. 그러면 두 posterior가 유사해지고 (실제 분포와 추정 분포), KLD가 0에 가까워진다. 결과적으로 ELBO를 최대화 하는 것은 KLD를 최소화 하는 것이고, 이것이 $\text{log}\,p(x)$ 즉 loglikelihood를 찾는 과정이 되는 것이다.

 

우리의 목표는 관찰된 데이터를 잘 설명하는 latent 구조를 학습하는 것이다. 다시 말해 $q_\phi(z|x)$의 최적의 파라미터를 찾아서 $p(z|x)$와 일치시켜야 한다. 그런데 우리는 $p(z|x)$에 접근할 수 없어서 KLD 항을 직접적으로 최소화 할 수 없다. 따라서 ELBO를 최대화하여 근사적으로 데이터의 likelihood를 추정하는 것이다.