[논문 요약] Auto-Encoding Variational Bayes

1. Introduction

Variational Bayesian method

구하고자 하는 조건부 확률 $P( Z | X )$ 가 intractable 할 때, 이것의 근삿값을 최적화하는 문제로 바꾸어 구하는 방법이다.

이 방법에서는 흔히 Mean-field Approach라고 불리는 기법을 사용하는데, $P(Z|X)$를 근사시킬 $Q(Z)$를 만들고 두 분포의 차이를 나타내는 쿨백-라이블리 발산을 최소화하는 $Q(Z)$를 구해서 계산에 사용하는 방법이다. 하지만 $Q(Z)$가 여전히 계산 불가능한 상황이 발생할 수 있다는 문제가 있다.

연구진은 SGVB – Reparameterization trick을 사용해 Mean-field Approach를 위와 같이 변형함으로써 문제를 해결했다.

Recognition model $q_{\varnothing}(z|x)$

데이터 x를 바탕으로 latent variable $z$를 만들어 내기 위한 확률 분포 $p_{\theta}(z|x)$를 모델링한 것

AEVB (Auto-Encoding Variational Bayesian)

• SGVB 기법을 사용해 Recognition model 을 학습시킨다
• 기존의 고비용 추론 모델 (ex : MCMC 알고리즘) 의 도움 없이 효율적으로 모델을 학습할 수 있다
• 학습된 모델은 인식, 노이즈 제거, 시각화에 쓰일 수 있다

VAE는 Auto Encoder와 동일한 구조를 갖지만, 오토인코더는 데이터를 압축하기 위한 앞단을 학습하기 위해 뒷단을 붙인 반면 Variational 오토 인코더는 데이터를 생성하기 위한 뒷단을 학습하기 위해 앞단을 붙이게 되었다.

2. Method

전제 조건

• 모든 데이터는 서로 독립적이고 동일한 확률분포를 갖는다 (i.i.d)
• 모든 데이터에 잠재 변수 (latent variable)가 있다
• Global Parameters ($\theta$, $\varnothing$)에 대해서는 ML 또는 MAP 추정을 시행하고, Latent Variables에 대해서는 Variational Inference를 시행한다
• Global Parameters에 대해서도 Variational Inference를 시행할 수 있지만 이것에 대한 실험은 나중의 작업에 맡긴다

문제 상황

• i.i.d를 만족하는 연속 or 이산 변수 N개로 구성된 데이터셋 $X = {x^{(i)}}^N_{i=1}$
• 데이터 $X$는 $𝑝_{\theta ∗} (𝑥 | 𝑧)$에 의해서, 잠재 변수 $Z$는 $𝑝_{\theta *} (𝑧)$ 에 의해서 랜덤하게 생성된다
• $𝑝_{\theta ∗} (𝑥 | 𝑧)$, $𝑝_{\theta *} (𝑧)$ 는 $𝑝_{\theta} (𝑥 | 𝑧)$, $𝑝_{\theta} (𝑧)$를 최적화한 확률 분포에서 오며, 두 분포의 PDF는 미분 가능하다
• $\theta *$, $z$는 알려져 있지 않다
• 계산 과정에서 확률을 간단하게 하는 가정을 만들지 않아 General한 알고리즘을 도출할 것이다

목표

• 효율적인 ML or MAP 추정으로 $\theta$를 구한다.
  이를 통해 Generative Model $𝑝_{\theta}(𝑧) 𝑝_{\theta} (𝑥 | 𝑧)$을 풀이해서 실제 같은 인공 데이터를 만들어 낼 수 있다.
• 데이터 $x$에 대한 잠재 변수 $z$의 사후 확률인 $𝑝_{\theta} (𝑧 | 𝑥)$을 추정한다.
  이를 통해 데이터를 그것의 특징을 담고 있는 잠재 변수들로 인코딩할 수 있다.
• 데이터 $x$의 확률분포 $P(x)$를 추정한다.
  노이즈 제거, 이미지 복원 등 $P(x)$가 필요한 여러 컴퓨터 비전에 쓰일 수 있다.

목표 달성을 위한 아이디어

• Intractable한 $𝑝_{\theta} (𝑧 | 𝑥)$에 대한 근사로 Recognition model $q_{\theta} (z | x )$를 도입한다. 이때, mean-field 방식과는 다르게 $\theta$가 항상 Closed-form하게 계산될 필요가 없다.
• 잠재 변수 $z$ 를 Code라고 하면, $q_{\theta} (z | x )$를 encoder, $𝑝_{\theta} (𝑥 | 𝑧)$를 decoder로 볼 수 있다.

The variational bound

Total marginal likelihood $p_{\theta}(x)$

                  $log p_{\theta} (x^{(1)}, \cdot \cdot \cdot, x^{(N)}) = \sum_{i=1}^{N}log p_{\theta}(x^{(i)})$

Marginal likelihood $p_{\theta}(x^{(i)})$

                  $log p_{\theta} (x^{(i)}) = D_{KL}(q_{\phi}(z|x^{(i)})||p_{\theta}(z|x^{(i)}))+L(\theta, \phi; x^{(i)})$

• $D_{KL}(\cdot\cdot\cdot)$ : 실제 Posterior와 그것의 근사치인 Recognition model간의 차이
• $L(\theta, \phi; x^{(i)})$ : Marginal likelihood의 Lower Bound

$log(p(x)) = \int log(p(x))q_\phi (z|x)dz\,\,\,\,\,\,\,\,\,\,\,\,\,\because \int q_\phi (z|x)dz=1$

$\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,=\int log({p(x,z)\over p(z|x)})q_\phi (z|x)dz\,\,\,\,\,\,\,\,\,\,\,\,\,\because p(x)={p(x,z)\over p(z|x)}$

$\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,=\int log({p(x,z)\over {q_\phi (z|x)}}\times {q_\phi(z|x)\over {p(z|x)}})q_\phi(z|x)dz$

$\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,=\int log({{p(x,z)}\over {q_\phi(z|x)}})q_\phi(z|x)dz+\int log({{q_\phi(z|x)}\over {p(z|x)}})q_\phi(z|x)dz$

$\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,=L(\theta,\phi;x)+D_{KL}(q_\phi (z|x)||p(z|x))$

위와 같은 유도를 통해 Total marginal likelihood를 두 개의 항으로 나눌 수 있다.

이때, 두번째 항 $D_{KL}(q_\phi (z|x)||p(z|x))$은 $p(z|x)$를 모르기 때문에 계산할 수 없다.

하지만 이 값은 항상 양수이기 때문에 첫번째 항을 Lower Bound로 보고 이것을 최대화하여 Total marginal likelihood에 가까워질 수 있다.

$log \,p_{\theta}(x^{(i)})\geq L(\theta, \phi; x^{(i)})=-D_{KL}(q_\phi(z|x^{(i)})||p_\theta(z))+E_{q_\phi(z|x^{(i)})}[log p_\theta(x^{(i)}|z)]$

목표는 $max_\phi max_\theta L(\theta, \phi; x^{(i)})$를 찾는 것인데, 일반적으로 많이 사용되는 Monte Carlo gradient estimator는 이 상황에서 굉장히 큰 Variance를 갖기 때문에 부적절하다. 또한 $q_\theta(z|x)$가 미분 불가능할 경우 gradient를 계산하기 어렵다.

이와 같은 문제를 해결하기 위해, 연구진은 Reparameterization trick이라는 이론을 만들었다.

자세한 내용은 뒤에서 다루고, 지금은 $q_\phi(z|x)$에서 생성된 잠재변수 z가 새로운 샘플링 함수 $g_\phi(\epsilon, x)$에 의해 생성된다고 가정하자.

$\widetilde{z}=g_\phi(\epsilon, x)\,\,with\,\,\epsilon \sim p(\epsilon)$

이때, 샘플링 함수 $g_\phi(\epsilon, x)$는 미분 가능해야 하며 노이즈 변수 $\epsilon$의 확률 분포로는 흔히 정규분포가 이용된다.

$E_{q_\phi(z|x^{(i)})}[f(z)]=E_{p(\epsilon)}[f(g_\phi(\epsilon, x^{(i)}))]\simeq {1\over L}\sum_{l=1}^Lf(g_\phi(\epsilon^{(l)}, x^{(i)}))\,\, where\,\,\epsilon^{(l)}\sim p(\epsilon)$

$z\sim q_\phi(z|x^{(i)})$

$\epsilon \sim p(\epsilon)$

Reparametrization trick을 이용하면 아래와 같이 몬테카를로 추정을 할 수 있다.

몬테카를로 추정

$E_{q_\theta(z|x^{(i)})}[f(z)]\simeq {1\over L}\sum^L_{l=1}f(g_\phi(\epsilon^{(l)}, x^{(i)})) \,\,where \,\,\epsilon^{(l)}\sim p(\epsilon)$

$p(\epsilon)$의 확률분포에서 랜덤하게 뽑아낸 L개의 샘플로 평균을 낸다

SGVB estimator (1st version)

$L(\theta, \phi; x^{(i)})=E_{q_\phi(z|x^{(i)})}[-log \,q_\phi(z|x^{(i)})+log\, p_\theta(x^{(i)}, z)]$

$\widetilde{L}^A(\theta, \phi; x^{(i)})={1\over L}\sum^L_{l=1}-log\,q_\phi(z^{(i,l)|x^{(i)}})+log\,p_\theta(x^{(i)}, z^{(i, l)})$

$\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,where \,\,z^{(i, l)}=g_\phi(\epsilon^{(i, l)}, x^{(i)})\,\, and\,\, \epsilon^{(l)}\sim p(\epsilon)$

$\widetilde{L}^A(\theta, \phi; x^{(i)})\simeq L(\theta, \phi; x^{(i)})$

이 방식은 Lower Bound 전체를 샘플링해 평균을 구해서 추정하는 방식이다.

SGVB estimator (2nd version)

$L(\theta, \phi; x^{(i)})=-D_{KL}(q_{\phi}(z|x^{(i)})||p_\theta(z))+E_{q_\phi(z|x^{(i)})}[log\, p_\theta(x^{(i)}|z)]$

$\widetilde{L}^B(\theta, \phi; x^{(i)})=-D_{KL}(q_{\phi}(z|x^{(i)})||p_\theta(z))+{1\over L}\sum^L_{l=1}(log\,p_\theta(x^{(i)}|z^{(i,l)}))$

$\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,where\,\,\,\,z^{(i,l)}=g_\phi(\epsilon^{(i,l)},x^{(i)})\,\,\,\,and\,\,\,\,\epsilon^{(l)}\sim p(\epsilon)$

$\widetilde{L}^B(\theta, \phi; x^{(i)})\simeq L(\theta, \phi; x^{(i)})$

이 방식은 $p_\theta(z)$와 $q_\phi(z|x)$가 둘 다 정규 분포를 따를 때 쿨백-라이블리 발산을 anaytic하게 계산할 수 있는 솔루션을 이용할 수 있다는 점을 적용했다.

따라서 쿨백-라이블리 발산은 샘플링 없이 계산 가능하고, 나머지 항만 샘플링해 추정한다.

전체 데이터셋의 lower bound의 marginal likelihood 계산 방법

$L(\theta, \phi; X)\simeq \widetilde{L}^M(\theta, \phi; X^M)={N\over M}\sum^M_{i=1}\widetilde{L}(\theta, \phi; x^{(i)})$

$\widetilde{L}^M(\theta, \phi; X^M)$ : Total SGVB estimator

$\widetilde{L}(\theta, \phi; x^{(i)})$ : SGVB estimator of one datapoint

전체 N개의 데이터 중 minibatch $X^M={x^{(i)}}^M_{i=1}$를 랜덤하게 골라 SGVB estimator 평균을 구한 뒤 N배 한다.

이후 추정한 estimator를 미분해 gradient를 구하고 SGD나 Adagrad로 최적화해 $\theta$와 $\phi$를 도출한다.

오토 인코더와의 연관성

Reparametrization trick

문제 상황

$q_\theta(z|x)$로 샘플을 생성한다면 $z\sim q_\theta(z|x)$는 미분이 불가능하기 때문에 Lower Bound의 Gradient 계산이 어려운 상황이 발생할 수 있따.

해결 방법

$z^{(i,l)} = g_\phi(\epsilon^{(i,l)}, x^{(i)})\,\,\,where\,\,\,\epsilon\sim p(\epsilon)$

$q_\theta(z|x)$의 분포를 따르는 연속 변수 $z$를 샘플링 함수 $g$를 이용해 deterministic하게 나타낸다. (이때, 샘플링 함수 $g$는 미분 가능해야 한다)

$Ex)$

$z\sim p(z|x)=\mathcal{N}(\mu, \sigma^2)$
z가 평균이 $\mu$, 표준편차가 $\sigma$인 정규분포를 따른다고 하면,

$z=\mu+\sigma \epsilon \,\,\,where \,\,\,\epsilon\sim \mathcal{N}(0,1)$
이와 같이 deterministic하게 표현할 수 있다.

$\therefore\,\mathbb{E}_{\mathcal{N}(z;\mu,\sigma^2)}[f(z)]=\mathbb{E}_{\mathcal{N}(\epsilon;0,1)}[f(\mu+\sigma\epsilon)]\simeq {1\over L}\sum^L_{l=1}f(\mu+\sigma\epsilon^{(l)})\,\,where\,\,\epsilon^{(l)}\sim\mathcal{N}(0,1)$
몬테카를로 추정을 통해 나온 최종식이 미분 가능하므로 gradient를 구할 수 있다.