[DL / Paper review] Auto-Encoding Variational Bayes

논문 링크(26527회 인용)

Summary

• probabilistic model에서 variational inference은 종종 사용되는데 closed form으로 업데이트를 하는 것이 아닌, elbow에 대한 gradient을 구해서 파라미터를 업데이트 하는 방법을 제시한다.
• niave monte carlo gradient estimator을 유도하고 이는 variance가 큼을 지적하며 reparameterization 기술을 통해 gradient estimator을 유도한다.
• 이는 objective function이 어떠한 형태이든, 적절한 reparameterization을 사용하여 미분 가능한 식으로 바꿔두면 어떠한 elbow에 대해서도 gradient을 구할 수 있기 때문에 꽤나 일반적이고 강력한 방법론이다.

Motivation

• variational inference은 mean-field 가정을 기반으로 variational parameters에 대해 식을 분해하고 closed-form으로 업데이트 식을 유도한다. 이는 marginal likelihood, posterior density, VB을 위한 integral 등이 intractable할 때 불가능하다는 단점이 있다. 따라서 intractable한 식이 나오는 모델을 가정했을 때에도 일반적으로 파라미터를 업데이트할 수 있는 방법이 필요하다.
• 데이터의 크기가 클 경우 minibatches 또는 하나의 데이터를 사용하여 파라미터 업데이트를 진행하는데 monte carlo em과 같은 샘플링 기반의 방법은 하나의 데이터당 샘플링하는 계산이 많이 요구되기 때문에 전체적인 computation burden이 늘어난다.
• 위 두가지 문제를 해결하기 위해 vae가 개발되었다.
• 데이터의 크기가 커도 효율적인 posterior inference을 제공하고 $\mathbf{x}$에 대한 prior가 필요한 모든 종류의 inference 문제도 해결할 수 있는 방법론이다.

Approach

Problem scenario

• $\mathbf{X} = \{ \mathbf{x}^{(i)} \}^N_{i=1}$: N개의 iid samples
• 관측되지 않은 연속형 확률 변수인 $\mathbf{z}$을 포함하는 어떤 random process에 의해 데이터가 생성된다고 가정한다.
• random process는 아래의 과정을 거친다.
  • $\mathbf{z}^{(i)} \sim p_{\boldsymbol{\theta^*}} (\mathbf{z})$
  • $p_{\boldsymbol{\theta^*}} (\mathbf{x} | \mathbf{z})$: likelihood function
  • 이 둘은 모든 $\boldsymbol{\theta}, \mathbf{z}$에 대해 미분 가능하다고 가정한다.
• 하지만 true parameters인 $\boldsymbol{\theta}^*$와 latent variables인 $\mathbf{z}^{(i)}$은 unknown이다.
• $q_\phi (\mathbf{z} | \mathbf{x})$: intractable true posterior인 $p_{\boldsymbol{\theta}} (\mathbf{z} | \mathbf{x})$에 대한 recognition model. 이전의 VI와 같이 파라미터에 대해 분해되거나 closed form solution으로 유도될 필요가 없다. 그 대신, 파라미터 $\phi$와 generative model 파라미터 $\boldsymbol{\theta}$을 동시에 학습할 수 있는 방법을 제안한다.

이쯤에서 이 논문의 제목에 왜 auto-encoding이 붙어있는지 짚고 넘어가자. 관측되지 않은 latent variable인 $\mathbf{z}$을 code로 해석하여 recognition model인 $q_\phi (\mathbf{z} | \mathbf{x})$을 probabilistic encoder라고 부른다. 이유는 데이터 $\mathbf{x}$가 주어졌을 때, $q_\phi$ 분포로부터 $\mathbf{z}$을 generate할 수 있기 때문이다. 간단하게 생각하면 데이터가 주어졌을 때 latent variables을 만드는 것이니까 encoder로 볼 수 있다. 비슷한 맥락으로 $p_{\boldsymbol{\theta}} (\mathbf{x} | \mathbf{z})$을 probabilistic decoder라고 부르는데 code $\mathbf{z}$가 주어졌을 때, $\mathbf{z}$에 대한 분포를 만들기 때문이다. 간단하게 생각하면 latent variable이 주어졌을 때 새로운 데이터를 생성하는 것이니까 decoder로 볼 수 있다.

The variational bound

marginal likelihood인 $\log p_{\theta} (\mathbf{x}^{(1)}, \cdots, \mathbf{x}^{(N)})$은 아래와 같이 분해할 수 있다.

관심있는 likelihood가 두 개의 항으로 분해되었다. 첫번째 항은 KL-Divergence로 true posterior와 recognition model과의 분포상 거리이다. 이는 정의상 양수이기 때문에 marginal likelihood은 최소 두번째 항의 값을 가진다. 즉, 다시 말해 아래의 관계가 성립한다.

오른쪽 항을 expected lower evidence bound (ELBO)라고 부른다.

ELBO인 $\mathcal{L}(\boldsymbol{\theta, \phi}; \mathbf{x}^{(i)})$를 \boldsymbol{\theta, \phi}에 대해 최적화하고자 하는 것이 목표이므로 이에 대한 gradient을 구하는 것이 일차적인 과제이다. 본 논문에서는 analytic closed form을 구하는 것을 가정하지 않으므로 일반적인 상황, 직접 gradient의 효율적인 estimator을 구하는 상황을 가정한다.

문제는, ELBO의 gradient estimator을 구하는게 쉽지 않다는 것이다. naive monte carlo gradient estimator는 아래와 같다.

선행 연구에 따르면, 이 estimator는 높은 variance을 가지는 것이 단점이다. 따라서 reparameterization을 통해 다른 형태의 gradient estimator을 제안한다.

The SGVB estimator and AEVB algorithm

reparameterization trick을 통해 확률 변수 $\tilde{z} \sim q_{\phi} (z | x)$을 미분 가능한 함수 $g_\phi (\epsilon, x)$와 auxiliary noise variable $\epsilon$을 통해 아래와 같이 재구성할 수 있다.

이를 이용하면 어떤 함수 $f(z)$의 $q_\phi (z|x)$에 대한 기댓값의 monte carlo estimates을 아래와 같이 나타낼 수 있다.

이를 ELBO에 적용하여 stochastic gradient variational bayes (SGVB)을 유도하면 아래와 같다.

Another interpretation of ELBO as regularization added objective function

ELBO을 아래와 같이 전개해볼 수 있다.

분포의 형태에 따라 KL Divergence 항이 analytical하게 적분이 된다고 한다. 이런 상황에서는 KL 항은 $\phi$을 정규화하는 역할을 하고 expected reconstruction error인 $E_{q_\phi} [ p( \mathbf{z} | \mathbf{x}^{(i)}) ]$에 대한 샘플링만 필요하게 된다. 이런 상황에서 두번째 형태의 SGVB을 아래와 같이 유도할 수 있다.

배치별로 M개의 데이터를 사용하는 minibatches 전략을 활용한다면 아래와 같이 ELBO을 추정하면 된다.

autoencoder과의 관계를 살펴보자.

• KL 항은 regularizer로, 두번째 항은 expected negative reconstruction error의 역할을 한다.
• $\mathbf{z}^{(i,l)} = g_\phi (\boldsymbol{\epsilon}^{(l)}, \mathbf{x}^{(i)})$이고 $\mathbf{z}^{(i,l)} \sim q_\phi (\mathbf{z} | \mathbf{x}^{(i)})$
• $\mathbf{z}^{(i,l)}$은 함수 $\log p_\theta (\mathbf{x}^{(i)} | \mathbf{z}^{(i,l)})$의 입력값으로 들어간다. 이는 $\mathbf{z}^{(i,l)}$ 하에 데이터 포인트에 대한 확률밀도함수인, generative model이다.
• latent variable 가정 하에 데이터의 분포를 생각하므로 decoder의 역할을 하는 것이다.

The reparameterization trick

목적은 ELBO의 monte carlo estimate을 $\phi$에 대해 미분 가능한 함수로 바꾸는 것이다.

• $\mathbf{z} \sim q_\phi (\mathbf{z} | \mathbf{x})$
• 확률변수 $\mathbf{z}$을  deterministic function인 $g_\phi$로 바꾼다: $\mathbf{z} = g_\phi (\boldsymbol{\epsilon}, \mathbf{x})$
• $\boldsymbol{\epsilon}$: auxiliary variable with independent marginal $p(\boldsymbol{\epsilon})$
• $g_\phi (\cdot)$: $\phi$를 파라미터로 가지는 vector-valued function

$g_\phi$가 주어졌다고 가정하면 $q_\phi (\mathbf{z} | \mathbf{x}) \prod_i dz_i = p(\boldsymbol{\epsilon}) \prod_i d \epsilon_i$이다. 따라서,

이를 적용하여 $E_{q_\phi (z|x)} [f(z)]$을 $\phi$에 대한 미분 가능한 monte carlo estimation으로 바꿔 쓸 수 있다.

그러면 $g_\phi$와 $\boldsymbol{\epsilon}$을 어떻게 선택해야할까? 이에 대한 내용은 논문의 내용을 그대로 첨부한다.

Experiments

mnist 데이터를 실험에 사용한다. 비교 대상 모델은 wake-sleep으로 논문에서 vae와 가장 유사하다고 언급되는 알고리즘이다.

• x축: training에 사용한 데이터의 수
• y축: elbo
• 모든 실험에서 vae가 더 빠르고 좋은 수렴 결과를 얻음을 확인할 수 있다.

• vae의 marginal likelihood가 더 높게 추정됨을 알 수 있다.
• mcem같은 경우, 데이터 수가 많을 때 수렴조차 쉽지 않은 것으로 보인다.

Conclusion

• cf에 vae을 적용한 논문 때문에 vae 논문을 찾아서 읽어보았는데 꽤나 강력한 방법을 제시했다는 생각이 들었다.
• 보통 vi을 하면 analytical form을 유도하는 경우가 많은데 이 논문은 그런 가정조차 하지 않고 좀 더 일반적인 상황에서 monte carlo estimation에 대한 gradient을 어떻게 유도할지 고민한다는 점에서 더 많은 상황을 커버할 수 있을 것 같다.
• 여기서 output distribution을 multinomial로 한 것이 cf로 생각하면 될 것 같다.
• generative model을 세우고 likelihood가 intractable 하다면 vae을 꺼내서 적용해보면 좋을 것 같다.