60.Hierarchical Variational Models (2015)

Hierarchical Variational Models ( 2015 )

Abstract

이 논문에서는 hierarchical variational models (HIERARCHICAL VM)를 통해, variational distribution을 (1) expressive하게 표현하면서도, (2) efficient하게 computation 할 수 있는 방법을 제안한다. 이름 “hierarchical”답게, parameter들에게 prior를 부여하고, 이러한 shared structure하에서 latent variable들은 서로 conditionally independent하게 된다.

1. Introduction

BBVI (Black Box Variational Inference)는, 연구자로 하여금 아무 probabilistic model을 specify할 수 있게끔 한다. 이는 주로 mean-field family와 함께 이루어지는데, 이는 BBVI 알고리즘을 efficient하게 만든다는 점이 있지만, 강한 factorization으로 dependency를 무시한다는 한계가 있다.

따라서, 이 논문에서는 mean-field를 넘어서는 BBVI 방법을 제안한다. 여기서 제안하는 방법은 tractable하면서도 richer family of variational distribution을 사용한다.

여기서 핵심은 다음과 같다.

“We treat the original variational family as a “model” of the latent variables, and then expand it hierarchically ~

여기서 제안하는 variational model은 다음과 같은 2-stage distribution이다.

• stage 1) prior로부터 variational parameter를 뽑는다
• stage 2) 위에서 뽑은 variational distribution으로부터 latent variable를 뽑는다.

Hierarchical VM에서는, mean-field parameter를 직접적으로 fitting시키지 않고, prior의 hyperparameter를 fitting한다.

한 줄 요약 : In this paper, we define HVM, and develop a general algorithm for fitting them in the context of black box inference

2. BBVI (Black Box Variational Inference)

ELBO : \(\mathscr{L}(\lambda)=\mathbb{E}_{q(z ; \lambda)}[\log p(\mathbf{x}, \mathbf{z})-\log q(\mathbf{z} ; \lambda)] .\).

Black box 방법은, variational approximation으로부터의 sample들을 통해 noisy gradient를 구함으로써 위의 ELBO를 최대화 한다. 이것은 model의 log-likelihood만 계산할 것을 요구하기 때문에, model-specific하지 않다. 이러한 것을 바탕으로, 이 논문은 더 richer한 variational family를 사용한다.

3. Variational Models

Black Box variational method들은 모든 probabilistic model에 적용 가능하지만, 어떠한 variational distribution을 특정할 것인지에 대해 의문이 남는다. Black box method는 아래의 3가지 사항을 만족시키는 approximating distribution을 찾아야 한다.

• 1) generating process for obtaining posterior samples \(\mathbf{z}\)
• 2) computation of its log density function \(\text{log}q(\mathbf{z})\)
• 3) proper support within the posterior

MFVM

이를 충족시키는 대표적인 것이 바로 Mean-Field Variational models이다.

\(q_{\mathrm{MF}}(\mathbf{z} ; \boldsymbol{\lambda})=\prod_{i=1}^{d} q\left(\mathbf{z}_{i} ; \lambda_{i}\right)\).

• \(\lambda_i\) : \(i\)번째 latent variable의 parameter

• 장) computationally feasible

  단) expressivity 포기

\(\rightarrow\) need more expressive model!

HVM

그래서 필요한 것이 바로 HVM, Hierarchical Variational models이다.

parameter \(\lambda\)에 prior (with hyperparameter \(\theta\) )를 부여한뒤, 이를 marginalize out한다.

\(q_{\mathrm{HVM}}(\mathbf{z} ; \theta)=\int\left[\prod_{i=1}^{d} q\left(\mathbf{z}_{i} \mid \lambda_{i}\right)\right] q(\lambda ; \theta) \mathrm{d} \lambda\).

• \(q(\lambda ; \theta)\) : variational prior

  ( = distribution over variational distribution )

  ( 이 구조가 주어진 하에서, \(\mathbf{z}\)들은 서로 conditionally independent하다 )

이 모델 하에서의 ELBO는 아래와 같다.

\(\mathscr{L}(\boldsymbol{\theta})=\mathbb{E}_{q_{\mathrm{HvM}}(z ; \theta)}\left[\log p(\mathrm{x}, \mathrm{z})-\log q_{\mathrm{HVM}}(\mathrm{z} ; \theta)\right]\).

• 1번째 term ) \(q\)로부터 sample만 할 수 있으면 tractable

• 2번째 term ) entropy ( 일반적으로 intractable )

  • \(r(\lambda \mid \mathbf{z} ; \phi)\)를 도입함으로써, 이 entropy에 대해 bound를 만들 수 있다.

    \(-\mathbb{E}_{q_{\mathrm{HVM}}}\left[\log q_{\mathrm{HVM}}(\mathbf{z})\right] \geq-\mathbb{E}_{q(\mathrm{z}, \lambda)}[\log q(\lambda)+\log q(\mathbf{z} \mid \lambda)-\log r(\lambda \mid \mathbf{z} ; \phi)]\).

  • 이 inequality는, \(r(\lambda \mid \mathbf{z} ; \phi)\)가 variational posterior \(q(\lambda \mid \mathbf{z} ; \boldsymbol{\theta})\)와 일치할 때 등식이 성립한다.

    ( 따라서 우리는 \(r\)를 recursive variational approximation으로 볼 수 있다 )

위의 entropy의 lower bound를 반영하여 새로 만든 Hierarchical ELBO는 아래와 같다.

\[\begin{array}{l} \widetilde{\mathscr{L}}(\theta, \phi)=\mathbb{E}_{q(z, \lambda ; \theta)}[\log p(\mathbf{x}, \mathbf{z})+\log r(\lambda \mid \mathbf{z} ; \phi) \left.-\sum_{i=1}^{d} \log q\left(\mathbf{z}_{i} \mid \lambda_{i}\right)-\log q(\lambda ; \theta)\right] \end{array}\]

• \(\mathscr{L}(\boldsymbol{\theta})=\mathbb{E}_{q_{\mathrm{HvM}}(z ; \theta)}\left[\log p(\mathrm{x}, \mathrm{z})-\log q_{\mathrm{HVM}}(\mathrm{z} ; \theta)\right]\).와
• \(-\mathbb{E}_{q_{\mathrm{HVM}}}\left[\log q_{\mathrm{HVM}}(\mathbf{z})\right] \geq-\mathbb{E}_{q(\mathrm{z}, \lambda)}[\log q(\lambda)+\log q(\mathbf{z} \mid \lambda)-\log r(\lambda \mid \mathbf{z} ; \phi)]\)를 결합!

이를 EM 알고리즘 관점에서 해석하면,

• E-step ) \(\theta\)를 optimize하는 것은 posterior approximation을 improve하는 것이고,
• M-step) \(\phi\)를 optimize하는 것은 upper bound를 타이트하게 하는 것이다.

위의 Hierarchical ELBO를, Mean-field lower bound ( = \(\mathscr{L}_{\mathrm{MF}}(\lambda)\) )의 관점에서 다시 작성하면,

\(\widetilde{\mathscr{L}}(\theta, \phi)=\mathbb{E}_{q}\left[\mathscr{L}_{\mathrm{MF}}(\lambda)\right]+\mathbb{E}_{q}[\log r(\lambda \mid \mathrm{z} ; \phi)-\log q(\lambda ; \theta)]\)이다.

• 1번째 term ) Bayesian model average of mean-field objectives

  ( weight = variational prior \(q(\lambda ; \theta)\) )

• 2번째 term ) correction term

4. Specifying the Hierarchical Variational Model

HVM을 specify하는데에는 다음과 같은 2가지 component를 요한다.

• 1) variational likelihood : \(q(\mathbf{z} \mid \lambda)\)
  • ex) mean-field
• 2) prior : \(q(\lambda ; \theta)\)
  • 2가지 요건을 충족해야함
    • 1) \(\lambda\) should not be independent
    • 2) \(q(\lambda ; \theta)\)가 discrete & continuous latent variable에 모두 적용 가능!

Variational prior의 ex)

• Mixture of Gaussian

  \(q(\lambda ; \theta)=\sum_{i=1}^{K} \pi_{k} \mathrm{~N}\left(\mu_{k}, \sigma_{k}\right) .\).

• Normalizing flow ( Bayesian Neural Network에 이와 관련된 포스트가 있으니, 이를 참고하길 바람 )

  $$\log r(\lambda \mid \mathbf{z})=\log r\left(\lambda_{0} \mid \mathbf{z}\right)+\sum_{k=1}^{K} \log \left(\left

  \operatorname{det}\left(\frac{\partial g_{k}^{-1}}{\partial \lambda_{k}}\right)\right

  \right)$$.

  \(r\left(\lambda_{0} \mid \mathbf{z}\right)=\prod_{i=1}^{d} r\left(\lambda_{0 i} \mid \mathbf{z}_{i}\right)\).

5. Optimizing the Hierarchical ELBO

(1) Stochastic gradient of ELBO

(a) score function estimator

ELBO의 gradient에 대한 estimator로써, score function estimator는 discrete / continuous한 경우에 모두 적용 가능하다.

( 이는 REINFORCE gradient로도 알려져 있다. )

\(\nabla_{\lambda}^{\text {score }} \mathscr{L}=\mathbb{E}_{q(\mathbf{z} \mid \lambda)}\left[\nabla_{\lambda} \log q(\mathbf{z} \mid \lambda)(\log p(\mathbf{x}, \mathbf{z})-\log q(\mathbf{z} \mid \lambda))\right]\).

위 식을 MC estimate로 근사하면 아래와 같다.

\(\begin{aligned} \frac{1}{S} \sum_{s=1}^{S} \nabla_{\lambda} \log q\left(\mathbf{z}^{S} \mid \lambda\right)\left(\log p\left(\mathbf{x}, \mathbf{z}^{s}\right)-\log q\left(\mathbf{z}^{s} \mid \lambda\right)\right) \\ \text { where } \mathbf{z}^{s} \sim q(\mathbf{z} \mid \lambda) \end{aligned}\).

하지만, score function estimator는 high variance의 문제를 종종 가진다.

(b) mean field ( + reparam trick )

따라서 Reparameterization trick을 사용하여 variance를 줄인다.

mean field model에서와, reparam trick을 적용했을 때의 gradient of ELBO는 아래와 같다.

\(\nabla_{\lambda_{i}} \mathscr{L}_{\mathrm{MF}}=\mathbb{E}_{q\left(\mathbf{z}_{i} ; \lambda_{i}\right)}\left[\nabla_{\lambda_{i}} \log q\left(\mathbf{z}_{i} ; \lambda_{i}\right)\right.\left.\left(\log p_{i}(\mathbf{x}, \mathbf{z})-\log q\left(\mathbf{z}_{i} ; \lambda_{i}\right)\right)\right]\).

\(\nabla_{\lambda}^{r e p} \mathscr{L}=\mathbb{E}_{s(\epsilon)}\left[\left(\nabla_{z} \log p(\mathbf{x}, \mathbf{z})-\nabla_{z} \log q(\mathbf{z})\right) \nabla_{\lambda} \mathbf{z}(\epsilon ; \lambda)\right]\).

(2) Stochastic gradient of Hierarchical ELBO

( focus on differentiable priors … ex. Normalizing Flows )

\(\lambda\) 를 reparameterization을 통해 \(\epsilon\) 과 \(\theta\) 에 대한 함수로 나타내자. ( = \(\lambda(\epsilon ; \theta)\) )

그런 다음, score function \(V\) 를 정의한다.

\(V=\nabla_{\lambda} \log q(\mathbf{z} \mid \lambda)\).

그러면, gradient of Hierarchical ELBO는 아래와 같다.

\[\begin{aligned} \nabla_{\theta} & \widetilde{L}(\theta, \phi)=\mathbb{E}_{s(\epsilon)}\left[\nabla_{\theta} \lambda(\epsilon) \nabla_{\lambda} \mathscr{L}_{\mathrm{MF}}(\lambda)\right] \\ &+\mathbb{E}_{s(\epsilon)}\left[\nabla_{\theta} \lambda(\epsilon) \nabla_{\lambda}[\log r(\lambda \mid \mathrm{z} ; \phi)-\log q(\lambda ; \theta)]\right] \\ &+\mathbb{E}_{s(\epsilon)}\left[\nabla_{\theta} \lambda(\epsilon) \mathbb{E}_{q(\mathrm{z} \mid \lambda)}[V \log r(\lambda \mid \mathrm{z} ; \phi)]\right] \end{aligned}\]

(3) Local Learning

위의 (2)의 local한 version이다.

\(V_{i}=\nabla_{\lambda} \log q\left(\mathbf{z}_{i} \mid \lambda_{i}\right)\).

그러면, 위 \(\nabla_{\theta} \widetilde{L}(\theta, \phi)\)에서 세번째 term은 다음으로 바뀌게 된다.

\[\begin{array}{l} \mathbb{E}_{s(\epsilon)}\left[\nabla_{\theta} \lambda(\epsilon ; \theta) \mathbb{E}_{q(\mathbf{z} \mid \lambda)}[V \log r(\lambda \mid \mathbf{z} ; \boldsymbol{\phi})]\right] \\ =\mathbb{E}_{s(\epsilon)}\left[\nabla_{\theta} \lambda(\epsilon ; \boldsymbol{\theta}) \mathbb{E}_{q(\mathbf{z} \mid \lambda)}\left[\sum_{i=1}^{d} V_{i} \log r_{i}(\boldsymbol{\lambda} \mid \mathbf{z} ; \boldsymbol{\phi})\right]\right] \end{array}\]

(4) Stochastic gradient w.r.t \(\phi\)

\(\nabla_{\phi} \widetilde{\mathscr{L}}=\mathbb{E}_{q(z, \lambda)}\left[\nabla_{\phi} \log r(\lambda \mid \mathbf{z}, \phi)\right]\).

Summary

Hierarchical VM을 사용한 Black box inference 알고리즘을 정리하면, 아래와 같다.