Tutorials › Statistical Methods for Single-Cell RNA Sequencing › EM Estimation With Latent Branch Labels

Branching and Lineage Models · Tutorial 418 of 700

EM Estimation With Latent Branch Labels

Learn how EM alternates between estimating latent branch memberships and updating branch models, and how to assess whether the algorithm has converged reliably.

Advanced 11 min read

What You'll Learn

  • Write the observed-data likelihood when branch labels are latent.
  • Derive posterior branch responsibilities from a specified mixture model.
  • Update branch proportions and branch-specific parameters using responsibility weights.
  • Verify the EM likelihood-monotonicity result and its conditions.
  • Read likelihood traces and distinguish numerical convergence from a trustworthy fit.
  • Recognize risks including empty branches, label switching, and variance collapse.

Why Estimate Branch Labels Rather Than Fix Them?

A cell observed near a branchpoint may be compatible with more than one lineage. Assigning it permanently to one branch discards that uncertainty; fitting every branch to every cell ignores the branching structure. Expectation-maximization (EM) provides a middle course: it represents each cell's branch membership as latent, computes a probability of membership in each branch, and uses those probabilities as weights when updating the model.

The model below treats the number and structure of branches as specified. This tutorial focuses on estimating latent branch assignments and branch-specific parameters, not on selecting a topology. As emphasized in the earlier material on identifiability, distinct parameter values can sometimes generate the same observed distribution; EM cannot resolve a lack of identifiability merely by iterating.

Definition: For cell \(i\), let \(y_i\) denote its observed data, \(z_i\in\{1,\ldots,B\}\) its unobserved branch label, and \(x_i\) any observed covariates or trajectory coordinate used by the model. Let \(\pi_b=P(z_i=b)\), where \(\pi_b\geq 0\) and \(\sum_{b=1}^{B}\pi_b=1\). If branch \(b\) has density or probability mass function \(f_b(y_i\mid x_i;\theta_b)\), then the observed-data contribution is \(\sum_{b=1}^{B}\pi_b f_b(y_i\mid x_i;\theta_b)\). The parameters \(\theta_b\) describe branch-specific distributions.

The basic likelihood assumes cells are conditionally independent given their covariates and the model parameters, and that each cell has one latent branch label. Thus \[ L(\pi,\theta)=\prod_{i=1}^{n}\sum_{b=1}^{B}\pi_b f_b(y_i\mid x_i;\theta_b), \qquad \ell(\pi,\theta)=\sum_{i=1}^{n}\log\left\{\sum_{b=1}^{B}\pi_b f_b(y_i\mid x_i;\theta_b)\right\}. \] Here \(\theta\) collects all branch-specific parameters. The likelihood is a sum over possible labels inside each logarithm. Direct maximization is possible in some models, but the latent labels make the expressions and optimization less convenient than the corresponding complete-data likelihood.

The EM Updates

If the labels were observed, the complete-data log-likelihood would be \[ \ell_c(\pi,\theta)=\sum_{i=1}^{n}\sum_{b=1}^{B} I(z_i=b)\left[\log(\pi_b)+\log f_b(y_i\mid x_i;\theta_b)\right], \] where \(I(z_i=b)\) is 1 when cell \(i\) belongs to branch \(b\), and 0 otherwise. EM replaces each indicator by its conditional expectation under the current parameter estimates.

Definition: At iteration \(t\), the responsibility of branch \(b\) for cell \(i\) is \[ r_{ib}^{(t)}=P(z_i=b\mid y_i,x_i;\pi^{(t)},\theta^{(t)}) =\frac{\pi_b^{(t)}f_b(y_i\mid x_i;\theta_b^{(t)})} {\sum_{k=1}^{B}\pi_k^{(t)}f_k(y_i\mid x_i;\theta_k^{(t)})}. \] Responsibilities are nonnegative and satisfy \(\sum_b r_{ib}^{(t)}=1\) for every cell.

The E-step calculates these responsibilities. The M-step maximizes the expected complete-data log-likelihood \[ Q(\pi,\theta\mid\pi^{(t)},\theta^{(t)}) =\sum_{i=1}^{n}\sum_{b=1}^{B}r_{ib}^{(t)} \left[\log(\pi_b)+\log f_b(y_i\mid x_i;\theta_b)\right] \] over the model parameters. When the mixing proportions are unrestricted, their update is \(\pi_b^{(t+1)}=n^{-1}\sum_i r_{ib}^{(t)}\). Each branch-specific parameter is estimated by maximizing its weighted log-likelihood, \(\sum_i r_{ib}^{(t)}\log f_b(y_i\mid x_i;\theta_b)\). A cell with responsibility 0.8 therefore contributes weight 0.8, not a fixed membership, to that branch's parameter update.

Why the Observed Likelihood Does Not Decrease

The key guarantee is conditional: exact EM updates do not decrease the observed-data likelihood, provided the model and likelihood are well-defined, the E-step uses the current posterior responsibilities, and the M-step maximizes \(Q\) (or at least increases it for a generalized EM update). This is a monotonicity result, not a guarantee that EM finds the global maximum.

1
Hold the current posterior fixed.
Set \(q_{ib}=r_{ib}^{(t)}\). These values sum to 1 across branches for each cell.
2
Apply the log-sum inequality.
For any candidate parameters \((\pi',\theta')\), let \(a'_{ib}=\pi'_b f_b(y_i\mid x_i;\theta'_b)\). Since \(\sum_{b:q_{ib}>0}q_{ib}=1\), the full likelihood sum is at least its sum over this support, and Jensen's inequality gives \(\log\sum_b a'_{ib}\geq\log\sum_{b:q_{ib}>0}a'_{ib}\geq\sum_{b:q_{ib}>0}q_{ib}\log(a'_{ib}/q_{ib})\). Terms with \(a'_{ib}=0\) are interpreted by a limit, giving \(-\infty\). At the current parameters, positive posterior responsibilities occur exactly for positive likelihood contributions, so the bound is exact there.
3
Compare candidate and current parameters.
Summing that inequality over cells gives \(\ell(\pi',\theta')-\ell(\pi^{(t)},\theta^{(t)})\geq Q(\pi',\theta'\mid\pi^{(t)},\theta^{(t)})-Q(\pi^{(t)},\theta^{(t)}\mid\pi^{(t)},\theta^{(t)})\). At the current parameters, \(q_{ib}\) is the posterior, so the bound is exact there. An M-step that increases \(Q\) therefore cannot decrease \(\ell\).

Worked Updates

Worked Example: One Gaussian EM Iteration

Consider three hypothetical scalar observations \(y=(0,1,4)\) and two branches with Gaussian densities of known variance 1. Initially, set \(\pi_1=\pi_2=0.5\), \(\mu_1=0\), and \(\mu_2=4\). The model assumptions are independent observations from the stated two-component mixture, a common known variance of 1, and freely estimated mixing proportions and means. The Gaussian density is \(\phi(y;\mu,1)=(2\pi)^{-1/2}\exp\{-(y-\mu)^2/2\}\).

For \(y=0\), the ratio of branch 1's density to branch 2's density is \(\exp(8)\), so \(r_{01}=1/(1+\exp(-8))\approx0.999665\). For \(y=1\), the density ratio is \(\exp(4)\), giving \(r_{11}=1/(1+\exp(-4))\approx0.982014\). For \(y=4\), the ratio is \(\exp(-8)\), giving \(r_{41}=1/(1+\exp(8))\approx0.000335\). The branch 2 responsibilities are the complements: approximately 0.000335, 0.017986, and 0.999665.

The total responsibility for branch 1 is \(0.999665+0.982014+0.000335=1.982014\); for branch 2 it is \(1.017986\). Therefore the updated mixing proportions are \(\pi_1^{(1)}=1.982014/3\approx0.660671\) and \(\pi_2^{(1)}\approx0.339329\). The weighted means are \[ \mu_1^{(1)}=\frac{0(0.999665)+1(0.982014)+4(0.000335)}{1.982014}\approx0.4961, \qquad \mu_2^{(1)}=\frac{0(0.000335)+1(0.017986)+4(0.999665)}{1.017986}\approx3.9457. \] These are responsibility-weighted means, not means of hard-assigned groups.

The initial observed log-likelihood is the sum of the logs of the three mixture densities: \[ \ell^{(0)}=\sum_{y\in\{0,1,4\}}\log\{0.5\phi(y;0,1)+0.5\phi(y;4,1)\} \approx\log(0.199538)+\log(0.123201)+\log(0.199538) \approx-5.317. \] Evaluating the same observed-data likelihood at the updated proportions and means gives approximately \(-4.906\), an increase as the EM monotonicity result requires. The observations are hypothetical, and rounding the displayed parameters may cause a small difference in a recalculation.

Worked Example: Updating Bernoulli Branch Models

Suppose four hypothetical cells have a binary marker \(x=(1,1,1,0)\). There are two branches with initial equal proportions. The marker probability is 0.8 in branch A and 0.2 in branch B, with independent Bernoulli observations conditional on branch. These assumptions specify the probabilities for both marker values: for example, \(P(x=0\mid A)=1-0.8=0.2\).

For each cell with \(x=1\), the E-step gives \[ r_{iA}=\frac{0.5(0.8)}{0.5(0.8)+0.5(0.2)}=0.8. \] For the cell with \(x=0\), it gives \[ r_{iA}=\frac{0.5(0.2)}{0.5(0.2)+0.5(0.8)}=0.2. \] Thus the branch A responsibilities are \(0.8,0.8,0.8,0.2\); the branch B responsibilities are their complements.

The responsibility totals are \(N_A=2.6\) and \(N_B=1.4\). The estimated branch proportion becomes \(\pi_A=2.6/4=0.65\), with \(\pi_B=0.35\). The weighted marker-success total for A is \(0.8+0.8+0.8+0=2.4\), so its updated Bernoulli probability is \(p_A=2.4/2.6\approx0.9231\). For B, the success weight is \(0.2+0.2+0.2+0=0.6\), giving \(p_B=0.6/1.4\approx0.4286\). The M-step uses expected branch-specific success and cell totals. It does not declare that the three positive cells truly belong to A.

Reading a Likelihood Trace

Worked Example: Deciding Whether Iteration Has Settled

Consider a hypothetical run whose observed log-likelihoods are \(-102.40,-101.98,-101.91,-101.901,-101.9006,-101.9003\) at iterations 0 through 5. The increments are \(0.42,0.07,0.009,0.0004,\) and \(0.0003\). They are nonnegative, as expected under exact EM, and the last two increments are both below a prespecified absolute threshold of 0.001 log-likelihood units. This trace supports stopping under a rule that requires two consecutive changes below that threshold.

A small change in likelihood is not by itself proof that all parameters are accurately estimated. A flat likelihood can occur far from a well-defined maximum, or along a weakly identified direction. In practice, inspect the parameter changes and responsibilities as well: check that mixing proportions remain valid, branch-specific parameters are finite, and branches have not become effectively empty. Multiple initial values can reveal sensitivity to local maxima or different labelings.

Diagnostics, Alternatives, and Limitations

A useful trace records the observed-data log-likelihood after every iteration, not just the value of \(Q\). Under the stated conditions, the observed log-likelihood should be nondecreasing up to numerical rounding. A material decrease suggests an incorrect E-step, an M-step that failed to increase \(Q\), inconsistent likelihood calculations, or numerical underflow. Compute log mixture densities with stable log-sum-exp arithmetic when component densities are very small; directly summing tiny floating-point values can produce zero and an undefined log-likelihood.

Other warning signs include a branch with \(\sum_i r_{ib}\) near zero, parameters drifting to extreme values, and markedly different solutions from different initializations. If variances are estimated in an unconstrained Gaussian mixture, a component can concentrate around one observation while its variance approaches zero; the likelihood may then become unbounded. A variance floor, a suitable constraint, or a justified penalty may be needed. Such modifications change the optimization problem, so the objective and convergence criterion must match the modified model.

EM is not the only approach. Direct numerical maximization of the observed likelihood optimizes the same target without alternating over latent labels, but mixture likelihoods can be difficult to optimize and still require good initialization and constraints. Hard assignment alternates between assigning each cell to its most likely branch and refitting; it discards uncertainty and is not the same as the soft weighted EM update. Bayesian inference instead represents uncertainty about parameters and assignments through a posterior distribution; that is a different inferential framework. No method can make indistinguishable branch models identifiable from the observed data alone.

Common pitfall: A cell with responsibility 0.51 for one branch is not known to belong to that branch, and rounding all responsibilities to hard labels before the M-step changes the procedure. Use the posterior weights in \(Q\) and preserve the observed-data likelihood trace. Under appropriate regularity conditions, a convergent EM parameter sequence may have a stationary limit point; a stable likelihood trace alone does not establish stationarity, branch-assignment certainty, or global optimality.
Key takeaway: EM alternates between posterior branch responsibilities and responsibility-weighted parameter updates. Its likelihood-monotonicity guarantee depends on valid model calculations and an E-step and M-step that match the stated likelihood; convergence and biological interpretability require separate diagnostics.

Check Your Understanding

Answer each question, then compare with the brief solutions below.

  1. Why must the responsibilities for one cell sum to 1?
  2. If a branch receives total responsibility 12 among 30 cells, what is its unrestricted mixing-proportion update?
  3. In a Bernoulli branch model, what weighted quantity determines the updated marker probability?
  4. A likelihood trace falls by 0.2 at one iteration. Is that compatible with exact EM under the assumptions given? Name one possible cause.
  5. Why does a stable likelihood trace not establish that branch assignments are certain?

Solutions. (1) The latent label must be one of the \(B\) branches, and the displayed responsibilities are conditional probabilities over this exhaustive set, so their sum is 1. (2) The update is \(12/30=0.4\). (3) Divide the responsibility-weighted number of marker successes by the branch's total responsibility. (4) No; exact EM should not decrease the observed likelihood. A mismatched likelihood calculation, numerical underflow, or a failed M-step could explain the drop. (5) The likelihood measures fit, not certainty in each latent label; cells can retain diffuse responsibilities even after parameter estimates have stabilized.

Clinical Biostats