[Loss, AI Math] Variational Inference


이전 포스트에서 MLE 와 MAP 를 살펴봤다. MLE 는 Likelihood 만 주어진 경우 이를 최대로 하는 parameter 를 찾는 일이고 MAP 는 Likelihood 와 사전 지식인 Prior 가 전부 주어진 경우 Posterior 를 최대로 하는 parameter 를 찾는 것이었다. 이런 방식으로 parameter $w$ 를 구하는 것은 $w$ 의 값을 확정짓는 Frequentist View 딥러닝으로, $w$ 의 분포를 구하는 Bayesian View 가 아니다.

Bayesian View 는 딥러닝에서 구하고자 하는 값인 $w$ 의 확률 분포, 즉 posterior($w \vert x$) 다. MAP 에서는 posterior 를 최대로 하는 $w$ 를 구할 뿐 posterior 자체를 정확히 구하지는 않는다.

그리고 바로 앞선 포스트에서는 MLE 와 KL-div 의 밀접한 관계를 살펴보면서, Generative Model 의 VAE 는 Maximum Likelihood 라는 목표를 위해 Variational Inference 를 수행함을 봤다. Variational Inference 는 대표적인 Bayesian Inference 로 데이터 $x$ 가 들어왔을 때 latent vector $z$ 에 대한 posterior $P_\theta (z \vert x)$ 를 구하고자 한다.

이 Variational Inference 에 대해서 이해해보자.

Posterior

  • Bayes Equation 을 다시 한 번 보자. $w$ 는 우리가 구하고자 하는 parameter 이고, $D$ 는 데이터이다.

    \[p(w \vert D) = \frac{p(D \vert w)p(w)}{\int p(D \vert w)p(w)dw}\]
  • 우리의 목표는 Neural Network 의 parameter 인 $w$ 를 찾는 것이다. 이를 찾기 위해 Loss function 을 설정하고 Gradient Descent 알고리즘을 돌린다. 이것이 바로 고정된 $w$ 를 찾는 Frequentist View 다.
  • 이 포스트에서는 $w$ 의 값만 알고 싶은 것이 아니라 $w$ 의 분포, 즉 Posterior 를 구하고자 한다. 이를 Bayesian View 라고 한다.
  • 그러나 Posterior 를 구하는 것은 아래의 섹션에서 다루는 이유들로 인해 쉽지 않다.

Evidence

  • 위 Bayes Equation 식을 뜯어보자.
  • 먼저 우리가 구하고자 하는 Posterior $p(w \vert D)$ 는 Train Data $D$ 가 주어졌을 때 이를 가장 잘 재현하기 위한 parameter $w$ 의 확률분포다.
  • Likelihood $p(D\vert w)$ 는 모델이 주어진 데이터를 얼마나 잘 설명하는지 나타내는 가능도이며 Gaussian 으로 가정하면 MSE, Bernoulli 로 가정하면 Cross Entropy 가 유도된다.
  • Prior $p(w)$ 는 우리의 사전 지식에 의해 설정되는 분포다. 사전 지식이 없다면 Prior 를 설정하여 우리가 원하는 방향의 $w$ 가 구해지게끔 할 수도 있다. Gaussian 을 설정하면 L2 Regularization(weight decay) 이 되고, Lapalace 를 설정하면 L1 Regularization 이 된다.
  • 여기까지는 해당 포스트에서 잘 정리한 내용이다.
  • 이처럼 Posterior 를 구하기 위한 모든 조건들은 전부 주어지거나 우리가 설정할 수 있다. 그런데도 불구하고 Posterior 를 구하지 못하는 이유는 뭘까?
  • 뜬금없게도 분모에 있는 Evidence 때문이다.

    \[p(D) = \int p(D\vert w)p(w) dw\]
  • Posterior 를 최대로 하는 $w$ 만 구하면 됐었던 MAP 에서는 $D$ 가 주어진 값이기 때문에 Evidence 가 상수가 되어 치환만 해주고 따로 구하지 않았었다.
  • 하지만 Posterior 분포를 정확히 알고 싶은 지금은 이 Evidence 를 구해야 하는데 $w$ 는 변수 하나로 적어 놔서 그렇지, 보통 굉장히 차원이 높은 편이고 Neural Net 의 Layer 가 깊어질수록 그 개수 또한 엄청나게 많아진다.
  • 이렇게 넓은 차원에 퍼져 있는 모든 가능한 $w$ 에 대하여 적분 연산을 한다는 것은 사실상 불가능에 가깝다.

Output

  • $w$ 에 대해 적분을 해내서 $w$ 의 Posterior $p(w \vert D)$ 를 구해 냈다고 해보자.
  • MAP 에서는 Posterior 가 최대가 되는 $w$ 값을 구해서 이 $w$ 를 가중치(parameter)로 가지는 Neural Network 를 통과시켜, 그대로 Output 을 얻기만 하면 됐었다.
  • 하지만 Posterior 를 정확히 알고 있는 지금, MAP 처럼 Posterior 가 최대가 되는 $w$ 를 찾아 쓰기만 하면 될까? 그러면 MAP 와 다를 것이 없기 때문에 Posterior 를 구할 이유가 없다.
  • Posterior 를 온전히 알고 있으면 우리는 $w$ 에 따른 Output 의 기댓값을 구할 수 있게 된다.
  • 이는 MAP 에 비해 Posterior 전체에 담긴 정보를 훨씬 많이 반영한 결과이기 때문에 더 신뢰할 수 있는 output 이라고 할 수 있다. 그리고 output 의 기댓값은 아래의 식을 통해 구할 수 있다.

    \[\mathbb{E}_{p(w \vert D)}[f(w)] = \int f(w)p(w \vert D) dw\]
  • 또 적분이다. 그것도 그냥 적분이 아니라 위에서 말한 것처럼 매우 고차원에다가 개수까지 많은 $w$ 에 대한 적분이다. 이 적분 역시 현실적으로 계산하는 것이 불가능하다.

Variational Inference

  • 위에서 살펴봤듯이, posterior $p(w \vert D)$ 를 직접 구하기는 힘들다. 그럼 어떻게 구할 수 있을까?
  • 딥러닝의 중요한 본질들 중 하나는 바로 우리가 모르는 함수 $f$ 가 있을 때, 비교적 적은 미지수를 가지고 있고 우리가 아는 함수 $g$ 를 정의하여 $g$ 가 $f$ 를 잘 흉내내도록 미지수들을 정해주는 것이다.
  • 즉 모델링하는 함수 혹은 분포를 가지고 현실 세계의 함수 혹은 분포에 근사시키는 것이다. 이와 똑같은 논리를 적용시켜보자.
  • 우리가 알고 비교적 적은 parameter $\theta$ 를 가지는 분포 $q(w\vert \theta)$ 를 정의하여, 이 $q(w \vert \theta)$ 가 $p(w\vert D)$ 를 가장 잘 흉내내도록 parameter $\theta$ 를 정해주는 것이다.
  • 이 때 $q(w\vert \theta)$ 가 $p(w \vert D)$ 를 얼마나 잘 흉내내고 있는지 그 척도를 정의해 줄 필요가 있다. 이 척도만 있다면 딥러닝에서의 Loss 처럼 써먹을 수 있을 것 이다.
  • 두 Distribution 이 얼마나 닮았는지를 비교하는 척도로는 아래의 Kullback-Leibler Divergence 가 있다.

    \[\mathbb{KL}[q(w\vert \theta) \Vert p(w\vert D)] = \int q(w\vert \theta) \log \frac{q(w\vert \theta)}{p(w \vert D)}dw\]
  • KL-div 의 정의와 식은 이 포스트이 포스트에서 자세히 다루었다.
  • 위 식에서 $w$ 를 연속으로 가정하고 있기 때문에 Integral 형태로 식을 작성된다. 이제 위 식에 Bayesian Equation을 적용시키면 아래와 같이 전개할 수 있다.

    \[\begin{aligned} \int q(w\vert \theta) \log \frac{q(w\vert \theta)}{p(w \vert D)}dw &= \int q(w\vert \theta) \log \frac{q(w\vert \theta)p(D)}{p(w)p(D \vert w)}dw \\ &= \int q(w \vert \theta)\left[\log \frac{q(w \vert \theta)}{p(w)} + \log p(D) -\log p(D\vert w)\right]dw \\ &= \int q(w \vert \theta) \log \frac{q(w \vert \theta)}{p(w)}dw + \mathbb{E}_q[\log p(D)] - \mathbb{E}_q[\log p(D\vert w)] \\ &= \mathbb{KL}[q(w\vert \theta) \Vert p(w)] + \mathbb{E}_q[\log p(D)] - \mathbb{E}_q[\log p(D\vert w)] \end{aligned}\]
  • 여기서 $\mathbb{E}_q[\log p(D)]$ 는 $q$ 와 무관하므로 제거하면, 최종적으로 우리는 다음 식을 Minimize 하는 $\theta$ 를 찾아야 한다.

    \[\color{red}{\mathbb{KL}[q(w\vert \theta) \Vert p(w)]} \color{blue}{- \mathbb{E}_q[\log p(D\vert w)]}\]
  • 위 식을 통해 $\theta$ 를 찾으면, 완벽히 같지는 않겠지만 $q(w \vert \theta)$ 를 $w$ 의 Posterior 라고 주장할 수 있게 된다.
  • 먼저 빨간색 term $\mathbb{KL}[q(w\vert \theta) \Vert p(w)]$ 은 $\mathbb{KL}[q(w\vert \theta) \Vert p(w \vert D)]$ 와는 달리 계산 가능한 값이다. $p(w)$ 도 알고 $q(w\vert \theta)$ 는 우리가 정의해주는 분포이기 때문이다.
  • 파란색 term 의 식을 보면 $\log p(D \vert w)$ 가 보이는데, 이는 Likelihood 다. 빨간색 term 은 계산 가능한 값이기 때문에 이 Likelihood 를 Maximize 시킨다.
  • 이는 Maximum Likelihood Estimation(MLE)에 해당하기 때문에 일반적인 Neural Network 를 학습시키는 일과 동치가 된다.
  • 그러나 단순히 Neural Network 를 학습시키는 것이 아니다. 그냥 Likelihood 가 아니라 $q(w\vert \theta)$ 에 대한 Likelihood 의 기대값이기 때문이다.
  • 여기에 우리는 Monte Carlo Method 를 응용해볼 수 있다. $\mathbb{E}_q[\log p(D \vert w)]$ 를 $q(w \vert \theta)$ 로부터 Sampling 된 $w$ 인 $\bar{w}$ 에 대하여 $\log p(D\vert \bar{w})$ 로 대체하는 것이다.
  • 뭔가 복잡하지만, 정리하면 다음과 같다.

    \[\color{red}{\mathbb{KL}[q(w\vert \theta) \Vert p(w)]} \color{blue}{- \mathbb{E}_q[\log p(D\vert \bar{w})]}\]
    • $q(w \vert \theta)$ 로부터 Sampling 된 $w(\bar{w})$ 에 대하여 얻은 output 과 label 사이의 loss, 그리고 $\mathbb{KL}[q(w \vert \theta) \Vert p(w)]$ 를 계산하여 얻은 식의 합을 최종적인 loss 로 하여 back propagation 을 진행한다.
    • 이 때 학습되는 parameter 는 $w$ 가 아니라 $q$ 의 parameter 인 $\theta$ 가 된다.

Practical Implementation

  • 이제 예를 들어서 Variational Inference 을 실제로 진행해보자.
  • 우리가 Bayesian Learning 으로 학습할 Neural Network 에는 $N$ 개의 parameter 들 $w_1, w_2, \ldots, w_N$ 이 있다.
  • 먼저 Prior $p(w_i)$ 를 정해줘야 할 텐데, $i$ 에 상관없이 Standard Normal Distribution 으로 가정해보자.

    \[p(w_i) = \mathcal{N}(0, 1^2) = \frac{1}{\sqrt{2\pi}}e^{-\frac{w^2_i}{2}}\]
  • 이제 우리가 만들어서 근사시킬 분포 $q$ 또한 Gaussian 으로 놓을 수 있다.

    \[q(w_i \vert \mu_i, \sigma_i) = \mathcal{N}(\mu_i, (\sigma_i)^2) = \frac{1}{\sqrt{2\pi}\sigma_i}e^{-\frac{(w_i - \mu_i)^2}{2\sigma^2_i}}\]
  • $w_i$ 는 $i$ 에 따라 다른 분포를 따를 수 있기 때문에 우리가 최적화시켜야 할 parameter 들은 $\mu_i, \sigma_i$ 로 총 $2N$ 개가 된다.
  • 여기서 헷갈리면 안되는 게, 우리는 기존의 딥러닝처럼 $w_i$ 를 업데이트하는 것이 아니다. $\mu_i$ 와 $\sigma_i$ 를 업데이트 할 것이고, 이로부터 $w_i$ 가 Sampling 되도록 할 것이다.
  • 이제 정의해야 할 것들을 정의했다. 그럼 먼저 우리가 계산할 수 있는 빨간색 term 을 계산하자.

    \[\mathbb{KL}[q(w \vert \theta_i) \Vert p(w)] = \mathbb{E}_{w \sim q(w \vert \theta)}[\log \frac{q(w \vert \theta)}{p(w)}]\]
  • 그 결과는 다음과 같게 된다.

    \[\mathbb{KL}[q(w_i \vert \theta_i) \Vert p(w_i)] = -\frac{1}{2}[(1 + \log (\sigma^2_i)) - \mu^2_i - \sigma^2_i]\]
  • 그리고 이를 Loss function 으로 이용하기 위하여 상수배 Scale 을 제거하고 모든 $i$ 에 대한 합을 구하면 다음과 같은 Loss 식이 완성된다.

    \[-\sum^N_{i=1}[(1 + \log(\sigma^2_i)) - \mu^2_i - \sigma^2_i]\]
  • 이제 파란색 term 인 $\mathbb{E}_q[\log p(D \vert w)]$ 을 해결해야 하는데, 위에서 말했듯이 MLE 에 해당하기 때문에 $w_i$ 들로 구성된 네트워크를 통과해서 나온 결과값 $y$ 와 실제 $t$ 사이의 loss 를 구해서 back propagation 하면 된다.
  • 간단하게 $L2$ Loss 로 표현해보자.

    \[L_2 = (y - t)^2\]
  • 이제 마지막 관문이 남았다. $w_i$ 를 어떻게 Sampling 할까?
  • 안타깝게도 단순히 $w_i$ 들을 각각의 distrbution 에 따라 random sampling 하여 Neural Network 에 집어 넣게 되면 $\mu_i$ 와 $\sigma_i$ 까지 gradient 가 도달하게 만들 수 없다. Sampling 과정이 미분 불가능하기 때문이다.
  • 그래서 다음과 같은 방법을 사용한다.

    \[w_i = \mu_i + \sigma_i \times \epsilon_i, \quad \epsilon_i \sim \mathcal{N}(0, 1^2)\]
  • 위 식을 통해 $w_i$ 를 Sampling 하면 $w_i$ 까지 도달한 gradient 를 $\mu_i$ 와 $\sigma_i$ 까지 propagate 할 수 있다.
  • 최종적으로 정리하자면, Variational Inference 를 이용하여 $w_i$ 들을 알아내는 과정은 다음과 같다.
    1. $\mu_i$ 와 $\sigma_i$ 를 전부 초기화 한다.
    2. $w_i = \mu_i + \sigma_i \times \epsilon_i$ 를 이용하여 학습시킬 Neural Network 의 parameter 들을 결정한다.
    3. $w_i$ 들이 결정 된 Neural Network 에 input $x$ 를 통과시켜 output $y$ 를 얻는다.
    4. $y$ 와 정답 $t$ 사이의 $L_2$ loss $(y-t)^2$ 를 계산한다.
    5. 최종 Loss 는 아래와 같이 되고, 이를 통해 backpropagation 한다. 그러면 $\mu_i$ 와 $\sigma_i$ 들이 업데이트 된다.

      \[\text{Loss} = -\sum^N_{i=1}[(1+\log(\sigma^2_i)) = \mu^2_i - \sigma^2_i] + (y-t)^2\]
    6. 2 ~ 5 를 반복한다.(학습)
  • 이런 과정을 통해 학습이 완료되면 우리는 최종적인 $\mu_i$ 와 $\sigma_i$ 들을 얻을 수 있고, 이렇게 얻은 값들을 parameter 로 하는 gaussian distribution $q(w_i \vert \mu_i, \sigma_i)$ 를 우리는 $w_i$ 의 Posterior 라고 말할 수 있다.

Variational Inference 와 딥러닝

  • 그러면 이러한 Variational Inference(VI) 는 딥러닝에서 어떻게 사용될까?
  • VI 는 확률적 모델의 잠재변수(latent variable)를 추론하거나 복잡한 분포를 근사하는 데 중요한 역할을 한다.
  • 이는 직접적인 계산이 어려운 확률분포를 계산 가능한 분포로 근사하는 기법으로, 주로 비선형, 고차원 데이터에서 유용하다.
  • 딥러닝에서 VI 는 특히 생성모델에서 자주 사용되며 대표적으로 Variational Autoencoder(VAE)에서 핵심적인 역할을 한다.
  • VAE
    • VAE 는 데이터의 잠재변수(혹은 잠재공간)를 학습하고 이를 통해 새로운 데이터를 생성하는 모델이다.
    • 이 때, 직접적인 잠재변수 $p(z \vert x)$ 의 분포를 알 수 없기 때문에 VI 를 사용해 근사한다.
    • 구체적으로 VAE 는 잠재변수의 사후분포를 근사하기 위해 KL-div 를 최소화하는데, 이 때 ELBO(Evidence Lower Bound)를 최적화한다.
    • 자세한 것은 이 포스트에서 정리해보자.
  • 또한 불확실성을 고려하는 Bayesian Neural Networks(BNN) 에서도 VI 가 사용된다.
    • BNN 은 bayesian inference 를 통해 모델의 출력값 이외에도 불확실성(uncertainty)도 출력한다. 이를 통해 딥러닝 기반 인식 시스템의 over-confidence 문제를 해결할 수 있다.
    • 즉 BNN 에서는 Frequentist View 와 같이 확률모델을 계속 학습하여 고정된 참 값의 확률을 찾아가는 것이 아니라 불확실성을 내포한 확률분포를 찾는 것이 목적이 된다.
    • 따라서 기존의 Neural Network 는 학습이 끝나면 weight 가 고정된 값을 가지게 되지만 BNN 은 각 weight 또한 불확실성을 가지는 확률분포로 나타나 진다는 점에서 차이가 있다. 각 weight 가 불확실성을 가지는 확률분포이므로 최종 출력인 $y$ 또한 불확실성을 가지는 확률분포를 가지게 된다.
    • weight 에 대한 분포를 구할 때 VI 가 사용된다.
    • 이를 통해 전통적인 신경망에 비해 불확실성을 고려한 모델링이 가능하게 하며, 더 안정적이고 견고한 예측을 가능하게 한다.

Reference

맨 위로 이동 ↑

댓글 남기기