Expectation maximization and Gaussian mixture models

It is natural to want a procedure that takes the shape of the clusters into account.

Gaussian mixture models, supervised¶

If, for data in $\mathbb{R}^d$, we are given a division into $k$ classes (labels), then for the subset of the data consisting of all points of one class we can compute the sample covariance and the sample expected value, and from them set up a $d$-variate normal distribution that models this subset.

This yields a model for the data as a whole: write $p_i$ for the number of points in class $i$ divided by the total number of points in the sample, and $f_i(x)$ for the estimated density of the normal distribution of class $i$; then the density is

This model, in turn, we can of course use to judge which class a point belongs to: for a point $x$ we compute the densities

and pick the $i$ for which this expression is largest.

One can picture the generative process like this: the data points were sampled from a (fictitious) Gaussian mixture model by first picking a class $j$ according to the discrete distribution of the $p_i$ and then drawing a sample from the density $f_j(x)$ – once for each data point.

Now if the clusters overlap, it can easily happen that the model classifies some points $x$ differently from their original classes – and in principle there is no remedy for that. Even so, in this way we can find the Gaussian mixture model that best fits the data, and use it to sample data that will look very much like the data at hand (provided the normality assumption was right).

One further observation: the Gaussian mixture model gives us more than just a new division into classes; for every point and every class, it gives us a probability that the point belongs to that class. That is a good deal more information, and it can be considerably more useful. For the points that are not correctly classified, above all, it gives us a way of quantifying how certain we are about the division into classes.

Gaussian mixture models, unsupervised¶

The procedure just described is not available if no division into classes is given to us. In that case the class is a latent variable: one we do not know, but which is essential to describing the model. What we can model well are the conditional probabilities given the value of the latent variable – that is, the distribution of a single class. The distribution of the latent variable itself is harder to describe when all we have is unlabelled data.

A two-stage procedure suggests itself: first find a division into classes with k-means (often not a particularly good one, if the homoscedasticity assumption fails) and then, on the basis of that division, set up a Gaussian mixture model – that is, obtain an estimate of the latent variable.

Now, in turn, the division into classes that the mixture model lets us compute will in general differ from the one we computed with k-means beforehand. This time, though, it is plausible that the mixture model achieves the better division, since now the shapes of the clusters are taken into account as well (approximated by ellipsoids, thanks to the normality assumption).

With this improved division into classes – that is, with the improved estimate of the latent variable – we can set up a Gaussian mixture model all over again, re-estimating the covariance matrices and the centres, much as in the iteration step of the k-means algorithm.

One more observation: if, in setting up the covariance matrices, we consider only the quality of the model for a single class (the best fit for that class alone), then we ignore the effect that change of model has on the other classes – and there usually is one. It would therefore be better to optimise all the parameters simultaneously in the iteration step, towards the target quantity – the mixture model with the best fit.

To construct a good estimator for an arbitrary distribution we have used the maximum likelihood method, and that is what we want to try for a mixture of normal distributions too.

Maximum likelihood in models with latent variables¶

The probability density in the Gaussian mixture model, where $\theta_j$ denotes the parameters $\mu_j,\Sigma_j$ of a single mixture component and $p_j$ the distribution of the components, is

and that is the likelihood (the probability of the data, computed in our model), writing $M$ for our model and $X$ for the data. So the model assumption is now: $X$ is distributed according to $f$. Here $\theta$ is the collection of the parameters $p_j$ together with the parameters $\mu_j,\Sigma_j$ hidden inside the $f_j$. Reading this quantity as a likelihood means reading it as a function of those parameters.

A tip: if you prefer to picture $\theta$ as a list indexed by $\theta_j$, simply put $\theta_0 := (p_1,\dots,p_k)$.

If we take seriously the modelling decision to express the division of the data points into classes by probabilities, we may as well model the division itself as a random variable $Y$: $Y_i = j$ if the $i$-th data point is of class $j$ – only we do not know the variable $Y$.

Our model, formulated with the latent variable $Y$, is

Finding the zeros of the derivative?¶

Finding the zeros of the derivative of $L(\theta)$ directly is not easy, and the popular trick of maximising $\ell(\theta) = \log L(\theta)$ instead does not help much here either, because

and the logarithm does not get us out of that sum business so easily.

Jensen to the rescue!¶

We can lower our sights a little and work with Jensen’s inequality (the logarithm is concave and $\sum_{j=1}^k p_j = 1$):

The charm of the expression on the right-hand side is that a $\theta$ which maximises the right-hand side cannot have a smaller left-hand side. In formulas: if we are given a $\theta^{(t)}$ and find a $\theta^{(t+1)}$ for which the right-hand side grows, then $\ell(\theta^{(t+1)}) \geq \ell(\theta^{(t)})$ as well.

Note that a fat $1 = \frac{q(j)}{q(j)}$ can also be multiplied into the expression for the contribution of a class, $p_j f_j$ – thus

The point of this becomes visible once Jensen’s inequality is applied, this time to the convex combination $\sum_{j=1}^k q(j) = 1$ – so from now on we assume that $q$ is a probability density:

So we have found a lower bound on $\ell(\theta)$ that we can write as the sum of an expression which is easier to compute and the entropy of the distribution with density $q$. Since the entropy does not depend on $\theta$, $L(\theta,q)$ can be maximised by maximising the other summand.

We shall come back to this observation (which is why we have given it a name, $L(\theta,q)$); for the moment, let us discuss a general approach to the solution.

Auxiliary Functions¶

If we cannot maximise a function $\ell(\theta)$, we can try a stopgap: pretend that $\theta$ is constant in some of the subexpressions of $\ell(\theta)$ and maximise, then hold the other subexpressions constant instead and maximise over the rest. To formalise this we introduce an auxiliary function $Q(\theta,\vartheta)$ that is easier to maximise in $\theta$ than $\ell$ is, but nevertheless satisfies $Q(\theta, \theta) = \ell(\theta)$ and, moreover, $Q(\theta,\vartheta) \leq \ell(\theta)$ for all $\vartheta$.

These assumptions are already enough to show that an iterative process leads to better and better approximations $\theta^{(t)}$ to the $\theta$ that maximises $\ell(\theta)$. Here $\theta^{(t+1)}$ denotes the $\theta$ that makes $Q(\theta, \theta^{(t)})$ maximal (which by assumption is easy to find).

All told, $\ell(\theta^{(t)}) \leq \ell(\theta^{(t+1)})$ – which is to say that at least our estimate never gets worse along the way!

If we want to take this approach, what remains is to choose a well-judged auxiliary function $Q(\theta,\vartheta)$.

The complete-data log-likelihood¶

The last ingredient for developing the EM algorithm for the Gaussian mixture model is the complete-data log-likelihood: the log-likelihood in the fictitious world in which we really do know the complete data – not only the values of $X$ but those of the latent variable $Y$ as well, that is, the actual classes of the data points.

The complete data are therefore the $(x_i,y_i)$, which we model with $(X,Y)$.

It is useful now to look at the variable $Y_i$ in one-hot encoding as well: $Y_{ij} = 1$ if $Y_i = j$, and $Y_{ij} = 0$ if $Y_i \neq j$. Some people write $[Y_i=j]$ for this (the Iverson bracket).

The log-likelihood of the data under the Gaussian mixture model is

where in the first step we inserted a fat 1: $Y_{ij}$ is 1 for just one factor of the product $\prod_{j=1}^k (...)^{Y_{ij}}$ and 0 otherwise, which makes the corresponding factor 1. The second step is just $\log \prod (...)^t = \sum t \log (...)$.

Note also that $\theta$ is now really only $\theta_1,\dots,\theta_k$, since the $p_i$ are encoded in the distribution of $Y$.

We can simplify further:

The expected value over the latent variable¶

If we assume a distribution $q(y)$ for the latent variable $Y$, then we can take the expected value over it:

The terms denoted $\tau_{ij}$ we call the responsibilities, since in our model they are responsible for the division of the points $x_i$ into classes:

If we hold the $\tau_{ij}$ fixed, then $\mathbb{E}_{Y \sim q} \ell_c(\theta)$ is a quantity we can maximise (even though the $\tau_{ij}$ do of course really depend on $\theta$). We still want to regard the $\tau_{ij}$ as a function of $\theta$, but from now on we call that argument $\vartheta$.

So we define

For $\theta=\vartheta$ this is the expected complete-data log-likelihood under our assumption $q$ about the distribution of the latent variable $Y$. Note, too, that the other term, $\log P(X=x_i, Y=j) = \left( \log p_{j} + \log f_{j}(x_i\ |\ \theta_j) \right)$, depends on $\theta$, not on $\vartheta$.

The difference between $L(\theta, q)$ and $Q(\theta, \theta)$¶

Let us recall the observation we reached with the help of Jensen’s inequality:

and note that the summand other than $H(q)$ is exactly the expected complete-data log-likelihood, that is

Since $H(q)$ does not depend on $\theta$, maximising $Q(\theta,\vartheta)$ in $\theta$ therefore also maximises $L(\theta, q)$ in $\theta$.

The difference between $\ell(\theta)$ and $L(\theta, q)$¶

Let us rearrange the formula for $L(\theta, q)$ once more, in a different way: we express the joint distribution $P(X=x, Y=j)$ through conditional distributions, $P(Y=j, X=x) = P(Y=j|X=x)P(X=x)$, and obtain:

The difference between $\ell(\theta)$ and $L(\theta,q)$ is therefore exactly the relative entropy of the estimated distribution $q$ to the actual distribution $P(Y \mid X=x, \theta)$ (in the notation of the chapter Relative entropy):

With that we also have $\ell(\theta) - Q(\theta,\theta) = H\big(q \,\Vert\, P(Y \mid X=x, \theta)\big) + H(q)$, so we now understand exactly in what sense $\ell(\theta)$ can be maximised by way of $Q(\theta,\theta)$.

The EM algorithm for fitting a Gaussian mixture model¶

The algorithm is initialised with a choice of $\theta^{(0)}$ (and, as we have already seen with k-means, the result generally depends very much on this initialisation). One possible choice for the $\mu_j$ (the cluster centres) is to run the k-means algorithm (with k-means++ initialisation) on the data and take the converged centroids as the $\mu_j$. For the covariance matrices one can start with identity matrices (which corresponds to a shifted standard normal distribution). The mixing ratio, that is the $p_i$, can be initialised to $p_i = \frac{1}{k}$.

Then the E-step and the M-step are carried out iteratively, alternating between the two:

• E-step: update the parameter $\vartheta$: set $\vartheta = \theta$ and compute the responsibilities $\tau_{ij} := \tau_{ij}(\theta)$, so as to have a correct formula for $Q(\theta,\theta)$ again – or rather a new estimate of $q$, and with it of $L(\theta, q)$. In this step we estimate the latent variable, aiming to minimise the relative entropy to its actual distribution (that is, to approximate it as well as possible). The better the estimate, the closer $L(\theta, q)$ lies to $\ell(\theta)$.

• M-step: update the parameter $\theta$: maximise the expected complete log-likelihood $Q(\theta,\vartheta)$ in $\theta$ (with $\vartheta$ held fixed, so that the responsibilities $\tau_{ij}$ are constant). This maximises $L(\theta, q)$, a lower bound on $\ell(\theta)$.

The procedure converges, since the approximations to a maximum of $\ell(\theta)$ never get worse (because of Jensen’s inequality, one might say).

So far we have not needed to know what form the density functions $f_j$ take – in fact everything discussed up to this point works for arbitrary distributions as well! With some care, the latent variable may even be continuous, as long as the sums from $j=1$ to $k$ are suitably replaced by the corresponding integrals. Seen that way, the EM algorithm is really not one algorithm but a template for developing many concrete EM algorithms.

More generally still, when the aim is not necessarily to maximise a likelihood in the presence of latent variables, one speaks of MM algorithms – for minorise–maximise (or, for minimisation problems, majorise–minimise) – and of optimisation transfer; a good introduction is Hunter & Lange, A Tutorial on MM Algorithms, The American Statistician 58 (2004). With these keywords you can track down the literature should you need a generalisation of the EM algorithm.

For the moment, though, we restrict ourselves to a discrete, finite latent variable $Y$ and head towards an implementation in Python.

Concrete formulas for the Gaussian mixture model¶

To implement the algorithm concretely we must of course use concrete density functions too, and that is what we do now – so let us recall that the multivariate normal distribution has this density:

To make life easy for ourselves, we look at the univariate case first, so that

At step $t$ of the iteration we are given the parameters $\theta^{(t)}$, that is, $\mu_j^{(t)}$ and $\Sigma_j^{(t)}$ together with the mixture parameters $p_j^{(t)}$. The superscript index $t$ counts the iterations; $i$ remains the sample index, as in $x_i$.

In the E-step we compute the responsibilities

In the M-step we maximise $Q(\theta,\vartheta)$ in $\theta$ by finding the zeros of the derivative with respect to $\theta$, that is, we solve the system of equations

First we look at the derivatives with respect to the $\mu_j$:

Multiplying through by $\sigma_j^2$ and solving for $\mu_j$ gives the closed formula

Next we look at the derivatives with respect to the $\sigma_j$:

Solving the equation $0 = \frac{d}{d\sigma_j} Q(\theta, \theta^{(t)})$ for $\sigma_j$ (multiplying through by $\sigma_j^3$) gives

Finally, in the same way, the zeros of the derivative with respect to $p_j$ give the formula that represents $p_j$ as the responsibility of $x_i$ for class $j$, averaged over the $x_i$:

With that, the new parameters $\theta^{(t+1)}$ have been found.

Special case: k-means¶

If we hold $p_j = \frac{1}{k}$ and $\Sigma_j = E_d$ fixed, we obtain the formulas

A real difference from the classical k-means algorithm is that here we have worked with soft assignments: for every possible class of a point $x_i$ we have determined a probability $\tau_{ij}$. Classical k-means, by contrast, commits itself – it makes a hard assignment. We can just as well set $\widetilde{\tau_{ij}} := 1$ when $\tau_{ij}$ is the largest of all the $\tau_{im}$, and $\widetilde{\tau_{ij}} := 0$ otherwise. With that we get

where $y_i$ is the class to which $x_i$ was hard-assigned. That is exactly the arithmetic mean of the $n_j$ points $x_i$ in class $j$. With the soft assignment instead, it is a weighted mean of the $x_i$ that gives less weight to the points whose assignment is less certain – another obvious improvement on classical k-means.

Implementing more general models¶

If you want to implement models more general than the ones discussed here, you have to adapt the part in which the explicit closed formulas for $p_j$, $\mu_j$, $\sigma_j$ are computed. We were able to find closed formulas because the $\log$ of a density that can be written as $e^{(...)}$ is especially simple. As a rule of thumb, then, every distribution of that same build can be treated the same way and used in place of the normal distribution. The precise statement is that every distribution from a so-called exponential family is suitable.

Implementation in Python¶

Univariate data¶

As with the k-means algorithm, we now want to watch the fit improve over the course of the iteration. First we knock together some test data:

Multivariate data¶

For the more general multivariate case we shall not labour through the derivation again, but go straight to the corresponding code:

Example applications¶

Whereabouts¶

data.csv contains a record of whereabouts. The data come, under the MIT licence, from a blog post which explains how to embed OpenStreetMaps – though that is not needed here.

The transformation to principal components is of little use with these data:

One interpretation would be that we have now identified the places where the longer stays happened.

Iris classification¶

We now use scikit-learn’s GMM (Gaussian mixture model) implementation to model the Iris dataset.

From the Gaussian mixture model we can also obtain the individual probabilities:

Here we look for the permutation that agrees best with the actual classes:

And finally we plot what the result looks like:

Handwritten digits¶

On the UCI ML hand-written digits dataset we look at how well such a GMM works and how it can be used generatively, to produce new samples of a class.

How many components should the model have? In the chapter on model selection we met tools for questions of this kind. For mixture models the Bayesian information criterion (BIC) is in common use: $\mathrm{BIC} = -2 \log L + m \log n$, with the log-likelihood $\log L$, the number $m$ of model parameters and the sample size $n$. Every additional component increases $m$ and has to earn that back first through a better likelihood. One picks the $k$ with the smallest BIC:

Here the BIC keeps falling right up to the largest $k$ tested: the digits evidently break up into very many small clusters (different ways of writing the same digit), and anyone who wants to find the minimum has to extend the range. In what follows we nevertheless try out various $k$ by hand, to see what the individual components learn.

These new digits were generated for us by the GMM. But why are the labels that go with them not the digits 0–9?

Because in training we never tried to have exactly one normal distribution per digit in the mixture. With the 100 individual normal distributions used here, generating new digits works all the same. If we additionally run the classification separately (with SVMs, say – or indeed directly, as a latent variable in the EM algorithm), we have a model that can also generate new digits of a particular class on demand (a 7, for instance).

To gain a little insight into our model, let us look at individual mixture components:

So we do see that each mixture component stays within a single class – there are simply several different “zeros”. With only 10 mixture components, the model looks worse:

To generate, say, a 7 within our paradigm and without additional algorithms, we can train directly on the subset of the data consisting of the 7s:

Fashion-MNIST¶

Simply because we can, let us look at exactly the same thing on the Fashion-MNIST dataset:

Let us see how long PCA.fit and EM (gm.fit) take on this somewhat larger dataset:

Looking back and looking ahead¶

With that, these notes come to an end. The road through the machine learning part of the course was a chain of ideas, each built on the last:

• k-means was our first algorithm: measure distances, assign points to centroids, recompute the centroids.

• In the chapter on statistical models we recognised k-means as a maximum likelihood estimator and introduced the notions of estimator, bias and likelihood.

• Principal component analysis and linear regression were the first methods to rest on the covariance matrix and on minimising squared deviations, respectively, with regularisation as the remedy for overfitting.

• The kernel trick showed that many of these methods need only distances or inner products and can therefore be carried over to nonlinear structures.

• Statistical tests and model selection dealt with how to assess results and models without letting the sample fool you: significance, cross-validation, systematic hyperparameter search.

• Expectation Maximization closes the circle: a Gaussian mixture model is a ‘soft’ k-means, one that follows from the maximum likelihood method and Jensen’s inequality, and with criteria such as the BIC even the number of components can be chosen on good grounds.

We have deliberately left out, among other things, neural networks and deep learning, time series analysis, reinforcement learning, and everything that comes after the model: deployment, monitoring and the maintenance of models in production. The foundations laid in these notes – probability, expected value, covariance, likelihood, regularisation and validation – turn up again in all of those subjects. The course Introduction to Deep Learning follows on directly from these notes; in particular it uses random variables and the expected value, cross-entropy as a loss function, linear regression as the simplest model, Numpy shapes and broadcasting, the train-test split with validation, and the lower bound $L(\theta,q)$ from this chapter, there called the ELBO. Anyone who would like to read further will find the mathematical depth in Friedman et al. (2001) and Deisenroth et al. (2020), and the practitioner’s perspective in Schutt & O'Neil (2013).