Gradient Descent and the EM Algorithm: Born of the Same Root, Part of the Same Lineage

PS: This article simply works out the relationship between gradient descent and the EM algorithm. Following a single unified line of reasoning, we derive ordinary gradient descent, the EM algorithm used in pLSA, and the EM algorithm used in K-Means, showing that they are all basically different facets of the same thing — as the saying goes, "seen from the side, a mountain range; seen from the peak, a single summit — its true face changes with each vantage point."

In machine learning, we usually formulate the problem we want to solve as a loss function with unknown parameters, such as mean squared error (MSE), and then try to find the minimum of this function in order to obtain the best parameter values and complete the modeling process. Since multiplying a function by -1 turns its maxima into minima, we can uniformly speak in terms of finding minima. In the field of machine learning, there are generally two major schools of thought on how to find the minimum of a function: 1) gradient descent; 2) the EM algorithm, i.e. the expectation-maximization algorithm, generally used for solving complicated maximum-likelihood problems.

In most tutorials these two methods are described as utterly different, as though they were two rival schools of thought, and the EM algorithm in particular tends to be described in an almost mystical way. But in fact these two methods are just different instances of the same idea — as the saying goes, they were "born from the same root," and they belong to one and the same lineage.

Let's start from the ancient method: Newton's method.

Newton's iteration method

Given some complicated nonlinear function $f(x)$, suppose we want to find its minimum. Generally, assuming it is sufficiently smooth, its minimum will be a critical point, satisfying $f'(x_0)=0$, so the problem reduces to finding the root of the equation $f'(x)=0$. We have Newton's method for finding roots of nonlinear equations, so

\begin{equation}x_{n+1} = x_{n} - \frac{f'(x_n)}{f''(x_n)}\end{equation}more

However, this approach strips away the geometric meaning, and doesn't let us peek into any deeper secrets. Instead, let's use the following line of reasoning: at the point $x=x_n$ of $y=f(x)$, we can use an approximating curve to locally mimic the original function, and if that approximating curve is easy to minimize, then we can use the minimum of the approximating curve as a stand-in for the minimum of the original curve:

Approximation-iterationApproximation-iteration

Clearly, the requirements on the approximating curve are:

1. It should match the true curve to some degree — generally we require at least first-order agreement;
2. It should have a minimum, and that minimum should be easy to solve for.

With this in mind, a very natural choice is the "tangent parabola":

\begin{equation}f(x)\approx g(x) = f(x_n)+f'(x_n)(x-x_n)+\frac{1}{2}f''(x_n)(x-x_n)^2\end{equation}

This parabola has second-order accuracy. For this parabola, the extremum point is

\begin{equation}x_n - \frac{f'(x_n)}{f''(x_n)}\end{equation}

so we recover the iteration formula of Newton's method:

\begin{equation}x_{n+1} = x_n - \frac{f'(x_n)}{f''(x_n)}\end{equation}

If $f(x)$ is sufficiently smooth and has only a single extremum globally, Newton's method converges quickly (with exponential speed). But real functions are rarely so well-behaved, and so its drawbacks become apparent:

1. It requires computing the second derivative, and for some functions the second derivative is already quite complicated;
2. Since the sign of $f''(x_n)$ is not fixed, the opening direction of $g(x)$ is not fixed either, so we cannot be sure whether the result we ultimately get is a maximum or a minimum.

Gradient descent

These two drawbacks are fatal for many problems. So, in order to fix them, we give up second-order accuracy — that is, we drop $f''(x_n)$ and replace it with a fixed positive constant $1/h$:

\begin{equation}g(x) = f(x_n)+f'(x_n)(x-x_n)+\frac{1}{2h}(x-x_n)^2\end{equation}

This approximating curve only has first-order accuracy, but in exchange it eliminates the need to compute the second derivative, and it guarantees an upward-opening parabola, so iterating with it is at least guaranteed to converge to some minimum (at least a local one). The minimum point of the above $g(x)$ is

\begin{equation}x_n - h f'(x_n)\end{equation}

so we get the iteration formula

\begin{equation}x_{n+1} = x_n - h f'(x_n)\end{equation}

and in higher dimensions

\begin{equation}\boldsymbol{x}_{n+1} = \boldsymbol{x}_n - h \nabla(\boldsymbol{x}_n)\end{equation}

This is the famous gradient descent method. Of course it has many issues of its own, but a great many improved algorithms — such as stochastic gradient descent — are built around it.

Here we've understood gradient descent as the result of approximating with a parabola. Seen this way, readers will naturally wonder: why should I necessarily use a parabola for the approximation? Couldn't I use some other curve? Of course you could — for many problems, gradient descent can actually make things harder, i.e., the parabola approximation breaks down, and in those cases we need to consider other forms of approximation. In fact, essentially all other approximation schemes are called "the EM algorithm" — as though gradient descent, despite being born from the same root, were somehow excluded — which is rather puzzling.

Maximum likelihood

When estimating probabilities, our usual optimization target is the likelihood function rather than MSE. For instance, when building a language model, we need to estimate the co-occurrence probability $p(x,y)$ of any two words $x,y$. Suppose that in a corpus of size $N$, the pair $x,y$ co-occurs $\#(x,y)$ times; then we can obtain the empirical estimate

\begin{equation}\tilde{p}(x,y)=\frac{\#(x,y)}{N}\end{equation}

This is of course the most basic result, but it suffers from sparsity issues, and moreover storing a result for every possible pair of words is already too much for memory to handle.

A better solution is to assume that $p(x,y)$ can be represented by a function $p(x,y;\theta)$ with unknown parameters $\theta$ ($\theta$ might be a vector) — for example, a neural network — and then simply optimize the parameters. So the question becomes: what should the optimization target be? Note that if we use MSE, the biggest problem is that there's no way to guarantee the resulting values are non-negative, whereas probabilities must always be non-negative.

For probability problems, statisticians proposed a more natural scheme — the maximum likelihood function. Philosophers often say "what exists is reasonable"; the idea behind the maximum likelihood function goes even further, saying "what exists is the most reasonable." Suppose the pair $x,y$ co-occurred $\#(x,y)$ times; since this event actually happened, it must be the most reasonable outcome, and so the probability function

\begin{equation}\prod_{x,y} p(x,y;\theta)^{\#(x,y)}\end{equation}

should attain its maximum. Taking the logarithm gives

\begin{equation}\sum_{x,y} \#(x,y)\log p(x,y;\theta)\end{equation}

and we should maximize this function — this is the maximum likelihood function. Clearly, $\#(x,y)$ can be replaced by the empirical frequency $\tilde{p}(x,y)$, giving an equivalent result:

\begin{equation}\sum_{x,y} \tilde{p}(x,y)\log p(x,y;\theta)\end{equation}

In fact, multiplying this by -1 gives

\begin{equation}S=-\sum_{x,y} \tilde{p}(x,y)\log p(x,y;\theta)\end{equation}

We give this a special name — cross-entropy — and it's one of the most common loss functions in machine learning. In other words, maximizing the likelihood function is equivalent to minimizing the cross-entropy. If instead of prespecifying a functional form for $p(x,y;\theta)$ we directly estimate $p(x,y;\theta)$, it's not hard to show that $p(x,y;\theta)=\tilde{p}(x,y)$, which is exactly what we would expect, and which also demonstrates the reasonableness of the maximum likelihood function as an optimization target.

The EM algorithm

For optimizing the cross-entropy, we would usually also try gradient descent. But in many cases gradient descent is ineffective, and we'd rather use the EM algorithm, which has been called "God's algorithm."

Continuing with the language model example, to improve the generalization ability of our estimate, we transform $p(x,y)$ into $p(x|y)p(y)$, and then factor $p(x|y)$ as $p(x|y)=\sum_z p(x|z)p(z|y)$. The meaning of this factorization was already discussed in the article SVD Decomposition (II): Why Does SVD Imply Clustering?: $z$ can be understood as a category, or a topic; $p(x|y)$ is the probability of $x$ following $y$; $p(z|y)$ is the probability that $y$ belongs to topic $z$; and $p(x|z)$ can be understood as the probability that $x$ appears within topic $z$. Generally, the number of $z$'s is much smaller than the number of $x,y$'s, which reduces the total number of parameters.

In this case, the cross-entropy becomes

\begin{equation}S=-\sum_{x,y} \tilde{p}(x,y)\log \sum_z p(x|z)p(z|y)p(y)\end{equation}

We can take $p(y)$ to be estimated exactly by $\tilde{p}(y)$, so this term only contributes a constant, and so the equivalent optimization target is:

\begin{equation}S=-\sum_{x,y} \tilde{p}(x,y)\log \sum_z p(x|z)p(z|y)\end{equation}

Here all the $p(x|z),p(z|y)$ are parameters to be solved for (ranging over all possible x, y, z combinations).

Its gradient is

\begin{equation}\begin{aligned}&\frac{\partial S}{\partial p(x|z)} = -\sum_{y} \frac{\tilde{p}(x,y)}{\sum_z p(x|z)p(z|y)}p(z|y)\\ &\frac{\partial S}{\partial p(z|y)} = -\sum_{x} \frac{\tilde{p}(x,y)}{\sum_z p(x|z)p(z|y)}p(x|z)\end{aligned}\end{equation}

Direct gradient descent won't work here, because the $p(x|z),p(z|y)$ are all non-negative and satisfy the constraint

\begin{equation}\sum_x p(x|z) = 1,\quad \sum_z p(z|y)=1\end{equation}

Gradient descent cannot guarantee this non-negativity constraint, which means it's fundamentally unable to solve this kind of problem effectively. Recall our earlier derivation of gradient descent — the idea of using an approximating curve in place of the original curve for iteration. Gradient descent uses a parabolic approximation; here, we will not use a parabolic approximation (since it cannot preserve the positivity constraint), but instead try to construct a new kind of approximation. Suppose we've already carried out $n$ iterations and obtained the estimate $p_n(x|z),p_n(z|y)$; then, according to the gradient formula, the gradient at this step is

\begin{equation}\begin{aligned}&\frac{\partial S}{\partial p_n(x|z)} = -\sum_{y} \frac{\tilde{p}(x,y)}{\sum_z p_n(x|z)p_n(z|y)}p_n(z|y)\\ &\frac{\partial S}{\partial p_n(z|y)} = -\sum_{x} \frac{\tilde{p}(x,y)}{\sum_z p_n(x|z)p_n(z|y)}p_n(x|z)\end{aligned}\end{equation}

The difficulty with the original problem is that $\log$ still contains a summation inside it — if we could move the summation to the outside, things would be much simpler. So, let's consider an approximating function of the form

\begin{equation}S_n'=-\sum_{x,y} \tilde{p}(x,y)\sum_z C_{x,y,z,n} \log p(x|z)p(z|y)\end{equation}

where $C$ is a constant (for this iteration). This way $S_n'$ also has a minimum, and it can be solved for exactly. Naturally, we want the gradient of $S'$ to match the gradient of the original $S$ (i.e., to have first-order accuracy). The gradient of $S'$ is

\begin{equation}\begin{aligned}&\frac{\partial S'}{\partial p_n(x|z)} = -\sum_{y} \frac{\tilde{p}(x,y)C_{x,y,z,n}}{p_n(x|z)}\\ &\frac{\partial S'}{\partial p_n(z|y)} = -\sum_{x} \frac{\tilde{p}(x,y)C_{x,y,z,n}}{p_n(z|y)}\end{aligned}\end{equation}

Comparing the two gradients, we obtain

\begin{equation}C_{x,y,z,n}=\frac{p_n(x|z)p_n(z|y)}{\sum_z p_n(x|z)p_n(z|y)}\end{equation}

In other words, once we have an initial set of parameters, we can substitute them into the above formula to get $C_{x,y,z,n}$, and then find the parameters that minimize $S_n'$ to serve as $p_{n+1}(x|z)$ and $p_{n+1}(z|y)$, and iterate in this manner.

This is basically the derivation used to solve the pLSA model; further details can be found in Natural Language Processing and PLSA. For our purposes here, the point is simply this: the EM algorithm, just like gradient descent, is born from the same root — both are based on approximating-curve techniques, and there is nothing mystical about it; moreover, the choice of this approximating function has a solid rational basis. Yet nearly every tutorial online just directly hands you the expression for $S'$ (commonly called the "Q function" in such tutorials), presenting it in a way that feels, to me, almost like mysticism — and that has always bothered me.

K-Means

K-Means clustering is easy to understand: given the coordinates $\boldsymbol{x}_i,\,i=1,\dots,N$ of $N$ points, we want to partition these points into $K$ clusters, each with a cluster center $\boldsymbol{c}_j,\,j=1,\dots,K$. Naturally, the category a point belongs to is determined by whichever cluster center $\boldsymbol{c}_j$ (representing that category) is closest to it, where distance is defined as Euclidean distance.

So the main task in K-Means clustering is to find the cluster centers $\boldsymbol{c}_j$. Naturally, we want each cluster center to sit right at the "center" of its category. Expressed as a function, we want the following function $L$ to be minimized:

\begin{equation}L=\sum_{i=1}^N \min\bigg\{|\boldsymbol{x}_i-\boldsymbol{c}_1|^2,|\boldsymbol{x}_i-\boldsymbol{c}_2|^2,\dots,|\boldsymbol{x}_i-\boldsymbol{c}_K|^2\bigg\}\end{equation}

where the $\min$ operation ensures that each point is assigned only to the category whose center is nearest to it.

If we try to optimize $L$ directly with gradient descent, we run into serious trouble — not because the $\min$ operation is hard to differentiate, but because this is an NP problem, and the theoretical convergence time grows exponentially with $N$. Here too we use the EM algorithm, which in this setting takes the following form:

1. Randomly select $K$ points as the initial cluster centers;
2. Given the current $K$ cluster centers, determine which category each point belongs to, and then use the average coordinates of all points within the same category as the new cluster center.

This method basically converges after just a few iterations — so where does the justification for it come from?

We again follow the idea of approximating-curve methods. But now the problem is: what do we do since $\min$ isn't differentiable? We can consider a smooth approximation, and then take a limit — the answer can be found in Looking for a Smooth Maximum Function. Take $M$ sufficiently large; then we may treat (the minimum equals negative the maximum of the negatives)

\begin{equation}\begin{aligned}&\min\bigg\{|\boldsymbol{x}_i-\boldsymbol{c}_1|^2,|\boldsymbol{x}_i-\boldsymbol{c}_2|^2,\dots,|\boldsymbol{x}_i-\boldsymbol{c}_K|^2\bigg\}\\ =&-\frac{1}{M}\ln\bigg(e^{-M|\boldsymbol{x}_i-\boldsymbol{c}_1|^2}+e^{-M|\boldsymbol{x}_i-\boldsymbol{c}_2|^2}+\dots+e^{-M|\boldsymbol{x}_i-\boldsymbol{c}_K|^2}\bigg)\end{aligned}\end{equation}

so that

\begin{equation}L=-\sum_{i=1}^N \frac{1}{M}\ln\bigg(e^{-M|\boldsymbol{x}_i-\boldsymbol{c}_1|^2}+e^{-M|\boldsymbol{x}_i-\boldsymbol{c}_2|^2}+\dots+e^{-M|\boldsymbol{x}_i-\boldsymbol{c}_K|^2}\bigg)\end{equation}

can now be differentiated:

\begin{equation}\frac{\partial L}{\boldsymbol{c}_j}=\sum_{i=1}^N \frac{2e^{-M|\boldsymbol{x}_i-\boldsymbol{c}_j|^2 } }{e^{-M|\boldsymbol{x}_i-\boldsymbol{c}_1|^2}+e^{-M|\boldsymbol{x}_i-\boldsymbol{c}_2|^2}+\dots+e^{-M|\boldsymbol{x}_i-\boldsymbol{c}_K|^2}}(\boldsymbol{c}_j-\boldsymbol{x}_i)\end{equation}

Let the result of the $n$-th iteration be $\boldsymbol{c}^{(n)}_j$; then the gradient at this round is:

\begin{equation}\frac{\partial L}{\boldsymbol{c}^{(n)}_j}=\sum_{i=1}^N \frac{2e^{-M|\boldsymbol{x}_i-\boldsymbol{c}^{(n)}_j|^2 } }{e^{-M|\boldsymbol{x}_i-\boldsymbol{c}^{(n)}_1|^2}+e^{-M|\boldsymbol{x}_i-\boldsymbol{c}^{(n)}_2|^2}+\dots+e^{-M|\boldsymbol{x}_i-\boldsymbol{c}_K^{(n)}|^2}}(\boldsymbol{c}^{(n)}_j-\boldsymbol{x}_i)\end{equation}

Based on the form of this expression, we can look for an approximating curve (really a hypersurface) of this form:

\begin{equation}L'=\sum_{i=1}^N \sum_{j=1}^K C^{(n)}_{i,j} |\boldsymbol{x}_i-\boldsymbol{c}_j|^2 \end{equation}

where $C^{(n)}_{i,j}$ is to be determined, and is a constant at each iteration, making this simply a quadratic function whose minimum is easy to find, at

\begin{equation}\boldsymbol{c}_j = \frac{\sum_{i=1}^N C^{(n)}_{i,j}\boldsymbol{x}_i}{\sum_{i=1}^N C^{(n)}_{i,j}}\end{equation}

that is, a weighted average of $\boldsymbol{x}_i$.

As before, we want this approximating curve to match the original function to at least first order, so we take its derivative:

\begin{equation}\frac{\partial L'}{\boldsymbol{c}_j}=\sum_{i=1}^N 2C^{(n)}_{i,j} (\boldsymbol{c}_j-\boldsymbol{x}_i)\end{equation}

Comparing this with the derivative of the original function, we readily obtain

\begin{equation}C^{(n)}_{i,j} = \frac{e^{-M|\boldsymbol{x}_i-\boldsymbol{c}^{(n)}_j|^2 } }{e^{-M|\boldsymbol{x}_i-\boldsymbol{c}^{(n)}_1|^2}+e^{-M|\boldsymbol{x}_i-\boldsymbol{c}^{(n)}_2|^2}+\dots+e^{-M|\boldsymbol{x}_i-\boldsymbol{c}_K^{(n)}|^2}}\end{equation}

which gives us the iteration formula:

\begin{equation}\boldsymbol{c}^{(n+1)}_j = \frac{\sum_{i=1}^N C^{(n)}_{i,j}\boldsymbol{x}_i}{\sum_{i=1}^N C^{(n)}_{i,j}}\end{equation}

At this point we've worked out every step of the derivation, but since we're still using a continuous approximation, in the end we need to take the limit as $M\to\infty$ goes to infinity, and taking that limit simplifies things considerably. From the formula above, we can derive:

\begin{equation}\lim_{M\to\infty} C^{(n)}_{i,j} = \Delta^{(n)}_{i,j} = \left\{\begin{aligned}&1,\text{for fixed i, min distance from j to i}\\ &0,\text{other cases}\end{aligned}\right.\end{equation}

In plain terms: think of $\Delta^{(n)}_{i,j}$ as a matrix with $N$ rows and $K$ columns; each row of this matrix can have only a single 1, with everything else being 0. If the $j$-th entry in row $i$ is 1, that means the cluster center nearest to $\boldsymbol{x}_i$ is $\boldsymbol{c}_j$. In that case, the iteration formula becomes

\begin{equation}\boldsymbol{c}^{(n+1)}_j = \frac{\sum_{i=1}^N \Delta^{(n)}_{i,j}\boldsymbol{x}_i}{\sum_{i=1}^N \Delta^{(n)}_{i,j}}\end{equation}

Given the meaning of $\Delta^{(n)}_{i,j}$, this simply says:

$\boldsymbol{c}^{(n+1)}_j$ is the average of all the points nearest to $\boldsymbol{c}^{(n)}_j$.

And this is exactly the iterative algorithm we normally use to solve K-Means, which is likewise called an instance of the EM algorithm.

Summary

As we can see, so-called "EM algorithm" isn't a single specific method but rather a whole family of methods, or perhaps better described as a strategy — and gradient descent turns out to be just a special case within this same family; they're really one and the same thing, and strictly speaking gradient descent shouldn't be excluded from it. What's called "God's algorithm" is really nothing more than an iterative method: through repeated self-updating it can approach perfection (the optimal solution), much like biological evolution converging on something exquisite, as though shaped by a careful designer — small wonder it earned the name "God's algorithm."

English translation of a post from 科学空间 | Scientific Spaces by 苏剑林. Original: https://kexue.fm/archives/4277
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.