Here's an example of Expectation Maximisation (EM) used to estimate the mean and standard deviation. The code is in Python, but it should be easy to follow even if you're not familiar with the language.
The motivation for EM
The red and blue points shown below are drawn from two different normal distributions, each with a particular mean and standard deviation:
To compute reasonable approximations of the "true" mean and standard deviation parameters for the red distribution, we could very easily look at the red points and record the position of each one, and then use the familiar formulae (and similarly for the blue group).
Now consider the case where we know that there are two groups of points, but we cannot see which point belongs to which group. In other words, the colours are hidden:
It's not at all obvious how to divide the points into two groups. We are now unable to just look at the positions and compute estimates for the parameters of the red distribution or the blue distribution.
This is where EM can be used to solve the problem.
Using EM to estimate parameters
Here is the code used to generate the points shown above. You can see the actual means and standard deviations of the normal distributions that the points were drawn from. The variables
blue hold the positions of each point in the red and blue groups respectively:
import numpy as np
from scipy import stats
np.random.seed(110) # for reproducible random results
# set parameters
red_mean = 3
red_std = 0.8
blue_mean = 7
blue_std = 2
# draw 20 samples from normal distributions with red/blue parameters
red = np.random.normal(red_mean, red_std, size=20)
blue = np.random.normal(blue_mean, blue_std, size=20)
both_colours = np.sort(np.concatenate((red, blue)))
If we could see the colour of each point, we would try and recover means and standard deviations using library functions:
But since the colours are hidden from us, we'll start the EM process...
First, we just guess at the values for the parameters of each group (step 1). These guesses don't have to be good:
# estimates for the mean
red_mean_guess = 1.1
blue_mean_guess = 9
# estimates for the standard deviation
red_std_guess = 2
blue_std_guess = 1.7
Pretty bad guesses - the means look like they are a long way from any "middle" of a group of points.
To continue with EM and improve these guesses, we compute the likelihood of each data point (regardless of its secret colour) appearing under these guesses for the mean and standard deviation (step 2).
both_colours holds each data point. The function
stats.norm computes the probability of the point under a normal distribution with the given parameters:
likelihood_of_red = stats.norm(red_mean_guess, red_std_guess).pdf(both_colours)
likelihood_of_blue = stats.norm(blue_mean_guess, blue_std_guess).pdf(both_colours)
This tells us, for example, that with our current guesses the data point at 1.761 is much more likely to be red (0.189) than blue (0.00003).
We can turn these two likelihood values into weights (step 3) so that they sum to 1 as follows:
likelihood_total = likelihood_of_red + likelihood_of_blue
red_weight = likelihood_of_red / likelihood_total
blue_weight = likelihood_of_blue / likelihood_total
With our current estimates and our newly-computed weights, we can now compute new, probably better, estimates for the parameters (step 4). We need a function for the mean and a function for the standard deviation:
def estimate_mean(data, weight):
return np.sum(data * weight) / np.sum(weight)
def estimate_std(data, weight, mean):
variance = np.sum(weight * (data - mean)**2) / np.sum(weight)
These look very similar to the usual functions to the mean and standard deviation of data. The difference is the use of a
weight parameter which assigns a weight to each data point.
This weighting is the key to EM. The greater the weight of a colour on a data point, the more the data point influences the next estimates for that colour's parameters. Ultimately, this has the effect of pulling each parameter in the right direction.
The new guesses are computed with these functions:
# new estimates for standard deviation
blue_std_guess = estimate_std(both_colours, blue_weight, blue_mean_guess)
red_std_guess = estimate_std(both_colours, red_weight, red_mean_guess)
# new estimates for mean
red_mean_guess = estimate_mean(both_colours, red_weight)
blue_mean_guess = estimate_mean(both_colours, blue_weight)
The EM process is then repeated with these new guesses from step 2 onward. We can repeat the steps for a given number of iterations (say 20), or until we see the parameters converge.
After five iterations, we see our initial bad guesses start to get better:
After 20 iterations, the EM process has more or less converged:
For comparison, here are the results of the EM process compared with the values computed where colour information is not hidden:
| EM guess | Actual
Red mean | 2.910 | 2.802
Red std | 0.854 | 0.871
Blue mean | 6.838 | 6.932
Blue std | 2.227 | 2.195
Note: this answer was adapted from my answer on Stack Overflow here.