53.Auto-Encoding Variational Bayes (2014)

Auto-Encoding Variational Bayes ( 2014 )

Abstract

이 논문 또한 앞선 논문들과 마찬가지로, posterior가 intractable한 경우의 문제를 풀기 위한 방법을 제안한다. 데이터셋의 크기가 클 때도 잘 작동하는(scales to large datasets) stochastic variational inference 방법을 제안한다.

그러기 위해, 아래와 같은 2가지의 contribution을 설명하고 있다.

• 1) reparameterization of the ELBO
• 2) posterior inference can be made especially efficient! ( by fitting an approximate inference model )
  

1. Introduction

VI의 대표적인 방법으로 널리 사용되는 MFVI는 analytical solution을 요구하지만, intractable한 경우가 많다. 이 논문은 ELBO에 대한 reparameterization을 통해, ELBO의 simple differentiable unbiased estimator를 제안한다. 이를 “SGVB (Stochastic Gradient Variational Bayes)” 라고 부른다.

또한, AEVB(Auto-Encoding Variational Bayes)를 제안하는데, 이는 우리가 흔히 알고 있는 VAE (Variational Autoencoder)를 만든 알고리즘이기도 하다.

2. Method

2-1. Problem Scenario

Notation을 다음과 같이 정의하겠다.

• \(\mathbf{X}=\left\{\mathbf{x}^{(i)}\right\}_{i=1}^{N}\) : continuous/discrete한 iid sample들
• \(\mathbf{z}\) : continuous random variable

우리는 data가 아래와 같은 step을 통해서 생성되었다고 가정한다.

• step 1) \(\mathbf{z}^{(i)}\) 는 어떠한 prior \(p_{\boldsymbol{\theta}^{*}}(\mathbf{z})\) 를 통해 sample 된다
• step 2) \(\mathbf{x}^{(i)}\) 는 \(p_{\boldsymbol{\theta}^{*}}(\mathbf{x} \mid \mathbf{z})\)에서 sample 된다.

( 여기서 , true parameter \(\theta^{*}\) 와 latent variable인 \(\mathbf{z}^{(i)}\)는 unknown인 상황이다 )

이 논문은, 아래의 두 상황에서도 잘 작동하는 general한 알고리즘을 제안한다

• (1) Intractability
• (2) Large dataset

위의 두 가지 상황 속에서도, 아래와 같은 solution을 제안한다.

Efficient approximate…

• 1) ML or MAP estimation for the paramters \(\theta\)
• 2) posterior inference of the latent variable \(\mathbf{z}\)
• 3) marginal inference of the variable \(\mathbf{x}\)

그러기 위해, recognition model ( \(q_{\phi}(\mathbf{z} \mid \mathbf{x}\) ) )를 도입한다. ( = 이는 intractable한 true posterior인 \(p_{\theta}(\mathbf{z} \mid \mathbf{x})\) 를 근사하는 함수이다. )

기존의 MFVI에서 true posterior를 근사하는 \(q\)와의 차이점은, factorize하지 않는다는 점과, parameter \(\phi\)가 closed-form expectation으로부터 계산되지 않는 다는 점이다. 대신에, recognition model의 parameter인 \(\phi\)를, generative model parameter인 \(\theta\)와 jointly하게 학습하는 방법을 제안한다.

앞으로는 조금 다른 관점으로 parameter들을 바라볼 것이다.

• \(\mathbf{z}\) : latent variable = code

  \(\rightarrow\) recognition model \(q_{\phi}(\mathbf{z} \mid \mathbf{x})\)는 “probabilistic ENCODER“이다!

• \(p_{\theta}(\mathbf{x} \mid \mathbf{z})\) : probabilistic DECODER이다.

2.2 The Variational Bound

다음을 각각의 individual datapoint의 marginal likelihood의 sum으로 나타낼 수 있다.

\(\log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(1)}, \cdots, \mathbf{x}^{(N)}\right)=\sum_{i=1}^{N} \log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)}\right),\).

위 식을, 아래와 같이 ELBO를 포함하도록 재정리할 수 있다.

\(\log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)}\right)=D_{K L}\left(q_{\boldsymbol{\phi}}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right) \| p_{\boldsymbol{\theta}}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right)\right)+\mathcal{L}\left(\boldsymbol{\theta}, \boldsymbol{\phi} ; \mathbf{x}^{(i)}\right)\).

• \(\log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)}\right) \geq \mathcal{L}\left(\boldsymbol{\theta}, \boldsymbol{\phi} ; \mathbf{x}^{(i)}\right)=\mathbb{E}_{q_{\phi}(\mathbf{z} \mid \mathbf{x})}\left[-\log q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})+\log p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z})\right]\).
• (ELBO) \(\mathcal{L}\left(\boldsymbol{\theta}, \boldsymbol{\phi} ; \mathbf{x}^{(i)}\right)=-D_{K L}\left(q_{\boldsymbol{\phi}}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right) \| p_{\boldsymbol{\theta}}(\mathbf{z})\right)+\mathbb{E}_{q_{\phi}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right)}\left[\log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)} \mid \mathbf{z}\right)\right]\)

위 식을, variational parameter인 \(\phi\)와, generative parameter인 \(\theta\)에 대해 최적화하고 싶다.

하지만, \(\phi\)에 대해서는 이것이 바로 풀리지 않기 때문에, MC gradient estimator를 통해 아래와 같이 풀 수 있다.

\(\nabla_{\phi} \mathbb{E}_{q_{\phi}(\mathbf{z})}[f(\mathbf{z})]=\mathbb{E}_{q_{\phi}(\mathbf{z})}\left[f(\mathbf{z}) \nabla_{q_{\phi}(\mathbf{z})} \log q_{\boldsymbol{\phi}}(\mathbf{z})\right] \simeq \frac{1}{L} \sum_{l=1}^{L} f(\mathbf{z}) \nabla_{q_{\phi}\left(\mathbf{z}^{(l)}\right)} \log q_{\phi}\left(\mathbf{z}^{(l)}\right)\).

​ where \(\mathbf{z}^{(l)} \sim q_{\phi}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right)\).

위 식의 문제점은, high variance를 가진다는 점이다. 이는 알고리즘 속도의 저하를 가져오기 때문에, 반드시 해결해야하는 문제이다!

2-3. The SGVB estimator and AEVB algorithm

ELBO와, 이것의 derivative에 대한 practical한 estimator를 제안한다.

이를 설명하기에 앞서서, reparameterization trick에 대한 이해가 필요하다.

우리는 \(\widetilde{\mathbf{z}} \sim q_{\phi}(\mathbf{z} \mid \mathbf{x})\)에서 \(\widetilde{\mathbf{z}}\)를 샘플해야 한다. 이를 아래와 같은 (auxiliary) noise variable \(\epsilon\)을 통해, reparameterize할 수 있다.

\(\widetilde{\mathbf{z}}=g_{\phi}(\boldsymbol{\epsilon}, \mathbf{x}) \quad \text { with } \quad \boldsymbol{\epsilon} \sim p(\boldsymbol{\epsilon})\).

\(\mathbb{E}_{q_{\phi}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right)}[f(\mathbf{z})]=\mathbb{E}_{p(\boldsymbol{\epsilon})}\left[f\left(g_{\boldsymbol{\phi}}\left(\boldsymbol{\epsilon}, \mathbf{x}^{(i)}\right)\right)\right] \simeq \frac{1}{L} \sum_{l=1}^{L} f\left(g_{\boldsymbol{\phi}}\left(\boldsymbol{\epsilon}^{(l)}, \mathbf{x}^{(i)}\right)\right) \quad \text { where } \quad \boldsymbol{\epsilon}^{(l)} \sim p(\boldsymbol{\epsilon})\).

위의 reparameterization trick을 우리의 ELBO에 적용하면 아래와 같이 나타낼 수 있고, 이렇게 해서 나온 estimator를 SGVB estimator 라고 한다.

이는 두 가지로 표현하여 서로 다른 해석을 할 수 있다.

(1 번째 표현) ELBO = \(\mathbb{E}_{q_{\phi}(\mathbf{z} \mid \mathbf{x})}\left[-\log q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})+\log p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z})\right]\)로 보는 경우

\(\widetilde{\mathcal{L}}^{A}\left(\boldsymbol{\theta}, \phi ; \mathbf{x}^{(i)}\right) \simeq \mathcal{L}\left(\boldsymbol{\theta}, \phi ; \mathbf{x}^{(i)}\right)\).

\(\begin{array}{l} \widetilde{\mathcal{L}}^{A}\left(\boldsymbol{\theta}, \phi ; \mathbf{x}^{(i)}\right)=\frac{1}{L} \sum_{l=1}^{L} \log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)}, \mathbf{z}^{(i, l)}\right)-\log q_{\boldsymbol{\phi}}\left(\mathbf{z}^{(i, l)} \mid \mathbf{x}^{(i)}\right) \\ \text { where } \mathbf{z}^{(i, l)}=g_{\boldsymbol{\phi}}\left(\boldsymbol{\epsilon}^{(i, l)}, \mathbf{x}^{(i)}\right) \text { and } \boldsymbol{\epsilon}^{(l)} \sim p(\boldsymbol{\epsilon}) \end{array}\).

(2 번째 표현) ELBO = \(-D_{K L}\left(q_{\boldsymbol{\phi}}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right) \| p_{\boldsymbol{\theta}}(\mathbf{z})\right)+\mathbb{E}_{q_{\phi}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right)}\left[\log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)} \mid \mathbf{z}\right)\right]\)로 보는 경우

\(\widetilde{\mathcal{L}}^{B}\left(\boldsymbol{\theta}, \phi ; \mathbf{x}^{(i)}\right) \simeq \mathcal{L}\left(\boldsymbol{\theta}, \phi ; \mathbf{x}^{(i)}\right)\).

\(\begin{array}{l} \widetilde{\mathcal{L}}^{B}\left(\boldsymbol{\theta}, \phi ; \mathbf{x}^{(i)}\right)=-D_{K L}\left(q_{\boldsymbol{\phi}}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right) \| p_{\boldsymbol{\theta}}(\mathbf{z})\right)+\frac{1}{L} \sum_{l=1}^{L}\left(\log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)} \mid \mathbf{z}^{(i, l)}\right)\right) \\ \text { where } \mathbf{z}^{(i, l)}=g_{\boldsymbol{\phi}}\left(\boldsymbol{\epsilon}^{(i, l)}, \mathbf{x}^{(i)}\right) \text { and } \boldsymbol{\epsilon}^{(l)} \sim p(\boldsymbol{\epsilon}) \end{array}\).

• 1번째 term 해석 ) regularizer
• 2번째 term 해석 ) expected negative reconstruction error

위의 방법을 통해 재표현한 ELBO를, 다음과 같이 minibatch 방법을 통해 나타낼 수 있다.

\(\mathcal{L}(\boldsymbol{\theta}, \boldsymbol{\phi} ; \mathbf{X}) \simeq \widetilde{\mathcal{L}}^{M}\left(\boldsymbol{\theta}, \boldsymbol{\phi} ; \mathbf{X}^{M}\right)=\frac{N}{M} \sum_{i=1}^{M} \widetilde{\mathcal{L}}\left(\boldsymbol{\theta}, \phi ; \mathbf{x}^{(i)}\right)\).

​ ( \(N\) datapoints, where minibatch \(\mathbf{X}^{M}=\left\{\mathbf{x}^{(i)}\right\}_{i=1}^{M})\)

2-4. The reparameterization trick

Reparameterization trick에 대해서는 위의 2-3에서 전부 다 설명했다.

요약하자면, 아래와 같이 나타낼 수 있다.

• \(\mathbf{z}\) : continuous random variable
• \(\mathbf{z} \sim\)\(q_{\phi}(\mathbf{z} \mid \mathbf{x}) .\)
• 이 \(\mathbf{z}\)를 \(\mathbf{z}=g_{\phi}(\boldsymbol{\epsilon}, \mathbf{x})\)로 재표현!
  • where \(\boldsymbol{\epsilon}\) is an auxiliary variable with independent marginal \(p(\boldsymbol{\epsilon}),\)
  • \(g_{\boldsymbol{\phi}}(.)\) is some vector-valued function parameterized by \(\phi .\)

그렇다면, 이 trick을 사용하는 이유는 무엇일까?

이는 expectation w.r.t \(q_{\phi}(\mathbf{z} \mid \mathbf{x})\)를 \(\phi\)에 대해 미분가능하도록 표현할 수 있기 때문이다!

3. Example : VAE (Variational Autoencdoer)

probabilistic encoder \(q_{\phi}(\mathbf{z} \mid \mathbf{x})\)를 Neural Net으로 설정하는 경우에 대해 설명할 것이다.

( 이는 곧 generative model \(p_{\theta}(\mathbf{x,z})\)의 posterior에 대한 approximation이다 )

이때, parameter \(\phi\)와 \(\theta\)는 “JOINTLY”하게 optimize된다.

이러한 알고리즘을, AEVB (Auto Encoder Variational Bayes) 라고 한다.

설정

• prior : \(p_{\theta}(\mathbf{z})=\mathcal{N}(\mathbf{z} ; \mathbf{0}, \mathbf{I})\)

• \(p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})\) : MVN (in case of real-valued data) or Bernoulli (in case of binary data)

  ( 이 때의 distribution parameter는 MLP를 통한 \(\mathbf{z}\)에 대한 함수로 정해진다 )

• \(p_{\boldsymbol{\theta}}(\mathbf{z} \mid \mathbf{x})\) 는 intractable하다

• \(\log q_{\phi}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right)=\log \mathcal{N}\left(\mathbf{z} ; \boldsymbol{\mu}^{(i)}, \boldsymbol{\sigma}^{2(i)} \mathbf{I}\right)\).

  • \(\mu^{(i)}\) & \(\sigma^{(i)}\) 는 encoding MLP의 output들이다.

    ( = nonlinear function of data \(x^{(i)}\) & variational parameter \(\phi\) )

위의 상황 속에서, 아래의 reparameterization trick을 통해

• \(\mathbf{z}^{(i, l)} \sim q_{\phi}\left(\mathbf{z} \mid \mathbf{x}^{(i)}\right)\).

• \(g_{\phi}\left(\mathbf{x}^{(i)}, \boldsymbol{\epsilon}^{(l)}\right)=\boldsymbol{\mu}^{(i)}+\boldsymbol{\sigma}^{(i)} \odot \boldsymbol{\epsilon}^{(l)}\).

  where \(\boldsymbol{\epsilon}^{(l)} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\).

ELBO를 아래와 같이 정리할 수 있다.

\(\begin{array}{l} \mathcal{L}\left(\boldsymbol{\theta}, \boldsymbol{\phi} ; \mathbf{x}^{(i)}\right) \simeq \frac{1}{2} \sum_{j=1}^{J}\left(1+\log \left(\left(\sigma_{j}^{(i)}\right)^{2}\right)-\left(\mu_{j}^{(i)}\right)^{2}-\left(\sigma_{j}^{(i)}\right)^{2}\right)+\frac{1}{L} \sum_{l=1}^{L} \log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)} \mid \mathbf{z}^{(i, l)}\right) \\ \text { where } \mathbf{z}^{(i, l)}=\boldsymbol{\mu}^{(i)}+\boldsymbol{\sigma}^{(i)} \odot \boldsymbol{\epsilon}^{(l)} \text { and } \boldsymbol{\epsilon}^{(l)} \sim \mathcal{N}(0, \mathbf{I}) \end{array}\).