Personal Perspective

Expectation Maximization

Published on 2018-06-25
Last updated on 2018-06-25


많은 종류의 기계 학습 모델은 likelihood를 최대화하는 방식으로 학습이 진행된다. 직접 관측할 수 없는 latent variable이 있는 경우 이는 다음과 같은 식으로 결정된다.

$$\arg\max_{\mathbf{\theta}} \log p(\mathbf{X}|\mathbf{\theta}) = \arg\max_{\mathbf{\theta}} \log \int_{\mathbf{Z}} p(\mathbf{X}, \mathbf{Z}|\mathbf{\theta}) d\mathbf{Z}$$

위의 식의 최적화는 gradient descent 기법으로도 가능하지만 $\mathbf{\theta}$에 제약조건이 있을 때에는 이를 문제 없이 처리하기 위한 노력이 필요하다. 가령 Gaussian Mixture Model (GMM)을 생각해보면, 클러스터에 속할 확률은 모두 더해서 1을 만족해야 하고, 각 가우스 분포의 공분산은 positive semi-definite 해야 하는 조건을 만족해야 하는데 단순하게 gardient descent를 수행하면 이런 제약 조건을 깨는 경우가 생길 수 있다.

반면 gradient descent 와는 다르게 latent variable model 에서의 학습을 다루는 방법으로 expectation maximization (EM) 기법이 있다. 대체로 EM은 위에서 설명한 문제가 없고 이후에 variational inference로 이어지기에 적합한 접근 방법이다.

Gaussian Mixture Model

EM 알고리즘의 소개에 앞서 GMM을 소개한다. GMM은 관측된 데이터 $\mathbf{x}$를 클러스터링 하는 기법이다. GMM은 데이터 생성 프로세스를 모델링하여 이를 토대로 클러스터를 역으로 추정하는 방식의 접근을 취한다.

하나의 데이터 $\mathbf{x}$가 생성되는 것은 다음의 과정을 거친다고 가정한다. 특정 클러스터 $k$ 가 $\pi_k$의 확률로 선택된다. 그리고는 해당 클러스터에서 데이터가 생성되는 규칙을 표현한 정규분포 $N(\mathbf{x}|\mathbf{\mu}_k, \mathbf{\Sigma}_k)$로부터 $\mathbf{x}$가 생성된다. 이를 수식으로 표현하면 다음과 같다.

$$p(\mathbf{x}) = \sum_{\mathbf{z}} p(\mathbf{z}) p(\mathbf{x}|\mathbf{z}) = \sum_{k=1}^K \pi_k \mathcal{N}(\mathbf{x}|\mathbf{\mu}_k, \mathbf{\Sigma}_k)$$

이로부터 학습 문제를 푸는 것을 생각해보자. 학습을 통해 궁극적으로 얻고자 하는 것은 모델 패러미터 $\mathbf{\theta}$이며, 이는 위의 식에서는 개별적인 $\pi_k$와 $\mu_k, \Sigma_k$이다. 흔히 적용하는 maximum likelihood 기법으로 접근하면 맨 위의 식과 같이

$$\arg\max_{\mathbf{\theta}} \log p(\mathbf{x}) = \arg\max_{\mathbf{\theta}} \log \sum_{\mathbf{z}} p(\mathbf{z}) p(\mathbf{x}|\mathbf{z})$$

를 풀어야 한다. 하지만 로그 안에 합이 들어가서 식이 간단하게 변하지 않으며 모든 $\mathbf{z}$에 대한 합을 고려해야 하는 상황이 생긴다.

Complete-data Log Likelihood

한 발자국 물러나서 $\mathbf{z}$가 주어진 상황을 생각해보자. 만약 $\mathbf{z}$가 주어졌다면 $\mathbf{x}$와 함께 필요한 데이터가 모두 주어진 상황이며, 이 때 $\log p(\mathbf{x}, \mathbf{z} | \mathbf{\theta})$ 를 계산하는 것은 간단하다. 이를 complete-data log likelihood라 한다. 이는 GMM의 예제에서 만약 데이터 $\mathbf{x}$가 $k$ 번째 클러스터로부터 생성되었다는 사실이 알려졌다면, $z = k$ 이고, $p(\mathbf{x}, z=k) = \pi_k \mathcal{N}(\mathbf{x}|\mathbf{\mu}_k, \mathbf{\Sigma}_k)$ 로 간단히 계산할 수 있다. 또한 이를 극대화 하는 $\pi_k$와 $\mathbf{\mu}_k, \mathbf{\Sigma}_k$ 를 찾을 수도 있을 것이다.

하지만 실제로 주어진 데이터에는 특정 $\mathbf{z}$ 가 포함되어 있지 않다. 그렇기에 이에 대한 가장 좋은 추정은 posterior distribution $p(\mathbf{z}|\mathbf{x})$ 로 나타낼 수 있다. 그리고 우리가 추정하려는 complete-data log likelihood 의 기대값을 사용하는 것이다.

$$\sum_{\mathbf{z}} p(\mathbf{z}|\mathbf{x}, \mathbf{\theta}') \log p(\mathbf{x}, \mathbf{z} | \mathbf{\theta})$$

즉, 특정 패러미터 $\mathbf{\theta}'$을 가정한 상황에서 latent variable $\mathbf{z}$ 의 posterior distribution 을 추정한다. 그런 뒤 이에 대한 complete-data log likelihood의 expectation을 계산하고, 이를 극대화하는 $\mathbf{\theta}$를 찾는 과정을 반복하는 것이다. 이를 expectation maximization (EM) 알고리즘이라 부른다.

Expectation Maximization Algorithm

  1. 모델 패러미터 $\mathbf{\theta}'$를 선택한다.
  2. (E-step) $p(\mathbf{Z}|\mathbf{X}, \mathbf{\theta}')$를 계산한다.
  3. (M-step) $\arg\max_{\mathbf{\theta}} Q(\mathbf{\theta}, \mathbf{\theta}')$ 를 계산한다. 이 때, $Q(\mathbf{\theta}, \mathbf{\theta}') = \sum_{\mathbf{z}} p(\mathbf{z}|\mathbf{x}, \mathbf{\theta}') \log p(\mathbf{x}, \mathbf{z}|\mathbf{\theta})$ 이다.
  4. 수렴을 확인한다. 수렴하지 않았으면 $\mathbf{\theta}' = \mathbf{\theta}$로 설정한 뒤 2로 돌아간다.

위와 같은 방식으로 업데이트를 계속 수행하면 log likelihood 가 감소하지는 않음을 어렵지 않게 증명할 수 있다.

EM for GMM

총 $N$ 개의 데이터가 있을 때 $K$ 개의 클러스터가 존재하는 GMM을 다루어보자. 우선 EM 알고리즘의 사용을 위해 complete-data log likelihood를 생각해보자.

$$p(\mathbf{x}, \mathbf{z}) = \prod_{k=1}^K \pi_{k}^{z_k} \mathcal{N}(\mathbf{x}|\mathbf{\mu}_{k}, \mathbf{C}_k)^{z_k}$$

여기에서 $z_k$는 클러스터 $k$가 선택되었는지 여부를 나타내는 0/1 값이다.

클러스터 $z_k$가 선택되었을 확률은 Bayes theorem에 의해 다음과 같다.

$$\begin{array}{ll} p(z_k = 1| \mathbf{x}) &= \frac{p(\mathbf{x} | z_k = 1)p(z_k = 1)}{p(\mathbf{x})} \\ &= \frac{\pi_k \mathcal{N}(\mathbf{x}|\mathbf{\mu}_k, \mathbf{C}_k)}{\sum_{j=1}^K \pi_j \mathcal{N}(\mathbf{x}|\mathbf{\mu}_j, \mathbf{C}_j)} \end{array}$$

이제 이를 사용해 전체 데이터에 대한 식을 세워보자.

$$\log p(\mathbf{x}, \mathbf{z}) = \sum_{n=1}^N \sum_{k=1}^K (z_k^{(n)} \log \pi_k + z_k^{(n)} \log \mathcal{N}(\mathbf{x}^{(n)} | \mathbf{\mu}_k, \mathbf{C}_k))$$

$Q$는 $p(\mathbf{z}|\mathbf{x})$에 대한 기대값으로 계산이 된다. 이는 $\mathbf{z}$를 제외한 나머지 부분은 일종의 상수가 된다고 생각하고 계산할 수 있다.

$$\begin{array}{ccl} E_{\mathbf{z} \sim p(\mathbf{z}|\mathbf{x}, \mathbf{\theta}')} [ \log p(\mathbf{x}, \mathbf{z}) ]&=&E_{\mathbf{z} \sim p(\mathbf{z}|\mathbf{x}, \mathbf{\theta}')} [\sum_{n=1}^N \sum_{k=1}^K (z_k^{(n)} \log \pi_k + z_k^{(n)} \log \mathcal{N}(\mathbf{x}^{(n)} | \mathbf{\mu}_k, \mathbf{C}_k)) ]\\ &=&E_{\mathbf{z} \sim p(\mathbf{z}|\mathbf{x}, \mathbf{\theta}')} [\sum_{n=1}^N \sum_{k=1}^K z_k^{(n)} \log \pi_k] + E_{\mathbf{z} \sim p(\mathbf{z}|\mathbf{x}, \mathbf{\theta}')} [\sum_{n=1}^N \sum_{k=1}^K z_k^{(n)} \log \mathcal{N}(\mathbf{x}^{(n)} | \mathbf{\mu}_k, \mathbf{C}_k) ] \\ &=&\sum_{n=1}^N \sum_{k=1}^K E_{\mathbf{z} \sim p(\mathbf{z}|\mathbf{x}, \mathbf{\theta}')} [z_k^{(n)}] \log \pi_k + \sum_{n=1}^N \sum_{k=1}^K E_{\mathbf{z} \sim p(\mathbf{z}|\mathbf{x}, \mathbf{\theta}')} [z_k^{(n)}] \log \mathcal{N}(\mathbf{x}^{(n)} | \mathbf{\mu}_k, \mathbf{C}_k) ] \\ \end{array}$$

위의 식을 보면 최종적으로 $E_{\mathbf{z} \sim p(\mathbf{z}|\mathbf{x}, \mathbf{\theta}')} [z_k^{(n)}]$ 계산만 남게 된다. 이는 흔히 responsibility 라고 부르며 위에서 계산한 $p(z_k = 1| \mathbf{x})$와 같다.

결과적으로 responsibility를 계산할 때에 E-step이 진행되며 $\mathbf{\theta}'$를 고정하게 된다. M-step에서는 위의 $Q(\mathbf{\theta}, \mathbf{\theta}')$를 극대화 하는데, responsibility의 고정된 $\mathbf{\theta}'$을 제외한 나머지 패러미터를 움직이게 된다. 이렇게 해서 얻은 새로운 패러미터가 다음 스텝의 새 $\mathbf{\theta}'$이 된다.

References

EM에 대한 더 자세한 소개는 Chris Bishop의 Pattern Recognition and Machine Learning의 해당 챕터를 참고하기 바란다.

Exponential family를 가정했을 때 EM 알고리즘은 매우 자연스러운데, 이에 대한 정당성을 제공하는 논의는 아래의 Mathematical monk의 강의에서 볼 수 있다.