[NOTES] Expectation Maximization

30 Mar, 2024

This is an extremely long and mathematically rigorous essay on Expectation Maximization (EM) Algorithm from the lens of a Computer Science graduate who has not taken a math heavy course in a while.

Some of the best long-form essays, be they in text or video form, have one common denominator: They come from individuals who have developed a deep obsession with the subject of the essay. Obsession has positive connotations in this context. However, I'm not obsessed with EM and it holds no special place in my heart (as you will see soon). It's just that I have set a 4-hour timer and am forcing myself to learn as much as I can about EM. Maybe a more general "proxy" for the quality of an essay is "hours spent meditating on the subject" (readers will be the judge). I'm sure that trained mathematicians will view this as blasphemy.

I'm not a big fan of the unnecessary amount of formalism that some math textbooks use. In fact, one of my controversial beliefs is that some (only some) mathematicians deliberately gate-keep knowledge by finding the most unidiomatic formalism for stating a simple property. I'm especially annoyed but this gate-keeping because you will rarely find this happening in software-engineering. For example, the following code would get me brutal PR reviews and maybe a call from my manager:

// Gatekeep version of checking if a number is even
bool isEvenUnidiomatically(int number) { 
    std::bitset<32> binaryRepresentation(number);
    return binaryRepresentation[0] == 0; 
}

Anyway, let's get into notations and get past the gatekeeping.

Background

EM is an optimization algorithm and is:

• Iterative
• Motivated by a notion of missing-ness
  • Conditional distribution of what is missing given what is observed
• Can reliably find the optimum through stable, uphill steps

Notation

• $Y=(X,Z)$: Y is complete data, X is observed data, Z is latent data. This may seem silly to write but note that Y is NOT a function of X and Z. Y is also not X + Z.
  • X = M(Y): this is a "many-to-fewer" mapping. Remember that X is a "part of" Y. Don't confuse this with the X we use in traditional ML techniques where X represents training dataset.
• $f_x(x|\theta )$: density function of observed data. Read as: "density of X given theta"
  • $f_x(x|\theta )={\int }_{y:M(y)=x},f_Y(y|\theta )dy$
    • Read the above equation as: the density of observed data X is integral of density of full data Y given theta. This effectively sums (or integrates) the probabilities of all the different ways in which the observed data $x$ could occur, according to the model.
• $f_y(y|\theta )$: density function of Y (complete data). Read as "density of Y given theta"
• $f_{Z|X}(z|x,\theta )=\frac{f_Y(y|\theta )}{f_X(x|\theta )}$: Conditional density of missing data
• $\theta$ is our parameter that we wanna estimate (aka free-parameters/coefficients)

Marginalization Model

Marginalization refers to the process of summing or integrating out one or more variables from a joint distribution to be able to study the distribution of the remaining variables.

$f_X(x|\theta )={\int }_{[y:M(y)=x]}{f}_Y(y|\theta ),dy$

• Read as "probability density of x given theta is given as integral of probability density of y given theta over all possible ways in which x can occur"
• OR "the density of observed data X is integral of density of full data Y given theta. This effectively sums (or integrates) the probabilities of all the different ways in which the observed data x could occur, according to the model."

If you see drop-downs like these, try to familiarize yourself with these terms to better grasp the math that will come.

EM Algorithm

EM is an optimization algorithm as stated above. It is used for estimating the parameters of probabilistic models. Typically, we deal with a (log-)likelihood function and want to find the set of parameters that are most likely to generated the sample data.

Suppose we want to maximize $L(\theta |x)$ with respect to $\theta$ we define the following function:

$Q(\theta |{\theta }^{(t)})=E\log L(\theta |Y)|x,{\theta }^{(t)}$

Read as: "Expectation of joint log likelihood of complete data conditional on observed data X = x" Can also be written as: $$Q(\theta |{\theta }^{(t)})=E\log f_Y(y|\theta )|x,{\theta }^{(t)}$$ OR $$Q(\theta |{\theta }^{(t)})=\int [\log f_Y(y|\theta )]f_{Z|X}(z|x,\theta ),dz$$ If this were to be written in a frequentist setting then we'd get a summation instead of an integral.

Now the EM algorithm boils down the the following two steps:

1. E-Step: Compute $Q(\theta |{\theta }^{(t)})$ (Expectation of the joint-likelihood function function given parameter (theta) estimate at iteration t)
2. M Step: Maximize $Q(\theta |{\theta }^{(t)})$ with respect to $\theta$. Set ${\theta }^{(t+1)}$ (theta at iteration t+1)

as the maximizer of $Q$.

At this point you will ask, what the fuck is E? And how the fuck did we go from E to an integral? Yes it looks intimidating if you have never had a background to Statistics, but I will explain what it means with an example because it's so silly.

When you have mastered numbers, you will in fact no longer be reading numbers, any more than you read words when reading books. You will be reading meanings._ (Harold Geneen, “Managing”)

If it were up to me, I'd reduce the formal definition for EM Algorithm and make it like three lines and avoid all math jargon cuz it did less than 1% to build my intuition. Let's just dive into the examples and get dirty.

Example 1.

This is going to be probably the only detailed example. Rest of the examples will assume that the reader has understood and internalised these steps.

$$Y_1,Y_2~\text{i.i.d. }\text{Exp}(\theta )\text{ with }{Y}_1=5\text{ observed but }{Y}_2\text{ missing.}$$

Let's break this down. Read the above statement as "Y1 and Y2 are independently and identically distributed (i.id) and follow (~ symbol) an exponential distribution".

Now the first step is to get the probability density function of an exponential distribution. The probability density function (PDF) for an exponential distribution is given by: $$ f(x|\lambda) = \lambda e^{-\lambda x} \quad \text{for } x \geq 0 $$And it's expectation is given by: $$E(X)=\frac{1}{\lambda }$$ The likelihood function $L(\theta )$ for a set of independent and identically distributed (i.i.d.) observations $x_1,x_2,\dots ,x_n$ is the product of the individual PDFs (literally just slap a big Product sign in front of the distribution):

$$L(\theta |x_1,x_2,\dots ,x_n)={\prod }_{i=1}^{n}f(x_i|\theta )$$

Substituting the exponential PDF into the likelihood function gives (this this is so dumb. We just added a product sign): $$ L(\theta | x_1, x_2, \ldots, x_n) = \prod_{i=1}^{n} \theta e^{-\theta x_i} = \theta^n e^{-\theta \sum_{i=1}^{n} x_i} $$ Taking the natural logarithm of the likelihood function gives the log-likelihood: $$\ln L(\theta |x_1,x_2,\dots ,x_n)=n\ln (\theta )-\theta {\sum }_{i=1}^n{x}_i$$

So now, coming back to our original function, can you guess what the log-likelihood function is? It's just this! $$\log L(\theta |y)=\log f_Y(y|\theta )=2\log \theta -\theta y_1-\theta y_2.$$

BOOM! (we just added the log-likelihoods for Y1 and Y2) Now, remember that the E-step involves computing the Q function. What is that? How do we do that?

Well, thats when you simply substitute y1 and y2 with their respective Expected value. Thus $$Q(\theta |{\theta }^{(t)})=2\log \theta -50\theta -\frac{\theta }{{\theta }^{(t)}}$$ since $E{Y}_2|y_1,{\theta }^{(t)}=E{Y}_2|{\theta }^{(t)}=\frac{1}{{\theta }^{(t)}}$ follows from independence.

Ok what about the M step?
The maximizer of $Q(\theta |{\theta }^{(t)})$ is the root of $\frac{2}{\theta }-5-\frac{1}{{\theta }^{(t)}}=0$. Thus ${\theta }^{(t+1)}=\frac{2{\theta }^{(t)}}{50{\theta }^{(t)}+1}.$ Converges quickly to $\hat{\theta }=0.2$. The code to implement the above thing looks as simple as follows:

import math
E = lambda y1, theta: 2 * math.log(theta) - y1 * theta - theta/theta
update = lambda theta: 2*theta/(5*theta + 1) # M step

tol = 2e-09
theta = 5.0 # initial guess
for i in range(100):
    new_theta = update(theta)
    err = (new_theta - theta)**2
    theta = new_theta
    if err <= tol: # optimal stopping
        break
    print(new_theta, E(5, theta)) 

We must note that since this is a minimization problem, with time, the value of the log-likelihood for ${\theta }^{(t)}$ decreases! Woohoo!

Example 2.

Okay, now the language from here on will get extremely terse, no bullshit, and only focusing on the most important findings in EM. Examples from this section onwards will be harder to keep up with. But will build a more general understanding of how EM works and how it can be applied.

This is an example covered in the Computational Statistics, by G. H. Givens and J. A. Hoeting (colostate.edu)

Wing color determined by a single gene with three possible alleles which we denote C, I, and T. C is dominant to I, and T is recessive to I. Thus genotypes CC, CI, and CT result in the carbonaria phenotype having solid black coloring. The TT genotype results in the typica phenotype with light-colored wings. The II and IT genotypes produce an intermediate insularia phenotype with mottled wings.

Key points:

• Genotypes are expresses as combination of two alleles. They cannot be directly observed
• The expression of Genotype is called Phenotype. Phenotype is the physical trait that can be directly observed in the nature.
• Carbonaria is dominant to Insularia. Insularia is dominant to Typica
• Hard Weinberg Equilibrium states that if two alleles match then their genotype frequency is represented as square of the allele frequency. If the two alleles are different then their genotype frequencies will be represented as 2 times the product of individual allele frequencies.
  • Allele frequency: $(p_c,p_i,p_t)$
  • genotype frequency: $(p_c^2,2p_c{p}_i,2p_c{p}_t,p_i^2,2p_i{p}_t,p_t^2)$

How is the problem modelled?

• Y: Genotype counts of the data $y=(n_{cc},n_{ci},n_{ct},n_{ii},n_{it},n_{tt})$
• X: Phenotype counts of the data (thins we have observed in the nature) $x=(n_C,n_I,n_T)$
• Notice how we haven't defined a model for Z. In EM, you do not need to model Z ever. You simply figure out a mapping between X and Y and the EM algorithm works.
• The said mapping is given as follows:

$$y:x=M(y)=(n_{cc}+n_{ci}+n_{ct},n_{ii}+n_{it},n_{tt})$$

• Here mapping simply refers to how corresponding terms are related. For example: $n_C=(n_{cc}+n_{CI}+n_{CT})$
• The Log-likelihood is a multinomial and is given as follows:

$$\log f_Y(y|𝐩)=n_{CC}\log 2p_C^2+n_{CI}\log 2p_C{p}_I+n_{CT}\log 2p_C{p}_T+n_{II}\log 2p_I^2+n_{IT}\log 2p_I{p}_T+n_{TT}\log 2p_T^2+\log \left(\frac{n_{CC}{n}_{CI}{n}_{CT}{n}_{II}{n}_{IT}{n}_{TT}}{n}\right)$$

• One must now ask, what are the parameters in this case?
  • Parameters will be "free". I.e. they will be those components in the log-likelihood that can be adjusted. OR parameters are all those values that ARE NOT observed values
  • In the log-likelihood function you will notice that we are not given $p=(p_c,p_t,p_i)$ and nor are we given the $M(y)$. Does that mean we estimate both the p and M(y)? Well no! Read the next part
  • This is where we exploit the Hardy Weinberg principle/equilibrium and relate our allele frequencies to genotype frequencies. We express each term in Y as a function of p.
• Now let's understand the Q function for this. Q function is simply an optimization transfer problem i.e. It converts the original log-likelihood function into what is called a "Minorizing function". Mathematically it is expressed as follows:

$$l(\theta |x)\ge Q(\theta |{\theta }^{(t)})+l({\theta }^{(t)}|x)-Q({\theta }^{(t)}|{\theta }^{(t)})=G(\theta |{\theta }^{(t)})$$

• Note that our Q function is going to be essentially the same as our Log-likelihood function but y is replaced with its corresponding expectation! The Q function for the above problem looks as follows:

$$Q(p|p^{(t)})=n_{CC}^{(t)}\log p_C^2+n_{CI}^{(t)}\log 2p_C{p}_I+n_{CT}^{(t)}\log 2p_C{p}_T+n_{II}^{(t)}\log p_I^2+n_{IT}^{(t)}\log 2p_I{p}_T+n_{TT}\log p_T^2+k(n_C,n_I,n_T,p^{(t)})$$

• Well this is the exact same at the log-likelihood function right? No
  • The (t) at the top represents the estimated value at iteration t
  • Recall that the Q function is expressed as $Q(\theta |{\theta }^{(t)})=E[\log (f_Y(y|x,{\theta }^{(t)}))]$. Expectation (E) of any function will replace the main input parameter of the function with its expected value (read average value).
  • Notice how Q function is NOT a function of y or x ! It's only a function parameters!
  • Ok then how do we know the expected value of Y?
• Getting Expectations is going to be the key learning for getting to the E step.
  • With help of an Example we can see that the E for n_CC is given as follows:

$$E[N_{CC}|n_C,n_I,n_T,p^{(t)}]=n_{CC}^{(t)}=\frac{n_C(p_C^{(t)}{)}^2}{(p_C^{(t)}{)}^2+2p_C^{(t)}{p}_I^{(t)}+2p_C^{(t)}{p}_T^{(t)}}$$

• How did we arrive at this? Well think of n_C as the total number of phenotypes with carbonaria dominant. The remaining part is simply a proportion. A thumb rule to use here is that to calculate the proportion, simple include the immediate genotype frequency in the neumerator and in the denominator use the sum of all the allele frequencies where Carbonaria is dominant
• Let's take another example. How would n_IT look?

$$E[N_{IT}|n_C,n_I,n_T,p^{(t)}]=n_{IT}^{(t)}=\frac{n_I*2<em>p_I^{(t)}{p}_T^{(t)}}{(p_I^{(t)}{)}^2+2</em>p_I^{(t)}{p}_T^{(t)}}$$ - Notice how the denominator only contains the sum of genotype frequencies where atleast one insularia and one typical allele is present. - One final example to seal the deal: Assume you now have another unknown population that you call $n_U$. All we know is that this sample is a mix of Insularia and Typica alleles. How would n_IT look now? - Answer: We start by writing a new mapping for Y and X

$$(n_C,n_I,n_T,n_U)=(n_{cc}+n_{ci}+n_{ct},n_{ii}+n_{it},n_{tt},N_{ii}+N_{it}+N_{tt})$$

$$E[N_{IT}|n_C,n_I,n_T,n_U,p^{(t)}]=n_{IT}^{(t)}=\frac{n_I*2*p_I^{(t)}{p}_T^{(t)}}{(p_I^{(t)}{)}^2}+\frac{n_U*2*p_I^{(t)}{p}_T^{(t)}}{(p_I^{(t)}{)}^2+2p_I^{(t)}{p}_T^{(t)}+(p_T^{(t)}{)}^2}$$ - Notice how for the second term, the proportion is estimated by including all genotype frequencies where Insularia and Typica occur.

• How do I get the M step? Well you simply take a derivative of the Q function w.r.t each parameter in p and then solve for the generation solution of p. One simple trick to write the update step is to just write it like this:

$$p_C^{(t+1)}=\frac{2n_{CC}^{(t)}+n_{CI}^{(t)}+n_{CT}^{(t)}}{2n}$$

• Well this is trivial to see because our allele frequencies HAVE to be the ratio sum of all corresponding genotypes to the total number of observations. Sometimes instead of doing the full derivation just common sense. M step is essentially your opportunity to update your parameter estimates. Update params can be used to compute the Q function.

Terence Tao views mathematics as climbing peaks, trying to reach various goals. "Some peaks are just out of your reach", he adds. I'm fond of this take as it implies that there exists a terrain of mathematics (perhaps too grand for the human-brain to comprehend). We are simply blind travellers on this vast and largely hidden super-structure. "Intuition" can be thought of as a viewport for this super-structure. Sadly, because humanity does not YET have a hive-mind, every new generation has to reconstruct this viewport over the course of several years. What hurts even more is that every person's viewport spawns at some random co-ordinates on this super-structure; some one might be cursed with a viewport that shows a flat terrain till the end of horizon (people who have poor math intuition) while some might have a viewport for a region that gives them a more zoomed-out or larger view of this super-structure. In my case, I was cursed with a view-port with a flat terrain however, because of the brown-family-peer-pressure I was forced to take higher level math. In just a couple of months, I had "an intelligence explosion" and I could suddenly "see" more of math.

I've been collecting a bunch of tricks to effectively expand my viewport. Before we get dirtier with EM math, I'll share one controversial trick. This trick has to do with the "sensation of understanding". It feels weird to type it out but I experience the moment of "I understand X" almost as a clear sensation similar to taste or touch. And one interesting personal revelation is that, the sensation of understanding is strongest when I'm able to successfully place a concept in my mind relative to other concepts. We'll come back to this but let me quickly tell you about the time I was trying to understand a complicated math proof and I was unable to tell how the author went from step 1 to step 2. Because I was running late for an exam, I just made up a shitty + vulgar mnemonic to remember the steps. Then the most interesting thing happened when I tried to study another proof a while later. I was able to follow the steps for this new unseen proof using the same shitty + vulgar mnemonic as a proxy for the real underlying intuition. This has a very similar sensation to that of understanding the move from step 1 to step 2.

So my biggest take away was that if you don't understand how some parts of math works, just make up vulgar assumptions (vulgar is easier to recall). There is an occlusion in your viewport and you gotta move on. I promise eventually your vulgar assumption get's replace with the real intuition (if you spend enough time playing around with math). Now for this next part, don't focus too much on reasoning or intuitive understanding. Just follow along and get used to it.

EM in Exponential Families

Steps:

• E Step: Compute the Expected values of the sufficient statistics for the complete data:

$$s^{(t)}=E[s(Y)|x,{\theta }^{(t)}]=\int s(y)f_{Z|X}(z|x,{\theta }^{(t)}),dz$$

• M step: Now with the newly acquired $s^{(t)}$ solve for ${\theta }^{(t+1)}$ by finding root of $E[s(Y)|\theta ]=s^{(t)}$. Notice how its the un-conditional expectation (no dependency on x) of sufficient statistic
• Run the above two steps till convergence

Missing Info Principle

We start with the following:

$$\log f_X(x|\theta )=Q(\theta |{\theta }^{(t)})-H(\theta |{\theta }^{(t)})$$

where $H(\theta |{\theta }^{(t)})=E[\log (f_{Z|X}(Z,x,{\theta }^{(t)}))|x,{\theta }^{(t)}]$. Read as "H is expected value of the log-likelihood of hidden/missing data given the parameter estimates at time t".

This intuitively makes sense because if Q function is the expectation of the full data then subtracting the expectation of missing data from it will simply give the expectation of the observed data.

Now taking the information of $\log f_X(x|\theta )$ we get:

$$-{\ell }^{″}(\theta |x)=-Q^{″}(\theta |\omega )|<em>\omega =\theta +H^{″}(\theta |\omega )|</em>\omega =\theta$$

Finally, the missing info principle is:

$$\hat{ı}x(\theta )=\hat{ı}y(\theta )-\hat{ı}z|x(\theta )$$

Supplemented EM (SEM) algorithm

This is used to find the variance of the parameter estimates.

Just remember this: $$\text{var}\hat{\theta }={\hat{ı}}_Y(\hat{\theta }{)}^{-1}(I+{\psi }^{\prime }(\hat{\theta }{)}^T(I-{\psi }^{\prime }(\hat{\theta }){)}^{-1})$$

The previous sections cover enough basics about EM. This part covers some hard-core EM application examples. This is an exercise in modelling: Taking real world problems and then casting it to a concrete model.

Example 3.

You are given a dataset consisting of $N$ data points assumed to be generated from a mixture of two Gaussian distributions with unknown means ${\mu }_1,{\mu }_2$ and known, equal variances ${\sigma }^2$. The mixing coefficient $\pi$ indicates the probability that a given data point comes from the first Gaussian distribution, and $(1-\pi )$ is the probability for the second distribution. The goal is to estimate the parameters $({\mu }_1,{\mu }_2)$, and $(\pi )$ using the Expectation-Maximization (EM) algorithm.

If you have followed the EM-algorithm steps so far, then it'd be be obvious to you on what needs to be done. To outline the steps roughly:

1. Start with the probability density function for GMM.
2. Put a $\prod$ symbol before it (cuz i.id) and boom this is our likelihood function.
3. Take a log-transform to get the log-likelihood function.
4. Compute the Q function as the expected value of the log-likelihood function.
5. Compute the M step by maximizing the Q function w.r.t the parameters. Seems easy right?

In case of Gaussian Mixture Models, we encounter a $\log {\sum }_{}^{}$ in step 3 which is truly disastrous. Because simplifying the log transform of a sum is pretty much not possible. We use some awesome-sauce tricks to by-pass this log-transform. Let me elaborate a little more:

Some important notations:

$$x_i\to \text{data}$$

$$z_i\to \text{cluster assignment for }i$$

$$\mu \to \text{center of cluster }k$$

$${\Sigma }_k\to \text{spread of cluster }k$$

$$w_k\to \text{proportion of data in cluster }k\text{ (mixture weights)}$$

The formula for a normal distribution looks as follows: $$p(X=x|\mu ,\Sigma )=\frac{1}{(2\pi {)}^{p/2}|\Sigma {|}^{1/2}}\exp (-\frac{1}{2}(X-\mu {)}^T{\Sigma }^{-1}(X-\mu )).$$

Likelihood for GMM

reference

The likelihood is given by: $$\text{likelihood}=P(X_1,\dots ,X_n=x_1,\dots ,x_n|w,\mu ,\Sigma ),$$

where

$$w=[w_1,\dots ,w_k],\mu =[{\mu }_1,\dots ,{\mu }_k],\Sigma =[{\Sigma }_1,\dots ,{\Sigma }_k].$$

$$\text{likelihood}(\theta )={\prod }_{i}P(X_i=x_i|\theta ),$$

$$={\prod }_i{\sum }_{k=1}^{K}P(X_i=x_i|z_i=k,\theta )P(z_i=k|\theta )(\text{law of total probability}),$$

$$={\prod }_i{\sum }_{k=1}^K𝒩(x_i;{\mu }_k,{\Sigma }_k)w_k.$$

On the second and third lines above, the sum is over possible cluster assignments (k) for point (i). Taking the log,

$$\log \text{likelihood}(\theta )=\log {\prod }_i{\sum }_{k}P(X_i=x_i|z_i=k,\theta )P(z_i=k|\theta ),$$

$$={\sum }_i\log {\sum }_{k}P(X_i=x_i|z_i=k,\theta )P(z_i=k|\theta ).$$

And there you go! You have the god-forsaken log of a sum that we cannot eliminate. BUT WE TOTALLY CAN ELIMINATE IT WITH SOME MATH WIZARDRY! USE JENSEN'S INEQUALITY MF.

Using Jensen's inequality, we get elegantly handle that summation symbol. I highly suggest you follow this source

$$\log \text{likelihood}(\theta )={\sum }_i\log {𝔼}_z\left[\frac{P(X_i=x_i,Z_i=k|\theta )}{P(Z_i=k|x_i,{\theta }_t)}\right],\backslash$$

$$\ge {\sum }_i{𝔼}_z\log \left[\frac{P(X_i=x_i,Z_i=k|\theta )}{P(Z_i=k|x_i,{\theta }_t)}\right](\text{Jensen's inequality}),$$

$$={\sum }_i{\sum }_{k}P(Z_i=k|x_i,{\theta }_t)\log \left[\frac{P(X_i=x_i,Z_i=k|\theta )}{P(Z_i=k|x_i,{\theta }_t)}\right]=:A(\theta ,{\theta }_t).$$

$(A(·,{\theta }_t))$ is called the auxiliary function. Now back to our original problem. The Auxiliary function for this problem is just the responsibility function.

E-step:

In the E-step, we compute the responsibility (\gamma (z_{i1})) that the first Gaussian distribution has for generating each data point (x_i). This is given by:

$$\gamma (z_{i1})=\frac{\pi 𝒩(x_i|{\mu }_1,{\sigma }^2)}{\pi 𝒩(x_i|{\mu }_1,{\sigma }^2)+(1-\pi )𝒩(x_i|{\mu }_2,{\sigma }^2)},$$

where $(𝒩(x|\mu ,{\sigma }^2))$ is the probability density function of the Gaussian distribution:

$$𝒩(x|\mu ,{\sigma }^2)=\frac{1}{\sqrt{2\pi {\sigma }^2}}\exp (-\frac{(x-\mu {)}^2}{2{\sigma }^2}).$$

M-step:

In the M-step, we update the parameters (\mu_1, \mu_2), and (\pi) using the responsibilities computed in the E-step.

1. Update $({\mu }_1)$ and $({\mu }_2)$:

The new estimates for $({\mu }_1)$ and $({\mu }_2)$ are calculated by taking the weighted average of all data points, weighted by their responsibilities: $${\mu }_1^{new}=\frac{{\sum }_{i=1}^N\gamma (z_{i1})x_i}{{\sum }_{i=1}^N\gamma (z_{i1})},$$

$${\mu }_2^{new}=\frac{{\sum }_{i=1}^N(1-\gamma (z_{i1}))x_i}{{\sum }_{i=1}^N(1-\gamma (z_{i1}))}.$$

1. Update $(\pi )$:

The new estimate for $(\pi )$ is the average responsibility that the first Gaussian distribution has for generating all data points:

$${\pi }^{new}=\frac{1}{N}{\sum }_{i=1}^N\gamma (z_{i1}).$$

The EM algorithm iteratively applies these E-step and M-step computations until the parameter estimates converge, that is, until the changes in the parameter estimates $(({\mu }_1,{\mu }_2,)$ and $(\pi ))$ between successive iterations are sufficiently small. Each iteration consists of using the current parameter estimates to compute the responsibilities in the E-step, and then updating the parameters based on these responsibilities in the M-step, thus maximizing the likelihood of the observed data given the model.

Example 4.

You are given a dataset from a series of experiments where two coins, A and B, are flipped, but the identity of the coin used in each experiment is unknown. Each experiment consists of a fixed number of coin flips, and only the number of heads observed is recorded.

Objective: Your task is to use the Expectation-Maximization (EM) algorithm to estimate the bias (probability of landing heads) of each coin, denoted as $({\theta }_A)$ and $({\theta }_B)$.

Given Data:

1. Number of experiments: $(M)$
2. Number of flips in the (i)-th experiment: $(N_i)$
3. Number of heads observed in the $(i)$-th experiment: $(X_i)$
4. Initial bias estimates: $({\theta }_A^{(0)})$ and $({\theta }_B^{(0)})$

Questions:

1. Derive the E-Step of the EM algorithm for this problem. Calculate the posterior probability that coin A was used in the $(i)$-th experiment, given the observed data and current bias estimates.
2. Derive the M-Step. Update the bias estimates for $({\theta }_A)$ and $({\theta }_B)$ based on the posterior probabilities computed in the E-Step.
3. Explain how the EM algorithm iteratively improves the estimates of $({\theta }_A)$ and $({\theta }_B)$ and discuss the convergence criterion for the algorithm.

Solution:

1. E-Step:

In the E-Step, we calculate the expected value of the latent variable $(Z_{iA})$, which indicates whether coin A was used in the $(i)$-th experiment. The posterior probability that coin A was used, given the observed data $(X_i)$ and the current parameter estimates, is calculated as follows: $$P(Z_{iA}=1|X_i;{\theta }_A,{\theta }_B)=\frac{{\theta }_A^{X_i}(1-{\theta }_A{)}^{N_i-X_i}}{{\theta }_A^{X_i}(1-{\theta }_A{)}^{N_i-X_i}+{\theta }_B^{X_i}(1-{\theta }_B{)}^{N_i-X_i}}$$

2. M-Step:

In the M-Step, we update the parameters $({\theta }_A)$ and $({\theta }_B)$ to maximize the expected complete-data log-likelihood. The updates are given by: $${\theta }_A^{(new)}=\frac{{\sum }_{i=1}^{M}P(Z_{iA}=1|X_i;{\theta }_A,{\theta }_B)X_i}{{\sum }_{i=1}^{M}P(Z_{iA}=1|X_i;{\theta }_A,{\theta }_B)N_i}$$ $${\theta }_B^{(new)}=\frac{{\sum }_{i=1}^M(1-P(Z_{iA}=1|X_i;{\theta }_A,{\theta }_B))X_i}{{\sum }_{i=1}^M(1-P(Z_{iA}=1|X_i;{\theta }_A,{\theta }_B))N_i}$$

3. Iterative Improvement and Convergence:

The EM algorithm iteratively applies the E-Step and M-Step until the changes in the parameter estimates $({\theta }_A)$ and $({\theta }_B)$ become negligible, indicating convergence. A common convergence criterion is the absolute change in the log-likelihood function, which is expected to increase with each iteration. When the increase falls below a predefined threshold, the algorithm is considered to have converged.

##optimization #expectation-maximization #machine-learning #statistics