Expectation-Maximization For Multi-Omics Factors

Why Expectation-Maximization Is Used

In a multi-omics factor model, the latent scores are not observed, but they explain dependence among measured features across omics views. If those scores were known, estimating loadings and residual variances would be relatively direct. If the parameters were known, estimating the scores would also be direct. The expectation-maximization (EM) algorithm alternates between these two calculations: it uses the current parameters to find conditional expectations of the latent scores, then uses those expectations to update the parameters.

This tutorial derives EM for a Gaussian model with shared factors, observed measurements, and feature-specific residual variances. The model is deliberately stated precisely: the validity of the updates depends on its distributional and independence assumptions. The method is an estimation procedure, not a way to establish that a fitted factor represents a biological mechanism.

Model and Assumptions

Suppose \(n\) samples have measurements in \(V\) omics views. View \(v\) contains \(p_v\) features, and there are \(K\) shared factors. Let \(y_{ivj}\) be the measurement for sample \(i\), view \(v\), and feature \(j\); let \(z_i\) be the \(K\)-dimensional latent score for sample \(i\); let \(\lambda_{vj}\) be the loading vector for that feature; and let \(\mu_{vj}\) be its mean.

Assume measurements are centered in advance or that the means \(\mu_{vj}\) are fixed at known values. Write \(y_i\) for the concatenated, mean-centered measurement vector across all views for sample \(i\), and \(\Lambda\) for the corresponding matrix of loading rows. The residual covariance matrix \(\Psi\) is diagonal: its entry for feature \(j\) in view \(v\) is \(\psi_{vj}>0\).

The standard normal score distribution fixes a scale for the latent scores, but it does not generally fix their orientation: rotations of the factors can leave the implied covariance unchanged. This is the rotational ambiguity covered in the earlier material on rotational invariance. An EM fit therefore need not yield a uniquely oriented loading matrix. Further interpretation may require constraints or an alignment convention.

The E-Step: Conditional Score Expectations

Let the current loading and residual parameters be \(\Lambda\) and \(\Psi\). The E-step finds the conditional distribution of \(z_i\) given \(y_i\) under these current parameters. Define \(C=(I_K+\Lambda^\mathsf{T}\Psi^{-1}\Lambda)^{-1}\) and \(m_i=C\Lambda^\mathsf{T}\Psi^{-1}y_i\). Then the posterior distribution is Gaussian:

The second moment is essential. It is not just \(m_i m_i^\mathsf{T}\), because the conditional score remains uncertain; its covariance \(C\) contributes to the expectation of \(z_i z_i^\mathsf{T}\).

To justify these expressions, the joint density as a function of \(z_i\) has log-kernel \(-\tfrac12[z_i^\mathsf{T}z_i+(y_i-\Lambda z_i)^\mathsf{T}\Psi^{-1}(y_i-\Lambda z_i)]\). Collecting the quadratic and linear terms in \(z_i\) gives precision matrix \(C^{-1}=I_K+\Lambda^\mathsf{T}\Psi^{-1}\Lambda\) and linear coefficient \(\Lambda^\mathsf{T}\Psi^{-1}y_i\). Completing the square gives the stated mean and covariance. Finally, for a random vector with mean \(m_i\) and covariance \(C\), its second moment is \(C+m_i m_i^\mathsf{T}\).

The M-Step: Updating Loadings and Residual Variances

The M-step maximizes the expected complete-data log-likelihood, where the expectation is taken over the E-step distributions of the scores. For compactness, define \(S=\sum_{i=1}^n E[z_i z_i^\mathsf{T}\mid y_i]\). For feature \(j\), let \(y_{\cdot j}\) be its \(n\)-vector of centered measurements and let \(\lambda_j\) be its loading vector. The expected squared-error term for this feature is \(\sum_i E[(y_{ij}-\lambda_j^\mathsf{T}z_i)^2\mid y_i]\).

Differentiating that term with respect to \(\lambda_j\), setting the derivative to zero, and solving gives the loading update. Then maximizing with respect to the residual variance gives the average expected squared error:

Here \(C\) is common across samples under the stated model. Apply the updates separately to each feature, using its own measurements and loading row; each feature has its own residual variance. The formula for \(\psi_j^{\text{new}}\) is an expected residual sum of squares divided by \(n\), so it has the units of that feature squared. In view-indexed notation, these are the updates for each \(j\) in each view \(v\).

For the loading update, the expected squared-error expression, up to terms not involving \(\lambda_j\), is \(\sum_i[-2y_{ij}\lambda_j^\mathsf{T}m_i+\lambda_j^\mathsf{T}E[z_i z_i^\mathsf{T}\mid y_i]\lambda_j]\). Its gradient is \(-2\sum_i y_{ij}m_i+2S\lambda_j\). Setting this to zero yields the update when \(S\) is invertible. For the variance, the expected Gaussian log-likelihood terms involving \(\psi_j\) are \(-\tfrac{n}{2}\log\psi_j-\tfrac{1}{2\psi_j}\sum_i E[(y_{ij}-\lambda_j^\mathsf{T}z_i)^2\mid y_i]\). Differentiating and setting the derivative to zero gives the expected squared error divided by \(n\).

Why the Iteration Increases Likelihood

Let \(\theta\) denote all model parameters and let \(q(z)\) be the conditional distribution of the latent scores under the current parameter value \(\theta^{\text{old}}\). The observed-data log-likelihood can be expressed as \(\log p(y\mid\theta)=\mathcal{L}(q,\theta)+\mathrm{KL}(q\,\|\,p(z\mid y,\theta))\), where \(\mathcal{L}(q,\theta)=E_q[\log p(y,z\mid\theta)]-E_q[\log q(z)]\), and \(\mathrm{KL}\) is the nonnegative Kullback–Leibler divergence. For the E-step choice \(q=p(z\mid y,\theta^{\text{old}})\), this divergence is zero at the old parameters. The M-step maximizes \(\mathcal{L}(q,\theta)\), so:

Therefore an exact EM iteration does not decrease the observed-data likelihood. This result does not imply convergence to the global maximum: the likelihood can have local optima, and different initializations can reach different solutions. Nor does it remove rotational ambiguity in the factor representation.

Diagnostics, Alternatives, and Limitations

EM is useful when the latent scores are missing variables in a likelihood model and the conditional score distribution has tractable moments. Its updates use all views jointly to estimate shared scores while keeping feature-specific residual variances. Compared with direct numerical maximization of the observed-data likelihood, EM often replaces one difficult optimization with simpler conditional-expectation and maximization steps. Direct optimization can be preferable when a desired constraint or model extension makes the M-step inconvenient, but it requires reliable handling of the observed likelihood and its gradients. Principal component analysis (PCA) is a different alternative: it gives a low-rank description of variation, but does not by itself provide this Gaussian latent-variable likelihood with feature-specific residual variances.

Check whether the observed-data log-likelihood is nondecreasing across exact EM iterations and whether fitted residual variances remain positive and plausible. A decrease may signal an incorrect update, numerical error, or an implementation that is not performing an exact M-step. Very small residual variances can make \(\Psi^{-1}\) unstable; a nearly singular \(S\) can make the loading update unstable. These issues may indicate inadequate data for the chosen factor count, problematic scaling, or numerical conditioning. They are not resolved just by continuing to iterate.

The assumptions also limit what this model can capture. Diagonal \(\Psi\) implies residual independence across features conditional on the scores; remaining structured dependence suggests that the model is missing structure or that the residual specification is inadequate. Gaussian errors may be inappropriate for counts or binary measurements, for which Gaussian updates do not apply. When some measurements are missing, the displayed formulas cannot simply be applied as though the missing entries were observed. Under an appropriate ignorable-missingness assumption, likelihood calculations can instead use each sample’s observed features, with corresponding posterior moments and updates; the exact observed-data pattern must be handled in both steps.

A common mistake is to use only \(m_i m_i^\mathsf{T}\) in the M-step and omit \(C\). That treats the estimated scores as known, even though they have conditional uncertainty. The correct second moment is \(C+m_i m_i^\mathsf{T}\); omitting \(C\) generally changes both loading and residual-variance updates. Another mistake is to treat likelihood increase as evidence that the factors are uniquely identified or biologically meaningful. EM guarantees neither: it optimizes within the specified model and may converge to a local solution with an orientation equivalent to other factor representations.