Variational Inference에서 ELBO설명

Variational Inference에서 ELBO란?

Variational Inference(VI)

Variational Inference(이하, VI)에서 inference는 모델의 예측 과정을 이야기 하는 것이 아니다. VI에서 모델의 인자를 찾는 과정이다. 관측한 데이터로부터 사후 확률(Posterior, $p(\theta|D) $) 분포를 직접적으로 계산하고 싶을 때 VI기법을 사용한다. 보통은 사후 확률 분포를 직접 계산하지 않고 풀 수있는 분야도 있다. 또한 사후 분포 확률은 계산하기 어려운 확률로 알려져있다. VI는 이 문제를 극복하기 위하여, 데이터와 비슷하지만 다루기쉬운 확률 분포도를 도입한다. 비록 딱 맞지는 않지만 쉬운 확률 분포로 관측 데이터($D$)를 근사하여, 최적의 인자를 계산하는 기법이 VI이다. 이에 반하여, MAP(Maximum a Posterior)는 관찰 데이터 $D$로부터 사후 확률이 최대화되는 지점의 $\theta$인자를 구하는 것입니다. 즉, 데이터를 가장 잘 설명하는 모델의 인자를 하나 찾는 것입니다. Bayesian 관점에서 필요한 진정한 사후 확률과 가장 비슷한 $\theta$의 분포를 구하는 것과는 차이가 있습니다(그림 1 참조). 한 점인가? 분포도 인가?로 볼 수 있을 듯 하다. 

$\theta$의 잠재 변수의 분포도를 가지고 있다는 것은, 많은 장점이 있다. 예를들면 데이터를 현실성있게 생성할 수 있거나, 손실 데이터를 채울 수 있다. 그리고 잠재 공간에서 두 점의 어느 중간 지점의 데이터를 예측해 볼 수 있다. 또한 비정상 탐지 또는 불확실성 예측도 가능할 수 있다.

보통 VI기반하여 모델로 예측된 결과가 큰 관점에서는 맞지만, 세밀한 부분에서 뭉개지는 경향이 발생할 수 있다. 근사화의 한계가 있는 것으로보아서 항상 만능 키는 없는 듯 하다. 그러나 이 VI 기법도 모든 가능성이 순식간의 계산 가능한 컴퓨터가 발견되면 쉽게 해결될 지도 모른다. 정보가 부족하거나 계산 기법이 느릴 수록 여러가지 가정을 도입한다. 그리고 더욱더 증명 과정으로 코드구현이 복잡해진다. 이에 반하여 가정으로부터 디테일을 손실하는 것이 발생한다. 하지만 단순하지만 충분한 데이터와 강력한 계산 파워가 있다면, 이러한 가정은 가능한 사라지고, 디테일도 살려내는 추세인 것은 확실해 보인다.

Bayesian Theorem의 어려운 점

왜? posterior 추정이 어려운지 이해가되면, VI의 필요성도 이해가 될 것 같다. 먼저 Bayesian theorem의 문제를 살펴보자.

Marginal Probability, Evidence (Bayesian Theorem)

잠재변수 $Z$를 소거하여, $p(X)$를 설명한다. 보통은 잠재 변수($Z$)를 도입하여, 관찰된 데이터($X$)를 설명하려 할때 사용한다. 보통 Bayesian theorem에서는 관찰 데이터의 marignal probabilty density, $p(X)$를 evidence라고 한다. 추가적으로 $p(z)$는 piror라고 한다.

$$p(X) = \sum_{z \in Z}{p(X,z)} = \sum_{z \in Z }{p(X|z)p(z)} $$

Bayesian Theorem에서 evidence $p(X)$계산이 어려운 점

Evidence에서 모델의 의존적인 적분항이 보통 문제가 발생한다. 직접적인 계산이 불가능하거나, 지수적인 계산시간이 필요하다. 이 문제를 극복하기 위해서 VI는 근사할 수 있는 효과적인 방법을 제공한다. VI는 marginal probability density(또는 evidence)의 lowr bound를 잘 관찰하여 문제를 해결하였다. Posterior를 계산할때 evidence의 어려운 점을 Bayesian 공식($\ref{bayesian}$)과 연결하여 자세히 설명해 보자.

$$\begin{aligned} p(Z|X) &= \frac{ p(X|Z)p(Z) } {p(X)} \\  &= \frac{ p(X|Z)p(Z) } {\sum_{z \in Z }{p(X|z)p(z)}}  \end{aligned} \label{bayesian}$$

결론적으로 사후확률(posterior probability, $p(Z|X)$ = likelihood * prior / partition function)를 잘 계산하기 위해서는 $\ref{bayesian}$에서 분모의 합(적분)형태(partition function)를 풀어야 한다. 하지만, 이것은 계산하기가 힘들다. 이유는 잠재변수 $Z=(z_1, z_2, z_3, \cdots )$와 연계하여 실제 관측 대상의 확률 분포를 계산하기가 어렵기 때문이다(적어도 나의 연구 기간내에서...). 그래서 일부만 자료를 사용하는 샘플링 기법을 사용해서 얻으려고 한다(Monte Carlo). 하지만, 만약 잠재 변수($z$)가 기하 급수적으로 크다면, 샘플링 기법도 비효율적이게 된다. 샘플링의 헛 발이 많다. 그러나 미래에는 양자컴퓨팅 기술이 이러한 문제를 쉽게 해결을 할 수 있을 수도 있다. 기다려 볼까? 절대 마법의 비법 소스도 없다. 현실적으로, 지금은 못하니 다른 대안을 소개해본다. 특히 본 블로그에서는 VI기법으로 해결하는 방법을 소개한다.

자 이제 Generative Model($z \rightarrow x$)을 설계를 해보자! 어떤 어려움이 있을까?

주의하자, 분류기는 관찰된 라벨링 변수($y$)가 있다. 여기서 $z$는 잠재변수를 칭하여, 향후 문맥 이해에 방해가 되어 확실히 분류기에서의 $y$와는 구분을 지어보려 한다.

Marginal Inference를 사용하여 관측 데이터를 설명한다.

관찰을 통한 데이터 분포인 $p(X)$만 덜컥 주어지면, 우리는 더 이상 할 수 있는 선택지가 별로 없다. 그냥 그러한 것이다. 그러나 보통은 우리의 요구사항은 앞으로의 데이터(보지 못한 데이터)까지 예상했으면 한다. 즉, 학습을 통한 예측을 하고 싶을 것이다. 그러기 위해서는 통상적으로 관측데이터로부터 모델링을 하고, 이를 잘 적용되도록 모델을 일반화를 시켜야 한다. 이때 도입하는 방법이 정확한 값을 알 수 없지만 있다고 가정하는 잠재 변수 $z$ (예, 벡터, 지능, 건강, 임베딩 값 등)과 파라미터($\theta$, 입력-잠재변수-출력사이에 연산에 필요한 추가 인자들등)도입이다. 잠재변수는 활용 용도에 따라서 다양하기 때문에, 명확히 규정하기가 힘들 수 있다. 특히 변수가 난무하면 이해가 어려우니, 본 블로그에서는 대부분은 수식에서 $\theta$는 생략하고 $z$와 $\theta$를 통합하여 한 개($z$)로 전개한다. 잠재변수 도입하여 관측 데이터의 분포($p(X)$)를 모델링을 할때 많이 사용하는 기법이 marginal inference(model prediction)기법이다. 즉, 

$$p(X) = \sum_{z \in Z}{p(X,z)} \label{marginal_probability}$$

로 해서 $p(X)$ 확률 분포를 근사할 수 있다. $z$로 $x$를 설명하게 된것이다. 결국 하고자 하는 것은 $z \rightarrow x$로서, 원인과 결과의 관계를 모델링하는 것이다. 

그런데 잠재변수 $z$의 도입하여 해를 해결하려고 하였다. 그러나 그것조차 모델이 복잡하다보니 적어도 나의 인생동안에는 계산이 어렵겠네... ?

모델이 복잡하다는 것은 $z$의 차원이 크거나, $z$사이에 서로 의존도가 복잡하거나, 적분이 어렵거나, 최적화 과정 중에서 어딘가에 계산량이 매우 비싼 경우가 있기 때문이다. 이때 우리는 근사한 접근 방법을 도입하여 효율적으로 문제에대한 해(즉, 모델의 인자)를 찾고 싶어한다. 특히 사후 확률(posterior)는 NP-hard problem으로 알려있다. 이 어려움 때문에 일반적인 최적화 모델에서는 MAP(likelihood * prior), MLE(likelihood probability)와 같은 최대화 기법을 대신 사용한다. 그러나 우리는 진짜 사후 확률($p(Z|D)$ or $p(\theta|D)$)을 얻고 싶은 것이다. MAP, MLE가 아니다. 특히 MAP은 사후 확률을 계산하는 것이 아님을 명심하다. 그래서 MAP, MLE를 사용할 수 없다. 결론적으로, $Z$의 분포 자체를 얻는 VI와 데이터를 가장 잘 설명하는 인자 $z$를 얻는 MAP은 서로 다른 해를 가진다고 할 수 있다(그림 1참고). 

[그림1] 왼쪽: Variational inference의 과정. 오른쪽: MAP과 variational inference의 차이. 여기서 $D$는 관측 데이터 집합이다. $\xi$는 $q$의 분포를 표시하는 variational 인자이다. [참조]

어찌되었든 VI는 사후 확률($p(Z|D)$) 자체를 근사하는 기법이 있다. 기본 원리는 다루기 쉬운 $q(z)$로 posterior proability를 근사 시키자는 것이다. 즉, $q(Z) \approx p(Z|X)$ 되는 $q(Z)$를 잘 찾자. 여기서 $p(Z|X)$를 사후 분포(posterior probability)라고 한다. 즉, 관측 데이터로부터 모델($Z$)을 연결한 것이다. 실제적으로 VI는 global optimization을 찾을 수는 없을 수 있으나, 상당히 실용적이고, 병렬처리와 확장성이 잘 설계된 stochastic gradient optimization를 사용할 수 있다[참조].

다루기 쉽고 계산하기 쉬운 최적의 $q^{*}{(Z)} \approx p(Z|X) $를 찾자!

이 VI 개념을 설명하기 위해서 보통은 $q(Z)$(surrogate)와 $p(Z|X)$라는 두 분포함수를 가지고 이야기 한다. 여기서 $q(Z)$는 다룰 수 있는(계산하기 쉬운) 분포 함수의 $Z$의 분포이다. 보통은 가우시안 모델을 사용한다. 그리고 $p(Z|X)$는 관측 한 데이터($X$)가 주어졌을때 모델의 인자의 $Z$의 분포이다. 즉 두 $Z$의 분포가 교차되는 지점을 잘 찾으면 된다. 

보통의 최적화의 대상은 잠재변수 $Z$를 찾는 것이다. 하지만, VI는 $q(Z)$를 찾는 점이 특이한 점이다. $q(Z)$를 잘 고르면, 관측 데이터를 가장 잘 설명하는 분포와 우리가 다룰 수 있는 분포가 가장 비슷한 시점을 고를 수 있다. 즉, $p(Z|X)$와 $q(Z)$가 가장 오차가 작도록 최적의 $q^{*}(Z)$를 찾으면 된다. 수식으로 표현하면 다음과 같다:

$$q^{*}(Z) = \underset{q(Z) \in Q }{argmin} D_{KL}[q(Z) || p(Z|X)] $$

두 분포의 오차(거리)를 측정하는 기법으로 KL-Divergence, $D_{KL}$를 많이 사용한다. 여기서 최적의 $q^{*}(Z)$를 찾는 것은 VI라고 할 수 있다. 즉, $q^{*}(Z) \approx P(Z|X)$ 가 된다.

보통 $q(Z)$를 variational distribution라고 이야기 한다[https://en.wikipedia.org/wiki/Variational_Bayesian_methods]. 필자의 뇌피셜인데, variational distribution으로 근사하는 $q(Z)$로 모델 추론하기 때문에 variational inference라고 하는 듯 하다.

어떻게 $q^{*}(Z)$를 찾았을까? ELBO 최적화!

KL Divergence, $D_{KL}[q(Z) || p(Z|X)] $ 에서 우리는 참 사후 확률(true posterior probability)인 $$p(Z|X) = \frac{\bar{p}(Z,X)}{p(X)} \label{z_normal}$$의 분포를 모른다(이미 모델을 할때 $z\rightarrow x$라고 가정한다. 그래서 수식 $\ref{z_normal}$로 표현이 가능하다). 더 자세히 말하면, $p(Z|X)$를 구체적으로 값을 평가할 수 없다. 즉 계산을 할 수 없다. 수식 $\ref{z_normal}$에서 보면 특히 문제의 항인 $p(X)$는 실제 데이터의 모든 가능성 있는 빈도의 값을 알아야 한다. 일부 관측된 X의 실체 모양을 알 수 없으니 계산할 수가 없다. 그러나 상수인 것은 확실 할 것이다. 여전히 이거참 난감한 상황이다. 하지만 그동안의 현명한 분들이 다 솔루션을 찾았다. 그 과정을 소개해본다. 자 그럼, 수집한 데이터 $X$가 있기 때문에, 즉 $\bar{p}(Z,X)$을 곰곰히 활용해 보자. 결론적으로 말하면, evidence $p(X)$때문에 KL Divergence를 직접 계산을 피하고, ELBO를 최적화하려는 기법인 VI를 활용한다. ELBO를 최적화 하여, $q^{*}(Z)$를 찾을 수 있다.

$D_{KL}[q|| \color{Blue}p]$와 $D_{KL}[q || \color{Red} \bar{p}]$ 사이의 관계는?

주로 이 사이트를 참조했음을 밝힌다. 가장 납득력이 있고 쉽게 설명하였다. KL Divergence에 $p(Z|X)$대신 $\bar{p}(Z,X)$를 넣고 $p(Z|X)$때와 비교를 분석을 해보는 것이 목적이다. $D_{KL}$은 확률 분포를 인자로 거리를 계산하는 것처럼, 비슷한 수식모형으로 형식으로 관측된 데이터로 거리를 표현하는 $J(q||\bar{p})$를 생각해보자.

$$J(q||\bar{p}) = \sum_{z \in Z}{q(z)\log \frac{q(z)}{\bar{p}(z,X)}} $$ 

여기서 $\bar{p}(z,X)$는 수식 $\ref{z_normal}$에있다.

기계적으로 풀어 보자.

$$\begin{aligned}J(q||\bar{p}) &=&  \sum_{z \in Z}{q(z)\log \frac{q(z)}{\bar{p}(z,X)}} \\ &=& \sum_{z \in Z}{q(z)(\log \frac{q(z)}{p(z|X)}} - \log p(X)) \\ &=& D_{KL}(q||p) - \sum_{z \in Z}{q(z)\log p(X)} \\ &=& D_{KL}(q||p) - \log p(X)\cdot 1\end{aligned}$$

여전히 $p(X)$와 $D_{KL}(q||p)$은 계산 불가이다(원래 모르는 것이다). 그럼에도 불구하고 이제 항을 정리해 놓고, $q$와 $p$와 $\bar{p}$의 관계를 잘 살펴보면 다음과 같다: 

$$\color{Red} \log p(X) = D_{KL}(q||p) - J(q||\bar{p}) \label{ELBO_EQ}$$

ELBO의 명칭 유래

두 분포의 차이를 나타내는 $D_{KL}(q||p) \geq 0$ 이므로, 수식 $\ref{ELBO_EQ}$으로부터 최초 비교해보려던 $q$와 $p$와 $\bar{p}$의 관계는 

$$\log p(X)  \geq - J(q||\bar{p}) = \color{Blue} ELBO \label{LB}$$ 

이다. 여기서 수식 $\ref{LB}$의 $-J(q||\bar{p})$를 variational lower bound 또는 ELBO(evidence lower bound)라고 한다. ELBO라고 붙인 이유는 아무래도 $\ref{LB}$의 부등식때문일 것이다. $\geq$을 기준으로 좌측에는 evidence $p(X)$가 있고, 우측에 하한이 있기 때문이다.

수식 $\ref{ELBO_EQ}$으로부터 ELBO($-J(q||\bar{p})$)와 variational distribution($q(z)$), true posterior probability($p(Z|X)$)의 고찰

$\log p(X)$ 는 상수항으로 고정되있기 때문에, $-J(q||\bar{p})$를 최대화하는 것은 $D_{KL}(q||p)$를 줄이는 것과 같다. 결과적으로 $q$가 ture posterior probability($p$)에 근사하게 되는 것이다. 

[                 $\log p(X)$                ] = -->[ $D_{KL}(q||p)$]<-- + [<-------  $-J(q||\bar{p})=ELBO(q(z))$  ------->]

$ELBO(q(z))$를 최대화로 이끌면, $q(z)$가 최적의 $q^{*}(z)$를 찾아 간다.  결국 $ELBO(q(z))$만 최적화를 하면 $q^{*}(Z)$를 찾을 수 있는 것이다. 특히 $ELBO(q(z))$에는 모두 계산이 가능한 항들만 있다. 

Expectation으로 $ELBO(q(z))$를 표현해보자.

먼저 mpatacchiola.github.io/blog 를 참조하였음을 밝혀둔다. 

$$\begin{aligned}-J(q||\bar{p}) &= -\sum_{z \in Z}{q(z)\log \frac{q(z)}{\bar{p}(z,X)}} \\ &= E_{z \sim q(Z)}[-\log\frac{q(z)}{\bar{p}(z,X)}] \\ &= E_{z \sim q(Z)}[\log{\bar{p}(z,X)} - \log{ q(z)}]\end{aligned}$$ 

$ELBO(q(z))$를 expectation($E_{p(x)}[f(x)]$)관점으로 정리를 하면, 다음과 같다(표현이 축약되어서 이해하기 쉬울 수가 있다): 

$$E_{z \sim q(Z)}[\log{\bar{p}(z,X)} - \log{ q(z)}] = \log p(X) - D_{KL}(q(Z)||p(Z|X)) \label{ELBO_EXP} $$

사실 수식 $\ref{ELBO_EXP}$의 오른쪽 항은 관심이 없다. 이유는 이제 ELBO의 원리를 알았기 때문이다. 최적화 설계할 때는 오른쪽항은 잊어 버리자. 그리고 ELBO만 자세히 살펴보자.

$$\begin{aligned} ELBO(q(z)) &= -J(q(z)||\bar{p}(z,X))) \\ &= E_{z \sim q(Z)}[\log{\bar{p}(z,X)} - \log{ q(z)}] \\ &= E_{z \sim q(Z)}[\log{\bar{p}(z,X)}] - E_{z \sim q(Z)}[\log{ q(z)}] \\ &= E_{z \sim q(Z)}[\log{\bar{p}(X|z)}\bar{p}(z)] - E_{z \sim q(Z)}[\log{ q(z)}] \\ &= E_{z \sim q(Z)}[\log{\bar{p}(X|z)}] + E_{z \sim q(Z)}[\bar{p}(z)] - E_{z \sim q(Z)}[\log{ q(z)}] \\ &= E_{z \sim q(Z)}[\log{\bar{p}(X|z)}] + E_{z \sim q(Z)}[\bar{p}(z) - \log{ q(z)}] \\ &= E_{z \sim q(Z)}[\log{\bar{p}(X|z)}] + E_{z \sim q(Z)}[\frac{\bar{p}(z)}{ q(z)}] \\ &= \color{Red} E_{z \sim q(Z)}[\log{\bar{p}(X|z)}] + D_{KL}(q||\bar{p}(z)) \\ \end{aligned}$$

여기서 $\bar{p}(X|z)$항은 모델의 잠재변수로 관측된 데이터를 예측하는 확률이다. 즉 여러개의 가능성있는 $q(z)$가 있을 경우, 기존 보다 더 좋은 것을 찾으면 된다. 특히 2번째 항은 $\bar{p}(z)$로서 잠재변수 $z$의 사전확률이므로 사람이 잘 잡아주면 된다. 이것은 $q(z)$가 $\bar{p}(z)$ 멀어지면 패널티(penalty)를 더 준다고 생각하면 된다.

결론적으로 관측한 데이터 $X$를 대상으로 기존보다 더 설명을 잘하는 $q(z)$를 반복적으로 잘 찾으면 된다. 사실 이렇게 반복적 찾는 방법은 다양하지 않을까? 너무 많아서 나열하기가 힘들 수도 있다.

$ELBO(q(z))$를 최적화 하는 방법은?

다음의 2개의 기법을 언급한다:

1. Gradient 기법으로
1. REINFORCE 기법(policy gradient), Score Function 기법
2. Pathwise 기법.
2. EM 기법이 있다.

이미 기존 것들은 많이 소개가 되었다. 그래서 생략을 한다. 하지만, 많이 소개되지 않은 ELBO의 최적화를 policy gradient로 연결하고, 최적화 하는 방법을 소개할 것이다.

REFERENCES

1. https://ermongroup.github.io/cs228-notes/inference/variational/
2. https://ratsgo.github.io/generative%20model/2017/12/19/vi/
3. https://pyro.ai/examples/intro_part_ii.html#A-Simple-Example
4. https://mpatacchiola.github.io/blog/2021/01/25/intro-variational-inference.html
5. https://www.cs.princeton.edu/courses/archive/fall11/cos597C/lectures/variational-inference-i.pdf
6. https://gregorygundersen.com/blog/2021/04/16/variational-inference/
7. https://fabiandablander.com/r/Variational-Inference.html
8. https://jonathan-hui.medium.com/machine-learning-summary-fundamental-6562ff7448a#ee5a
10. Roeder, Geoffrey, Yuhuai Wu, and David Duvenaud. "Sticking the landing: Simple, lower-variance gradient estimators for variational inference." arXiv preprint arXiv:1703.09194 (2017).