How can I derive the EM algorithm for a mixture of two Bernoulli distributions?

Question

How can I derive the E-step and M-step in the EM algorithm for a mixture of two Bernoulli distributions? Note that I am aware that there are several notes online that explain how to do this for the more general case of $K$ components, but I am only interested in the case of two components.

Answer 1

The following derivation was adapted from the references:

• "What is the expectation maximization algorithm?" by Chuong B Do & Serafim Batzoglou (2008)
• "A Gentle Tutorial of the EM Algorithm and its Application to Parameter Estimation for Gaussian Mixture and Hidden Markov Models" by Jeff Bilmes (1998)
• Section 9.3.3 in Christopher Bishop's "Pattern Recognition and Machine Learning" (2006)

Preliminaries

A mixture of two Bernoulli distributions can model the following random experiment: Given two coins, suppose one of these coins is chosen at random. Then, suppose that the randomly chosen coin is flipped. We want to estimate the probability of choosing each coin and the probability of observing a heads for each coin. If we repeat this random experiment several times, and we knew which flip belonged to which coin, this would be fairly straightforward. However, suppose that your friend conducts this random experiment several times behind a curtain, and they only inform you of the outcomes of the coin flips without revealing which coin each flip belongs to. This is where the EM algorithm comes in.

To make things more precise, suppose that your friend conducts this random experiment $N$ times and informs you only of the outcomes of the coin flips. Let $\mathcal{D}_X$ be the set of these outcomes, such that $$ \mathcal{D}_X = \{X_1 = x_1,X_2 = x_2,\dots,X_N = x_N\}, $$ where $X_i = 1$ if the $i^\text{th}$ coin flip was heads, and $X_i = 0$ if the $i^\text{th}$ coin flip was tails. For example, if your friend conducts this experiment $N = 4$ times, the outcomes could be $$ \mathcal{D}_X = \{X_1 = 1,X_2 = 0,X_3 = 0, X_4 = 1\}, $$ which corresponds to the sequence: heads, tails, tails, heads. So far, we do not know which coin generated which flip. All we know is the sequence of coin flips. Because of this, we call $\mathcal{D}_X$ the incomplete dataset. The complete dataset, which includes knowledge of which coin generated which flip, would then be $$ \mathcal{D}_{XZ} = \{(X_1 = x_1,Z_1),(X_2 = x_2,Z_2),\dots,(X_N = x_N,Z_N)\}, $$ where $Z_i = 0$ if the $i^\text{th}$ coin flip was generated by the first coin and $Z_i = 1$ if the $i^\text{th}$ coin flip was generated by the second coin. Note that no values were assigned to any of the $Z_i$'s in $\mathcal{D}_{XZ}$, since we do not know which coin generated which flip. In other words, each $Z_i$ is assumed to be random with a probability distribution $p(Z_i = z_i)$. For simplicity, we assume that each pair $(X_i = x_i, Z_i)$ in $\mathcal{D}_{XZ}$ is independent and identically distributed (i.i.d.). We are therefore interested in estimating,

• the probability of choosing the first coin, $p(Z_i = 0)$,
• the probability of choosing the second coin, $p(Z_i = 1)$,
• the probability of observing a heads when the first coin is chosen, $p(X_i = 1 \mid Z_i = 0)$,
• the probability of observing a tails when the first coin is chosen, $p(X_i = 0 \mid Z_i = 0)$,
• the probability of observing a heads when the second coin is chosen, $p(X_i = 1 \mid Z_i = 1)$,
• and the probability of observing a tails when the second coin is chosen, $p(X_i = 0 \mid Z_i = 1)$.

Because each $X_i$ and $Z_i$ are Bernoulli random variables, we only need to estimate $p(Z_i = 1), p(X_i = 1 \mid Z_i = 0),$ and $p(X_i = 1 \mid Z_i = 1)$, which can be used to determine the other probabilities.

The complete data log-likelihood

We estimate these probabilities via maximum likelihood estimation. First, we describe what form these probability mass functions take: $$ \begin{align} p(X_i = x_i \mid Z_i = 0) &= \begin{cases} p(X_i = 1 \mid Z_i = 0), &x_i = 1 \\ 1 - p(X_i = 1 \mid Z_i = 0), &x_i = 0\end{cases} \\ &= p(X_i = 1 \mid Z_i = 0)^{x_i} \cdot (1 - p(X_i = 1 \mid Z_i = 0))^{1-x_i} \\ p(X_i = x_i \mid Z_i = 1) &= \begin{cases} p(X_i = 1 \mid Z_i = 1), &x_i = 1 \\ 1 - p(X_i = 1 \mid Z_i = 1), &x_i = 0\end{cases} \\ &= p(X_i = 1 \mid Z_i = 1)^{x_i} \cdot (1 - p(X_i = 1 \mid Z_i = 1))^{1-x_i} \\ p(Z_i = z_i) &= \begin{cases} p(Z_i = 1), &z_i = 1 \\ 1 - p(Z_i = 1), &z_i = 0\end{cases} \\ &= p(Z_i = 1)^{z_i} \cdot (1 - p(Z_i = 1))^{1-z_i} \\ \end{align} $$ Next, we form the complete data likelihood function as follows. First, let $$ \theta = \{p(Z_i = 1), p(X_i = 1 \mid Z_i = 0),p(X_i = 1 \mid Z_i = 1)\}. $$ then, the complete data likelihood function is $$ \begin{align} p(\mathcal{D}_{XZ};\theta) &= \prod_{i=1}^N p(X_i = x_i,Z_i) \\ &= \prod_{i=1}^N p(X_i = x_i \mid Z_i) \cdot p(Z_i) \end{align} $$ where $$ \begin{align} p(X_i = x_i \mid Z_i) &= \begin{cases} p(X_i = x_i \mid Z_i = 0), &Z_i = 0 \\ p(X_i = x_i \mid Z_i = 1), &Z_i = 1 \end{cases} \\ &= p(X_i = x_i \mid Z_i = 1)^{Z_i} \cdot p(X_i = x_i \mid Z_i = 0)^{1 - Z_i} \\ &= \left[p(X_i = 1 \mid Z_i = 1)^{x_i} \cdot (1 - p(X_i = 1 \mid Z_i = 1))^{1-x_i}\right]^{Z_i} \\ &\cdot \left[p(X_i = 1 \mid Z_i = 0)^{x_i} \cdot (1 - p(X_i = 1 \mid Z_i = 0))^{1-x_i}\right]^{1 - Z_i} \end{align} $$ and $$ \begin{align} p(Z_i) &= \begin{cases} p(Z_i = 1), &Z_i = 1 \\ 1 - p(Z_i = 1), &Z_i = 0 \end{cases} \\ &= p(Z_i = 1)^{Z_i} \cdot (1 - p(Z_i = 1))^{1 - Z_i} \end{align} $$ are random functions. The randomness of these functions arises from the randomness of each $Z_i$. Consequently, this means that the complete data likelihood function $p(\mathcal{D}_{XZ};\theta)$ is also random. Before attempting to make $p(\mathcal{D}_{XZ};\theta)$ deterministic, we first apply a $\log$ transformation to it to obtain $$ \begin{align} \log p(\mathcal{D}_{XZ};\theta) &= \sum_{i=1}^N \log p(X_i = x_i,Z_i) \\ &= \sum_{i=1}^N \log p(X_i = x_i \mid Z_i) + \log p(Z_i) \end{align} $$ where $$ \begin{alignat}{2} \log p(X_i = x_i \mid Z_i) &= \log \biggl\{\left[p(X_i = 1 \mid Z_i = 1)^{x_i} \cdot (1 - p(X_i = 1 \mid Z_i = 1))^{1-x_i}\right]^{Z_i} \cdot \left[p(X_i = 1 \mid Z_i = 0)^{x_i} \cdot (1 - p(X_i = 1 \mid Z_i = 0))^{1-x_i}\right]^{1 - Z_i} \biggr\} \\ &= Z_i \left[x_i \log p(X_i = 1 \mid Z_i = 1) + (1 - x_i) \log (1 - p(X_i = 1 \mid Z_i = 1))\right] + (1 - Z_i) \left[x_i \log p(X_i = 1 \mid Z_i = 0) + (1 - x_i) \log (1 - p(X_i = 1 \mid Z_i = 0))\right] \\ \log p(Z_i) &= Z_i \log p(Z_i = 1) + (1 - Z_i) \log (1 - p(Z_i = 1)) \end{alignat} $$ To make the complete data log-likelihood function $\log p(\mathcal{D}_{XZ};\theta)$ deterministic, we "average out" its randomness by computing its conditional expectation with respect to $\mathbf{Z} = (Z_1,\dots,Z_N)$ given the complete dataset $\mathcal{D}_X$ (see this answer for some motivation of the "averaging out" process): $$ \begin{align} \mathbb{E}_{\mathbf{Z}}\left[\log p(\mathcal{D}_{XZ};\theta) \mid \mathcal{D}_X\right] &= \sum_{i=1}^N \mathbb{E}_{\mathbf{Z}}\left[\log p(X_i = x_i \mid Z_i) \mid \mathcal{D}_X\right] + \mathbb{E}_{\mathbf{Z}}\left[\log p(Z_i) \mid \mathcal{D}_X\right] \\ &= \sum_{i=1}^N \mathbb{E}_{Z_i}\left[\log p(X_i = x_i \mid Z_i) \mid \mathcal{D}_X\right] + \mathbb{E}_{Z_i}\left[\log p(Z_i) \mid \mathcal{D}_X\right] \end{align} $$ To simplify the expectations $\mathbb{E}_{Z_i}[\log p(X_i = x_i \mid Z_i) \mid \mathcal{D}_X]$ and $\mathbb{E}_{Z_i}[\log p(Z_i) \mid \mathcal{D}_X]$, note that $$ \begin{align} p(Z_i = z_i \mid \mathcal{D}_X) &= \frac{p(Z_i = z_i,\mathcal{D}_X)}{p(\mathcal{D}_X)} \\ &= \frac{\sum_{Z_j} p(Z_i = z_i,\{Z_j\}_{j\neq i},\mathcal{D}_X)}{\sum_{Z_i,Z_j} p(Z_i = z_i,\{Z_j\}_{j\neq i},\mathcal{D}_X)}\\ &= \frac{\sum_{Z_j} \prod_{k=1}^N p(X_k = x_k,Z_k = z_k)}{\sum_{Z_i,Z_j} \prod_{k=1}^N p(X_k = x_k,Z_k = z_k)} \\ &= \frac{p(X_j = x_j) \prod_{k \neq j} p(X_k = x_k,Z_k = z_k)}{\prod_{k=1}^N p(X_k = x_k)} \\ &= p(Z_i = z_i \mid X_i = x_i) \end{align} $$ Therefore, $$ \begin{align} \label{eq:acdlf} \mathbb{E}_{\mathbf{Z}}\left[\log p(\mathcal{D}_{XZ};\theta) \mid \mathcal{D}_X\right] &= \sum_{i=1}^N \mathbb{E}_{Z_i}\left[\log p(X_i = x_i \mid Z_i) \mid X_i = x_i\right] + \mathbb{E}_{Z_i}\left[\log p(Z_i) \mid X_i = x_i\right] \tag{1} \end{align} $$ where $$ \begin{align} \mathbb{E}_{Z_i}[\log p(X_i = x_i \mid Z_i) \mid X_i = x_i] &= \mathbb{E}_{Z_i}[Z_i \mid X_i = x_i] \left[x_i \log p(X_i = 1 \mid Z_i = 1) + (1 - x_i) \log (1 - p(X_i = 1 \mid Z_i = 1))\right] \\ &+ (1 - \mathbb{E}_{Z_i}[Z_i \mid X_i = x_i]) \left[x_i \log p(X_i = 1 \mid Z_i = 0) + (1 - x_i) \log (1 - p(X_i = 1 \mid Z_i = 0))\right] \\ &= p(Z_i = 1 \mid X_i = x_i) \left[x_i \log p(X_i = 1 \mid Z_i = 1) + (1 - x_i) \log (1 - p(X_i = 1 \mid Z_i = 1))\right] \\ &+ (1 - p(Z_i = 1 \mid X_i = x_i)) \left[x_i \log p(X_i = 1 \mid Z_i = 0) + (1 - x_i) \log (1 - p(X_i = 1 \mid Z_i = 0))\right] \label{eq:ll_avg} \tag{2} \end{align} $$ and $$ \begin{align} \mathbb{E}_{Z_i}[\log p(Z_i) \mid X_i = x_i] &= \mathbb{E}_{Z_i}[Z_i \mid X_i = x_i] \cdot \log p(Z_i = 1) \\ &+ (1 - \mathbb{E}_{Z_i}[Z_i \mid X_i = x_i]) \cdot \log (1 - p(Z_i = 1)) \\ &= p(Z_i = 1 \mid X_i = x_i) \cdot \log p(Z_i = 1) \\ &+ (1 - p(Z_i = 1 \mid X_i = x_i)) \cdot \log (1 - p(Z_i = 1)) \label{eq:prior_avg} \tag{3} \end{align} $$ Note that $\mathbb{E}_{Z_i}[Z_i \mid X_i = x_i]$ was replaced with $p(Z_i = 1 \mid X_i = x_i)$ because $Z_i$ is a Bernoulli random variable.

The EM Algorithm

Recall that $\theta = \{p(Z_i = 1), p(X_i = 1 \mid Z_i = 0),p(X_i = 1 \mid Z_i = 1)\}$. Our objective is to determine the value of $\theta$ that maximizes the "averaged" complete data log-likelihood function (ACDLF) $\mathbb{E}_{\mathbf{Z}}\left[\log p(\mathcal{D}_{XZ};\theta) \mid \mathcal{D}_X\right]$ in $\eqref{eq:acdlf}$. From $\eqref{eq:ll_avg}$ and $\eqref{eq:prior_avg}$ above, we see that that the ACDLF is a function of $\theta$, the set of parameters that we want to estimate, and $p(Z_i = 1 \mid X_i = 0)$ and $p(Z_i = 1 \mid X_i = 1)$, the "other" parameters in the ACDLF. We are trying to maximize the ACDLF with respect to $\theta$, and to do so, we should be able to evaluate the ACDLF. However, we do not know the values of $p(Z_i = 1 \mid X_i = 0)$ and $p(Z_i = 1 \mid X_i = 1)$ to be able to evaluate the ACDLF. Therefore, we proceed as follows:

1. Guess initial values for the "other" parameters, $p(Z_i = 1 \mid X_i = 0)$ and $p(Z_i = 1 \mid X_i = 1)$, so that we are able to evaluate the ACDLF.
2. Determine a value of $\theta$ that maximizes the ACDLF, called $\theta^*$. (M step)
3. Since both $p(Z_i = 1 \mid X_i = 0)$ and $p(Z_i = 1 \mid X_i = 1)$ depend on the parameters in $\theta$ via Bayes' rule, we compute new values for them using $\theta^*$. (E step)
4. Repeat steps 2 and 3 until convergence.

Note: in the literature of the EM algorithm, step 1 above would be slightly different. Instead of guessing values for $p(Z_i = 1 \mid X_i = 0)$ and $p(Z_i = 1 \mid X_i = 1)$ directly, we would instead guess values for the parameters in $\theta$. Then, using these guesses, we compute values for $p(Z_i = 1 \mid X_i = 0)$ and $p(Z_i = 1 \mid X_i = 1)$ using Bayes' rule. Therefore, "step 0" would be to guess initial values for the parameters in $\theta$ first. However, the current wording of step 1 was chosen to avoid any confusion with steps 2 and 3.

Implementation

We now derive expressions for the M step and E step in steps 2 and 3 above, respectively. First, for brevity, let the $k^\text{th}$ estimate for $\theta$ and the the "other" parameters be $$ \begin{align} \theta^{[k]} &= \left\{p_k(Z_i = 1),p_k(X_i = 1 \mid Z_i = 0),p_k(X_i = 1 \mid Z_i = 1)\right\} \\ &= \left\{p_Z^{[k]},p_{XZ_0}^{[k]},p_{XZ_1}^{[k]}\right\} \\ p_k(Z_i = 1 \mid X_i = 0) &= p_{ZX_0}^{[k]} \\ p_k(Z_i = 1 \mid X_i = 1) &= p_{ZX_1}^{[k]} \end{align} $$ We first guess initial values the parameters in $\theta$. Let these initial guesses be $\theta^{[0]} = \left\{p_Z^{[0]},p_{XZ_0}^{[0]},p_{XZ_1}^{[0]}\right\}$. Next, we compute initial guess for $p(Z_i = 1 \mid X_i = 0)$ and $p(Z_i = 1 \mid X_i = 1)$ using Bayes' rule: $$ \begin{align} p_{ZX_0}^{[0]} &= \frac{\left(1 - p_{XZ_1}^{[0]}\right) \cdot p_Z^{[0]}}{\left(1 - p_{XZ_1}^{[0]}\right) \cdot p_Z^{[0]} + \left(1 - p_{XZ_0}^{[0]}\right) \cdot \left(1-p_Z^{[0]}\right)} \\ p_{ZX_1}^{[0]} &= \frac{p_{XZ_1}^{[0]} \cdot p_Z^{[0]}}{p_{XZ_1}^{[0]} \cdot p_Z^{[0]} + p_{XZ_0}^{[0]} \cdot \left(1-p_Z^{[0]}\right)} \end{align} $$

Next, we plug these guesses into the ACDLF in $\eqref{eq:acdlf}$ such that $\eqref{eq:ll_avg}$ and $\eqref{eq:prior_avg}$ become

$$ \begin{align} \mathbb{E}_{Z_i}[\log p(X_i = x_i \mid Z_i) \mid X_i = x_i] &= p_{ZX_{x_i}}^{[0]} \left[x_i \log p(X_i = 1 \mid Z_i = 1) + (1 - x_i) \log (1 - p(X_i = 1 \mid Z_i = 1))\right] \\ &+ (1 - p_{ZX_{x_i}}^{[0]}) \left[x_i \log p(X_i = 1 \mid Z_i = 0) + (1 - x_i) \log (1 - p(X_i = 1 \mid Z_i = 0))\right] \label{eq:ll_avg_v2} \tag{4} \\ \mathbb{E}_{Z_i}[\log p(Z_i) \mid X_i = x_i] &= p_{ZX_{x_i}}^{[0]} \cdot \log p(Z_i = 1) + (1 - p_{ZX_{x_i}}^{[0]}) \cdot \log (1 - p(Z_i = 1)) \label{eq:prior_avg_v2} \tag{5} \end{align} $$ This completes step 1 in the section "The EM Algorithm" above. We now move on to step 2, which is to determine a value of $\theta$ that maximizes the ACDLF. To do this, we differentiate the ACDLF in $\eqref{eq:acdlf}$ with respect to $p(Z_i = 1), p(X_i = 1 \mid Z_i = 0),$ and $p(X_i = 1 \mid Z_i = 1)$, equate each of the resulting expressions to $0$, and then solve for each of these parameters. Note that $$ \begin{align} \frac{\partial}{\partial p(Z_i = 1)} \mathbb{E}_{\mathbf{Z}}\left[\log p(\mathcal{D}_{XZ};\theta) \mid \mathcal{D}_X\right] &= \sum_{i=1}^N \frac{\partial}{\partial p(Z_i = 1)} \mathbb{E}_{Z_i}\left[\log p(Z_i) \mid X_i = x_i\right] \label{eq:pz_deriv_acdlf} \tag{6} \\ \frac{\partial}{\partial p(X_i = 1 \mid Z_i = 0)} \mathbb{E}_{\mathbf{Z}}\left[\log p(\mathcal{D}_{XZ};\theta) \mid \mathcal{D}_X\right] &= \sum_{i=1}^N \frac{\partial}{\partial p(X_i = 1 \mid Z_i = 0)} \mathbb{E}_{Z_i}\left[\log p(X_i = x_i \mid Z_i) \mid X_i = x_i\right] \label{eq:pxz0_deriv_acdlf} \tag{7} \\ \frac{\partial}{\partial p(X_i = 1 \mid Z_i = 1)} \mathbb{E}_{\mathbf{Z}}\left[\log p(\mathcal{D}_{XZ};\theta) \mid \mathcal{D}_X\right] &= \sum_{i=1}^N \frac{\partial}{\partial p(X_i = 1 \mid Z_i = 1)} \mathbb{E}_{Z_i}\left[\log p(X_i = x_i \mid Z_i) \mid X_i = x_i\right] \label{eq:pxz1_deriv_acdlf} \tag{8} \end{align} $$ where $$ \begin{align} \frac{\partial}{\partial p(Z_i = 1)} \mathbb{E}_{Z_i}\left[\log p(Z_i) \mid X_i = x_i\right] &= \frac{\partial}{\partial p(Z_i = 1)} \left\{p_{ZX_{x_i}}^{[0]} \cdot \log p(Z_i = 1) + (1 - p_{ZX_{x_i}}^{[0]}) \cdot \log (1 - p(Z_i = 1)) \right\} \\ &= \frac{p_{ZX_{x_i}}^{[0]}}{p(Z_i = 1)} - \frac{1 - p_{ZX_{x_i}}^{[0]}}{1 - p(Z_i = 1)} \\ \frac{\partial}{\partial p(X_i = 1 \mid Z_i = 0)} \mathbb{E}_{Z_i}\left[\log p(X_i = x_i \mid Z_i) \mid X_i = x_i\right] &= \frac{\partial}{\partial p(X_i = 1 \mid Z_i = 0)} \biggl\{ p_{ZX_{x_i}}^{[0]} \left[x_i \log p(X_i = 1 \mid Z_i = 1) + (1 - x_i) \log (1 - p(X_i = 1 \mid Z_i = 1))\right] \\ &+ (1 - p_{ZX_{x_i}}^{[0]}) \left[x_i \log p(X_i = 1 \mid Z_i = 0) + (1 - x_i) \log (1 - p(X_i = 1 \mid Z_i = 0))\right] \biggr\} \\ &= (1 - p_{ZX_{x_i}}^{[0]}) \left[\frac{x_i}{p(X_i = 1 \mid Z_i = 0)} - \frac{1 - x_i}{1 - p(X_i = 1 \mid Z_i = 0)}\right] \\ \frac{\partial}{\partial p(X_i = 1 \mid Z_i = 1)} \mathbb{E}_{Z_i}\left[\log p(X_i = x_i \mid Z_i) \mid X_i = x_i\right] &= \frac{\partial}{\partial p(X_i = 1 \mid Z_i = 1)} \biggl\{ p_{ZX_{x_i}}^{[0]} \left[x_i \log p(X_i = 1 \mid Z_i = 1) + (1 - x_i) \log (1 - p(X_i = 1 \mid Z_i = 1))\right] \\ &+ (1 - p_{ZX_{x_i}}^{[0]}) \left[x_i \log p(X_i = 1 \mid Z_i = 0) + (1 - x_i) \log (1 - p(X_i = 1 \mid Z_i = 0))\right] \biggr\} \\ &= p_{ZX_{x_i}}^{[0]} \left[\frac{x_i}{p(X_i = 1 \mid Z_i = 1)} - \frac{1 - x_i}{1 - p(X_i = 1 \mid Z_i = 1)}\right] \end{align} $$

Note that, when either $X_i$ or $Z_i$ take on more than 2 values, the aforementioned maximizations must be constrained such that the respective probability mass functions sum to $1$. Possible methods for enforcing these constraints is the use of Lagrange multipliers or the use of the softmax function. However, these methods are outside the scope of this derivation.

We then equate each of $\eqref{eq:pz_deriv_acdlf},\eqref{eq:pxz0_deriv_acdlf},$ and $\eqref{eq:pxz1_deriv_acdlf}$ to $0$ and solve for the corresponding parameter in $\theta$. We first equate $\eqref{eq:pz_deriv_acdlf}$ to $0$ and solve for $p(Z_i = 1)$: $$ \begin{align} \sum_{i=1}^N \frac{p_{ZX_{x_i}}^{[0]}}{p(Z_i = 1)} - \frac{1 - p_{ZX_{x_i}}^{[0]}}{1 - p(Z_i = 1)} &= 0 \\ \sum_{i=1}^N \frac{p_{ZX_{x_i}}^{[0]}\left(1 - p(Z_i = 1)\right) - p(Z_i = 1)\left(1 - p_{ZX_{x_i}}^{[0]}\right)}{p(Z_i = 1)\left(1 - p(Z_i = 1)\right)} &= 0 \end{align} $$ Because it was assumed that each pair $(X_i,Z_i)$ is i.i.d., then $p(Z_i = 1)$ is the same for every $i$. So, $$ \begin{align} \sum_{i=1}^N \frac{p_{ZX_{x_i}}^{[0]}\left(1 - p(Z_i = 1)\right) - p(Z_i = 1)\left(1 - p_{ZX_{x_i}}^{[0]}\right)}{p(Z_i = 1)\left(1 - p(Z_i = 1)\right)} &= 0 \\ \frac{1}{p(Z_i = 1)\left(1 - p(Z_i = 1)\right)} \left[\sum_{i=1}^N p_{ZX_{x_i}}^{[0]}\left(1 - p(Z_i = 1)\right) - p(Z_i = 1)\left(1 - p_{ZX_{x_i}}^{[0]}\right)\right] &= 0 \\ \frac{1}{p(Z_i = 1)\left(1 - p(Z_i = 1)\right)} \left[\sum_{i=1}^N p_{ZX_{x_i}}^{[0]} - p_{ZX_{x_i}}^{[0]}p(Z_i = 1) - p(Z_i = 1) + p_{ZX_{x_i}}^{[0]}p(Z_i = 1)\right] &= 0 \\ \frac{1}{p(Z_i = 1)\left(1 - p(Z_i = 1)\right)} \left[\sum_{i=1}^N p_{ZX_{x_i}}^{[0]} - p(Z_i = 1)\right] &= 0 \\ \frac{1}{p(Z_i = 1)\left(1 - p(Z_i = 1)\right)} \left[\left(\sum_{i=1}^N p_{ZX_{x_i}}^{[0]}\right) - Np(Z_i = 1)\right] &= 0 \\ \frac{1}{p(Z_i = 1)\left(1 - p(Z_i = 1)\right)} \sum_{i=1}^N p_{ZX_{x_i}}^{[0]} - \frac{N}{1 - p(Z_i = 1)} &= 0 \\ \frac{1}{p(Z_i = 1)} \sum_{i=1}^N p_{ZX_{x_i}}^{[0]} - N &= 0 \\ \frac{1}{p(Z_i = 1)} \sum_{i=1}^N p_{ZX_{x_i}}^{[0]} &= N \\ p(Z_i = 1) &= \frac{1}{N} \sum_{i=1}^N p_{ZX_{x_i}}^{[0]} \end{align} $$ Next, we equate $\eqref{eq:pxz0_deriv_acdlf}$ to $0$ and solve for $p(X_i = 1 \mid Z_i = 0)$: $$ \begin{align} \sum_{i=1}^N (1 - p_{ZX_{x_i}}^{[0]}) \left[\frac{x_i}{p(X_i = 1 \mid Z_i = 0)} - \frac{1 - x_i}{1 - p(X_i = 1 \mid Z_i = 0)}\right] &= 0 \\ \sum_{i=1}^N \frac{x_i(1 - p_{ZX_{x_i}}^{[0]})}{p(X_i = 1 \mid Z_i = 0)} - \sum_{i=1}^N \frac{(1 - x_i)\left(1 - p_{ZX_{x_i}}^{[0]}\right)}{1 - p(X_i = 1 \mid Z_i = 0)} &= 0 \\ \end{align} $$ Because it was assumed that each pair $(X_i,Z_i)$ is i.i.d., then $p(X_i = 1 \mid Z_i = 0)$ is the same for every $i$. So, $$ \begin{align} \sum_{i=1}^N \frac{x_i\left(1 - p_{ZX_{x_i}}^{[0]}\right)}{p(X_i = 1 \mid Z_i = 0)} - \sum_{i=1}^N \frac{(1 - x_i)\left(1 - p_{ZX_{x_i}}^{[0]}\right)}{1 - p(X_i = 1 \mid Z_i = 0)} &= 0 \\ \frac{1}{p(X_i = 1 \mid Z_i = 0)} \sum_{i=1}^N x_i\left(1 - p_{ZX_{x_i}}^{[0]}\right) - \frac{1}{1 - p(X_i = 1 \mid Z_i = 0)} \sum_{i=1}^N (1 - x_i)\left(1 - p_{ZX_{x_i}}^{[0]}\right) &= 0 \\ \left(1 - p(X_i = 1 \mid Z_i = 0)\right) \sum_{i=1}^N x_i\left(1 - p_{ZX_{x_i}}^{[0]}\right) - p(X_i = 1 \mid Z_i = 0) \sum_{i=1}^N (1 - x_i)\left(1 - p_{ZX_{x_i}}^{[0]}\right) &= 0 \\ \sum_{i=1}^N x_i\left(1 - p_{ZX_{x_i}}^{[0]}\right) - p(X_i = 1 \mid Z_i = 0) \sum_{i=1}^N x_i\left(1 - p_{ZX_{x_i}}^{[0]}\right) - p(X_i = 1 \mid Z_i = 0) \sum_{i=1}^N (1 - x_i)\left(1 - p_{ZX_{x_i}}^{[0]}\right) &= 0 \\ \sum_{i=1}^N x_i\left(1 - p_{ZX_{x_i}}^{[0]}\right) - p(X_i = 1 \mid Z_i = 0) \left[\sum_{i=1}^N x_i\left(1 - p_{ZX_{x_i}}^{[0]}\right) + \sum_{i=1}^N (1 - x_i)\left(1 - p_{ZX_{x_i}}^{[0]}\right)\right] &= 0 \\ \sum_{i=1}^N x_i\left(1 - p_{ZX_{x_i}}^{[0]}\right) - p(X_i = 1 \mid Z_i = 0) \sum_{i=1}^N \left(1 - p_{ZX_{x_i}}^{[0]}\right) &= 0 \\ p(X_i = 1 \mid Z_i = 0) &= \frac{\sum_{i=1}^N x_i\left(1 - p_{ZX_{x_i}}^{[0]}\right)}{\sum_{i=1}^N \left(1 - p_{ZX_{x_i}}^{[0]}\right)} \end{align} $$

Finally, we equate $\eqref{eq:pxz1_deriv_acdlf}$ to $0$ and solve for $p(X_i = 1 \mid Z_i = 1)$. This derivation follows a very similar approach to the one used to derive an expression for $p(X_i = 1 \mid Z_i = 0)$ above. So, we simply state the result of this derivation: $$ p(X_i = 1 \mid Z_i = 1) = \frac{\sum_{i=1}^N x_i \cdot p_{ZX_{x_i}}^{[0]}}{\sum_{i=1}^N p_{ZX_{x_i}}^{[0]}} $$ To summarize, the update rules for the parameters in $\theta$ are: $$ \begin{align} p_Z^{[k+1]} &= \frac{1}{N} \sum_{i=1}^N p_{ZX_{x_i}}^{[k]} \label{eq:pz_update} \tag{9} \\ p_{XZ_0}^{[k+1]} &= \frac{\sum_{i=1}^N x_i\left(1 - p_{ZX_{x_i}}^{[k]}\right)}{\sum_{i=1}^N \left(1 - p_{ZX_{x_i}}^{[k]}\right)} \label{eq:pxz0_update} \tag{10} \\ p_{XZ_1}^{[k+1]} &= \frac{\sum_{i=1}^N x_i \cdot p_{ZX_{x_i}}^{[k]}}{\sum_{i=1}^N p_{ZX_{x_i}}^{[k]}} \label{eq:pxz1_update} \tag{11} \end{align} $$ This completes step 2 in the section "The EM Algorithm" above. We now move on to step 3, which is much simpler than step 2. Using Bayes' rule, the update rules for the "other" parameters are: $$ \begin{align} p_{ZX_0}^{[k+1]} &= \frac{\left(1 - p_{XZ_1}^{[k]}\right) \cdot p_Z^{[k]}}{\left(1 - p_{XZ_1}^{[k]}\right) \cdot p_Z^{[k]} + \left(1 - p_{XZ_0}^{[k]}\right) \cdot \left(1-p_Z^{[k]}\right)} \label{eq:pzx0_update} \tag{12} \\ p_{ZX_1}^{[k+1]} &= \frac{p_{XZ_1}^{[k]} \cdot p_Z^{[k]}}{p_{XZ_1}^{[k]} \cdot p_Z^{[k]} + p_{XZ_0}^{[k]} \cdot \left(1-p_Z^{[k]}\right)} \label{eq:pzx1_update} \tag{13} \end{align} $$ This completes step 3 in the section "The EM Algorithm" above. Finally, we move on to step 4. Step 4 consists of iterating between the update rules in $\eqref{eq:pz_update},\eqref{eq:pxz0_update},\eqref{eq:pxz1_update}$ and the update rules in $\eqref{eq:pzx0_update},\eqref{eq:pzx1_update}$ repeatedly until convergence. That is, until $|\theta^{[k+1]} - \theta^{[k]}| < \varepsilon$, where $\varepsilon$ is a small positive number.

Here is a Python implementation of steps 1 - 4 in the section "The EM algorithm" above, which includes all the aforementioned update rules.

import numpy as np

# computes the "averaged" complete data log-likelihood function (ACDLF)
def ACDLF(X,pXZ0,pXZ1,pZ):
    
    # first, compute posterior probabilities using Bayes' rule
    pZX0 = (1 - pXZ1) * pZ / ((1 - pXZ1) * pZ + (1 - pXZ0) * (1 - pZ))
    pZX1 = pXZ1 * pZ / (pXZ1 * pZ + pXZ0 * (1 - pZ))
    
    # stack posterior probabilities for broadcasting later
    pZX = np.array([pZX0,pZX1])
    
    # first term in the ACDLF. Same as eq (2) in derivation above.
    first_term = pZX[X] * (X * np.log(pXZ1) + (1 - X) * np.log(1 - pXZ1)) + \
                 (1 - pZX[X]) * (X * np.log(pXZ0) + (1 - X) * np.log(1 - pXZ0))
    first_term = np.sum(first_term)
    
    # second term in the ACDLF. Same as eq (3) in derivation above.
    second_term = pZX[X] * np.log(pZ) + \
                  (1 - pZX[X]) * np.log(1 - pZ)
    second_term = np.sum(second_term)
    
    acdlf = first_term + second_term
    
    return acdlf

rng = np.random.default_rng(seed = 42)

# number of samples
N = 10000

# ground truth probabilities used to generate the complete dataset D_{XZ}
pXZ0 = 0.2 # p(X_i = 1 | Z_i = 0)
pXZ1 = 0.9 # p(X_i = 1 | Z_i = 1)
pZ = 0.75 # p(Z_i = 1)

# hidden dataset that contains which coin was chosen. Z = 0 corresponds to the
# first coin, while Z = 1 corresponds to the second coin. The second coin is
# chosen with probability pZ.
Z = rng.binomial(1,pZ,size = N)

# stack pXZ0 and pXZ1 into an np array to make use of broadcasting later
pXZ = np.array([pXZ0,pXZ1])

# given the hidden dataset, generate the incomplete dataset that consists only
# of coin flips. X = 0 corresponds to tails, while X = 1 corresponds to heads
X = rng.binomial(1,pXZ[Z])

# compute maximum likelihood estimates for the probabilities to compare
# with later
pXZ0_ML = X[Z == 0].mean()
pXZ1_ML = X[Z == 1].mean()
pZ_ML = Z.mean()

# step 1 in the section "The EM Algorithm":
# guess initial values for the parameters in \theta, then use these
# guesses to generate initial values for the "other" parameters.
# change the value of "sigma" below to change how far away the initial
# guesses are from the ground truth probabilities. sigma = 0 results in the
# ground truth probabilities.
# NOTE: if sigma is too big, such that any one of these probabilities is
# negative or exceeds 1, a runtime error will be raised
sigma = 0.05
pXZ0_hat = pXZ0 + rng.normal(scale = sigma)
pXZ1_hat = pXZ1 + rng.normal(scale = sigma)
pZ_hat = pZ + rng.normal(scale = sigma)

# compute guesses for posterior probabilities using Bayes' rule
pZX0_hat = (1 - pXZ1_hat) * pZ_hat / ((1 - pXZ1_hat) * pZ_hat + (1 - pXZ0_hat) * (1 - pZ_hat))
pZX1_hat = pXZ1_hat * pZ_hat / (pXZ1_hat * pZ_hat + pXZ0_hat * (1 - pZ_hat))

# stack posterior probabilities for broadcasting later
pZX_hat = np.array([pZX0_hat,pZX1_hat])

# number of iterations
num_iter = 20

for i in range(num_iter):
    
    # check that the ACDLF is non-decreasing with every iteration
    print("Iteration {}, ACDLF = {:.3f}".format(i+1,
                                                ACDLF(X,
                                                      pXZ0_hat,
                                                      pXZ1_hat,
                                                      pZ_hat)))
    
    # step 2 in the section "The EM Algorithm":
    # compute estimates for parameters in theta using posterior probabilities
    pZ_hat = pZX_hat[X].mean()
    pXZ0_hat = np.sum(X * (1 - pZX_hat[X])) / np.sum(1 - pZX_hat[X])
    pXZ1_hat = np.sum(X * pZX_hat[X]) / np.sum(pZX_hat[X])

    # step 3 in the section "The EM Algorithm":    
    # compute estimates for posterior probabilities using parameters in theta
    pZX0_hat = (1 - pXZ1_hat) * pZ_hat / ((1 - pXZ1_hat) * pZ_hat + (1 - pXZ0_hat) * (1 - pZ_hat))
    pZX1_hat = pXZ1_hat * pZ_hat / (pXZ1_hat * pZ_hat + pXZ0_hat * (1 - pZ_hat))
    
    # stack posterior probabilities for broadcasting later
    pZX_hat = np.array([pZX0_hat,pZX1_hat])

# show ground truth probabilities, probabilities estimated using maximum
# likelihood, and probabilities estimated using the EM algorithm
print("""
p(X_i = 1 | Z_i = 0) (true, ML, EM):
({:.4},{:.4},{:.4})

p(X_i = 1 | Z_i = 1) (true, ML, EM):
({:.4},{:.4},{:.4})

p(Z_i = 1) (true, ML, EM):
({:.4},{:.4},{:.4})\
""".format(pXZ0,pXZ0_ML,pXZ0_hat,
           pXZ1,pXZ1_ML,pXZ1_hat,
           pZ,pZ_ML,pZ_hat))