Part IX The EM algorithm - Stanford University

1 CS229 Lecture notesTengyu Ma and Andrew NgMay 13, 2019 Part IXThe EM algorithmIn the previous set of notes, we talked about the EM algorithmas applied tofitting a mixture of Gaussians. In this set of notes, we give a broader viewof the EM algorithm , and show how it can be applied to a large family ofestimation problems with latent variables. We begin our discussion with avery useful result calledJensen s inequality1 Jensen s inequalityLetfbe a function whose domain is the set of real numbers. Recall thatfis a convex function iff (x) 0 (for allx R). In the case offtakingvector-valued inputs, this is generalized to the conditionthat its hessianHis positive semi-definite (H 0). Iff (x)>0 for allx, then we sayfisstrictlyconvex (in the vector-valued case, the corresponding statement isthatHmust be positive definite, writtenH >0). Jensen s inequality canthen be stated as a convex function, and letXbe a random :E[f(X)] f(EX).

2 Moreover, iffis strictly convex, then E[f(X)] =f(EX) holds true if andonly ifX= E[X] with probability 1 ( , ifXis a constant).Recall our convention of occasionally dropping the parentheses when writ-ing expectations, so in the theorem above,f(EX) =f(E[X]).For an interpretation of the theorem, consider the figure [X]bf(a)f(b)f(EX)E[f(X)]fHere,fis a convex function shown by the solid line. Also,Xis a randomvariable that has a chance of taking the valuea, and a chance oftaking the valueb(indicated on thex-axis). Thus, the expected value ofXis given by the midpoint also see the valuesf(a),f(b) andf(E[X]) indicated on , the value E[f(X)] is now the midpoint on they-axis betweenf(a)andf(b). From our example, we see that becausefis convex, it must be thecase that E[f(X)] f(EX).Incidentally, quite a lot of people have trouble remembering which waythe inequality goes, and remembering a picture like this is agood way toquickly figure out the thatfis [strictly] concave if and only if fis [strictly]convex ( ,f (x) 0 orH 0).

3 Jensen s inequality also holds for concavefunctionsf, but with the direction of all the inequalities reversed (E[f(X)] f(EX), etc.).2 The EM algorithmSuppose we have an estimation problem in which we have a training set{x(1),..,x(n)}consisting ofnindependent examples. We have a latent vari-able modelp(x,z; ) withzbeing the latent variable (which for simplicity isassumed to take finite number of values). The density forxcan be obtainedby marginalized over the latent variablez:p(x; ) = zp(x,z; )(1)3We wish to fit the parameters by maximizing the log-likelihood of thedata, defined by ( ) =n i=1logp(x(i); )(2)We can rewrite the objective in terms of the joint densityp(x,z; ) by ( ) =n i=1logp(x(i); )(3)=n i=1log z(i)p(x(i),z(i); ).(4)But, explicitly finding the maximum likelihood estimates ofthe parameters may be hard since it will result in difficult non-convex optimization , thez(i) s are the latent random variables; and it is often the casethat if thez(i) s were observed, then maximum likelihood estimation wouldbe such a setting, the EM algorithm gives an efficient method for max-imum likelihood estimation.

4 Maximizing ( ) explicitly might be difficult,and our strategy will be to instead repeatedly construct a lower-bound on (E-step), and then optimize that lower-bound (M-step).2It turns out that the summation ni=1is not essential here, and towards asimpler exposition of the EM algorithm , we will first consider optimizing thethe likelihood logp(x) fora single examplex. After we derive the algorithmfor optimizing logp(x), we will convert it to an algorithm that works fornexamples by adding back the sum to each of the relevant equations. Thus,now we aim to optimize logp(x; ) which can be rewritten aslogp(x; ) = log zp(x,z; )(5)1It s mostly an empirical observation that the optimizationproblem is difficult to , the E-step and M-step can often be computed more efficiently than op-timizing the function ( ) directly. However, it doesn t necessarily mean that alternatingthe two steps can always converge to the global optimum of ( ).

5 Even for mixture ofGaussians, the EM algorithm can either converge to a global optimum or get stuck, de-pending on the properties of the training data. Empirically, for real-world data, often EMcan converge to a solution with relatively high likelihood (if not the optimum), and thetheory behind it is still largely not a distribution over the possible values ofz. That is, zQ(z) = 1,Q(z) 0).Consider the following:3logp(x; ) = log zp(x,z; )= log zQ(z)p(x,z; )Q(z)(6) zQ(z) logp(x,z; )Q(z)(7)The last step of this derivation used Jensen s inequality. Specifically,f(x) = logxis a concave function, sincef (x) = 1/x2<0 over its domainx R+. Also, the term zQ(z)[p(x,z; )Q(z)]in the summation is just an expectation of the quantity [p(x,z; )/Q(z)] withrespect tozdrawn according to the distribution given Jensen sinequality, we havef(Ez Q[p(x,z; )Q(z)]) Ez Q[f(p(x,z; )Q(z))],where the z Q subscripts above indicate that the expectations are withrespect tozdrawn fromQ.

6 This allowed us to go from Equation (6) toEquation (7).Now, foranydistributionQ, the formula (7) gives a lower-bound onlogp(x; ). There are many possible choices for theQ s. Which should wechoose? Well, if we have some current guess of the parameters, it seemsnatural to try to make the lower-bound tight at that value of . , we willmake the inequality above hold with equality at our particular value of .To make the bound tight for a particular value of , we need for the stepinvolving Jensen s inequality in our derivation above to hold with continuous, thenQwould be a density, and the summations overzin ourdiscussion are replaced with integrals note that the notionp(x,z; )Q(z)only makes sense ifQ(z)6= 0 wheneverp(x, z; )6= we implicitly assume that we only consider thoseQwith such a this to be true, we know it is sufficient that the expectation be takenover a constant -valued random variable.

7 , we requirethatp(x,z; )Q(z)=cfor some constantcthat does not depend onz. This is easily accomplishedby choosingQ(z) p(x,z; ).Actually, since we know zQ(z) = 1 (because it is a distribution), thisfurther tells us thatQ(z) =p(x,z; ) zp(x,z; )=p(x,z; )p(x; )=p(z|x; )(8)Thus, we simply set theQ s to be the posterior distribution of thez s givenxand the setting of the parameters .Indeed, we can directly verify that whenQ(z) =p(z|x; ), then equa-tion (7) is an equality because zQ(z) logp(x,z; )Q(z)= zp(z|x; ) logp(x,z; )p(z|x; )= zp(z|x; ) logp(z|x; )p(x; )p(z|x; )= zp(z|x; ) logp(x; )= logp(x; ) zp(z|x; )= logp(x; )(because zp(z|x; ) = 1)For convenience, we call the expression in Equation (7) theevidencelower bound(ELBO) and we denote it byELBO(x;Q, ) = zQ(z) logp(x,z; )Q(z)(9)6 With this equation, we can re-write equation (7) as Q, ,x,logp(x; ) ELBO(x;Q, )(10)Intuitively, the EM algorithm alternatively updatesQand by a) settingQ(z) =p(z|x; ) following Equation (8) so that ELBO(x;Q, ) = logp(x; )forxand the current , and b) maximizing ELBO(x;Q, ) while fixingthe choice that all the discussion above was under the assumption that weaim to optimize the log-likelihood logp(x; ) for a single examplex.

8 It turnsout that with multiple training examples, the basic idea is the same and weonly needs to take a sum over examples at relevant places. Next, we willbuild the evidence lower bound for multiple training examples and make theEM algorithm we have a training set{x(1),..,x(n)}. Note that the optimal choiceofQisp(z|x; ), and it depends on the particular examplex. Therefore herewe will introducendistributionsQ1,..,Qn, one for each examplex(i). Foreach examplex(i), we can build the evidence lower boundlogp(x(i); ) ELBO(x(i);Qi, ) = z(i)Qi(z(i)) logp(x(i),z(i); )Qi(z(i))Taking sum over all the examples, we obtain a lower bound for the log-likelihood ( ) iELBO(x(i);Qi, )(11)= i z(i)Qi(z(i)) logp(x(i),z(i); )Qi(z(i))Foranyset of distributionsQ1,..,Qn, the formula (11) gives a lower-bound on ( ), and analogous to the argument around equation (8), theQithat attains equality satisfiesQi(z(i)) =p(z(i)|x(i); )Thus, we simply set theQi s to be the posterior distribution of thez(i) sgivenx(i)with the current setting of the parameters.

9 Now, for this choice of theQi s, Equation (11) gives a lower-bound onthe loglikelihood that we re trying to maximize. This is the E-step. Inthe M-step of the algorithm , we then maximize our formula in Equation (11)with respect to the parameters to obtain a new setting of the s. Repeatedlycarrying out these two steps gives us the EM algorithm , whichis as follows:7 Repeat until convergence{(E-step) For eachi, setQi(z(i)) :=p(z(i)|x(i); ).(M-step) Set := arg max n i=1 ELBO(x(i);Qi, )= arg max i z(i)Qi(z(i)) logp(x(i),z(i); )Qi(z(i)).(12)}How do we know if this algorithm will converge? Well, suppose (t)and (t+1)are the parameters from two successive iterations of EM. We will nowprove that ( (t)) ( (t+1)), which shows EM always monotonically im-proves the log-likelihood. The key to showing this result lies in our choice oftheQi s. Specifically, on the iteration of EM in which the parameters hadstarted out as (t), we would have chosenQ(t)i(z(i)) :=p(z(i)|x(i); (t)).

10 Wesaw earlier that this choice ensures that Jensen s inequality, as applied to getEquation (11), holds with equality, and hence ( (t)) =n i=1 ELBO(x(i);Q(t)i, (t))(13)The parameters (t+1)are then obtained by maximizing the right hand sideof the equation above. Thus, ( (t+1)) n i=1 ELBO(x(i);Q(t)i, (t+1))(because ineqaulity (11) holds for allQand ) n i=1 ELBO(x(i);Q(t)i, (t))(see reason below)= ( (t))(by equation (13))where the last inequality follows from that (t+1)is chosen explicitly to bearg max n i=1 ELBO(x(i);Q(t)i, )8 Hence, EM causes the likelihood to converge monotonically. In our de-scription of the EM algorithm , we said we d run it until convergence. Giventhe result that we just showed, one reasonable convergence test would beto check if the increase in ( ) between successive iterations is smaller thansome tolerance parameter, and to declare convergence if EM is improving ( ) too we define (by overloading ELBO( ))ELBO(Q, ) =n i=1 ELBO(x(i);Qi, ) = i z(i)Qi(z(i)) logp(x(i),z(i); )Qi(z(i))(14)then we know ( ) ELBO(Q, ) from our previous derivation.