Expectation Maximization Correctness of Problem Formulation

29 Views Asked by At

Suppose I draw $n$ iid samples from a Poisson$(\lambda)$ distribution, with $\lambda$ unknown. Now, I artificially turn every 3 I draw into a 1, so that the probability of observing any particular non-negative integer $x$ is given by:

$P(Y_i = x | \lambda) = \begin{cases} f_a(x; \lambda) = e^{-\lambda}\lambda^x/x! & \text{if $x \neq 3, x \neq 1$} \\ f_b(x; \lambda) = e^{-\lambda}\lambda^1/1! + e^{-\lambda}\lambda^3/3! & \text{if $x = 1$} \\ f_c(x; \lambda) = 0 & \text{if $x = 3$} \\ \end{cases}$

Now, I'm looking to estimate $\lambda$ using the EM algorithm and my set of (modified) samples. My idea right now is to suppose that there are 3 populations, corresponding to $x \neq 3, x \neq 1$, $x = 1$ and $x = 3$, respectively. Then, supposing some latent variable $Z_i \sim \text{Multinomial} (1, p_a, p_b, p_c)$, where $p_a, p_b, p_c$ are the true population fractions, and $p_a + p_b + p_c = 1$, we can do the E-step pretty easily; i.e. $E[z_{ia} | y_i, \lambda^{(k)}, p_a^{(k)}, p_b^{(k)}, p_c^{(k)}] = \frac{p_a^{(k)}f_a(y_i)}{p_a^{(k)}f_a(y_i) + p_b^{(k)}f_b(y_i)}$.

Next, we differentiate the E-step with respect to $\lambda, p_a, p_b, p_c$ and obtain the appropriate updates for the parameters.

Is this the correct formulation for estimating $\lambda$ via the EM algorithm? Is there perhaps an easier construction for the EM algorithm, or a better way to do inference here?

EDIT: I've realized that this approach won't work...the hint I've been given is that I should consider the complete data as the true counts. So now I'm thinking of some transformation which yields the 1-to-3 switching mechanism, but I'm not really clear on how to compute the conditional expectation anymore.