Machine Learning, Deep Learning

Variational autoencoder

n.han 2021. 4. 16. 22:09

1. Intuition
Variational autoencoder는 generative model의 일종입니다. VAE는 autoencoder와 구조적으로 굉장히 유사한데, 목적 함수의 수학적 식은 다릅니다. VAE는 확률론적 그래프 모델로, 데이터가 어떻게 생성되는지를 배웁니다.

먼저 직관적으로 모델을 이해해보죠. autoencoder를 활용하면 입력 이미지의 latent representation을 얻을 수 있는데요.

latent rerpresentation with autoencoder

autoencoder에서는 discrete value이기 때문에 latent space로부터 새로운 이미지를 생성할 수 없습니다. 이를 우리가 상대적으로 편하게 다룰 수 있는 정규 분포로 가정하고, 정규 분포의 평균과 표준편차가 autoencoder의 diescre value라고 생각할 수 있습니다.

latent representation as probabilitic distribution

이렇게 학습이 완료되면, 가운데 확률 분포로부터 각 latent feature들을 샘플링하여 새로운 이미지를 생성할 수 있습니다.

sample from latent distributions

2. Backgrounds
2.1 Maximum Likelihood Estimation
Maximum Likelihood Estimation은 모델이 정해져 있을 때, Model의 파라미터들을 조절해가면서 해당 분포에서 주어진 데이터가 발생할 확률인 likelihood들의 곱을 최대화하는 것을 말합니다.

Maximum Likelihood Estimation. 출처: https://www.youtube.com/watch?v=XepXtl9YKwc
Maximum Likelihood Estimation. 출처: https://towardsdatascience.com/mse-is-cross-entropy-at-heart-maximum-likelihood-estimation-explained-181a29450a0b


로그는 단조 함수이고, likelihood의 곱보다 덧셈이 여러모로 계산하기 쉬우므로 log likelihood를 쓸 수 있습니다. 더하여 대부분 모델은 손실 함수가 감소되는 방향으로 업데이트 되므로 negative log likelihood를 많이 쓰게 됩니다.

2.2 MLE for Neural Network
Neural Network는 입력 x와 출력 y와의 관계식을 approximation하는데, 각 노드의 출력을 노드의 실제 값이 따르는 정규 분포의 평균이라고 생각하면 결국 Neural Network를 학습한다는 것은 negative log likelihood를 최소화하는 것과 같게 됩니다.

Log-likelihood for Neural Network

2.3 Marginal probability
X, Y로 구성된 joint distribution이 있을 때, 이를 통해 p(X)를 구하거나 p(Y)를 구하는 확률을 Marginal probability라고 합니다.

Marginal probability. 출처: https://en.wikipedia.org/wiki/Marginal_distribution
Marginal probability


3.VAE Architecture
3.1 ELBO

우리는 아래와 같이 z를 입력으로 하여 새로운 데이터 x를 생성하는 Generator를 만드는데 목표가 있습니다.

negative log likelihood 최소화하도록 네트워크의 파라미터를 업데이트해야 하는데, 이를 입력 z와 함께 marginal probability로 표현하면 다음과 같습니다.
$$
log(g_\theta (x))=log(\int g_{\theta}(x|z)p(z)dz))
$$
p(z)는 우리가 편하게 다룰 수 있는 정규 분포와 같은 분포로 설정할 수 있지만, x를 잘 생성할 수 있는 정규 분포의 파라미터를 알지 못합니다. 이를 위해 x가 주어진 상태에서 p(z)를 approximate하는 q(z|x)를 뉴럴 네트워크를 앞단에 붙입니다. 이렇게 모르는 p(z)를 추정(Variational Inference)한다고 하여 모델 이름에 Variational 라는 접두사가 붙었습니다.

Inference & Generation Networks

위의 모델의 Log-likelihood를 정리하면 아래와 같습니다.
$$
log(g_\theta (x))=log(\int g_{\theta}(x|z)p(z)dz)) = log(\int g_{\theta}(x|z){p(z) \over q_{\phi}(z|x)} q_{\phi}(z|x) dz)) \\
\ge \int log(g_{\theta}(x|z){p(z) \over q_{\phi}(z|x)})*q_{\phi}(z|x) dz (Jensen's Inequality)\\ 
\ge \int (log(g_{\theta}(x|z))-log({q_{\phi}(z|x) \over p(z)}))*q_{\phi}(z|x) dz \\
\ge \int log(g_\theta(x|z))q_\phi(z|x)dz-\int log({q_{\phi}(z|x) \over p(z)})q_\phi(z|x)dz
$$

최종 부등식을 ELBO(Evidence Lower BOund)라고 합니다. 즉 Log-likelihood를 Maximize하는 것은 ELBO를 Maximize하는 것과 같게 생각할 수 있습니다.

ELBO의 구성

ELBO는 Reconstruction Error와 Approximate posterior close to prior(Posterior와 prior간 KL Divergence)로 나뉠 수 있습니다. 먼저 Reconstruction Error 부분을 이해해보죠.

3.2 Reconstruction Error
Reconstruction Error는 Monte-carlo Approximation에 의해 간단히 표현될 수 있습니다.

Monte-carlo Approximation. 출처: https://hugrypiggykim.com/2018/09/07/variational-autoencoder%EC%99%80-elboevidence-lower-bound/

L개의 샘플을 뽑을 수 있지만, 보통 1개만 추출해서 Error를 계산합니다.
$$
\int log(g_\theta(x|z))q_\phi(z|x)dz \approx {1 \over L}\sum_{l} log(g_\theta(x|z^{l})) \approx log(g_\theta(x|z))
$$
우리의 출력이 1, 0과 같은 베르누이 분포를 따른다고 가정하면, Reconstruction Error는 다음과 같습니다:
$$
log(g_\theta(x|z)) \\
= \sum_{j=1}^D log(g_\theta(x_j|z)) \\
= \sum_{j=1}^D log(p_j^{x_j}(1-p_j)^{1-x_j}) \\
= \sum_{j=1}^D x_jlog(p_j) + (1-x_j)log(1-p_j) \\
$$
이를 코드로 구현하면 아래와 같습니다:

reconstruction_loss = tf.reduce_mean(
	tf.reduce_sum(x * y_hat + (1-x) * (1-y_hat), 1)
)

우리의 출력이 가우시안 분포를 따른다고 가정하면, Reconstruction Error는 다음과 같습니다:
$$
log(g_\theta(x|z)) \\
= \sum_{j=1}^D log(g_\theta(x_j|z)) \\
= \sum_{j=1}^D log({1 \over \sqrt{2\pi}}exp(-{(x_j-\mu_j)^2 \over 2})) \\
\propto \sum_{j=1}^D (x_j - \mu_j)^2
$$
이를 코드로 구현하면 아래와 같습니다:

reconstruction_loss = tf.reduce_mean(
	-tf.reduce_sum(
    	tf.math.square(x - y_hat), 1
    )
)

 

3.3 Approximate posterior close to prior
Error의 두번째 Term은 posterior와 prior 간 KL Divergence와 같습니다.
$$
\int log({q_{\phi}(z|x) \over p(z)})q_\phi(z|x)dz = D_{KL}(q_{\phi}(z|x) || p(z))
$$
prior는 보통 표준정규분포가 되므로, 전체 posterior 분포로 샘플링되는 z가 reconstruction error를 줄이도록 업데이트 하되 가능한 표준정규분포를 따르도록 만듦으로 Regularization 역할을 하여 Regularization term이라고도 표현합니다.

두 정규 분포간의 KL Divergence은 다음과 같이 유도됩니다.

위의 두 번째 식은 다음과 같이 전개됩니다:

세 번째 식은 다음과 같이 전개됩니다:

전개된 식을 정리하고, 두 번째 분포가 표준정규분포를 따른다고 두면 다음과 같습니다.

이 식을 풀면 다음과 같습니다.

따라서 posterior와 prior간 KL Divergence은 다음과 같이 정리할 수 있습니다:
$$
D_{KL}(q_{\phi}(z|x)||p(z)) = D_{KL}(N((\mu_1, \mu_2, ... , \mu_k), (\sigma_1, \sigma_2, ... , \sigma_k)) || N(0, 1)) \\
= {1 \over 2}\sum_{i=1}^k(\mu_i^2+\sigma_i^2-ln(\sigma_i^2)-1)
$$
코드로 이를 다음과 같이 구현할 수 있습니다:

kl_loss = tf.reduce.mean(
	0.5 * tf.reduce_sum(1 - tf.math.log(tf.math.square(sigma) + 1e-8) + tf.math.square(mu) + tf.math.square(sigma), 1)
)

 

3.4 Reparameterization trick

Decoder의 입력은 KL Divergence로 prior에 가깝게 근사하려는 posterior로 샘플링된 z이기 때문에 backpropagation을 하지 못합니다. 이를 개선하기 위해 reparameterization trick를 사용합니다. 방식은 아래와 같습니다.

Reparameterization trick. 출처: https://stats.stackexchange.com/questions/429315/why-is-reparameterization-trick-necessary-for-variational-autoencoders

코드로 이를 다음과 같이 구현할 수 있습니다:

def reparameterize(self, mu, sigma):
	eps = tf.random.normal(shape=mu.shape)
    return eps * sigma + mu


References

1. www.slideshare.net/NaverEngineering/ss-96581209
2. www.jeremyjordan.me/variational-autoencoders/  
3. hugrypiggykim.com/2018/09/07/variational-autoencoder%EC%99%80-elboevidence-lower-bound/