Contents

Score-based Generative Model

Generative Modeling by Estimating Gradients of the Data Distribution

학습데이터의 data distribution을 학습하여 “실제 있을법한” 데이터를 만들어낼 수 있는 모델을 생성모델(Generative model)이라고하며 이를 달성하기 위해 다양한 방법이 제시되었다. GAN, VAE, Normalizing Flow는 모두 생성모델을 위한 framework으로 각각 발전하며 더 자연스럽고 사실같은 데이터를 생성하기 위해 연구되고 있다. 이 글은 Yang Song이 1저자로 참여한 논문인 “Generative Modeling by Estimating Gradients of the Data Distribution"을 기반으로 한다. 다만 논문리뷰의 형식으로 다루는 것이 아니라 생성모델에 대해 정리하면서 제시된 score-based generative model의 배경과 내용을 전반적으로 파악하는데 집중하고자 한다. 이 글은 Yang Song의 지도교수인 Stefano Ermon교수의 세미나를 기반으로 하였다.

Benefits of Accessing Probability Mass/Density Function

우선 Ian Goodfellow가 NIPS 2016 Tutorial에서 발표한 “NIPS 2016 Tutorial: Generative Adversarial Networks”에서 사용된 분류체계를 살펴보자.

/score-based-models/taxonomy.png
From Ian Goodfellow's NIPS 2016 Tutorial

GAN은 implicit density로 분류되어 있다. GAN의 구조를 생각해보면 임의의 노이즈에서 시작해 GAN generator를 통과하면서 sample을 생성하게 된다. 이후에도 강조되겠지만 확률을 모델링할 때 중요한 것은 확률의 기본적인 성질, 즉 합이 1이되어야하며 어떤 곳에서도 음수가 될 수는 없다는 성질을 만족해야 한다는 것이다. 하지만 다루는 확률공간의 크기를 생각할 때, 합이 1이 되도록 해주는 것은 결코 쉽지 않다. 1이 되도록 normalize해주는 것을 parition function이라고 하는데 이 함수 자체를 찾는게 쉽지가 않다. 확률분포를 정의할 수 있다면 이 확률의 likelihood를 직접 높이면 참 좋겠으나 결과적으로 이 과정이 어려운 것이다. GAN은 generative model과 discriminative model을 적대적으로 학습시킴으로써 (two-player minimax game) 확률분포를 직접모델링하는 어려움을 피할 수 있었다. 이는 GAN의 학습과정에서 어디에도 likelihood를 사용하지 않는 것에서 확인할 수 있다. 결과물의 생성에 있어서 성공적이었으나 아쉬움이 남는 부분은 학습의 어려움은 차치하더라도 바로 확률분포에 대한 추정이 학습과정에서 explicit하게 드러나지 않았기 때문에 확률분포를 사용하는데 제약이 있다는 것이다.

VAE는 비록 확률을 직접 최적화하지는 못하지만 evidence lower bound를 구해 이를 최적화함으로써 간접적으로 확률분포를 추정하므로 GAN보다는 확률분포에 더 접근할 수 있는 framework이다. 물론 Gaussian prior 사용이 강제되고 최종 결과물의 품질이 GAN이 만들어내는 것보다는 낮다는 문제가 있었으나 이러한 단점은 VQ-VAE와 같은 방법의 등장으로 빠르게 개선되고있다.

Normalizing flow는 tractable한 layer를 거쳐 확률분포를 모델링하는 방법으로 가장 강력한 장점은 바로 확률분포를 tractable하게 추정하므로 확률분포에 대한 log-likelihood를 직접 최적화할 수 있다는 것이다. 하지만 tractability를 보장하기 위해서 invertible해야하며 매 단계마다 가하는 변환의 determinant를 쉽게 계산할 수 있어야한다는 제한이 붙게 된다.

여기서 가장 핵심적인 문제는 “확률분포에 직접 접근하기가 어렵다"는 것이다. 확률 질량/밀도 함수에 직접 접근할 수 있으면 이 분포가 우리가 찾고 싶은 데이터의 분포 그 자체이므로 이 분포에 가깝도록 최적화를 진행할 수 있으나 확률분포가 가지는 성질을 만족해야하기 때문에 모델링과정에서 formulate하기가 어렵고 따라서 maximum likelihood 접근을 사용하기 어렵다는 문제가 있다. Trade-off관계는 확률분포를 단순하게 가정하면 tractable하고 직접 최적화가 가능해진다는 장점이 있지만 단순한 확률분포를 가정하게되면 데이터가 가지는 표현성(expressivity)를 크게 제한하게 되고 복잡한 형태의 데이터를 모델링할 수 없게 된다는 것이다.

Score

Yang Song은 논문에서 기존과는 다른 새로운 framework을 제안한다. 제시된 접근방법은 score-based generative modeling으로 불리며 단계적으로 알아가보자.

우선 score에 대해 이해할 필요가 있다. 현재 우리가 다루는 생성모델 문맥에서의 score는 다음을 의미한다.

Score

데이터의 확률분포의 log값에 대한 gradient를 score라고 한다. $$\nabla_x \log p(\mathbf{x})$$

Score는 모델 parameter에 대한 gradient가 아닌 input 데이터에 대한 gradient임에 유의하자.

다음과 같이 $(-1.5, 1.5)$와 $(1.0, 1.0)$의 평균을 갖는 두 개의 Gaussian mixture로 구성된 확률분포가 있다고 해보자.

gradient_plot.png
Gradient Plot

여기서 score는 gradient의 vector field로 화살표로 확인할 수 있다. Score는 각 지점에서 어떻게 움직여야 likelihood가 높은 지점으로 갈 수 있는지를 알려주는 역할을 한다. 따라서 위 분포에서는 화살표를 따라가면 각각의 Gaussian의 mean vector에 도착하게 된다.

Score-based model의 핵심아이디어는 확률분포함수를 사용하는 대신에 바로 이 score를 활용해보자는 것이다.

Deep Energy-Based Models

Score를 사용하는 모델이 여기서 처음 제시되는 것은 아니다. Energy based model은 네트워크로 하여금 input $\mathbf{x}$에 대해 임의의 실수를 출력하도록 하는 network를 만들고 이 값을 다음의 식에 넣어 확률분포를 추정하고자 하였다.

$$f_{\theta}(\mathbf{x}) \in \mathbb{R}$$ $$p_{\theta}(\mathbf{x}) = \frac{e^{-f_{\theta}(\mathbf{x})}}{Z_{\theta}}$$

네트워크는 임의의 실수 $f_{\theta}(\mathbf{x})$를 출력하므로 음수가 될 수 있다. Exponential을 취해 양수로 만들어주어 음수가 되면 안된다는 조건은 해결할 수 있는데 문제는 역시 확률분포의 전체 합이 1이 되어야 한다는 것이다. 즉 위 식에서 $Z_{\theta}$를 구하는 과정이 intractable하다. Log-likelihood기반으로 학습하려면 다음의 NLL loss를 구해야 한다. $$ \mathbb{E}_{p_{\text{data}}}[-\log p_{\theta}(\mathbf{x})] = \mathbb{E}_{p_{\text{data}}}[\log f_{\theta}(\mathbf{x}) - \log Z_{\theta}] $$ 위 식에서 $Z_{\theta}$가 intractable하기 때문에 MLE를 사용할 수 없게 된다.

Learning Deep Energy-Based Models using Scores

Score는 partition function에 영향을 주지 않으므로 다음 식에서 partition function에 대한 gradient를 0으로 처리할 수 있다.

$$\nabla_{\mathbf{x}} \log p_{\theta}(\mathbf{x}) = - \nabla_{\mathbf{x}} f_{\theta}(x) - \nabla_{\mathbf{x}} \log Z_{\theta}$$

여기서 중요한 아이디어는 확률분포를 직접 다루려고하면 partition function을 처리해야 하지만 확률분포가 아닌 그 gradient인 score를 추정하는 것은 partition function 없이도 할 수 있다는 것이다.

실제 데이터의 score function은 $\nabla_{\mathbf{x}} \log p_{\text{data}}(\mathbf{x})$라고 하자. 이 때 parameter $\theta$를 갖는 network를 사용해 $\nabla_{\mathbf{x}} \log p_{\theta}(\mathbf{x})$가 $\nabla_{\mathbf{x}} \log p_{\text{data}}(\mathbf{x})$이 되도록 훈련시키자는 것이 핵심 아이디어이다.

Score Estimation

문제의 setting을 정리해보자. 우리가 하고 싶은 것은 data distribution $p_{\text{data}(\mathbf{x})}$를 추정하는 것이며 직접계산하는 것은 partition function때문에 할 수 없다. 따라서 score function을 통해서 접근하려고 한다.

/score-based-models/score_estimation.png
From Stefano Ermon's Seminar

하지만 위 그림에서 보듯 우리는 실제분포를 모르고 다만 그 분포에 존재하는 training data에만 접근할 수 있다. 이 training data가 갖는 score function을 추정해야하는데 training data는 그저 있을 뿐이지 gradient가 얼마인지에 대한 내용을 들고 있지는 않다. 이제 문제는 보다 구체적이 되었다. 어떻게 하면 $\nabla_{\mathbf{x}} \log p_{\text{data}}(\mathbf{x})$를 추정하는 score fucntion $s_{\theta}(\mathbf{x}): \mathbb{R}^D \rightarrow \mathbb{R}^{D}$를 training data로부터 구할 수 있을지를 찾아야 한다. 실제 데이터의 score들이 표현되는 vector field를 찾아야 하는 것이다. 이 때 $D$는 데이터의 차원이다.

문제 formulation은 다음과 같다.

  • Given: i.i.d. samples $\{ \mathbf{x}_1, \mathbf{x}_2, \cdots, \mathbf{x}_{N} \} \stackrel{\text{i.i.d.}}{\sim} p_{\text{data}} (\mathbf{x}) = p(\mathbf{x})$
  • Task: Estimating the score $\nabla_{\mathbf{x}} \log p(\mathbf{x})$
  • Score Model: A vector-valued function $s_{\theta}(\mathbf{x}): \mathbb{R}^{D} \rightarrow \mathbb{R}^{D}$
  • Objective: How to compare two vector fields of scores?

Score Matching

$\theta$에대해 parameterized된 함수 $s_{\theta}$가 $\nabla_{\mathbf{x}} \log p(\mathbf{x})$와 같아지면되고 거리를 Euclidean distance로 정의하면 다음을 loss function으로 설정해 볼 수 있다.

Fisher Divergence
$$ \frac{1}{2} \mathbb{E}_{p(\mathbf{x})} \left[ \lVert \nabla_{\mathbf{x}} \log p(\mathbf{x}) - s_{\theta}(\mathbf{x}) \rVert_2^2 \right] $$ 왼쪽이 true score, 오른쪽이 이를 추정하고자 하는 network라고 보면된다. 위 식을 Fisher divergence라고 한다. Fisher divergence가 0이면 두 분포는 같다.

하지만 여전히 문제는 우리는 $p(\mathbf{x})$를 모르므로 이 값의 score도 알 수 없어 왼쪽 true score term을 계산할 수 없다. 위 식을 integration by parts를 사용하면 다음과 같이 바꿀 수 있다.

$$ \mathbb{E}_{p(\mathbf{x})}\left[\frac{1}{2} \lVert s_{\theta}(\mathbf{x}) \rVert_{2}^{2}+\operatorname{tr}(\underbrace{\nabla_{\mathbf{x}} s_{\theta}(\mathbf{x})}_{\text {Jacobian of } s_{\theta}(\mathbf{x})})\right] $$

이렇게 바꿀 때 얻을 수 있는 이점은 $p(\mathbf{x})$를 식에서 제거할 수 있다는 것이다. 이 과정을 Score Matching (Hyvarinen, 2005)라고 한다. 이제 $p(\mathbf{x})$는 expectation에 대해서만 남아있으므로 training data를 통해 다음과 같이 추정이 가능하다.

$$ \frac{1}{N} \sum_{i=1}^{N}\left[\frac{1}{2} \lVert s_{\theta}\left(\mathbf{x}_{i}\right) \rVert_{2}^{2}+\operatorname{tr}\left(\nabla_{\mathbf{x}} s_{\theta}\left(\mathbf{x}_{i}\right)\right)\right], \quad \{ \mathbf{x}_{1}, \mathbf{x}_{2}, \ldots, \mathbf{x}\_{N} \} \stackrel{\text{i.i.d.}}{\sim} p(\mathbf{x}) $$

Derivation

위의 score matching과정을 이해하기 위해 1차원 데이터에 대한 예시를 살펴보자.

$$ \begin{aligned} & \frac{1}{2} \mathbb{E}_{p(x)}\left[\left(\nabla_{x} \log p(x)-s_{\theta}(x)\right)^{2}\right] \cr =& \frac{1}{2} \int p(x)\left(\nabla_{x} \log p(x)-s_{\theta}(x)\right)^{2} \mathrm{~d} x \cr =& \frac{1}{2} \int p(x)\left(\nabla_{x} \log p(x)\right)^{2} \mathrm{~d} x+\frac{1}{2} \int p(x) s_{\theta}(x)^{2} \mathrm{~d} x \cr &-\int p(x) \nabla_{x} \log p(x) s_{\theta}(x) \mathrm{d} x \end{aligned} $$

위 전개에서 첫 번째 term은 $\theta$에 대해 independent하다. 두 번째 term은 쉽게 expectation으로 묶을 수 있는 항이다. 마지막 term이 문제인데 마지막 term은 다음과 같이 풀어쓸 수 있다.

$$ \begin{aligned} &-\int p(x) \nabla_{x} \log p(x) s_{\theta}(x) \mathrm{d} x \cr =&-\int p(x) \frac{\nabla_{x} p(x)}{p(x)} s_{\theta}(x) \mathrm{d} x=-\int \nabla_{x} p(x) s_{\theta}(x) \mathrm{d} x \cr =&-\left.p(x) s_{\theta}(x)\right|_{x=-\infty} ^{\infty}+\int p(x) \nabla_{x} s_{\theta}(x) \mathrm{d} x \cr =& \int p(x) \nabla_{x} s_{\theta}(x) \mathrm{d} x \end{aligned} $$

세 번째 줄의 boundary term은 $\pm \infty$에 대해서 0이 될 것이므로 결과적으로 expectation으로 묶을 수 있는 형태로 정리된다.

다시 쓰면 1차원 $D$에 대해서는 다음과 같이 표현할 수 있다.

$$ \begin{aligned} & \frac{1}{2} \mathbb{E}_{p(x)}\left[\left(\nabla_{x} \log p(x)-s_{\theta}(x)\right)^{2}\right] \cr =& \frac{1}{2} \int p(x)\left(\nabla_{x} \log p(x)\right)^{2} \mathrm{~d} x+\frac{1}{2} \int p(x) s_{\theta}(x)^{2} \mathrm{~d} x \cr &+\int p(x) \nabla_{x} s_{\theta}(x) \mathrm{d} x \cr =& \text { const }+\frac{1}{2} \mathbb{E}_{p(x)}\left[s_{\theta}(x)^{2}\right]+\mathbb{E}_{p(x)}\left[\nabla_{x} s_{\theta}(x)\right] \end{aligned} $$

같은 원리로 $D \geq 1$일때는 위에서 표기한대로 유도된다.

Score Matching is not Scalable

Score matching을 통해 Monte Carlo방식으로 cost function을 줄이면 score function을 찾을 수 있고 이로써 우리가 풀고싶었던 true score를 추정하는 network를 학습시킬 수 있다. 하지만 아직도 문제는 남아있다. score matching 방식이 scalable하지 않다는 것이다. 이 문제를 살펴보자.

예를들어 데이터가 세 개의 변수로 표현된다고 생각해보자. 그럼 input data는 $\mathbf{x} \in \mathbb{R}^3$의 꼴을 가지며 score function $s_{\theta}(\mathbf{x})$는 $\mathbb{R}^{D} \rightarrow \mathbb{R}^{D}$의 함수이므로 출력도 세개의 값을 갖는다.

$$ \mathbb{E}_{p(\mathbf{x})}\left[\frac{1}{2} \lVert s_{\theta}(\mathbf{x}) \rVert_{2}^{2}+\operatorname{tr}(\underbrace{\nabla_{\mathbf{x}} s_{\theta}(\mathbf{x})}_{\text {Jacobian of } s_{\theta}(\mathbf{x})})\right] $$

Cost function이 위와 같이 생겼으므로 계산해야 하는 건 $\lVert s_{\theta} (\mathbf{x}) \rVert_2^2$와 $\operatorname{tr}(\nabla_{\mathbf{x}} s_{\theta}(\mathbf{x}))$ 두 가지 이다. 이 때 $\lVert s_{\theta} (\mathbf{x}) \rVert_2^2$의 계산은 어렵지 않으나 $\operatorname{tr}(\nabla_{\mathbf{x}} s_{\theta}(\mathbf{x}))$는 간단하지 않다. Trace하기 전의 score function에 대한 gradient는 전개하면 다음과 같다.

$$ \nabla_{\mathbf{x}} s_{\theta}(\mathbf{x})=\left(\begin{array}{lll} \frac{\partial s_{\theta, 1}(\mathbf{x})}{\partial x_{1}} & \frac{\partial s_{\theta, 1}(\mathbf{x})}{\partial x_{2}} & \frac{\partial s_{\theta, 1}(\mathbf{x})}{\partial x_{3}} \cr \frac{\partial s_{\theta, 2}(\mathbf{x})}{\partial x_{1}} & \frac{\partial s_{\theta, 2}(\mathbf{x})}{\partial x_{2}} & \frac{\partial s_{\theta, 2}(\mathbf{x})}{\partial x_{3}} \cr \frac{\partial s_{\theta, 3}(\mathbf{x})}{\partial x_{1}} & \frac{\partial s_{\theta, 3}(\mathbf{x})}{\partial x_{2}} & \frac{\partial s_{\theta, 3}(\mathbf{x})}{\partial x_{3}} \end{array}\right) $$

Trace이므로 위 행렬의 대각성분의 합을 구하면 된다. 만약 변수가 $D$개라면 $O(D)$번의 Backprop이 필요하다. 이 부분에서 score matching이 scalable하지 않다고 하는 것이다.

Scalable Score Matching: Sliced Score Matching

Scalable하게 만드는 방법으로 sliced score matching을 제시하는데 아이디어는 간단하다. 모든 차원에 대해서 gradient를 구해 더하기가 어려우니 random direction의 set을 준비해서 해당 direction들로 projection을 한 값이 같아지도록 만드는 것이다.

/score-based-models/sliced_score_matching.png
From Stefano Ermon's Seminar

위 그림에서 왼쪽처럼 모든 방향에서 같아지게 하는 것이 목적이나 이는 scalable하지 않으므로 점선으로 표시된 특정 방향을 잡고 해당 방향으로의 projection이 같아지게 하는 문제로 바꾸어주게 된다. 이를 적용하면 objective는 다음과 같이 쓸 수 있다.

Sliced Fisher Divergence
Sliced Fisher Divergence는 모든 방향이 아닌 임의의 방향에 대한 projection에 대한 fisher divergence이다. 식을 보면 내적을 통해 projection하는 것을 확인할 수 있다. $$ \frac{1}{2} \mathbb{E}_{p_{\mathbf{v}}} \mathbb{E}_{p_{\text {data }}}\left[\left(\mathbf{v}^{\top} \nabla_{\mathbf{x}} \log p_{\theta}(\mathbf{x})-\mathbf{v}^{\top} \nabla_{\mathbf{x}} \log p_{\text {data }}(\mathbf{x})\right)^{2}\right] $$

Integration by parts에 의해 위와 마찬가지로 sliced fisher divergence는 다음과 같이 쓸 수 있다.

$$ \mathbb{E}_{p_{\mathbf{v}}} \mathbb{E}_{p_{\text {data }}}\left[\mathbf{v}^{\top} \nabla_{\mathbf{x}}^{2} \log p_{\theta}(\mathbf{x}) \mathbf{v}+\frac{1}{2}\left(\mathbf{v}^{\top} \nabla_{\mathbf{x}} \log p_{\theta}(\mathbf{x})\right)^{2}\right]+\text { const } $$

Sliced score matching은 score matching보다 빠르다는 이점이 있다. 위 식으로 학습을 하면 finite-sample estimator가 되므로 training data로 계산이 가능하며 random direction $p_{\mathbf{v}}$는 multivariate standard normal이나 multivarate Rademacher distribution과 같은 분포에서 sampling할 수 있다.

From Score Estimation to Sample Generation

Score matching을 통해서 주어진 sampling을 통해 score를 추정할 수 있게 되었다. 이제 score로 구성된 vector field를 추정할 수 있으니 gradient를 따라가게 되면 확률분포에서 high density region에서 데이터를 sample할 수 있다. 하지만 단순하게 gradient를 따라간다고 해서 좋은 sample이 나오지는 않는데 이는 확률분포에서 local optima로 도착할 수 있기 때문이다. 하지만 MCMC방법의 일종인 Langevin dynamics를 사용해 local optima에 갇히지 않고 분포를 탐색해 실제 underground distribution에서 sampling을 할 수 있다.

Langevin Dynamics Sampling

Langevin dynamics를 사용한 sampling은 다음의 과정을 따른다. 여기서 목적은 $\nabla_{\mathbf{x}} \log p(\mathbf{x})$만을 사용해서 $p(\mathbf{x})$로부터 sampling을 하는 것이다.

  1. Initialize $\tilde{\mathbf{x}}_{0} \sim \pi(\mathbf{x})$
  2. Repeat for $t \leftarrow 1, 2, \cdots, T$: $$\begin{aligned} \mathbf{z}_{t} & \sim \mathcal{N}(0, I) \cr \tilde{\mathbf{x}}_{t} & \leftarrow \tilde{\mathbf{x}}_{t-1}+\frac{\epsilon}{2} \nabla_{\mathbf{x}} \log p\left(\tilde{\mathbf{x}}_{t-1}\right)+\sqrt{\epsilon} \mathbf{z}_{t} \end{aligned}$$

우선 초기값을 뽑아낸 이후 gradient를 따라가는데 그냥 따라가는 것이 아닌 perturbed gradient를 따라간다는 점에서 차이가 있다. 식을 보면 gradient방향을 따라가면서 마지막에 gaussian noise에 의한 perturbation이 있음을 볼 수 있다. 마지막 항으로 인해 local optima에 갇히는 것을 억제해 줄 수 있다.

Score-Based Generative Modeling

지금까지 다룬 내용을 바탕으로 score-based generative model이 어떻게 sample을 만들어내는지를 정리할 수 있다.

/score-based-models/sbgm.jpg
From Yang Song's Blog

  1. 학습 데이터 $\{ \mathbf{x}_1, \mathbf{x}_2, \cdots, \mathbf{x}_{N} \} \stackrel{\text{i.i.d.}}{\sim} p_{\text{data}} (\mathbf{x}) $를 구성한다.
  2. Score matching을 사용해서 score function을 학습한다.
  3. 학습된 score function을 통해서 $\nabla_{\mathbf{x}} \log p_{\text{data}}(\mathbf{x})$를 추정할 수 있으므로 Langevin dynamics를 사용해 sample을 생성한다.

Pitfalls

Score-based generative model을 위한 이론적인 토대는 모두 마련되었다. 하지만 실제 생성을 위해서는 다음과 같은 점을 추가로 신경써주어야 한다.

Pitfall 1: Manifold Hypothesis

Data score가 well-defined되어 있지 않다. 실제 데이터는 Gaussian처럼 예쁘게 수렴하는 형태로 구성되지 않는다. 예를들어 반지형태로 실제 데이터분포가 있다고 생각해보면 반지 주변부에서 score는 몰리게되는데 이 분포가 심지어 좁은구간이라면 주변 gradient는 큰 값을 가지게 되고 이런 성질은 gradient-based 학습에서는 최적화하기 어려운 형태가 된다. 실제로, CIFAR-10과 같은 이미지 데이터에 대해 sliced score matching을 적용해 학습을 시킬 때 loss함수가 일정하게 떨어지는게 아니라 매우 불안정한 형태로 등락하는 그래프를 보여주었다고 한다.

Pitfall 2: Inaccurate Score Estimation in Low Data-Density Regions

앞서 사용한 gradient plot을 다시 살펴보자.

gradient_plot.png
Gradient Plot
당연하게도 data generating distribution의 밀도가 높은 곳은 sample이 많이 존재하기 때문에 이러한 영역에서의 score estimation을 비교적 정확하게 이루어진다. 하지만 반대로 확률밀도가 낮은 2사분면이나 4사분면과 같은 영역에서는 데이터 자체가 적게 존재하고 따라서 해당영역에서의 score estimation은 매우 부정확하다는 문제가 있다. Score estimation이 부정확하면 결국 Langevin dynamics과정도 부정확해지게 될 것이다.

Pitfall 3: Slow Mixing of Langevin Dynamics between Data Modes

다른 문제로는 mixture비율을 알 수 없다는 문제가 있다. 예를들어 다음과 같은 두 분포의 mixture로 구성된 분포가 있다고 해보자. $$p_{\text{data}}(\mathbf{x}) = \pi p_{1}(\mathbf{x}) + (1-\pi) p_{2}(\mathbf{x})$$ $p_{1}$ 근처에서는 $\nabla_{\mathbf{x}} \log p_{1}(\mathbf{x})$의 score가 추정될 것이고 $p_{2}$의 근처에서는 $\nabla_{\mathbf{x}} \log p_{2}(\mathbf{x})$에 가까운 score가 추정될 것이다. 문제는 score를 추정할 때 실제 분포에서는 중요한 역할을 끼친 각 분포의 가중치 $\pi$가 반영되지 않는다는 것이다.

Gaussian Perturbation

위의 문제를 다루는 방식으로 Gaussian perturbation이 유효했다고 한다. 아래 그래프는 sliced score matching에 아주 적은 Gaussian noise를 추가했을 때 학습이 훨씬 안정적으로 바뀌는 것을 보여준다.

/score-based-models/loss_fn.png
From Stefano Ermon's Seminar
또한 noise를 추가하게 되면 low density region에서의 sample을 더 얻을 수 있으므로 두 번째 pitfall이 개선되었다고 한다. 물론 noise를 너무 크게 주면 학습을 방해할 수가 있으므로 다양한 noise level에 대해 학습을 시키는게 효과적이었다고 한다.
/score-based-models/noise_step.png
From Stefano Ermon's Seminar
이와 같이 Langevin Dynamics과정에서 sampling을 할 때 처음에는 큰 $\sigma$값을 사용하고 점점 $\sigma$를 줄여가면서 하는 방법을 annealed langevin dynamics라고 한다. 결과적으로 annealed 방식을 사용하면 high density 근방에서 집중적으로 sampling되는 것을 줄여줌으로써 앞서 언급한 pitfall들을 완화할 수 있게 된다.

Joint Score Estimation

그렇다면 sampling은 annealed Langevin dynamics를 사용한다고 하더라도 score estimator도 이러한 noised sample에 대해 학습을 해야 올바른 score를 예측할 수 있을 것이다. 다양한 scale의 noise에 대해 학습시키기위해 sample에 대해 isotropic Gaussian noise $\mathcal{N}(0, \sigma_{i}^{2} I), i=1,2,\cdots,L$를 더하는 방식으로 접근한다. $\sigma$는 $L$로 갈 수록 커지는 형태로 $\sigma_{1} < \sigma_{2} < \cdots < \sigma_{L}$이다. $$p_{\sigma_{i}}(\mathbf{x})=\int p(\mathbf{y}) \mathcal{N}\left(\mathbf{x} ; \mathbf{y}, \sigma_{i}^{2} I\right) \mathrm{d} \mathbf{y}$$ Sample은 $\mathbf{x} \sim p(\mathbf{x})$에 noise $\mathbf{z} \sim \mathcal{N}(0, I)$를 더해 $\mathbf{x} + \sigma_{i} \mathbf{z}$로 만들면 된다. 다양한 noise level에 대해 score function을 학습시키기 위해 noise conditional score-based model $s_{\theta}(\mathbf{x}, i) \approx \nabla_{\mathbf{x}} \log p_{\sigma_{i}}(\mathbf{x})$를 학습시키게 되고 이 때 $\sigma_{i}$는 각각 다른 noise level을 의미하게 된다. 식을 보면 여타 conditioning model과 같이 $\sigma$에 conditioned 되어 있음을 볼 수 있다. 네트워크 구조에서는 $D$개의 input에 추가적으로 noise 수준을 결정하는 $\sigma$가 추가적인 입력으로 들어가게 된다. 따라서 $s_{\theta}(\mathbf{x}, i): \mathbb{R}^{D+1} \rightarrow \mathbb{R}^{D}$이 된다. $s_{\theta}(\mathbf{x}, i)$의 목적함수는 Fisher divergence의 noise scale에 대한 weighted sum으로 다음과 같이 정의한다.

$$ \sum_{i=1}^{L} \lambda(i) \mathbb{E}_{p_{\sigma_{i}}(\mathbf{x})}\left[ \lVert \nabla_{\mathbf{x}} \log p_{\sigma_{i}}(\mathbf{x})-\mathbf{s}_{\theta}(\mathbf{x}, i) \rVert _{2}^{2}\right] $$

$\lambda(i) \in \mathbb{R}_{>0}$이며 $\lambda(i) = \sigma_{i}^2$을 주로 사용한다. 이외 부분은 앞서 설명한 바와 마찬가지로 score matching을 기반으로 학습하면 된다. 이렇게 noise-conditional score-based model $s_{\theta}(\mathbf{x}, i)$를 학습하였다면 Langevin dynamics를 $i=L, L-1, \cdots, 1$의 순서로 실행해 sample을 생성할 수 있게 된다.

Tuning Score-based Models

다음은 Yang Song이 score-based generative model을 multiple noise scale로 학습할 때 추천한 사항이다.

  • $\sigma_{1} < \sigma_{2} < \cdots < \sigma_{L}$은 기하수열로 구성한다. $\sigma_{1}$은 충분히 작은 값에서 십작하여야 하며 $\sigma_{L}$은 데이터의 maximum pairwise distance 수준으로 결정되는 것이 좋다. $L$은 주로 수백 또는 수천 수준의 크기를 갖는다.
  • Score-based model $s_{\theta}(\mathbf{x}, i)$는 U-Net skip connection과 같은 모델을 사용한다.
  • Test time에는 score-based model의 가중치를 exponential moving average를 적용한다.

Concluding Remarks

아래는 score-based generative model로 생성한 이미지로 충분히 좋은 결과를 보여주고 있다.

/score-based-models/cifar10_large.gif
From Yang Song's Blog
/score-based-models/ncsnv2.jpg
Improved Techniques for Training Score-Based Generative Models

이 글에서는 score-based generative model이 근간으로하는 이론에 대해 알아보았다. Score-based generative model은 기존 GAN, VAE, Normalizing Flow에 이어 새로운 framework을 사용하는 generative model로 score-based를 기반으로 한 생성모델에 대한 논문이 증가하고 있는 추세이다. 이론적으로도 굉장히 흥미로운 접근이고 최근에 나오는 논문들은 기존 GAN이 이루어낸 상당한 수준의 결과물을 앞지른 것으로 보여진다. 또한 multi-modal leraning에서도 주목할 만한 결과를 보여주면서 생성모델의 가능성을 확대하고 있다.

Reference