EM(Expectation-Maximization) 알고리즘의 완벽한 기초

오늘은 다양한 곳에서 사용되는 bound optimization의 하나인 Expectation-Maximization(EM) 설명을 진행해볼까 한다. EM 알고리즘에서 나오는 ELBO에 대한 개념은 VAE(Variational Auto Encoder), 그리고 이미지 계열에서 핫한 diffusion 수식을 이해하는 데에도 도움이 될 것이다. 

해당 내용은 책 머피의 머신러닝 1의 8.7에 해당하는 내용을 기반으로 하며 전개할 예정이다. 

 

참고 : 머피의 머신러닝 1

 

---

필수 사전지식

본론으로 들어가기 전 알아두면 좋을 개념들이다.

bound optimization

 많은 머신러닝 문제에서 optimization이라 함은 적절한 objective function(목적 함수)을 설계해서 이를 최적화하며 모델을 학습시키는 과정이다. 주로 목적함수는 loss function(손실 함수)를 일컫는다. 하지만 손실함수를 최적화 할 수 없는 문제가 존재할 수 있다. 이를테면 log likelihood를 최대화해야 하는 모델이 주어졌다고 가정해 보자. 이 log likelihood는 l(θ)로 나타내겠다.

• 달성하고자 하는 문제 : maximize log likelihood l(θ)

 

 대부분의 경우 l(θ)를 최대화하자니 이 함수가 복잡해서 최대화하기 어려웠던 것이다. 그래서 l(θ)에 근사하며 l(θ)이 가질 수 있는 최소한의 bound를 새로운 함수로 정의해 이 bound를 최대화하기로 한다. 이 bound를 surrogate function이라고 부른다.  아래 그림의 Q는 surrogate function를 의미한다. 빨간색 line은 우리가 구하고자 하는 '최대화하고 싶은' 함수이다. 빨간색 함수와 최대한 가까워지기 위해 Q가 maximize하는 방향으로 최적화되고 있다. log likelihood를 최대화(maximize) 하기 위해 그 최소한의 영역(minorize)으로 지정한 함수(surrogate function)를 최적화하는 것이다.

 

lower bound Q

 

그렇기 때문에 이 방법을 bound optimization, 또는 MM 알고리즘으로도 불린다. MM은 Majorize-minimize, 혹은 Miniorize-maximize라고 해석될 수 있는데 위에서 든 예시에서는 Miniorize-maximize 형태에 해당한다. 위 그림과 같이 l와 surrogate function Q는 다음과 같이 학습 스텝(t)이 커질수록 값이 커지는 식을 만족해야 한다. 

surrogate function

• 풀고자 하는 문제 : maximize surrogate function Q, (Q == minorized)

 

조금 가벼운 사전지식

MLE, MAP

파라미터 가 있을 때 prior, posterior, likelihood의 정의이다.

• prior : P(θ)
• posterior : P(θ|D)
• likelihood : P(D|θ)

θ의 최적값을 찾기 위해서 사용되는 것 중 대표적인 것은 optimize function, 주로 loss function을 optimize 하는 것이다. 즉 모델을 피팅하기 위해서 우리가 하고자 하는 일은 loss(L)를 최소화할 수 있는 상태의, 최적의 θ를 찾는 것이다. 

• θ = argmin L(θ)

머신러닝 model fitting의 목적인 최적의 θ를 찾는 다른 방법들이 있으며, Likelihood를 사용해 θ를 최적화하는 방법이 MLE, posterior를 사용해 θ를 최적화하는 방법이 MAP이다.  

 

Baysian

P(A|B) = P(A,B) / P(B)로 나타낼 수 있다. 다르게 쓰면 아래와 같이도 변형된다.

• P(A,B) = P(A|B) * P(B) = P(B|A) * P(A) 

 

KL Divergence

KL은 항상 0보다 같거나 큰 성질이 있는데, 이는 ELBO계산에서 제법 유용하다. KL 식은 아래와 같다. 

• KL( q || p ) = q log q - q log p

참고로 KL == 0을 만족할 땐 q == p를 만족해야 한다. 

 

Jensen's Inequality

 아래 식은 convex function에서의 Jensens's inequality를 나타낸다.  

Jensens's inequality (convex function)

 

EM 알고리즘

bound optimization의 종류에 EM 알고리즘이 있다. bound optimization을 사용하는데 해당 확률모델에서 데이터의 유실이 있거나 hidden variable이 있는 경우가 여기에 속한다. EM의 이름을 풀어보면 다음과 같다.

• Expectation 기댓값
• Maximization 최대화

이름과 같이 간단하게 EM을 한 줄로 설명하면 유실된 데이터 혹은 variable 기댓값을 찾고, log likelihood를 최대화하는 알고리즘이다. 이 방법에도 surrogate function에 해당하는 lower bound를 찾아야 한다. 바로 ELBO이다. 이제부터 ELBO를 찾은 후 E step과 M step을 거쳐 EM의 흐름이 어떻게 되는지 확인해 보자. 

 

Lower bound -ELBO 찾기

최적화 알고리즘의 목표는 목적 함수를 최적화하는 것이며, EM에서는 log likelihood를 최대화하는 것이 목표가 된다. 1차적인 목표는 다음과 같이 log p(y|θ)를 최대화하는 것이 된다. 하나 다른 점이 있다면 EM의 전제인 hidden variable z를 고려해야 한다는 것이다. 아래 식의 N은 모든 데이터를 의미한다. 

log likelihood

문제는 이 log 함수의 합을 구하는게 제법 까다롭다는 것이다. 그렇기 때문에  bound optimization에서의 surrogate 함수를 구했던 것과 같이 l(θ)의 lower bound를 찾아야 한다. 

 

log likelihood 함수를 아래 식과 같이 변형할 수 있다. 분모와 분자에 각각 q(z)를 곱해주는데, q(z)는  hidden variable z에 대한 임의의 분포이다. 아직까지는 log likelihood 수식과 완전히 동일하다. 

log likelihood with arbitary distribution of z

여기서 log 함수는 concave(==upper convex)함수이기 때문에 Jensen's inequality를 적용할 수 있다. 이를 적용하기 위해 위 l(θ)에서 p(y,z|θ)/q(z)를 x로, log를 f로, q(z)를 λ로 생각하면 된다. 

concav function의 jensens' inequality

log함수 내부를 이제 풀어보면 KL divergence 형태로, ∑_zn (q * log p - q * log q)로 튀어나온다. 오른쪽 항인 - ∑_zn (q * log q)는 q의 엔트로피를 의미하기 때문에 H(q)로 쓸 수 있다. 왼쪽 항은 크로스 엔트로피 (cross entropy), 즉 q에 대한  log p 확률분포의 기댓값이 된다. 

모든 n에 대한 L(θ,q)를 더한 값이 최종적인 ELBO가 된다.

ELBO 과정 정리

ELBO를 정리하자면 다음과 같다. 

 

1. log likelihood 구하기 

log likelihood l(θ)를 구했다. hidden value z가 포함되어 있다. 

2. Jensen's Inequality 적용하기

l(θ)의 lower bound를 구하기 위해 함수를 변형할 필요가 있다. 그래서 Jensen's Inequality를 적용했다.

3. log 식 전개하기

log 분모와 분자를 쪼갠다. 

4. 엔트로피로 묶기 

새로운 함수 L을 구했다.

5. ELBO 구하기

모든 데이터 N에 대한 L들의 합이 ELBO다. 

 

이제 우리는 log likelihood가 아닌 ELBO를 maximize할 것이다. 

 

 

E step

Expectation step에서는 우리가 모르는 값인 z에 대해 근사하길 원하며, 이는 z의 임의 분포인 q(z)를 어떻게든 처리해야 함을 의미한다. 그러기 위해선 ELBO를 조금 다듬어야 한다. ELBO 정리 수식의 2번에 해당하는 식에서부터 시작해서 L(θ,q)를 다르게 풀어볼 것이다. 

 

식을 아래와 같은 순서로 정리해보자.

 

1. 다시 정리해 보자  

ELBO수식 정리본의 2번째 열에 해당한다. 여기서 q(z)를 어떻게 할건지 고민해야 한다. 

 

2. bayes 이론 적용하기

분모에 bayes 이론을 적용한다. 여기서 p(y, z) = p(z|y) * p(y) 꼴로 변형이 됐다.

 

3. log 풀어쓰기

log함수 내 p(y|θ)를 쪼개서 표현한다. 그래서 우항은 ∑_zn (q(zn) * logp) 꼴이 되었다. 

 

4. ∑q =1

우항을 계속 보면 logp는 ∑와 상관없는 변수이기 때문에 밖으로 나올 수 있다. q는 z에 대한 분포기 때문에 ∑q =1이다. 

5. KL divergence로 나타내기

좌항은 q log p - q log q 로 - KL이다. 

 

 

5단계를 거쳐서 L(θ,q)를 다시 정리해 아래와 같은 최종 L 을 구했다.

최종 L(θ,q)

거듭 강조하지만 우리의 목표는 L(θ)를 최대화하는 것이다. 그러려면 D_KL의 값은 작아야 한다. KL divergence의 가장 작은 값은 0이다. KL이 0이 되기 위한 전제는 q == p를 만족해야 한다. 즉 hidden variable z의 분포 q(z)와 p(z|y,θ)가 같아야 L(θ)를 최대화할 수 있단 것이다!

 

이 말은 즉, ELBO를 최대화하기 위해서는 q(z) 가 p(z|y,θ)에 근사해야 한다는 것으로, 임의로 정했던 분포 q(z)를 더 이상 고려하지 않고 p로 대체할 수 있단 것이다. 그리고 KL이 0으로 사라지기 때문에 L(θ,q)는 아래와 같이 정리된다. 

 

M step

이제 L을 Maximize하며 θ를 학습할 차례다. ELBO 설명의 4번째 수식을 다시 데려와 보자. θ 관점으로 본다면 H(q)는 constant로 취급하기 때문에 고려할 필요가 없다. 

entropy로 식 정리. ELBO.

따라서, t번째로 E step이 진행되었을 때 최적화해야 할 함수는 아래와 같이 정리된다. 

정리 2

 

위 수식을 maximize하는 θ를 고르는 것이 EM 알고리즘의 마지막 스텝이 되겠다. 

EM 알고리즘의 끝.

이 수식을 통해 추가적으로 알 수 있는 건 z의 분포 q가 최적화를 위해 굳이 필요하지 않다는 거다. 알아야 하는 부분은 log p(y,z) 분포의 합일뿐인 거다.

---

 

오늘은 EM 수식을 보며 대한 전반적인 이해도를 높였다.

• 임의의 변수 z가 있을 때 사용할 수 있는 bound optimization 방법론인 EM에 대해 다뤘다.
• EM에서 log likelihood를 대체할 ELBO를 구할 수 있다.
• E step에서는 z의 임의의 분포 q(z)를 p(z|y,θ)로 대체할 수 있음을 알았다. 그리고 수식 L(θ,q)가 log p(y|θ)로 정리됨을 보였다.
• M step에서는 ELBO를 정리해서 ∑ log p(y,z|θ)를 최대화하는 θ를 구하면 됨을 알았다.  z의 분포인 q(z)를 정확하게 알 필요가 없다.

 

ELBO를 정리하면서 엔트로피에 대해 자세히 한 번 다뤄봐도 좋겠다는 생각이 들었다. 이미 세상에 양질의 관련 문서가 있어서 이 부분을 진행할진 잘 모르겠다. 내 글솜씨가 조금 더 나아진다면 해보고 싶다. 

 

다음 글에서는 EM 을 Gaussian Mixture Model (GMM)에 적용하는 예시를 다뤄보고자 한다.