이 내용은 Energy-based models에 대한 기본적인 개념 및 training 방법을 알고 있다는 전제하에 작성되었습니다. 이에 대한 내용을 모르시는 분들은 아래 "Energy-Based Models (EBMs)이란? (1)"편에 자세히 적어두었으니, 먼저 읽으시고 이번 내용을 보시면 더 수월하게 이해하실 수 있을 겁니다.
Energy-Based Models (EBMs)이란? (1) - 개념, training, sampling - 웅 공부방
Energy-based generative Models이 나오기 이전에도 Autoregressive models, Normalizing flow models, Variational autoencoders(VAE), GANs 등 많은 generative mode이 존재했다. 다만, 이 모델들은 모두 치명적인 단점을 갖고 있다. A
woo-ng.tistory.com
Energy-based models은 $Z({\theta})$가 intractable 하기 때문에 구하는 것이 쉽지 않다. 하지만 다행인 점은 $Z({\theta})$를 구하지 않고도 training을 할 수 있다는 것이다. 이 과정에서 $p_{\theta}$에서 sampling 된 데이터가 필요하고, $p_{\theta}$ 또한 intractable 하기에 MH-MCMC, unadjusted Langevin MCMC와 같은 방법을 이용해야 한다. 이 방법들의 치명적인 단점은 1개의 sample을 얻는데도 시간이 오래 걸린다는 것이다. 따라서 전체 training 또한 오래걸린다는 단점을 갖고 있다. 그래서 이를 극복하기 위해 sampling을 하지 않는 여러 방법론이 제시되었으며, 본 포스팅에서는 3가지 방법론에 대해 설명할 것이다.
Score Matching
Score Matching 방법론은 말 그대로 score를 서로 동일하게 혹은 비슷하게 맞추어준다는 뜻이다. 여기서 score는 정확히 말하자면 score function을 의미한다. 그렇다면 score function은 무엇일까?
Score Function
score function은 어떤 함수 혹은 분포 $p_{\theta}$의 $\log$의 gradient 값이다. 이를 수식으로 표현하면 아래와 같다.
$$ \mathbf{s}_{\theta}(\mathbf{x}) = \nabla_{\mathbf{x}} \log p_{\theta}(\mathbf{x}) $$
예를 들어 $p_{\theta}(\mathbf{x})$를 2-D gaussian mixture model이라고 가정을 했을때, $p_{\theta}(\mathbf{x})$와 $\mathbf{s}_{\theta}(\mathbf{x})$를 그래프에 그려보면 다음과 같다.
Figure 1에서 contour는 $p_{\theta}(\mathbf{x})$를 나타낸 것이고 빨간색 화살표는 $\mathbf{s}_{\theta}(\mathbf{x})$를 나타낸 것이다. 2개의 gaussian model 모두 평균에서 가까운 $\mathbf{x}$ 일수록 score function의 방향이 평균을 가리키는 것을 확인할 수 있다. 또한 평균에서의 값이 최대치이기 때문에 평균에서 score의 화살표가 없는 것을 확인할 수 있다.
Fisher divergence
Fisher divergence는 KL-divergence처럼 두 확률분포(pdf, pmf)의 차이를 측정하는 지표이다. 하지만 다른 점은 KL-divergence는 두 확률분포의 값을 직접 비교하지만, Fisher divergence는 두 확률분포의 score function 값을 비교한다는 데에 있다. Fisher divergence를 수식으로 표현하면 다음과 같다.
$$ D_{F}(p, q) = \frac{1}{2} \mathbf{E}_{\mathbf{x} \sim p} \left[||\nabla_{\mathbf{x}} \log p(\mathbf{x}) - \nabla_{\mathbf{x}} \log q(\mathbf{x})||_{2}^{2} \right] $$
위의 식을 보면 두 분포 $p$, $q$의 score 값의 $L_{2}$ norm 값이 작을수록 divergence의 값이 작아지는 것을 확인할 수 있다. 그러므로 Fisher divergence는 $p$, $q$의 score가 서로 비슷하다면, 이들 분포도 비슷할 것이다는 것이다는 아이디어를 기반으로 하고 있음을 알 수 있다.
Score matching
Fisher divergence를 $p_{\text{data}}(\mathbf{x})$, $p_{\theta}(\mathbf{x})$에 대해 적용하면 Energy-based models을 training하기 위한 loss를 만들 수 있다. $p$에 $p_{\text{data}}(\mathbf{x})$, $q$에 $p_{\theta}(\mathbf{x})$를 적용하면 아래와 같이 나타낼 수 있다.
$$ \frac{1}{2} \mathbf{E}_{\mathbf{x} \sim p_{\text{data}}} \left[||\nabla_{\mathbf{x}} \log p_{\text{data}}(\mathbf{x}) - \nabla_{\mathbf{x}} \log p_{\theta}(\mathbf{x})||_{2}^{2} \right] $$
우리가 갖고 있는 모든 training데이터에 대해 Fisher divergence를 계산하여 parameter를 업데이트하면 된다. 하지만 여기서 MLE에서 발생했던 문제점이 똑같이 발생한다. $p_{\text{data}}$를 알지 못하는데 $\nabla_{\mathbf{x}} \log p_{\text{data}}(\mathbf{x})$를 어떻게 알 수 있을까? 우선 Fisher divergence 파트를 전개해 보자. 논의를 쉽게 하기 위해 $\mathbf{x}$를 1-d라고 가정하면, Fisher divergence를 아래와 같이 전개할 수 있다.
\begin{align*} &\frac{1}{2} \mathbf{E}_{x \sim p_{\text{data}}} \left[(\nabla_{x} \log p_{\text{data}}(x) - \nabla_{x} \log p_{\theta}(x))^{2} \right] \\ =& \frac{1}{2} \int p_{\text{data}}(x) \left[(\nabla_{x} \log p_{\text{data}}(x) - \nabla_{x} \log p_{\theta}(x))^{2} \right] \\ =& \frac{1}{2} \int p_{\text{data}}(x)(\nabla_{x} \log p_{\text{data}}(x))^{2}dx + \frac{1}{2} \int p_{\text{data}}(x) (\nabla_{x} \log p_{\theta}(x))^{2}dx \\ -& \int p_{\text{data}}(x) \nabla_{x} \log p_{\text{data}}(x) \nabla_{x} \log p_{\theta}(x)dx \end{align*}
위 식을 보았을 때, Fisher divergence를 전개하면 총 3개의 term이 생기는 것을 확인할 수 있다. 첫 번째 term은 직접 구할 수는 없지만 모델 parameter $\theta$의 변화와 무관한 상수이므로 무시할 수 있다. 두 번째 term은 모델의 score function이므로 쉽게 구할 수 있다. 문제는 마지막 term이다. 문제가 되는 이유는 $p_{\text{data}}$의 score 값을 직접 구할 수 없기 때문이다. 그렇기에 추가적인 전개가 더 필요하다. 마지막 term을 계산하여 $p_{\text{data}}$의 score를 제거한 식은 다음과 같다.
\begin{align*} & - \int p_{\text{data}}(x) \nabla_{x} \log p_{\text{data}}(x) \nabla_{x} \log p_{\theta}(x)dx \\ = & - \int p_{\text{data}}(x) \frac{1}{p_{\text{data}}(x)} \nabla_{x} p_{\text{data}}(x) \nabla_{x} \log p_{\theta}(x)dx \\ = & - \int \nabla_{x} p_{\text{data}}(x) \nabla_{x} \log p_{\theta}(x)dx \\ = & \;\; -[p_{\text{data}}(x)\nabla_{x} \log p_{\theta}(x)]_{-\infty}^{\infty} + \int p_{\text{data}}(x) \nabla_{x}^{2} \log p_{\theta}(x)dx \\ = & \int p_{\text{data}}(x) \nabla_{x}^{2} \log p_{\theta}(x)dx \end{align*}
$p_{\text{data}}$가 ${-\infty}$, ${\infty}$에서 0의 값을 갖는다고 가정하면, 최종적으로 위의 식을 얻을 수 있다. 위 식을 보면 $p_{\text{data}}$의 score의 값을 직접 구하지 않고 $p_{\theta}(\mathbf{x})$의 score 값의 미분 형태를 통해 3번째 term을 구할 수 있다는 것을 확인할 수 있다. 최종적으로 $\mathbf{x}$가 1-d일 때, Fisher divergence(score matching loss)를 두 번째, 세 번째 term만 살려서 다음과 같이 쓸 수 있다.
$$ \frac{1}{2} \mathbf{E}_{x \sim p_{\text{data}}} \left[(\nabla_{x} \log p_{\text{data}}(x) - \nabla_{x} \log p_{\theta}(x))^{2} \right] \\ \approx \mathbf{E}_{x \sim p_{\text{data}}} \left[\frac{1}{2} (\nabla_{x} \log p_{\theta}(x))^{2} + \nabla_{x}^{2} \log p_{\theta}(x) \right] $$
우리는 $p_{\text{data}}$의 score 값을 몰라도 $p_{\text{data}}$를 따르는 training data를 통해 Fisher divergence를 간접적으로 계산할 수 있다! 즉, training시 정확한 Fisher divergence의 값을 evaluation 하면서 Fisher divergence을 minimization 하는 것은 불가능하지만, training 목표가 정확한 값을 얻는 것이 아닌 단순히 Fisher divergence을 minimization하는 것이기 때문에 위 식을 minimization 하면 된다. 지금 예시에서 $\mathbf{x}$를 다차원으로 확장하면 다음과 같이 쓸 수 있다.
$$ \frac{1}{2} \mathbf{E}_{\mathbf{x} \sim p_{\text{data}}} \left[||\nabla_{\mathbf{x}} \log p_{\text{data}}(\mathbf{x}) - \nabla_{\mathbf{x}} \log p_{\theta}(\mathbf{x})||_{2}^{2} \right] \\ = \mathbf{E}_{\mathbf{x} \sim p_{\text{data}}} \left[\frac{1}{2} ||\nabla_{\mathbf{x}} \log p_{\theta}(\mathbf{x})||_{2}^{2} + \text{tr}(\nabla_{\mathbf{x}}^{2} \log p_{\theta}(\mathbf{x})) \right]$$
이제 $p_{\theta}$는 Energy-based models, mini-batch data samples $\mathbf{x}_{1}, \dots\,, \mathbf{x}_{n} \sim P_{\text{data}}$라고 가정하자. 이를 이용하여 loss를 계산한 식은 아래와 같다.
$$ \frac{1}{n} \sum_{i=1}^{n} \left[\frac{1}{2} ||\nabla_{\mathbf{x}} \log p_{\theta}(\mathbf{x}_{i})||_{2}^{2} + \text{tr}(\nabla_{\mathbf{x}}^{2} \log p_{\theta}(\mathbf{x}_{i})) \right] \\ = \frac{1}{n} \sum_{i=1}^{n} \left[\frac{1}{2} ||\nabla_{\mathbf{x}} f_{\theta}(\mathbf{x}_{i})||_{2}^{2} + \text{tr}(\nabla_{\mathbf{x}}^{2} f_{\theta}(\mathbf{x}_{i})) \right] $$
최종적으로 우리가 최적화를 해야 하는 식은 위의 식이다. 우리는 이전까지 Energy-based models에서 MLE를 이용하여 $p_{\theta}$를 maximize 하기 위해 $p_{\theta}$에서 sampling 한 데이터가 필요했다. 하지만 score function을 이용한 Fisher divergence를 통해 얻은 score matching loss인 위의 식을 보았을 때, 더 이상 training을 위해 $p_{\theta}$에서 sampling을 하는 것이 필요 없다는 것을 알 수 있다. 그렇기에 training 속도 측면에서 sampling 하는 것보다 훨씬 더 빠른 속도를 기대할 수 있다. 하지만 이 방법론에도 문제가 있다. 바로 $\nabla_{\mathbf{x}}^{2} f_{\theta}(\mathbf{x}_{i})$(Hessian)를 구하는 과정이 computationally expensive하다는 것이다. 어떤 함수의 Hessian을 구하면 변수가 $n$차원일 때, $(n \times n)$ matrix가 나온다. 즉, 차원수가 늘어날수록 계산량이 증가하며 이는 모델의 크기를 늘려 capacity를 증가시키는데 있어 치명적인 단점이 된다. 이를 극복하여 model을 더욱 scalable하게 하는 방법으로는 대표적으로 Denoising score matching이 있다.
Noise Contrastive Estmiation (NCE)
Noise Contrastive Estimation (NCE)는 Score Matching과 같이 $p_{\theta}$에서 sampling하는 과정을 없애는 방법이다. NCE는 GAN처럼 Discriminator를 학습시키는 방법을 이용하여 $p_{\theta}$와 $p_{\text{data}}$를 implicit 하게 근사한다. NCE와 GAN이 다른 점은 NCE는 Discriminator만 학습시키기에 GAN처럼 minmax problem이 아닌 max problem이다.
NCE는 noise distribution $p_{n}(\mathbf{x})$가 필요하다. $p_{n}(\mathbf{x})$는 간단한 분포 이어야 한다. 즉, density 값을 알기 쉬워야 하며(tractable), sampling 하는 것 또한 쉬워야 한다. 따라서 gaussian distribution 같은 분포들을 noise distribution에 이용할 수 있다. 또 Discriminator를 학습시켜야 하므로 $D_{\theta} \in [0, 1]$가 필요하다. 이제 GAN의 Discriminator를 학습할 때 이용하는 loss를 그대로 이용하면 다음과 같이 쓸 수 있다.
$$ \max_{\theta} \; \mathbf{E}_{\mathbf{x} \sim p_{\text{data}}} \left[\log D_{\theta}(\mathbf{x}) \right] + \mathbf{E}_{\mathbf{x} \sim p_{n} } \left[\log (1 - D_{\theta}(\mathbf{x})) \right] $$
이때 Discriminator가 완벽히 학습되어 parameter $\theta^{\star}$에 최댓값을 갖는다고 할 때, Discriminator의 값은 다음과 같다.
$$ D_{\theta^{\star}}(\mathbf{x}) = \frac{p_{\text{data}}(\mathbf{x})}{p_{\text{data}}(\mathbf{x})+ p_{n}(\mathbf{x})} $$
NCE에서는 $p_{\theta}$가 $p_{\text{data}}$를 implicit 하게 배우기 위해 Discriminator를 다음과 같이 모델링한다.
$$ D_{\theta}(\mathbf{x}) = \frac{p_{\theta}(\mathbf{x})}{p_{\theta}(\mathbf{x})+ p_{n}(\mathbf{x})} $$
이렇게 했을 때, optimal parameter $\theta^{\star}$에서 Discriminator의 값을 나타내면 다음과 같다.
$$ D_{\theta^{\star}}(\mathbf{x}) = \frac{p_{\theta^{\star}}(\mathbf{x})}{p_{\theta^{\star}}(\mathbf{x})+ p_{n}(\mathbf{x})} = \frac{p_{\text{data}}(\mathbf{x})}{p_{\text{data}}(\mathbf{x})+ p_{n}(\mathbf{x})}$$
위 식을 보면 Discriminator를 학습시키면 $p_{\theta^{\star}}$와 $p_{\text{data}}$가 implicit 하게 같아지는 것을 기대할 수 있다. 즉, MLE처럼 $p_{\text{data}}$와 $p_{\theta}$의 차이가 얼마나 작아졌는지 explicit 하게 확인할 수는 없지만 학습이 끝난 후에는 $p_{\theta^{\star}}$가 $p_{\text{data}}$를 잘 묘사할 것임을 알 수 있다. 이제 $p_{\theta}$를 Energy-based model로 대체하여 학습을 진행하면 된다. 하지만 여기서 늘 그랬듯 $Z({\theta})$를 계산할 수 없다는 문제가 발생한다. 그래서 NCE는 $Z({\theta})$를 $\theta$처럼 새로운 learnable parameter $Z$로 간주한다. 즉, 이제는 $Z$ 값이 Energy-based models에서 normalizing constant의 역할을 수행하는 것이 아닌 자유로운 parameter가 되는 것이다. 이를 반영하여 Discriminator를 다시 나타내면 다음과 같다.
$$ D_{\theta, Z}(\mathbf{x}) = \frac{\frac{\exp (f_{\theta}(\mathbf{x}))}{Z}}{\frac{\exp (f_{\theta}(\mathbf{x}))}{Z} + p_{n}(\mathbf{x})} $$
이를 반영하여 NCE objective를 나타내면 다음과 같다.
$$ \max_{\theta, Z} \; \mathbf{E}_{\mathbf{x} \sim p_{\text{data}}} \left[\log D_{\theta, Z}(\mathbf{x}) \right] + \mathbf{E}_{\mathbf{x} \sim p_{n} } \left[\log (1 - D_{\theta, Z}(\mathbf{x})) \right] $$
mini-batch data samples $\mathbf{x}_{1}, \dots\,, \mathbf{x}_{n} \sim p_{\text{data}}$이고, noise samples $\mathbf{y}_{1}, \dots\,, \mathbf{y}_{n} \sim p_{n}$라고 가정하고 위 objective를 전개하면 아래와 같이 나타낼 수 있다.
\begin{align*} \frac{1}{n} \sum_{i=1}^{n} \left[ f_{\theta}(\mathbf{x}_{i}) - \log (\exp(f_{\theta}(\mathbf{x}_{i}) + Zp_{n}(\mathbf{x}_{i})) \\ \quad\quad \quad\quad\quad\quad+ \log Zp_{n}(\mathbf{y}_{i}) - \log (\exp(f_{\theta}(\mathbf{y}_{i}) + Zp_{n}(\mathbf{y}_{i}) \right] \end{align*}
이 식을 maximize 하는 $\theta^{\star}$, $Z^{\star}$를 찾았다고 가정했을 때, 우리는 $p_{\theta^{\star}, Z^{\star}} = \frac{\exp (f_{\theta^{\star}}(\mathbf{x}))}{Z^{\star}}$가 $p_{\text{data}}$를 근사하길 바란다. 하지만 NCE는 $Z$ 값을 자유로운 learnable parameter로 두었기 때문에 유한한 training data로 학습했을 때, $Z^{\star}$가 normalizing constant의 역할을 하지 못할 가능성이 있다. 그럼에도 불구하고 NCE 또한 $p_{\theta}$ 에서의 sampling 과정을 없앴기에 학습 속도 관점에서는 장점을 갖고 있다고 할 수 있다.
Flow Contrastive Estmiation (FCE)
FCE는 NCE의 연장선이라고 생각하면 된다. NCE에서 문제점은 noise distribution $p_{n}(\mathbf{x})$에 있다. 만약 $p_{n}(\mathbf{x})$이 gaussian distribution처럼 단순한 모델로 놓는다면, Discriminator 입장에서 데이터 $\mathbf{x}$가 noise인지 아닌지 classification 하기 매우 쉬워진다. 이 말은 Discriminator를 학습하여 얻은 우리 모델의 분포 $p_{\theta}$ 역시 단순한 분포가 되기 쉽다는 것이다. 왜냐하면 $p_{\theta}$ 단순한 gaussian 분포를 가져도 Discriminator 입장에서는 충분히 우리가 가진 training data를 classification이 가능하기 때문이다. 예시를 통해 이해해 보자.
위 그림에서 볼 수 있듯 $\mathbf{x}$를 1차원 data $x$로 가정하고, $p_{n}(x) = \mathcal{N}(0, 1)$라고 생각해보자. 분홍색으로 표시된 점이 우리가 가지고 있는 training data이고 보라색으로 표시된 점이 noise data라고 할 때, Discriminator는 $p_{\theta}$를 noise distribution과 비슷한 단순한 분포인 gaussian 분포 $\mathcal{N}(3, 1)$여도 우리가 갖고 있는 데이터에 대해 꽤나 좋은 성능을 낸다.
NCE는 Discriminator를 학습한 후 얻은 $p_{\theta}$를 $p_{\text{data}}$로 간주하기 때문에 이 예제에서는 $p_{\text{data}} = \mathcal{N}(3, 1)$가 된다. 하지만 실제 $p_{\theta}$인 파란색 분포를 보면 그렇지 않다는 것을 알 수 있다. 이렇듯 noise distribution이 단순하다면 classification 문제가 쉬워지고, 결과적으로 $p_{\theta}$ 또한 단순한 분포가 될 확률이 높다.
이 문제는 간단한 noise distribution에 의해 발생한 것이므로 noise distribution을 복잡하게 해서 classification이 어렵게 만들면 된다. 즉, noise distribution을 $p_{\text{data}}$와 비슷하게 만들면 된다. 그래서 FCE에선 noise distribution을 normalizing flow model을 이용한다. normalizing flow model은 여러 번의 합성과정을 통해 복잡한 분포를 만들 수 있고, likelihood evaluation도 쉽기 때문에 sampling 또한 쉽다. Normalizing flow model의 parameter를 $\phi$로 설정하고 Discriminator를 나타내면 다음과 같다.
$$D_{\theta, Z, \phi}(\mathbf{x}) = \frac{\frac{\exp (f_{\theta}(\mathbf{x}))}{Z}}{\frac{\exp (f_{\theta}(\mathbf{x}))}{Z} + p_{n, \phi}(\mathbf{x})}$$
FCE의 objective function은 다음과 같다.
$$ \min_{\phi} \max_{\theta, Z} \; \mathbf{E}_{\mathbf{x} \sim p_{\text{data}}} \left[\log D_{\theta, Z, \phi}(\mathbf{x}) \right] + \mathbf{E}_{\mathbf{x} \sim p_{n, \phi} } \left[\log (1 - D_{\theta, Z, \phi}(\mathbf{x})) \right] $$
noise distribution에서 나온 data가 noise인지 아닌지 구분하기 어렵게 만들어야 하므로, 원래 NCE에서의 max problem이 minmax problem으로 바뀌었다. 이 objective function을 최적화한다면 단순한 noise distribution을 이용했을때 보다 classification 문제가 어려워졌기에 $p_{\theta}$가 더 복잡한 모양을 가질 수 있다.
Summary
Energy-based models을 훈련하는 과정에서의 bottleneck이라고 할 수 있는 sampling과정을 제거한 training 전략 3가지를 보았다. 3가지 모두 sampling 하는 과정을 없애 훈련 속도 측면에서 상당한 이점을 갖고 있지만, 각자의 단점도 갖고 있다. Score Matching의 경우 모델의 Hessian을 구해야 해서 모델을 scaling 하는데 제약이 있고, NCE의 경우는 학습 후 얻게 되는 $Z$ 값이 normalizing constant 역할을 하지 못할 가능성이 존재하고 noise가 단순할 경우 모델의 분포가 실제 분포와 전혀 맞지 않을 가능성 또한 존재한다.
Reference
[1] Stanford CS236, lecture 12
'Generative AI' 카테고리의 다른 글
| Energy-Based Models (EBMs)이란? (1) - 개념, training, sampling - 웅 공부방 (1) | 2025.01.05 |
|---|