Variational Inference (VAE)

JoongHyun's Blog·2024년 3월 30일

variational autoencoder(VAE)를 2가지 관점에서 살펴보았다.
1. Explicit generative model 관점
2. Energybased model 관점

Explicit Generative Model

먼저, 첫째 관점은 어떤 샘플 space로 부터 무언가를 생성하고 싶다는것이다. 고양이 이미지로 예시를 들어 생각해보자. 어떤 고양이의 분포를 가진 이미지 포인트 X가 존재한다고 가정하고 X같은 포인트를 뽑을 확률을 높이는 게임을 해보자. 그러나 그런 이미지 분포를 가진 데이터 포인트가 존재한다고 해도 그걸 바로 샘플링하기 어려운 상황이다. 그래서 우리는 최고의 function approximator인 neural network를 사용해서 어떤 latent space Z에서부터 샘플링을 해서 f(z)를 통해 고양이를 만들것이다. 이걸 수식으로 생각해보면(단 z는 continous variable이다.)

variational approximation of marginalization over Z

P(X) = \int_Z P(X,Z)dZ

하지만 여기서 Z는 무한대의 sample을 가지기 때문에 intractable 하다. 그래서 우리는 여기에 variational inference라는 걸 사용해서 P(X)를 최대한 근사화 시키는 분포를 찾을것이다. variational variable Q를 도입하면 다음과 같아진다.

\log \int_Z P(X,z)dz = \log \int_Z \frac{Q(Z|X)}{Q(Z|X)}P(X,Z)dZ

여기에 Jensen's inequality를 적용하면 log는 concave한 함수이므로,

ELBO

Jensen's inequality

\geq \int_Z Q(Z|X) \log \frac{P(X,Z)}{Q(Z|X)}dZ

위처럼 나타낼 수 있고 이걸 Evidence Lower Bound(ELBO)라고 부른다. 이 ElBO를 maximize한다면 하한을 최대화하도록 optimize하는것과 같다.조금더 변현시키면 Expectation termr과 KL divergence term으로 나눠진다.

= \int_Z Q(Z|X) \log P(X|Z)dZ + \int_Z Q(Z|X) \log \frac{P(Z)}{Q(Z|X)}dZ \\ = E_{Z \ Q(Z|X)}[\log P(X|Z)] -D_{KL}(Q(Z|X)||P(Z))

Reconstruction (Expectation term)

E_{Z \ Q(Z|X)}[\log P(X|Z)]

첫번째 term은 reconstruction term으로 생각할 수 있고, 첫번째 문제점으로 말한 intractable한 문제를 해결할 수 있다. 왜냐하면 우리는 이제 무한한 Z space가 아닌 variational variable로 도입한 Q에 의해서 Z의 공간을 제한하고 있기때문이다. 그래서 이제 Q를 Monte-carlo sampling을 통해서 Q의 pdf를 알 수 있고 그를 통해 근사시킬 수 있다.

Regularizer (KL divergence)

-D_{KL}(Q(Z|X)||P(Z))

여기에 우리는 P(Z)는 N(0,I)을 따를것이라는 explicit한 prior를 준다. 또한, 계산의 편리성을 위해 Q분포 역시 P와 conjugacy function 중에서 선택한다. 그중에서 multivariate gaussian Distribution을 따른다고 가정하고, covariance matrix의 offdiagonal term은 0이라고 가정한다. 즉 분산값만 가지는 diagonal Matrix이다. 그래서 다차원 Z를 시각화 해서 볼 수 있다면 어떤 평균 값에서 각 평균에서 동일하게 떨어진 분포가 Z에서 샘플링 될것이다.

Energy Based Model

에너지 관점도 위 상황을 똑같이 가정한다. 다만 다른점은 바라보는 시각이다. minimize free engery를 목표로 잡아보자. 그걸 F로 표현하겠다.

열역학의 partition function 형태로 나타내면 다음과 같다.

E(x,z) = C(x,DEC(z)) \\ F(x) = -\frac{1}{\beta} \log \int_Z \exp^{-\beta E(x,z)}

아까와 마찬가지로 위 적분은 numerically 게산 할 수 없다(intractable). 의미를 보면 평균 에너지값을 objective를 세웠다. 그리고 똑똑한 사람들이 variational approximation을 도입했기에 마찬가지로 똑같이 나타내보자.

F(x) = -\frac{1}{\beta} \log \int_Z Q(z|x)\frac{\exp^{-\beta E(x,z)}}{Q (z|x)}

그리고 Q는 Gibbs distribution 형태로 나타내면

Q(z|x) = \frac{\exp^{q(z|x)}}{\int_{x^\cdot} \exp^{q(z|x^{\cdot})}}

variational approximation of marginalization over Z

F(x) \leq \tilde F(x) = \int_Z Q(z|x) [-\frac{1}{\beta} \log \frac{\exp^{-\beta E(x,z)}}{Q (z|x)}] \\ \tilde F(x) = <E> - TS \\ L(w) = \tilde F_w(x)

즉 Objective function은 에너지의 평균을 낮추도록 만든다.
이러한 두 관점에서 VAE는 데이터 생성과 자유 에너지 최소화의 목표를 달성하기 위해 variational inference, variational approximation을 사용하여 근사할 수 있었다. 모델 아키텍쳐는 똑같더라도 바라보는 관점에 따라 고스란히 모델에 녹여서 생각하게 된다고 느꼈다.

Reference

variational inference 설명
https://mbernste.github.io/posts/variational_inference/
kaist 온라인 강좌 Implicit Deep Generative model
https://www.youtube.com/watch?v=j3JBjaaJF5k&list=PLzZ7PPT4KK5qt6LirBzs5EbtmSnwS8x3l
lecunn의 energy based 해석 영상
https://www.youtube.com/watch?v=bdebHVF__mo
PRML 정리
https://norman3.github.io/prml/
Lagrangian을 통한 multinomial distribution MLE
https://blog.jakuba.net/maximum-likelihood-for-multinomial-distribution/
VQ-VAE code
https://github.com/MishaLaskin/vqvae

Appendix : VAE 구체적인 설명

Object Function

MLE = \Pi_d^D P(X_d) = \Pi_d^D \sum_i^N P(X_d,Z_i)

D : # of data

N : # of latent variable state

NLL = - \sum_d^D \log \sum_i^N P(X_d,Z_i) = - \sum_d^D \log \sum_i^N P(X_d,Z_i) \frac{Q(Z_i|X_d)}{Q(Z_i|X_d)}

Jensen's Inequality

\leq - \sum_d^D \sum_i^N Q(Z_i|X_d) \log \frac{P(X_d,Z_i) }{Q(Z_i|X_d)}

= - \sum_d^D \sum_i^N Q(Z_i|X_d) \log \frac{P(X_d|Z_i) P(Z_i) }{Q(Z_i|X_d)}

= \sum_d^D [ D_{KL}(Q(Z_i|X_d)|| P(Z_i)) - \sum_i^N Q(Z_i|X_d) \log P(X_d|Z_i) ]

= \sum_d^D D_{KL}(Q(Z_i|X_d)|| P(Z_i)) - E_Z[\log P(X_d|Z_i)]

KL- divergence term (Regularizer)

$Z_i$ is continous, so we assume Z as Gaussain.

D_{KL}(Q(Z_i|X_d)|| P(Z_i)) = -\int_{Z_i} Q(Z_i|X_d) \log P(Z_i)) + \int_{Z_i} Q(Z_i|X_d) \log [Q(Z_i|X_d) ]

multivariate Gaussian Distribution

if basis is appropriately adjusted, then we can assume that sigma is diagonal matrix which means that offdiagonal term is 0. so covariance matrix $\sum$ is written as $\sigma^2I$

Q(Z|X) \sim N(\mu,\sigma^2I)

D_{KL} = (\mu^2 + \sigma^2 -1)/2 - \log \sigma

Q() can be reparameterize to differentiable, determistic variable

P(Z_i) \sim N(0,I) \\ Z_i \sim Q(Z_i|X_d) = Q_\phi (\epsilon,X_d) \\ \epsilon \sim p(\epsilon) \\

applying amortized inference Q

Q(Z_i|X_d;\phi) = \Pi_i Q(Z_i;NN_i(X_d;\phi_i))

Reparameterization Trick

Q_\phi (\epsilon,X_k) = \mu + \epsilon * \sigma

Monte-carlo sampling

where $\epsilon^{(d)} \sim p(\epsilon)$

$\epsilon \sim \mathcal{N}(0, I)$ .

Expectation term (Reconstruction)

$E_i^{(d)}$ is reconstructed data and $E_i$ is original data

Depending upon the assumption of $E_i$ and $p_\theta\left(E \mid E^{(d)}\right)$

If $E_i$ is continuous and $p_\theta\left(E_i \mid E_i^{(d)}\right) \sim \mathrm{N}\left(E_i^{(d)}, 1^2\right)$

\log p_\theta\left(E \mid E^{(d)}\right)=C \times \log \exp \left(-\frac{\left(E_i-E_i^{(d)}\right)^2}{2}\right)=-C \times\left(E_i-E_i^{(d)}\right)^2

$\cdot$ If $E_i$ is discrete and $p_\theta\left(E_i \mid E_i^{(d)}\right) \sim \operatorname{Bern}\left(E_i^{(d)}\right)$

\log p_{\theta}(E \mid E^{(d)})

=\log E^{(d)^{K_i}}(1-E^{(d)})^{1-E_i}

=E_i \log E^{(d)}+(1-E_i) \log (1-E^{(d)})

JoongHyun's Blog

AI와 수학, 일상 기록들

이전 포스트

Variational Inference (VAE)

Explicit Generative Model

variational approximation of marginalization over Z

ELBO

Jensen's inequality

Reconstruction (Expectation term)

Regularizer (KL divergence)

Energy Based Model

variational approximation of marginalization over Z

Reference

Appendix : VAE 구체적인 설명

Object Function

Jensen's Inequality

KL- divergence term (Regularizer)

multivariate Gaussian Distribution

Reparameterization Trick

Monte-carlo sampling

Expectation term (Reconstruction)

LoRA를 통한 CLIP 모델 PEFT

0개의 댓글