Learning Mixture of Markov Models with EM

Here I present my Matlab implementation for learning Markov mixtures with EM algorithm. David Barber’s book explains the model and the EM derivation very nicely, and also provides a Matlab code package for it. Both the book and codes are open, you’re recommended to have look at them. However, Barber’s code is written for readability and not for performance. I’ve tried to write the code as efficient as possible. Additionally, if you are uncomfortable with the Expectation-Maximization algorithm, you can read this wonderful tutorial. It’s very detailed, and contains additional side information to help you to understand the method fully.

Here I explain the model and the EM derivations in my own terms. I have also added a Matlab code which is optimised for speed.

The Model

We are going to work with discrete-state Markov chains so let’s call the cardinality of our discrete alphabet \(D\). A Markov model is defined by an initial state probabilities \(\pi\) and a state transition probabilities \(A\):
\[
\begin{align}
\pi(i) &= p(x(1)=i)\\
A(i,j) &= p(x(t)=i \mid x(t-1)=j)
\end{align}
\]
The likelihood of observing a sequence \(x\) is written as
\[
\begin{align}
p(x \mid \pi,A) &= \prod_{i=1}^D \pi(i)^{[x(1)=i]} \prod_{t=2}^{T} \prod_{i=1}^D \prod_{j=1}^D A(i,j)^{[x(t)=i][x(t-1)=j]}
\end{align}
\]

where \([]\) is the indicator function takes the value 1 if the expression inside it is true and 0 otherwise. I’m happier when I re-formulate this equation by getting rid of the product over the sequence length:
\[
\begin{align}
p(x \mid \pi, A) & \prod_{i=1}^D \pi(i)^{[x(1)=i]} \prod_{i=1}^D \prod_{j=1}^D A(i,j)^{\sum_{t=1}^{T}[x(t)=i][x(t-1)=j]} \\
&= \prod_{i=1}^D \pi(i)^{[x(1)=i]} \prod_{i=1}^D \prod_{j=1}^D A(i,j)^{s(i,j)}
\end{align}
\]
where \(s(i,j)\) is the transition counts from state \(i\) to state \(j\). The \(s\) together with the initial state \(x(1)\) forms the sufficient statistics for the sequence \(x\). Once we collect them in the beginning of the computation, we will have shorter likelihood calculations.

Our observations \(x=\{x_1,x_2,\ldots,x_N\}\) are generated by a mixture of Markov models. Meaning that each sequence \(x_n\) is generated by one of the \(K\) Markov models. The probability that \(x_n\) is generated by the \(k^{th}\) Markov model is denoted as
\[
p(z_n=k) = \alpha(k)
\]
From now on, we use the subscripts as enumerators, and the numbers in the parenthesis as array indices. So, \(x_n\) is the \(n^{th}\) sequence and \(x_n(t)\) is its \(t^{th}\) state. If there’s no subscript on a symbol, it denotes the whole set: \( x = \{x_1,\ldots,x_N\} \). Furthermore, we introduce \(\theta_k = \{\pi_k, A_k\}\) to denote the set of a Markov model parameters (for ease of notation).

In the mixture model, we calculate the likelihood of a sequence \(x\) given model parameters as
\[
\begin{align}
p(x|\theta) &= \sum_{k=1}^K p(z=k) p(x|z=k,\theta_k) \\
&= \sum_{k=1}^K \alpha(k) \left ( \prod_{i=1}^D \pi_k(i)^{[x(1)=i]} \prod_{i=1}^D \prod_{j=1}^D A_k(i,j)^{s(i,j)} \right)
\end{align}
\]

Expectation-Maximization Derivation

Given the observations \(x=\{x_1,x_2,\ldots,x_N\}\), we want to learn the parameters of Markov models \(\theta=\{\theta_1,\theta_2,\ldots,\theta_K\}\). Maximum likelihood estimation for this is
\[
\arg\!\max_\theta \log p(x \mid \theta) = \arg\!\max_\theta \log \sum_z p(z \mid \theta) p(x \mid z,\theta)
\]
This is not easy to maximize since there’s a summation over all possible latent state assignments (\(K^N\) possibilities) inside a logarithm. Therefore we need to maximize a tight lower bound for this:
\[
\begin{align}
L(\theta \mid \theta^{(t)}) &= \underbrace{\sum_z p(z \mid x,\theta^{(t)}) \log p(x,z \mid \theta)}_{E(\theta \mid \theta^{(t)})} – \underbrace{\sum_z p(z \mid x,\theta^{(t)}) \log p(z \mid x,\theta^{(t)})}_{H(\theta^{(t)})}
\end{align}
\]
where \(E(\theta \mid \theta^{(t)})\) is the energy part, and \(H(\theta^{(t)})\) is the entropy part and \(\theta^{(t)}\) is the maximum likelihood estimation of \(\theta\) at the \(t^{th}\) iteration. We will be able to maximize this bound easily to get the next estimate of the model parameters:
\[
\theta^{(t+1)} = \arg\!\max_{\theta} L(\theta \mid \theta^{(t)})
\]
The bound will be redefined by using the new expectation \(p(z \mid x,\theta^{(t+1)})\). And the process will continue until the likelihood converges.

Expectation Step

The expectation \(p(z \mid x,\theta_t)\) can be factorized as \(\prod_{n} p(z_n \mid x_n,\theta_t)\) since the observations are independent given \(\theta_t\). We calculate each expectation of \(z_n\) as
\[
\begin{align}
p(z_n \mid x_n,\theta^{(t)}) &= \frac{p(x_n \mid z_n,\theta^{(t)})p(z_n)}{p(x \mid \theta^{(t)})} \propto p(z_n) p(x_n \mid z_n, \theta^{(t)})
\end{align}
\]
In the Matlab code, we calculate \(p(z_n=k \mid x_n, \theta^{(t)})\) for each \(k=1\ldots K\), and normalize to get the distribution \(p(z_n \mid x_n, \theta^{(t)})\).

Maximization Step

We are going to maximize the bound \(L(\theta \mid \theta^{(t)})\) wrt. \(\pi_k\)’s and \(A_k(i,j)\)’s. The entropy term can be omitted, since it’s only a function of \(\theta^{(t)}\), therefore constant. So, we only need to maximize the energy term:
\[
\begin{align}
\theta^{(t+1)} &= \arg\!\max_{\theta} E(\theta \mid \theta^{(t)}) \\
&= \arg\!\max_{\theta} \sum_z p(z \mid x,\theta^{(t)}) \log p(x,z \mid \theta)
\end{align}
\]
Let’s write the energy explicitly. First the joint likelihood \(p(x,z \mid \theta)\).
\[
\begin{align}
\log p(x,z \mid \theta) &= \log p(x \mid z, \theta) + \log p(z \mid \theta) \\
&= \log \prod_{n=1}^N p(x_n \mid z_n, \theta) + \log \prod_{n=1}^N p(z_n \mid \theta) \\
&= \log \prod_{n=1}^N p(x_n \mid z_n, \theta) + \log \prod_{n=1}^N p(z_n \mid \theta) \\
&= \sum_{n=1}^N \log p(x_n \mid z_n, \theta) + \sum_{n=1}^N \log p(z_n \mid \theta) \\
&= \sum_{n=1}^N \sum_{k=1}^K [z_n=k] \log \left( \prod_{i=1}^D \pi_k(i)^{[x_n(1)=i]} \prod_{i=1}^D \prod_{j=1}^D A_k(i,j)^{s_n(i,j)} \right) \\
& \quad + \sum_{n=1}^N \sum_{k=1}^K [z_n=k] \log \alpha(k) \\
&= \sum_{n=1}^N \sum_{k=1}^K [z_n=k] \left ( \sum_{i=1}^D [x_n(1)=i] \log \pi_k(i) + \sum_{i=1}^D \sum_{j=1}^D s_n(i,j) \log A_k(i,j) \right)\\
& \quad +\sum_{n=1}^N \sum_{k=1}^K [z_n=k] \log \alpha(k)
\end{align}
\]
Now, let’s write the expectation of this joint likelihood.
\[
\begin{align}
E(\theta \mid \theta^{(t)}) &= \sum_z p(z \mid x,\theta^{(t)}) \log p(x,z \mid \theta) \\
&= \sum_{n=1}^N \sum_{k=1}^K p(z_n=k \mid x_n, \theta^{(t)}) \left ( \sum_{i=1}^D [x_n(1)=i] \log \pi_k(i) + \sum_{i=1}^D \sum_{j=1}^D s_n(i,j) \log A_k(i,j) \right) \\
& \quad + \sum_{n=1}^N \sum_{k=1}^K p(z_n=k \mid x_n, \theta^{(t)}) \log \alpha(k)
\end{align}
\]
How did we come to that? It’s not easy to see immediately. Keep in mind that \(\sum_z\) is an abuse of notation:
\[
\begin{align}
\sum_z = \sum_{z_1=1}^{K}\sum_{z_2=1}^{K}\ldots \sum_{z_N=1}^{K}
\end{align}
\]
If you do the algebra, you will arrive the above equation. Now, let’s maximize the energy wrt. \(\pi_k(i)\) to find \(\pi_k(i)^{(t+1)}\) (by using Lagrange multipliers):
\[
\begin{align}
\frac{\partial}{\partial \pi_k(i)} \left ( E(\theta \mid \theta^{(t)}) + \lambda( \sum_{j=1}^D \pi_k(j)-1) \right ) &=0 \\
\frac{\partial}{\partial \pi_k(i)} \left ( \sum_{n=1}^N p(z_n=k \mid x_n,\theta^{(t)}) [x_n(1)=i] \log \pi_k(i) + \lambda \pi_k(i) \right ) &=0 \\
\frac{1}{\pi_k(i)} \left( \sum_{n=1}^N p(z_n=k \mid x_n,\theta^{(t)}) [x_n(1)=i] \right) + \lambda &=0 \\
-\frac{1}{\lambda} \left( \sum_{n=1}^N p(z_n=k \mid x_n,\theta^{(t)}) [x_n(1)=i] \right) &= \pi_k(i)
\end{align}
\]
We calculate \(\lambda\) as:
\[
\begin{align}
1 &= \sum_{j=1}^D \pi_k(j) \\
&= \sum_{j=1}^D -\frac{1}{\lambda} \left( \sum_{i=n}^N p(z_n=k \mid x_n,\theta^{(t)})[x_n(1)=j] \right) \\
\lambda &= – \sum_{j=1}^D \sum_{n=1}^N p(z_n=k \mid x_n,\theta^{(t)}) [x_n(1)=j]
\end{align}
\]
Finally we arrive at:
\[
\pi_k^{(t+1)}(i) = \frac{ \sum_{n=1}^N p(z_n=k \mid x_n,\theta^{(t)})[x_n(1)=i] }{\sum_{j=1}^D \sum_{n=1}^N p(z_n=k \mid x_n,\theta^{(t)} )[x_n(1)=j]}
\]
This is easy to compute: we will calculate the numerator for all \(i=\{1,\ldots,D\}\) and normalize. Similarly, we can derive the following formulas:
\[
\begin{align}
A_k^{(t+1)}(i,j) &= \frac{\sum_{n=1}^N p(z_n=k \mid x_n,\theta^{(t)}) s_n(i,j) }{\sum_{m=1}^D \sum_{n=1}^N p(z_n=k \mid x_n,\theta^{(t)}) s_n(m,j)} \\
\alpha^{(t+1)}(k) &= \frac{\sum_{n=1}^N p(z_n=k \mid x_n,\theta^{(t)})}{ \sum_{w=1}^K \sum_{n=1}^N p(z_n=w \mid x_n,\theta^{(t)}) }
\end{align}
\]

Calculating the likelihood

After every iteration, we should check the likelihood \(p(x \mid \theta)\) to see whether the process has converged to a solution. This likelihood has to be non-decreasing. If you a decrease on the likelihood, it’s 100% you did something wrong.
We can calculate the likelihood at the \(t^{th}\) step as
\[
p(x \mid \theta^{(t)}) = \prod_{n=1}^N p(x_n \mid \theta^{(t)}) = \prod_{n=1}^N \sum_{k=1}^K p(x_n \mid z_n=k, \theta^{(t)}) p(z_n=k \mid \theta^{(t)})
\]
This should be equal to \(L(\theta^{(t)} \mid \theta^{(t)})\). That is
\[
L(\theta^{(t)} \mid \theta^{(t)}) = \sum_z p(z \mid x,\theta^{(t)}) \log p(x,z \mid \theta^{(t)}) – \sum_z p(z \mid x,\theta^{(t)}) \log p(z \mid x,\theta^{(t)})
\]

The code and the demo

Here I present my Matlab implementation for the EM algoerithm, which I believe runs quite fast. Instead of processing each sequence for calculating the likelihoods at each iteration, I calculate the sufficient statistics for each sequence in the beginning. Therefore, every sequence, independent of its length is presented by a \(D\times D\) matrix where \(D\) is the dimension of the input space and each likelihood is calculated by \(D^2\) multiplications. For calculating the sufficient statistics, I avoid for loops in the code by using accumarray function. Also in the E and M steps, I avoid for loops by making matrix multiplications.

Sample Data: sample_data.mat
Code: mixMarkovEM.m

Executing the code with the sample data is straightforward. There’s a nice help at the top of the code.

Bon Apetite.