Score-based Generative Model 파헤치기!

 

 

목차

• Generative Modeling Techniques
• Score Function, Score-based Models, and Score Matching
• Langevin Dynamics
• Score-based Generative Modeling
  • Naive (Unconditional) Score-based Generative Modeling
  • Score-based Generative Modeling with Multiple Noise Perturbations
• Score-based generative modeling with Stochastic Differential Equations (SDEs)
  • Perturbing data with an SDE (Training)
  • Reversing the SDE for sample generation

 

 

이전까지 대표적인 생성모델인 AutoRegressive Models(ARMs), Flow-based models, Variational Auto-Encoders (VAEs), Generative Adversarial Networks(GANs)를 살펴보았다. 이번 글에서는 최근에 diffusion-based models와 함께 관심을 끌고 있는 score-based model에 대해 알아보려 한다.

사실 score-based model이 처음 제안된 시기는 2019년(Generative Modeling by Estimating Gradients of the Data Distribution, NeurIPS), 2021년(Score-based Generative Modeling through Stochastic Differential Equations, ICLR)으로 diffusion(2015년 Deep Unsupervised Learning using Nonequilibrium Thermodynamics, ICML, 2020년 Denoising Diffusion Probabilistic Models, NeurIPS)보다 늦지만, 개인적으로 diffusion에 가장 관심이 있기도 하고, diffusion을 이해하는 방법 중에 score-based model을 이해해야 하는 부분이 있어 먼저 다뤄보고, 마지막으로 diffusion을 다뤄볼 것이다.

 

참고한 블로그 글은 아래와 같다.

https://yang-song.net/blog/2021/score/

Generative Modeling by Estimating Gradients of the Data Distribution | Yang Song

 

 

 

 

 

Generative Modeling Techniques

 

Fig 1. Representative deep generative models

 

이제까지 알아본 generative model 중, 최근까지도 자주 사용되는 모델은 Fig 1과 같다. 생성모델을 크게 나누면 likelihood를 사용하는지, 아닌지에 따라 다음과 같이 나눠볼 수 있다.

 

Likelihood-based models

Fig 2. Example of likelihood-based models

 

Likelihood-based models는 Data의 distribution을 직접 모델링하여 학습한다. 이에는 ARMs(AutoRegressive Models), Normalizing flows(Flow-based models), EBMs(Energy-Based Models), 그리고 VAEs(Variational AutoEncoders)가 포함된다.

이러한 생성모델의 한계점은 엄격한 restriction이 필요하거나, likelihood를 근사하기 위한 objective에 의존적이라는 것이다.

 

Implicit models

Fig 3. Example of implicit models

 

Implicit models는 모델이 샘플링 과정에서 data의 probability distribution을 implicit하게 표현한다. 이에는 GANs (Generative Adversarial Networks)가 대표적이다.

이 모델은 Adversarial training으로 인해 학습이 불안정하고, mode collapse 문제가 발생하기 쉽다.

 

Score-based models

Fig 4. Score function (vector field) and density funciton (contours) of a mixture of two Gaussians

위 두 유형과 달리, 이번 글에서 다룰 score-based model은 score function을 modeling 및 estimating하면서 여러 generative task에 적용한다.

Likelihood 기반 모델과 달리 계산 가능한 normalizing constant가 필요 없으며, adversarial loss에 비해 안정적인 score matching이라는 방법으로 학습한다. 또한, Normalizing flow model과 연관이 있어 정확한 likelihood 계산과 representation learning이 가능하다.

 

 

 

Score Function, Score-based Models, and Score Matching

 

주어진 데이터셋 $\{{\mathbf{x}}_1,{\mathbf{x}}_2,\dots ,{\mathbf{x}}_N\}$에 대한 data distribution이 $p(\mathbf{x})$라고 가정하면, generative modeling은 모델의 distribution을 $p(\mathbf{x})$에 맞추어 새로운 data point를 샘플링하는 것이다.

이때 probability density function (pdf)는 다음과 같이 정의한다.

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

• Unnormalized probabilistic model(energy-based model) $f_{\theta }(\mathbf{x})\in \mathbb{R}$ : Learnable parameter $\theta$로 parameterize한 함수
• $Z_{\theta }>0$ : $\theta$에 의존하는 Normalizing constant

 

그리고 이러한 모델은 maximum log-likelihood로 학습한다.

$\underset{\theta }{max}\sum _{i=1}^N\log p_{\theta }({\mathbf{x}}_i)$

 

이때, pdf curve 아래 면적의 합은 항상 1$\int p_{\theta }(\mathbf{x})d\mathbf{x}=1$)이어야 하므로, $p_{\theta }(\mathbf{x})$가 normalized pdf여야 하는데, 일반적인 $f_{\theta }(\mathbf{x})$에 대해 normalizing constant $Z_{\theta }$는 계산이 불가능하다. Likelihood-based model에서는 $Z_{\theta }$를 계산 가능하게 만들기 위해 모델 아키텍쳐를 제한하거나(ARM에서의 causal convolution, Normalizing flow에서의 invertible network), $Z_{\theta }$를 근사(VAE에서의 variational inference)한다. 이 과정에서 계산량이 많아진다. (자세한 과정은 각 링크를 참조하자.)

Score-based model $s_{\theta }(\mathbf{x})$에서는 density function 대신 probability distribution $p(\mathbf{x})$의 score function을 정의한다. (log pdf의 gradient 개념)

${\mathrm{\nabla }}_{\mathbf{x}}\log p(\mathbf{x})$

이 모델의 objective는 $s_{\theta }(\mathbf{x})\approx {\mathrm{\nabla }}_{\mathbf{x}}\log p(\mathbf{x})$, 즉 모델이 data distribution의 (log의) gradient를 예측하는 것이다.

PDF를 정의할 때 등장했던 energy-based model $f_{\theta }(\mathbf{x})$로 score-based model을 parameterize할 수 있다.

$s_{\theta }(\mathbf{x})={\mathrm{\nabla }}_{\mathbf{x}}\log p_{\theta }(\mathbf{x})=-{\mathrm{\nabla }}_{\mathbf{x}}{f}_{\theta }(\mathbf{x})-\cancel{{\mathrm{\nabla }}_{\mathbf{x}}\log Z_{\theta }}=-{\mathrm{\nabla }}_{\mathbf{x}}{f}_{\theta }(\mathbf{x})$

이 과정에서 constant가 미분 과정에서 소거되므로, VAE 등에서 계산이 불가능했던 원흉인 normalizing constant $Z_{\theta }$를 고려할 필요가 없다.

 

Score-based model을 학습할 때에는 model과 data distribution 사이의 Fisher divergence를 활용한다.

${\mathbb{E}}_{p(\mathbf{x})}[\Vert {\mathrm{\nabla }}_{\mathbf{x}}\log p(\mathbf{x})-s_{\theta }(\mathbf{x}){\Vert }_2^2]$

Fisher divergence는 GT data score와 score-based model의 예측 값 간의 L2 거리 개념이다. $s_{\theta }(\mathbf{x})$의 형태가 자유로우므로, 모델링이 flexible하다는 장점이 있다.

Data score ${\mathrm{\nabla }}_{\mathbf{x}}\log p(\mathbf{x})$를 모르므로 직접 계산을 할 수는 없으나, score matching이라는 방법으로 GT data score를 몰라도 위 식을 최소화할 수 있도록 한다.

Score matching object는 dataset 기반으로 추정이 가능하며 stochastic graident descent로 최적화할 수 있다. 이는 normalizing constant를 알고 있는 상태에서의 log-likelihood based model들의 objective와 같다.

 

 

 

 

 

Langevin Dynamics

 

Langevin Dynamics는 Fig 5와 같이 학습된 score-based model $s_{\theta }(\mathbf{x})\approx {\mathrm{\nabla }}_{\mathbf{x}}\log p(\mathbf{x})$를 활용하여 샘플링하는 방법이다.

 

Fig 5. Using Langevin dynamics to sample from a mixture of two Gaussians

 

이 방법은 iterative하게 작동하며, Markov Chain Monte Carlo (MCMC) procedure를 사용한다.

${\mathbf{x}}_{i+1}\leftarrow {\mathbf{x}}_i+ϵ{\mathrm{\nabla }}_{\mathbf{x}}\log p(\mathbf{x})+\sqrt{2ϵ}{\mathbf{z}}_i,i=,1,\cdots ,K$

• Initialization ${\mathbf{x}}_0\sim \pi (\mathbf{x})$ : 임이의 prior distribution에서 샘플링
• ${\mathbf{z}}_i\sim \mathcal{N}(\mathbf{0},\mathbf{I})$

이 과정에서 Regularization을 적용하면, 과정이 반복될수록($ϵ\to 0,K\to \mathrm{\infty }$) ${\mathbf{x}}_K$는 $p(\mathbf{x})$의 샘플에 수렴한다. 이를 통해 학습된 $s_{\theta }(\mathbf{x})$를 식에 대입하고 반복하여 샘플링을 수행한다.

 

 

 

Score-based Generative Modeling

 

 

Naive (Unconditional) Score-based Generative Modeling

 

Fig 6. Learning and sampling of naive score-based model

 

이 모델은 Fig 6과 같이 score matching으로 학습을, Langevin dynamics를 활용하여 샘플링을 수행한다.

그러나 한계점이 있는데, low density 지역(data point가 많이 없는 지역)에서 추정한 score function이 부정확하다는 것이다.

 

Fig 7. Estimated scores for a mixture of two Gaussians

 

Objective를 보면 그 이유를 알 수 있다.

${\mathbb{E}}_{p(\mathbf{x})}[\Vert {\mathrm{\nabla }}_{\mathbf{x}}\log p(\mathbf{x})-s_{\theta }(\mathbf{x}){\Vert }_2^2]=\int p(\mathbf{x})\Vert {\mathrm{\nabla }}_{\mathbf{x}}\log p(\mathbf{x})-s_{\theta }(\mathbf{x}){\Vert }_2^{2}d\mathbf{x}$

 

우변을 보면, L2 difference에 $p(\mathbf{x})$를 weighting하는 개념으로, data point 자체가 적을 경우 $p(\mathbf{x})$가 작으므로 값이 무시된다. 그리고 Langevin dynamics로 샘플링할 때 또한 첫 샘플링은 low density 지역에서 하는 경우가 많을텐데, score가 부정확하면 당연히 생성한 데이터의 quality가 매우 낮을 것이다.

이 문제를 해결하기 위해 Noise perturbation을 사용한다.

 

 

Score-based Generative Modeling with Multiple Noise Perturbations

 

Noise perturbation이란, data point에 perturbation을 적용한 후 score-based model을 학습하는 방법이다. 이 방법을 사용한 score-based model을 noise conditional score-based model $s_{\theta }(\mathbf{x},i)$이라 한다. NCSN(Noise Conditional Score Network)으로 parameterize하며, score function은 \(\nabla_\mathbf{x} \log p_{\sigma_i} ( \mathbf{x})\)이다.

 

Fig 8. Estimated scores for a mixture of two Gaussians perturbed by additional Gaussian noise

 

Fig 8에서와 같이, perturbation을 적용하면 data point의 분포가 더 퍼지게 되어 기존의 문제를 해결할 수 있게 된다. 하지만, perturbation을 얼마나 줄 것인가(noise의 scale)가 문제인데, 여러 scale의 noise를 동시에 적용하는 Multiple noise perturbation을 사용한다. (${\sigma }_1<{\sigma }_2<\cdots <{\sigma }_L$)

$p_{{\sigma }_i}(\mathbf{x})=\int p(\mathbf{y})\mathcal{N}(\mathbf{x};\mathbf{y},{\sigma }_i^2\mathbf{I})d\mathbf{y}$

• Multiple Gaussian noise $\mathcal{N}(\mathbf{0},{\sigma }_i^2\mathbb{I}),i=1,2,\cdots ,L$

 

샘플링 할 때는 $p(\mathbf{x})$에서 data point $\mathbf{x}$를 샘플링한 후, $\mathbf{x}+{\sigma }_i\mathbf{z}$를 계산한다. 이때 $\mathbf{z}$는 standard Gaussian $\mathcal{N}(\mathbf{0},\mathbf{I})$에서 샘플링한 값이다.

 

Score matching (Training) of noise conditional score-based model

이 모델의 training objective는 모든 noise${\sigma }_1,{\sigma }_2,\dots ,{\sigma }_L$에 대해 Fisher divergence의 weighted sum을 최소화하는 것이다.

$\sum _{i=1}^L\lambda (i){\mathbb{E}}_{p_{{\sigma }_i}(\mathbf{x})}[\Vert {\mathrm{\nabla }}_{\mathbf{x}}\log p_{{\sigma }_i}(\mathbf{x})-s_{\theta }(\mathbf{x},i){\Vert }_2^2]$

• $\lambda (i)\in {\mathbb{R}}_{>0}$ : (positive) weighting function으로, 보통 $\lambda (i)={\sigma }_i^2$로 사용한다.

 

Fig 9. Multiple scales of Gaussian noise (top) and jointly estimated score functions (bottom)

 

Annealed Langevin dynamics (Sampling) of noise conditional score-based model

샘플링 할 때에는 noise-conditional score-based model $s_{\theta }(\mathbf{x},i)$를 활용하여 Langevin dynamics를 $i=L,L-1,\cdots ,1$ 순서로 실행한다.

 

Fig 10. Annealed Langevin dynamics (sampling)

 

Fig 10과 같이, 샘플링을 진행할수록 점점 noise가 사라진다. (이는 diffusion에서 time step에 따라 denoising되는 것과 비슷하다!)

 

NCSN Implementation

• ${\sigma }_i$를 geometric progression(자세한 내용은 링크를 참조하자.)방법으로 정한다.
  • ${\sigma }_1$은 충분히 작게, ${\sigma }_L$은 모든 training data points 중에 가장 거리가 먼 두 data point 정도로 정한다.
  • $L$은 몇 백 ~ 몇 천 정도로 정한다.
• Score-based model $s_{\theta }(\mathbf{x},i)$는 U-Net 아키텍쳐를 사용하여 구현한다.
• Test 시 model의 weight에 exponential moving average를 적용한다.

 

 

 

Score-based generative modeling with Stochastic Differential Equations (SDEs)

 

Multiple noise를 일반화하여 scale 개수를 무한히 늘린다. 이러한 모델을 time-dependent score-based model $s_{\theta }(\mathbf{x},t)$라 한다. 이때 학습 과정과 샘플링 과정을 모델링하기 위해 stochastic differential equation(SDE)을 활용한다.

이를 통해 샘플 퀄리티를 높이고, 정확한 log-likelihood를 계산할 수 있으며, controllable generation이 가능하다. 이는 diffusion의 forward process 및 reverse process와 같은 개념이 된다. Diffusion 내용은다음 글에서 따로 다룰 것이다.

 

 

Perturbing data with an SDE (Training)

 

Fig 11. Perturbing data

 

Data에 무수히 많은 noise를 가하는 과정을 SDE process라 하고, SDE의 해를 구하는 과정이다. 이는 Fig 11과 같이 diffusion process를 연상시킨다.

$d\mathbf{x}=f(\mathbf{x},t)dt+g(t)d\mathbf{w}$

• $f(\cdot ,t):{\mathbb{R}}^d\to {\mathbb{R}}^d$ : Drift coefficient (벡터함수)
• $g(t)\in \mathbb{R}$ : Diffusion coefficient (실수함수)
• $\mathbf{w}$ : Standard Brownian motion ($d\mathbf{w}$ : noise의 극소(infinitesimal)량)
• 해 : Random variables의 continuous collection $\{\mathbf{x}(t){\}}_{t\in [0,T]}$
  • Random variable $\mathbf{x}(t)$는 사진에 보이는 경로의 점을 의미한다.

 

SDE는 다양하게 설정할 수 있다. 예를 들어, $d\mathbf{x}=e^{t}d\mathbf{w}$로 설정할 경우, 평균($f$)이 0이고, exponential하게 증가하는 variance($g$) 사용하는 것을 의미한다.

 

Training objective는 다음과 같이 Fisher divergence에 time step $t$에 관한 항이 추가된다.

${\mathbb{E}}_{t\in \mathcal{U}(0,T)}{\mathbb{E}}_{p_t(\mathbf{x})}[\lambda (t)\Vert {\mathrm{\nabla }}_{\mathbf{x}}\log p_t(\mathbf{x})-s_{\theta }(\mathbf{x},t){\Vert }_2^2]$

• $\mathcal{U}(0,T)$ : Time interval $[0,T]$에 대한 uniform distribution
• $\lambda :\mathbb{R}\to \mathbb{R}>0$ : Positive weighting function
  • 보통 $\lambda (t)\propto 1/\mathbb{E}[\Vert \mathrm{\nabla }\mathbf{x}(t)\log p(\mathbf{x}(t)|\mathbf{x}(0)){\Vert }_2^2]$ 으로 설정하며, time step에 따라 달라지는 score matching loss의 balance를 맞춰준다.

 

$\lambda (t)=g^2(t)$일 때 $\lambda$를 likelihood weighting function이라 하고, regularity condition을 추가하면 objective는 $p_0$와 $p_{\theta }$ 사이의 KL divergence 개념과 연결된다. 즉, KL divergence를 minimize하는 게 곧 likelihood를 maximize하는 것(likelihood-based model을 학습하는 것)과 같은 개념이다.

$\begin{array}{rl}& D_{\text{KL}}(p_0(\mathbf{x})\Vert p_{\theta }(\mathbf{x}))\le \\ & \frac{T}{2}{\mathbb{E}}_{t\in \mathcal{U}(0,T)}{\mathbb{E}}_{p_t(\mathbf{x})}[\lambda (t)\Vert {\mathrm{\nabla }}_{\mathbf{x}}\log p_t(\mathbf{x})-s_{\theta }(\mathbf{x},t){\Vert }_2^2]+D_{\text{KL}}(p_T\Vert \pi )\end{array}$

 

이전에 알아본 noise scale 개수가 유한할 때와 비교해보자면, $t$는 Noise scale 개수가 유한할 때의 $i$와, $p_t(\mathbf{x})$는 $p_{{\sigma }_i}(\mathbf{x})$와 같다.

$t=0$일 때는 $p_0(\mathbf{x})=p(\mathbf{x})$로, perturbation이 적용되지 않은 (original) data distribution이며, $t=T$일 때에는 terminal distribution $p_T(\mathbf{x})=\pi (\mathbf{x})$으로, 계산 가능한 noise distribution (prior distribution)을 사용한다.

 

 

Reversing the SDE for sample generation

 

Fig 12. Denoising data

 

샘플링 과정은 Fig 12에서 보는 바와 같이, diffusion 모델의 reverse process와 같다. SDE도 time step $t$의 역순으로 푼다.

$d\mathbf{x}=[f(\mathbf{x},t)-g^2(t){\mathrm{\nabla }}_{\mathbf{x}}\log p_t(\mathbf{x})]dt+g(t)d\mathbf{w}$

• $dt$ : Negative infinitesimal time step
• ${\mathrm{\nabla }}_{\mathbf{x}}\log p_t(\mathbf{x})$ : $p_t(\mathbf{x})$의 score function

 

학습된 score-based model $s_{\theta }(\mathbf{x},t)\approx {\mathrm{\nabla }}_{\mathbf{x}}\log p_t(\mathbf{x})$를 위 reverse SDE 식에 대입하여 샘플링을 진행한다. $\mathbf{x}(T)\sim \pi$에서 시작하여 $\mathbf{x}(0)$까지 계산할 수 있다. Reverse SDE를 푸는 자세한 과정은 링크를 참조하자.

Score-based model $s_{\theta }(\mathbf{x},t)$이 잘 학습되었다면 $p_{\theta }\approx p_0$을 만족한다. 즉, data distribution $p_0$에서의 샘플과 계산을 통해 얻은 $\mathbf{x}(0)$이 유사하다.

 

Forward SDE, reverse SDE 전체 과정은 다음과 같이 그림으로 나타낼 수 있다.

 

Fig 13. Forward SDE, reverse SDE

 

Fig 13에서 흰색 선은 ODE(ordinary differential equation, 상미분방정식)를 활용해서, 즉 forward process와 reverse process를 deterministic하게 모델링한 경우이다. Continuous한 normalizing flow로 해석해볼 수 있는데, 여기서는 자세히 다루지 않겠다. 더 자세한 내용은 링크를 참조하자.