[Bayesian] Evidence lower bound (ELBO) and EM-algorithm

 

Evidence lower bound (ELBO)

파라미터가 $\theta$이고 latent variable이 $z$인 확률모델을 생각해보자. $z$를 적분하여 marginal을 구할 수 있다.

\[ p(x; \theta) = \int p(x, z; \theta) dz \]

 

non-Bayesian modeling에서는 log-likelihood를 최대로 만드는 $\theta^*$를 찾는데 관심을 갖는다. 즉

\begin{align} \theta^* &= \underset{\theta}{\mathrm{argmax}} \log p(x;\theta) \\ &= \underset{\theta}{\mathrm{argmax}} \log \int p(x, z; \theta) dz \end{align}

 

그러나 만일 적분이 불가능하다면? (what if the integration is not tractable?)

 

Evidence lower-bound

Jensen's inequality를 이용하여 위 적분값의 하한값을 구할 수 있다.

\begin{align} \log p(x;\theta) &= \log \int p(x, z;\theta) dz \\ &= \log \int q(z) \cfrac{p(x, z; \theta)}{q(z)} dz \\ &\ge \int q(z) \log \cfrac{p(x, z; \theta)}{q(z)} dz \end{align}

 

이 lower-bound를 보통 Evidence Lower BOund (ELBO)라고 부른다. 

그렇다면 어떤 분포 $q(z)$가 ELBO를 가장 tight하게 만들까?

 

Euler-Lagrange

함수 $J$를 다음과 같이 정의한다.

\[ J[f] = \int_a^b L(x, f(x), f'(x)) dx \]

$J$가 $f$에서 extreme이면 다음이 성립한다. 

\[ \cfrac{\partial L}{\partial f} - \cfrac{d}{dx} \cfrac{\partial L}{\partial f'} = 0 \]

 

Optimizing ELBO

\begin{align} J[q] &= \int q(z) \log \cfrac{p(x, z; \theta)}{q(z)} dz \\ &= \mathbb{E}_{q(z)}[\log p(x, z; \theta)] + \mathcal{H}[q] \end{align}

 

optimizing problem은 다음과 같다.

\[ \underset{q(z)}{\max} J[q] \text{ such that } \int q(z) dz = 1 \]

 

이는 Lagrange multiplier $\lambda$를 이용하여 풀 수 있다.

\[ J'[q, \lambda] = J[q] + \lambda \left( 1 - \int q(z) dz \right) \]

 

$q(z)$가 optimal 할 때 다음을 항상 만족해야한다.

\[ \log \cfrac{p(x, z; \theta)}{q(z)} - 1 - \lambda = 0 \Rightarrow q(z) = \exp \left( \log p(x, z; \theta) - 1 - \lambda \right) \]

 

다음의 사실 $\int q(z) dz = \cfrac{p(x; \theta)}{\exp (\lambda + 1)} = 1$ 을 이용하면

\[ q(z) = p(z|x; \theta) \]

를 얻는다. 즉 $q(z)$의 optimal distribution은 posterior distribution인 $p(z|x; \theta)$ 이다.

 

Expectation-Maximization algorithm

training data $X = \{ x_i \}_{i=1}^{n}$에 대하여 다음의 maximum likelihood parameter를 찾아보자.

\[ \theta^* = \underset{\theta}{\text{argmax}} \sum_{i=1}^{n} \log p(x_i; \theta) \]

 

위의 log-likelihood를 직접 계산하는 대신, lower-bound를 최적화한다.

\[ \sum_{i=1}^{n} \left\{ \mathbb{E}_{q(z_i)} [\log p(x_i, z_i)] + \mathcal{H}[q(z_i)] \right\} \]

 

E-step(Expectation)과 M-step(Maximization)을 번갈아가며 반복하여 local optimum을 찾는다.

 

initialize $\theta^{(1)}$ randomly and set $t=1$
while not converged do
    E-step) for $i=1, \dots, n$, 다음을 최대화하는 $q^{(t)}$를 찾는다.
\[ J_{\text{E-step}}[q, \theta^{(t)}] = \mathbb{E}_{q(z_i)}[\log p(x_i, z_i; \theta^{(t)})] + \mathcal{H}[q(z_i)] \]
    이때 $q^{(t)}(z_i) = p(z_i|x_i; \theta^{(t)})$가 된다. 

    M-step) 다음을 최대화 하는 $\theta^{(t)}$를 찾는다.
\[ J_{\text{M-step}}[q^{(t)}, \theta] = \sum_{i=1}^{n} \mathbb{E}_{q^{(t)}(z_i)}[\log p(x_i, z_i; \theta)] \]
    $\theta^* \leftarrow \theta^{(t+1)}$이고 $t \leftarrow t+1$로 한다.
end while
Return $\theta^*$