EM algorithm for mixture of categorical distributions instantly stabilizes

Question

Brief Summary of Question

I'm trying to fit a mixture model of categorical distributions (see https://en.wikipedia.org/wiki/Categorical_distribution). The expectation at the second time step is always equal to the first, causing the algorithm to halt. This leads me to believe I have made an error in my derivation in either the E or M step (or worse, I don't understand something deeper).

Problem Setup

I have data points $Y = \{y_i\}_{i=1}^N$, where each $y_i$ takes on one of $S$ possible values - that is $y_i \in \mathcal{S} = \{1, \dots S\}$. The $y_i$ are being generated by $K$ different distributions over the labels $\mathcal{S}$. In other words, I have a mixture model of $K$ categorical distrubutions:

$$ P(y_i | \Theta) = \sum_{k=1}^K \alpha_k p(y_i;\theta_k) $$ where the collection $\{\alpha_k\}_{k=1}^K$ are the mixing terms, each $\theta_k = (\theta_{k,1}, \dots \theta_{k,S})$ is a probability distribution over the labels $\mathcal{S}$ (i.e. $\sum_{s=1}^S \theta_{k,s} = 1$), and $p(\cdot;\theta_k)$ is a categorical distribution on $\mathcal{S}$ using $\theta_k$. For clarity, if $y_i = s$, $p(y_i;\theta_k) = \theta_{k,s}$.

Given the data $Y$ are iid from this distribution, I seek to fit (as best I can) the parameters $\Theta = \{\alpha_k, \theta_k\}_{k=1}^K$.

What I tried

From what I understand, mixture models are perfect for expectation maximization (EM). It feels like someone has should have clearly written down the update steps somewhere, but since I'm new to this community I was unable to find them (if someone has them, please send me a link because this is where I believe I have done something wrong). Doing some basic Lagrange multipliers led me to the following M step (the E step was just plugging information in).

E-Step

Given parameters $\Theta^{(t)}$, for $i = 1, \dots , N$ and $k = 1, \dots K$, compute $$ q_{i,k}^{(t+1)} = \frac{\alpha_k^{(t)} p(y_i;\theta_k^{(t)})}{\sum_{l=1}^K \alpha_l^{(t)} p(y_i;\theta_l^{(t)})}. $$

M-Step

The M-Step requires some notational clarity. Each $y_i$ is just a whole number in $\mathcal{S} = \{1, 2, \dots , S\}$. We can thus represent $y_i$ as a $S$-long vector of 0's with a single 1 in the slot representing the value. E.g. if $y_i = 2$, then we could write $y _i = (0,1,0, \dots, 0)$. Let $y_i = (y_{i,1}, y_{i,2}, \dots y_{i,S})$ denote the generic 0-1 representation of $y_i$.

Using this notation, we now have a nice form for the categorical distribution: $$ p(y_i;\theta_k) = \prod_{s=1}^S \left(\theta_{k,s} \right)^{y_{i,s}}. $$

Given $\{q_{i,k}^{(t)} \}_{i=1, k=1}^{N,K}$, for $k=1, \dots , K$ and $s = 1, \dots , S$ update the model parameters by $$ \alpha_k^{(t+1)} =\frac{1}{N} \sum_{i=1}^N q_{i,k}^{(t)} $$ and $$ \theta_{k,s}^{(t+1)} = \frac{\sum_{i=1}^N q_{i,k}^{(t)} y_{i,s}}{\sum_{i=1}^N q_{i,k}^{(t)}} $$ where $y_{i,s}$ is as described above.

The Issue

After I initialize $\Theta^{(0)}$, I compute the $\{q_{i,k}^{(1)} \}_{i=1, k=1}^{N,K}$, then $\Theta^{(1)}$. Now, $\Theta^{(0)} \neq \Theta^{(1)}$, however when I compute the next E-step, I find $q_{i,k}^{(2)} = q_{i,k}^{(1)}$ for all $i,k$.

I have done the laborious EM computations by hand with some examples just to ensure that there wasn't an error with the code I wrote. This means there might be an error with the update steps. However, everything is straightforward calculus, so I'm not sure that is the issue (steps can be provided if needed). This leads me to believe there is a fundamental deeper flaw I don't understand. Any guidance would be appreciated. Thanks!

Answer 1

The question here is, the mixture of the categorical model is equivalent to a single categorical model. Unlike GMM, introducing the mixture structure to categorical distributions does not enhance the model capacity. In your formulation, your model is equivalent to a single categorical model with parameter $\sum_{k=1}^K\alpha_k\theta_k$.

Answer 2

Without looking at your exact code and initial conditions it is hard to answer. However, did you consider that you converged after one iteration? In this case the parameter update is so small that without careful check the parameters stay constant. In fact, I observed very fast convergence (1 or 2 iteration) when I implemented that myself.

Now, to continue with @Tao Li answers: This model is not identifiable meaning that different mixtures will produce the exact same distribution for the observation Y. One extreme case being the example of @Tao Li where all the weight is concentrated in one mixture $k=1$, $P(Z=k=1)$, for example. (Keep in mind that other choices work). My guess is that this make the solution landscape crazy and the algorithm will converge to some solution most probably not relevant for your problem in a few iterations.

To understand the "identifiable problem", look at the marginal distribution of Y, $P(Y=i)$ for $i$ in $1$ to $S$. It needs $n_{Marginal} = S-1$ parameters to be fully specified (the last one is found by summing everything to 1).

Similarly, the mixture needs $n_{Mix} =(K-1) + K(S-1)$ which is always larger for $K>1$ than $n_{Marginal}$.

To solve this problem you need to change things a bit to get $n_{Mix}\leq n_{Marginal}$. You have at least two options:

• If you imagine you have $K$ dices with $S$ faces. Instead to choose one randomly i.e. $Z=k$ and then throw it once $P(Y=i|Z=k)$ you could throw it several times. If you throw it enough time before shuffling the dices then the problem becomes identifiable. In that case you have in fact a mixture of Multinomial distributions.

• Other option, after piking which $k$ you throw once, not one but multiple dices (they can have different categorical distributions). This is in fact very similar to what is done for D dimensional Bernoulli mixture used to recognized numbers. There, $D=28$ (pixels) and $n_{Mix} =(K-1) + K(S-1)D \leq n_{Mar} = S^D-1 $ is easily satisfied.