###### tags: `sampling` `em` `vi` `monte carlo` `one-offs`

# Interacting Particle Systems for Expectation-Maximisation

**Overview**: In this note, I will describe the problem of parameter estimation in latent variable models, and its resolution via the EM algorithm. I will then describe some particular implementations of the EM idea involving Monte Carlo approximations, particle methods, and other variations.

## Latent Variable Models

Classical statistical models are often 'fully-observed': there is some parameter of interest $\theta$, an associated generative model $P ( x | \theta )$ which generates data $\{ x_i \}_{i = 1}^N$, and we are tasked with finding an estimate of the $\theta$ which generated our data. This is often handled via the principle of *maximum likelihood*, wherein one estimates

\begin{align}
\theta \approx \hat{\theta} = \arg\max_\theta \sum_{i = 1}^N \log P(x_i | \theta).
\end{align}

When we can write down $P(x | \theta)$ in closed form, one can then apply standard methods of numerical optimisation to compute $\hat{\theta}$.

However, there are a number of statistical models for which this is *not* the case, due to the presence of unobserved *latent* (also 'hidden') variables. That is, for each observation $x_i$, there is some unobserved quantity $z_i$ such that the joint probability $P(x_i, z_i | \theta)$ is available, even though the marginal probability $P(x_i | \theta)$ may not be. Some standard examples include Mixture Models, Topic Models, Matrix Factorisation problems, and Hidden Markov (Chain) Models.

From a modeling point of view, latent variable models of this form are often desirable, as they allow for richer descriptions of the data-generating process, accounting for heterogeneity, heavy tails, and other phenomena. However, this comes at a cost, in the sense that parameter estimation becomes more challenging. One could attempt a principle of 'maximum joint likelihood'. i.e.

\begin{align}
\left( \hat{\theta}, \hat{z}_{1:N} \right) = \arg\max_{\theta, z_{1:N}} \sum_{i = 1}^N \log P(x_i, z_i | \theta),
\end{align}

but this is typically a bad idea, as the freedom to optimise over the $z_i$ can lead to overfitting - the flexibility of including these additional variables in the model is essentially over-exploited.

As such, it is more common to persist with the original principle of maximum (marginal) likelihood, and to focus on alternative strategies for solving the MLE problem.

## Surrogate Approaches to Maximum Likelihood Estimation

A useful strategy in optimisation is to replace the direct optimisation of a hard function by the indirect optimisation of an easier function which admits the same optimisers. For instance, suppose that we are interested in maximising a complicated function $f : \mathcal{X} \to \mathbf{R}$, but we have access to the representation

\begin{align}
f(x) = \max_{y \in \mathcal{Y}} F(x, y),
\end{align}

where $F$ is a simpler function. One can then attempt to instead solve the problem of maximising $F$, and recover the solution to the original problem as a byproduct.

This principle is useful in the context of latent variable models. Consider our likelihood function

\begin{align}
\log P(x|\theta) = \log \left( \int P(x, z| \theta) dz \right).
\end{align}

Using a change of measure, we can equally write this as

\begin{align}
\log P(x|\theta) = \log \left( \int Q(z) \cdot \frac{P(x, z| \theta)}{Q(z)} dz \right),
\end{align}

and using the concavity of $\log$ to apply Jensen's inequality demonstrates that

\begin{align}
\log P(x|\theta) &= \log \left( \int Q(z) \cdot \frac{P(x, z| \theta)}{Q(z)} dz \right) \\
&\geqslant \int Q(z) \log \left( \frac{P(x, z| \theta)}{Q(z)} \right) dz.
\end{align}

In fact, this expression can be written as

\begin{align}
\int Q(z) \log \left( \frac{P(x, z| \theta)}{Q(z)} \right) dz &= \log P (x | \theta) - \text{KL} \left( Q(z), P(z | x, \theta) \right),
\end{align}

from which the inequality follows immediately. Moreover, it furnishes us with the representation

\begin{align}
\log P (x | \theta) &= \max_Q \int Q(z) \log \left( \frac{P(x, z| \theta)}{Q(z)} \right) dz \\
&= \max_Q \,\mathbf{E}_Q \left[ \log \left( \frac{P(x, z| \theta)}{Q(z)} \right) \right],
\end{align}

which might be useful, in accordance with our earlier construction. Crucially, for many statistical problems of interest, the expectation $\mathbf{E}_Q \left[ \log P(x, z| \theta) \right]$ can be computed easily (as a function of $\theta$).

With this in mind, one can then rewrite the MLE problem as

\begin{align}
\max_\theta \sum_{i = 1}^N \log P(x_i | \theta) &= \max_\theta \sum_{i = 1}^N \max_{Q_i} \mathbf{E}_{Q_i (z_i)} \left[ \log \left( \frac{P(x_i, z_i| \theta)}{Q_i (z_i)} \right) \right],
\end{align}

and so by defining

\begin{align}
\mathcal{F} \left( \theta, \{ Q_i \}_{i = 1}^N \right) &= \sum_{i = 1}^N \mathbf{E}_{Q_i (z_i)} \left[ \log \left( \frac{P(x_i, z_i| \theta)}{Q_i (z_i)} \right) \right],
\end{align}

one arrives at a joint maximisation problem which is often tractable, at least in the sense of being able to write down many of the terms in closed form.

## Algorithms for Maximising $\mathcal{F}$

### Standard EM Algorithm

Given the task of maximising $\mathcal{F}$, a fruitful strategy is that of alternating maximisation, i.e. iterating

1. Fixing $\{ Q_i \}_{i = 1}^N$, solve

\begin{align}
\max_\theta \sum_{i = 1}^N \mathbf{E}_{Q_i (z_i)} \left[ \log P(x_i, z_i| \theta) \right],
\end{align}

either in closed form or numerically.

2. Fixing $\theta$, solve

\begin{align}
\max_{\{ Q_i \}_{i = 1}^N} \sum_{i = 1}^N \mathbf{E}_{Q_i (z_i)} \left[ \log \left( \frac{P(x_i, z_i| \theta)}{Q_i (z_i)} \right) \right],
\end{align}

by setting $Q_i (z_i) = P(z_i | x_i, \theta)$.

This scheme is known as the *Expectation-Maximisation* (EM) Algorithm, and has been highly influential since its inception in 1977.

### Monte Carlo EM

A drawback of the standard EM algorithm is that it presupposes that $\mathbf{E}_{Q_i (z_i)} \left[ \log P(x_i, z_i| \theta) \right]$ can be evaluated easily as a function of $\theta$. In practice, this need not be the case, and so one might instead draw some samples from $Q_i$, and then maximise the resulting empirical approximation with respect to $\theta$.

A slight weakness of this approach is that it requires the user to choose the number of samples to draw from $Q_i$. There are also strategies wherein one draws only a small number of samples from $Q_i$ at each step, but carries out the maximisation step with respect to a weighted sum of all such past samples; this is essentially the *stochastic approximation EM algorithm*. There are also variations in which one only samples from the $Q_i$ approximately, using e.g. Markov chain Monte Carlo methods.

### Variational EM

An alternative approach is to restrict a priori to using classes of $Q_i$ such that the expectation $\mathbf{E}_{Q_i (z_i)} \left[ \log P(x_i, z_i| \theta) \right]$ can be handled easily. The first step of the EM algorithm then remains the same, but the second step is replaced by solving

\begin{align}
Q_i^* &= \arg\min_{Q_i \in \mathcal{Q}} \text{KL} \left( Q_i (z_i), P(z_i | x_i, \theta) \right),
\end{align}

where $\mathcal{Q}$ is some sufficiently nice class of approximating measures.

One can also imagine fusing this method with Monte Carlo EM, for situations in which the expectation remains intractable, but sampling from $Q_i \in \mathcal{Q}$ is much easier than sampling from the globally-optimal $Q_i$. Equally, one can imagine deploying such a variational approximation within a Monte Carlo EM strategy which uses importance sampling.

### Gradient EM

When the first step of the EM algorithm is not available in closed form, one may simply take a gradient step with respect to $\theta$, i.e. move in the direction of

\begin{align}
\nabla_\theta \left( \sum_{i = 1}^N \mathbf{E}_{Q_i (z_i)} \left[ \log P(x_i, z_i| \theta) \right] \right) &= \sum_{i = 1}^N \mathbf{E}_{Q_i (z_i)} \left[ \nabla_\theta \log P(x_i, z_i| \theta) \right],
\end{align}

either analytically, or by constructing an unbiased stochastic gradient estimate.

There are also many other variants of EM which I will not address here.

## Expectation-Maximisation with Particle Systems

Return now to our objective function:

\begin{align}
\mathcal{F} \left( \theta, \{ Q_i \}_{i = 1}^N \right) &= \sum_{i = 1}^N \mathbf{E}_{Q_i (z_i)} \left[ \log \left( \frac{P(x_i, z_i| \theta)}{Q_i (z_i)} \right) \right].
\end{align}

Instead of optimising $Q_i$ analytically in an alternating fashion, one can consider approximating $Q_i$ with a particle system

\begin{align}
Q_i (dz_i) \approx Q_t^i (dz_i) = \frac{1}{M} \sum_{a = 1}^M \delta (z_t^{i, a}, dz_i)
\end{align}

and then optimising $\mathcal{F}$ nonparametrically, using ideas from Wasserstein gradient flows.

Indeed, for fixed $\theta$, one can write that

\begin{align}
\mathcal{F} \left( \theta, \{ Q_i \}_{i = 1}^N \right) &= \sum_{i = 1}^N \text{KL} \left( Q_i (z_i), P(z_i | x_i, \theta) \right),
\end{align}

and so under the evolution of the appropriate Wasserstein gradient flow (with respect to the $Q_i$), the particles corresponding to $Q_i$ could be taken to follow the motion of an overdamped Langevin diffusion with invariant measure $P(z_i | x_i, \theta)$, i.e.

\begin{align}
\text{for}\, a \in [M], \quad dz_t^{i, a} = \nabla_z \log P \left( z_t^{i, a} | x_i, \theta \right) dt + \sqrt{2} dW_t^{i, a}.
\end{align}

Note that strictly speaking, from the WGF point of view, one could simulate using only a single particle, as there are no interactions in the above. However, when it comes to optimising $\mathcal{F}$ with respect to $\theta$, the presence of multiple particles should be helpful.

In this direction, assume that $\theta$ follows an (ascending) gradient flow with respect to $\mathcal{F}$, i.e.

\begin{align}
\frac{d \theta_t}{dt} &= \frac{\partial \mathcal{F}}{\partial \theta} \left( \theta_t, \{ Q_t^i \}_{i = 1}^N\right) \\
&= \sum_{i = 1}^N \mathbf{E}_{Q_t^i (z_i)} \left[ \nabla_\theta \log P(x_i, z_i| \theta_t ) \right] \\
&= \frac{1}{M} \sum_{i = 1}^N \sum_{a = 1}^M \nabla_\theta \log P(x_i, z_t^{i, a}| \theta_t ).
\end{align}

One can then hope that $\theta_t$ will converge to a local maximum of $\log P (x_{1:N} | \theta)$, and that the $z_t^i$ will converge in distribution to the measure $P(z_i | x_i, \theta)$. Some of these results may hold at finite $M$, others only asymptotically.

There are of course also many questions of implementation: how to discretise these dynamics, whether interactions between the particles may be beneficial, whether information-geometric preconditioning could improve convergence, and so on. These could collectively amount to a large difference in practical performance.

While this approach is presented as nominally being distinct to the preceding EM approaches, there are of course key similarities. The evolution of the $\theta$ is based on gradient information, as in gradient EM, but also on an ensemble approximation to the objective function, as in Monte Carlo EM.

## Conclusion

In this note, I have provided a brief derivation of the EM algorithm, outlined some of its conventional implementations, and sketched out a separate conceptual implementation based on gradient flows and particle methods. The extent to which one would expect the modifications included in this new method to provide practical benefits remains to be seen.