Markov Chain Monte Carlo and Variational Inference ; Briding the Gap ( 2015 )

Abstract

Stochastic Gradient Variational Bayes (SGVB) 덕분에, variational inference에서 auxiliary r.v를 사용하는 것이 가능해졌다. 이로 인해, variational approximation 단계에서 한 개 이상의 MCMC step을 추가하는 것이 가능해졌다. 이로 인해, rich class of inference algorithm이 가능해졌다.


1. MCMC and VI

Bayesian에서 posterior의 intractability를 푸는 데에는 크게 2가지 방법인 (1) VI와 (2) MCMC가 있다. (1) VI는 explicit한 objective를 maximize하는 문제로, 더 빠르다는 장점이 있지만, (2) MCMC의 경우네는 non-parametric하고 (computation 능력만 충분하다면) asymptotically exact하다라는 장점이 있다. 여기서는 이 두 방법이 서로 잘 combine될 수 있음을 보여준다.

1.1 Variational Inference

posterior approximation인 \(q_{\theta}(z \mid x)\)를 구하는 문제로, 다음의 Lower Bound를 최대화 하는 문제로 바꿔서 풀 수 있다.

\(\begin{aligned} \log p(x) & \geq \log p(x)-D_{K L}\left(q_{\theta}(z \mid x) \| p(z \mid x)\right) \\ &=\mathbb{E}_{q_{\theta}(z \mid x)}\left[\log p(x, z)-\log q_{\theta}(z \mid x)\right]=\mathcal{L} \end{aligned}\).

이는 곧 아래의 KL-divergence를 minimize하는 것과 같다.

\(D_{K L}\left(q_{\theta}(z \mid x) \| p(z \mid x)\right)\).

( 즉, \(p(z \mid x)\)와 \(q_{\theta}(z \mid x)\)가 일치할 때 optimal하다)


1.2 MCMC and Auxiliary Variables

VI와 함께 대표적인 또 다른 방법이 바로 MCMC이다.

MCMC도 VI와 마찬가지로, initial distribution \(q(z_0)\)에서 \(z_0\)를 sample하는 것으로 시작한다.

하지만, VI와 다르게 해당 분포를 optimize하는 것이 아니라, 아래와 같이 stochastic transition operator를 순차적으로 적용해서 sampling을 해나간다.

\(z_{t} \sim q\left(z_{t} \mid z_{t-1}, x\right)\).

해당 transition을 잘 고르면, 이는 exact posterior \(p(z\mid x)\)로 수렴하게 된다. 다만, 이에 필요한 시간이 매우 길 수 있다는 점이 해당 방법의 단점이다.

이 논문에서 제시하는 바는, 아래와 같다.

We can interpret the stochastic Markov chain \(q(z \mid x)=\) \(q\left(z_{0} \mid x\right) \prod_{t=1}^{T} q\left(z_{t} \mid z_{t-1}, x\right)\) as a variational approximation in an expanded space by considering \(y=z_{0}, z_{1}, \ldots, z_{t-1}\) to be a set of auxiliary random variables.

즉, MCMC에서 샘플되는 \(z\)들의 모음을 일종의 Variational Approximation에서 사용하는 auxiliary variable로 보는 것이다!

위의 auxiliary variable를 ELBO안으로 integrate 할 경우, 아래와 같은 ELBO를 구할 수 있다.

\[\begin{aligned} \mathcal{L}_{\text {aux }} &=\mathbb{E}_{q\left(y, z_{T} \mid x\right)}\left[\log \left[p\left(x, z_{T}\right) r\left(y \mid x, z_{T}\right)\right]-\log q\left(y, z_{T} \mid x\right)\right.]\\ &=\mathcal{L}-\mathbb{E}_{q\left(z_{T} \mid x\right)}\left\{D_{K L}\left[q\left(y \mid z_{T}, x\right) \| r\left(y \mid z_{T}, x\right)\right]\right\} \\ & \leq \mathcal{L} \\ &\leq \log [p(x)] \end{aligned}\]
  • 여기서 \(r\left(y \mid x, z_{T}\right)\)는 auxiliary inference distribution이다.
  • \(q\left(z_{T} \mid x\right)=\int q\left(y, z_{T} \mid x\right) \mathrm{d} y\).
    • 해석 : mixture of distributions of the form \(q\left(z_{T} \mid x, y\right)\)
    • rich class of distribution!

위의 \(\mathcal{L}_{\text {aux }}\)를 maximize하는 조건으로는, \(r\left(y \mid x, z_{T}\right)=q\left(y \mid x, z_{T}\right)\)이 될 것이다. 하지만 이는 주로 intractable하기 때문에, \(q\left(y \mid x, z_{T}\right)\)를 잘 근사하는 \(r\left(y \mid x, z_{T}\right)\)로 설정하는 경우가 많다.

이 논문에서는 위의 auxiliary inference distribution ( = \(r\left(y \mid x, z_{T}\right)\) )를 posterior distribution과 마찬가지로 아래와 같이 Markov structure를 가진다고 가정한다.

\(r\left(z_{0}, \ldots, z_{t-1} \mid x, z_{T}\right)=\prod_{t=1}^{T} r_{t}\left(z_{t-1} \mid x, z_{t}\right)\).

따라서 ELBO를 다시 아래와 같이 재정리할 수 있다.

\(\begin{aligned} \mathcal{L}_{\text {aux }} &=\mathbb{E}_{q\left(y, z_{T} \mid x\right)}\left[\log \left[p\left(x, z_{T}\right) r\left(y \mid x, z_{T}\right)\right]-\log q\left(y, z_{T} \mid x\right)\right.]\\ &=\mathbb{E}_{q\left(y, z_{T} \mid x\right)}\left[\log p\left(x, z_{T}\right)+\text{log} r\left(y \mid x, z_{T}\right)-\log q\left(y, z_{T} \mid x\right)\right.]\\ &=\mathbb{E}_{q}\left[\log p\left(x, z_{T}\right)-\log q\left(z_{0}, \ldots, z_{T} \mid x\right)\right.\left.+\log r\left(z_{0}, \ldots, z_{t-1} \mid x, z_{T}\right)\right]\\ &=\mathbb{E}_{q}[\log [p(x, z_{T}) / q(z_{0} \mid x)]+\sum_{t=1}^{T} \log [r_{t}(z_{t-1} \mid x, z_{t}) / q_{t}(z_{t} \mid x, z_{t-1})]]\\ &\leq \log p(x) \end{aligned}\)

위의 transition operator \(q_t\)와, inverse model \(r_t\)를 어떻게 flexible한 parametric form으로 지정하느냐에 따라 ELBO를 optimize할 수 있다.


2. Optimizing the lower bound

위의 transition operator \(q_t\)와, inverse model \(r_t\)를 설정 해도, 위의 ELBO는 대부분 analytically하게 풀리지 않는다.

하지만, 적어도 \(q_t\)로 부터 sampling을 할 수 있고, 해당 샘플을 \(r_t\)에 대해서 evaluate할 수 있다면, 우리는 위의 ELBO를 잘 근사할 수 있다!

( with 아래의 알고리즘 )

figure2

SGVI (Stochastic Gradient Variational Inference)의 주요 특징 중 하나는, 모든 step이 \(\theta\)에 대해 (위의 \(q\)와 \(r\)에 대해) differentiable하다면, 이로 인한 결과물인 \(L\) 또한 differentiable하다. \(L\)은 unbiased estimate of ELBO이기 때문에, 이의 derivative 또한 unbiased estimate이다. 따라서 우리는 이에 stochastic optimization에 사용할 수 있다.

위에서 \(z_t \sim q_t(z_t \mid x, z_{t-1})\)의 단계는, 아래의 2 step을 통해 실행할 수 있다

  • 1) \(u_t \sim p(u_t)\)
  • 2) \(z_t = g_{\theta}(u_t,x)\).

위를 통해 stochastic estimate of the gradient of ELBO w.r.t \(\theta\)를 구했다. 우리는 이를 true posterior \(p(z \mid x)\)에 대한 stochastic gradient-based optimization에 사용할 수 있다.

해당 알고리즘은 아래와 같다.

figure2


2-1. ex) bivariate Gaussian

아래의 Gaussian distribution에서 sampling하는 예시를 보자.

\(p\left(z^{1}, z^{2}\right) \propto \exp \left[-\frac{1}{2 \sigma_{1}^{2}}\left(z^{1}-z^{2}\right)^{2}-\frac{1}{2 \sigma_{2}^{2}}\left(z^{1}+z^{2}\right)^{2}\right]\).

위의 분포를 사용해서 \(z^1\)와 \(z^2\)를 교대로 update한다. 이를 다음과 같은 2가지 방법으로 실행할수 있다.

1) Gibbs sampling

  • sample from Gaussian full conditional distributions.
  • \(p\left(z^{i} \mid z^{-i}\right)=N\left(\mu_{i}, \sigma_{i}^{2}\right)\).

2) Over-relaxation method

  • \(z^i\)를 다음 식을 통해서 update한다.
  • \(q\left(z_{t}^{i} \mid z_{t-1}\right)=N\left[\mu_{i}+\alpha\left(z_{t-1}^{i}-\mu_{i}\right), \sigma_{i}^{2}\left(1-\alpha^{2}\right)\right]\).
  • 위 식에서
    • \(\alpha=0\) : Gibbs sampler와 똑같다
    • \(\alpha \neq 0\) :
    • 2) 방법이 Gibbs보다 빠를 수 있다.

( + inverse model (\(r\left(z_{t-1} \mid z_{t}\right)\) = Gaussian )


3. Hamiltonian Variational Inference

HMC의 핵심은, momentum variable인 \(v\)를 (auxiliary variable로써) 사용한다는 점이다.

Hamiltonian dynamics는 ( exact log posterior의 gradient의 도움을 받아 ) posterior distribution을 매우 효과적으로 explore한다.

이러한 auxiliary variable은 다음과 같은 분포에서 sample된다.

\(v_{t}^{\prime} \sim q\left(v_{t}^{\prime} \mid x, z_{t-1}\right)\).

  • notation:
    • from) \(v_{t}^{\prime}, z_{t-1}\)
    • to) \(v_{t}, z_{t}\)

HMC에서 transition은 deterministic, invertible, volume preserving하다.

즉, 아래의 식이 성립한다.

\(\begin{array}{c} q\left(v_{t}, z_{t} \mid z_{t-1}, x\right)=q\left(v_{t}, z_{t}, z_{t-1} \mid x\right) / q\left(z_{t-1} \mid x\right) \\ =q\left(v_{t}^{\prime}, z_{t-1} \mid x\right) / q\left(z_{t-1} \mid x\right)=q\left(v_{t}^{\prime} \mid z_{t-1}, x\right) \end{array}\).

\(r\left(v_{t}^{\prime}, z_{t-1} \mid z_{t}, x\right)=r\left(v_{t} \mid z_{t}, x\right)\).
이와 같은

  • transition operator \(q_{t}\left(v_{t}, z_{t} \mid z_{t-1}, x\right)\)와
  • inverse model \(r_{t}\left(v_{t}^{\prime}, z_{t-1} \mid z_{t}, x\right)\)를 통해,

우리는 아래의 알고리즘과 처럼 stochastic approximation of the log marginal likelihood lower bound를 구할 수 있다.

figure2