Black Box Variational Inference ( 2013 )

Abstract

이 논문은 “Black box” VI 알고리즘을 제안한다.

핵심 : “Stochastic Optimization”에 기반한다!

( ELBO의 “noisy gradient가 MC sample”로 부터 구해진다 )

이 뿐만 아니라, stochastic gradient에 대한 다양한 “variance reduction” 테크닉을 제안한다.

이 알고리즘은 기존의 샘플링 방법들보다 더 빠르고 나은 성능을 보인다.

1. Introduction

빠른 복습!

Posterior를 정확히 계산하는 것은 불가능하기 때문에, 근사(approximate)하는 방법을 사용하는데, 이것의 대표적인 방법이 바로 Variational Inference이다. **간단한 probability distribution **\(q\)를 사용하여, 구하고자하는 true posterior에 근사한다.

하지만, 대부분은 closed-form형태로 존재하지 않는다. 이를 해결하기 위해, model-specific한 알고리즘들이 많이 개발되었다. 하지만, 이 논문에서 제안하는 것은“black box” variational inference 알고리즘은, 거의 모든 모델에 적용 가능(general)하다!

이 방법은, ELBO의 gradient를 다루기 쉬운 함수인 \(f\)를 사용하여 나타낸다.

그런 뒤, 대략적으로 아래와 같은 순서로 진행한다.

1) variational distribution으로부터 sampling을 한다
2) sample로부터 \(f\)를 계산한다

gradient에 대한 MC estimate를 구한다
3) 이러한 noisy gradient를 사용하여 variational parameter를 optimize한다.

이 논문은, 두 가지의 방법을 통해 gradient의 variance를 줄인다.

1) Rao-Blackwellization
2) Control Variates

이 두 방법은, 특정 모델의 형태를 요구하지 않기 떄문에, 위에서 제시한 ‘black-box inference’에 적용하는 데에 무리가 없다. 보다 자세한 내용은 뒤에서 보다 자세히 다루겠다.

마지막으로, 이 알고리즘을 scale up 하고 speed up하기 위한 recent innovation를 소개한다

1) adaptive learning rates
2) generic stochastic variational inference with subsampling

2. Black Box Variational Inference

notation 소개

\(x\) : observation
\(z\) : latent variable
\(\lambda\) : free parameters of variational distribution ( = \(q(z \mid \lambda)\) )

목표 : \(\lambda\)를 update하여 \(p(z\mid x)\)에 근사하는 것!

그러기 위해…

ELBO를 maximize ( = KL divergence를 minimize )

( ELBO : \(\mathcal{L}(\lambda) \triangleq \mathrm{E}_{q_{\lambda}(z)}[\log p(x, z)-\log q(z)]\) )
ELBO의 해석
- 첫번째 term) 잘 fitting하도록
- 두번째 term) entropic하도록 ( overfitting 방지 효과)

이전의 많은 방법들은, ELBO를 최대화하기 위해 coordinate-ascent update방법을 사용해왔다. 하지만 이것은 conjugate family model에서만 closed-form의 형태로 존재를 했었다.

따라서 이 논문은 “stochastic optimization” 한 방법으로 ELBO를 극대화 한다

( = 즉, noisy estimates of the gradient에 대한 함수를 maximize한다 )

Stochastic Optimization

( = ELBO를 gradient ascent 방법을 이용해서 optimize하는 것을 의미 )

notation

\(f(x)\) : 최대화 시키고 싶은 함수
\(H(x)\) : 기대값이 \(f(x)\)의 gradient인 함수
\(h_t(x)\) : \(H(x)\)의 realization
\(\rho_t\) : learning rate (non-negative)

Stochastic Optimization은 아래와 같은 식으로 parameter들을 update해나간다.

\(x_{t+1} \leftarrow x_{t}+\rho_{t} h_{t}\left(x_{t}\right)\).

Robbins-Monro condition하에서는, 위 updating equation은 수렴하게 되어있다.

\(\sum_{t=1}^{\infty} \rho_{t}=\infty\).
\(\sum_{t=1}^{\infty} \rho_{t}^{2}<\infty\).

Noisy gradient of the ELBO

우리는 gradient에 대한 unbiased estimator를 구해야 한다.

그러기 위해 , 우리는 ELBO의 gradient를 아래와 같이 re-write할 수 있다.

\(\nabla_{\lambda} \mathcal{L}=\mathrm{E}_{q}\left[\nabla_{\lambda} \log q(z \mid \lambda)(\log p(x, z)-\log q(z \mid \lambda))\right]\).

여기서 \(\nabla_{\lambda} \log q(z \mid \lambda)\) 는 score function이다

위 식을 MC sample를 통해 근사하면 아래와 같다.

\(\nabla_{\lambda} \mathcal{L} \approx \frac{1}{S} \sum_{s=1}^{S} \nabla_{\lambda} \log q\left(z_{s} \mid \lambda\right)\left(\log p\left(x, z_{s}\right)-\log q\left(z_{s} \mid \lambda\right)\right)\).

where \(z_{s} \sim q(z \mid \lambda)\).

이 식을 사용하여 우리는 stochastic optimization을 진행하면 된다.

지금까지 설명한 알고리즘을 정리하면 아래와 같다

( 여기서 알아야할 점은, score function과 sampling algorithm이 variational distribution에만 의존할 뿐, 우리의 underlying model에는 의존하지 않는다는 점이다 )

3. Controlling the Variance

방금 위에서 소개한 알고리즘은 ELBO를 최대화하는데 사용된다. 하지만 우리가 놓친 점이 있다. 위 알고리즘대로 진행할 경우, estimator of the gradient의 variance가 너무 클 수 있다는 것이다. 이럴 경우, 수렴 속도가 느려지기 때문에, 우리는 다음과 같은 2가지 방법을 통해 variance를 reduce한다.

1) Rao-Blackwellization
2) control variates

( 이전 포스트에서도 다뤘지만, 다시 한번 짚고 넘어가겠다 )

3-1. Rao-Blackwellization

핵심은 다음이다.

“Replace it with its conditional expectation w.r.t a subset of the variables”

( 다음 설명을 통해 쉽게 이해할 수 있을 것이다 )

두 개의 r.v \(X\)와 \(Y\)가 있다고 하자. 목표는 \(X\) 와\(Y\)에 대해 \(\mathrm{E}[J(X, Y)]\)를 최대화하는 것이다.

또한, 다음과 같이 정의를 하자.

\(\hat{J}(X)=\mathrm{E}[J(X, Y) \mid X]\).

이것의 expectation은, \(\mathrm{E}[\hat{J}(X)]=\mathrm{E}[J(X, Y)]\)가 된다.

이것이 의미하는 바는, \(\hat{J}(X)\)를 \(J(X, Y)\)를 대신해서 사용할 수 있다는 점이다. 이것을 이와 같이 대신해서 사용하면 좋은 점은, variance가 줄어든다는 점이다! (아래의 식 참고)

\(\operatorname{Var}(\hat{J}(X))=\operatorname{Var}(J(X, Y))-\mathrm{E}\left[(J(X, Y)-\hat{J}(X))^{2}\right]\).

다시 이전의 ELBO의 gradient를 estimate하는 문제로 돌아와보자. 총 \(n\)개의 latent variable이 있다고 가정하면, 아래와 같이 factorize할 수 있다. (mean-field approximation)

\(q(z \mid \lambda)=\prod_{i=1}^{n} q\left(z_{i} \mid \lambda_{i}\right)\).

따라서, 우리의 ELBO의 gradient는 iterated conditional expectation을 사용하여 아래와 같이 표현할 수 있다.

\(\begin{array}{l} \nabla_{\lambda_{i}} \mathcal{L}= E_{q_{(i)}}\left[\nabla_{\lambda_{i}} \log q\left(z_{i} \mid \lambda_{i}\right)\left(\log p_{i}\left(x, z_{(i)}\right)-\log q\left(z_{i} \mid \lambda_{i}\right)\right)\right] \end{array}\).

위 식을 MC estimation을 통해 approximate할 경우, 아래와 같다.

\(\begin{array}{c} \frac{1}{S} \sum_{s=1}^{S} \nabla_{\lambda_{i}} \log q_{i}\left(z_{s} \mid \lambda_{i}\right)\left(\log p_{i}\left(x, z_{s}\right)-\log q_{i}\left(z_{s} \mid \lambda_{i}\right)\right) \end{array}\).

where \(z_{s} \sim q_{(i)}(z \mid \lambda)\).

3-2. Control Variates

위 3-1 Rao-Blackwellization을 통해서 봤듯, 우리의 function을 expectation값은 같지만. variance가 더 작은 함수로 replace할 수 있다.

control variate는, expectation이 같은 function들의 family로,아래와 같이 나타낼 수 있다.

\(\hat{f}(z) \triangleq f(z)-a(h(z)-E[h(z)])\).

expectation은 같다 ( \(E_{q}[\hat{f}]=\mathrm{E}_{q}[f]\) )
variance는 줄어든다 ( \(\operatorname{Var}(\hat{f})=\operatorname{Var}(f)+a^{2} \operatorname{Var}(h)-2 a \operatorname{Cov}(f, h)\) )

위 variance식을 최소화하기 위해, \(a\)를 다음과 같이 설정하면 된다.

\(a^{*}=\operatorname{Cov}(f, h) / \operatorname{Var}(h)\).

그런 뒤, replace할 함수 \(h\)를 score function으로 설정하면,

\(d\)번째 entry에서의 gradient는 아래와 같이 나오게 된다.

\[\begin{array}{l} f_{d}(z)=\nabla_{\lambda_{d}} \log q\left(z \mid \lambda_{i}\right)\left(\log p_{i}(x, z)-\log q_{i}(x, z)\right) \\ h_{d}(z)=\nabla_{\lambda_{d}} \log q\left(z \mid \lambda_{i}\right) \end{array}\]

Rao-Blackwellization과 control variate를 모두 사용한, 최종적인 MC estimate of the gradient는 다음과 같다.

\(\hat{\nabla}_{\lambda_{d}} \mathcal{L} \triangleq \frac{1}{S} \sum_{s=1}^{S} \nabla_{\lambda_{d}} \log q_{i}\left(z_{s} \mid \lambda_{i}\right)\left(\log p_{i}\left(x, z_{s}\right)-\log q_{i}\left(z_{s}\right)-\hat{a_{d}}\right)\).

3-3. Black Box Variational Inference (2)

(1) noisy gradient

(2) Rao-Blackwellization

(3) control variates

를 모두 사용한 최종적인 알고리즘은, 아래와 같다.

4. Extension

이 알고리즘을 크게 다음과 같은 2개의 방법으로 확장한다

1) Adagrad

( address the difficulty of setting the step size schedule )
2) Stochastic Inference in Hierarchical Bayesian Models

( scalability by subsampling observations )

4-1. Adagrad

Learning rate를 설정하는 것은, stochastic optimization을 푸는 문제에서 주로 겪는 challenge이다.

직관적으로, 우리는 gradient의 variance가 클 때 learning rate가 작길 바란다 (vise versa)

이를 해결하기 위한 것이 AdaGrad로, 아래와 같이 learning rate를 정의한다.

\(\rho_{t}=\eta \operatorname{diag}\left(G_{t}\right)^{-1 / 2}\).

\(G_{t}\) : sum across the first \(t\) iterations of the outer products of the gradient
Adagrad는 \(G_{t}\)의 대각 요소들만 사용하기 때문에, 연산량이 많이 요구되지는 않는다!

4-2. Stochastic Inference in Hierarchical Bayesian Models

Stochastic optimization의 basic idea는 “noisy gradient”를 계산하기 위해 subsample을 한다는 것이다. 이를 hierarchical Bayesian model에 적용해보자.