Variational Inference For Multi-Omics Factors

Why Approximate the Posterior?

A multi-omics factor model represents aligned samples using latent scores and view-specific feature loadings. Even with a simple Gaussian observation model, the posterior distribution of all scores and loadings is generally dependent and can be costly to calculate exactly. Variational inference replaces that posterior with a tractable distribution chosen to approximate it. The choice of approximation determines both the computation and which posterior dependencies can be represented.

The central objective is the evidence lower bound, or ELBO. Maximizing it improves the chosen approximation to the posterior while also giving a lower bound on the marginal log probability of the observed data. Coordinate updates make this optimization practical when the number of samples and features is large.

A Gaussian Multi-Omics Model and Variational Family

For a precise derivation, consider \(V\) aligned omics views. In view \(v\), let \(Y_v\) be an \(n \times p_v\) data matrix, \(Z\) an \(n \times K\) matrix of sample scores, and \(W_v\) a \(p_v \times K\) matrix of feature loadings. Write \(z_i\) for the \(K\)-vector of scores for sample \(i\), and \(w_{vj}\) for the \(K\)-vector of loadings for feature \(j\) in view \(v\). For observed entries, use the model

Assume entries are conditionally independent given the scores and loadings. The precision \(\tau_{vj}\) is treated as known and positive. To make the example model fully specified, use independent standard-normal priors \(z_i\sim N(0,I_K)\) and \(w_{vj}\sim N(0,I_K)\). This is a convenient baseline model, not a claim that these prior choices suit every analysis. Let \(\Omega_v\) denote the set of observed sample-feature pairs in view \(v\); missing entries contribute no likelihood term.

This factorization is an approximation, not a property of the model’s exact posterior. In particular, it removes posterior dependence between a score and its loadings, and between different score or loading vectors. The approximation retains uncertainty within each vector through its covariance matrix.

The ELBO and Its Coordinate Updates

Let \(Y_{\mathrm{obs}}\) denote all observed entries. For any normalized \(q\) supported wherever the posterior is supported, the ELBO is

The bound follows by writing the marginal likelihood as an expectation under \(q\) and applying Jensen’s inequality:

Equivalently, the gap is a Kullback–Leibler divergence: \(\log p(Y_{\mathrm{obs}})-\mathcal{L}(q)=\operatorname{KL}\{q(Z,W)\,\|\,p(Z,W\mid Y_{\mathrm{obs}})\}\), which is nonnegative. Thus maximizing the ELBO over a chosen family minimizes this divergence within that family. A larger ELBO does not by itself establish that the family captures all posterior structure.

For this Gaussian model, define the second moments \(R_i=S_i+m_im_i^\mathsf{T}\) and \(Q_{vj}=C_{vj}+a_{vj}a_{vj}^\mathsf{T}\). The mean-field coordinate optimum for one factor is proportional to the exponential of the expected log joint density, where the expectation is over all other factors. Applying this result gives

The identity matrices come from the standard-normal priors. The loading second moments enter the score precision; the score second moments enter the loading precision. The means use the other factor’s current mean. These are coordinate updates: update one block using the current expected moments of the other block, then alternate. Each exact coordinate maximization cannot decrease the ELBO, although it need not find the global maximum.

Calculating and Interpreting the Bound

For the Gaussian model, the expected squared residual for an observed entry is

This identity uses independence of \(q(z_i)\) and \(q(w_{vj})\), so that the expected squared inner product is \(\operatorname{tr}(R_iQ_{vj})\). The likelihood contribution to the ELBO is the sum, over observed entries, of \(-\frac12\{\log(2\pi/\tau_{vj})+\tau_{vj}\mathbb{E}_q[(y_{vij}-z_i^\mathsf{T}w_{vj})^2]\}\). The remaining terms are expected log priors minus the entropy of \(q\), where entropy is \(-\mathbb{E}_q[\log q]\). Including these terms matters: fitting expected squared residuals alone is not ELBO optimization.

Computation, Diagnostics, and Alternatives

The coordinate formulas explain the computational appeal. Each score update only requires the observed features for that sample; each loading update only requires the observed samples for that feature. If entries are missing, the sums use observed pairs only. Accumulating the relevant \(K \times K\) second moments avoids forming or storing a full joint covariance over every score and loading. With dense data, full covariance factors can still be expensive as \(K\) grows; diagonal covariance factors reduce memory and computation, at the cost of a more restrictive approximation.

Check that the ELBO is finite and nondecreasing across exact coordinate updates, and inspect expected residuals and predictive performance rather than relying on the bound alone. A falling ELBO may indicate an implementation error, inconsistent updates, or an inaccurate numerical calculation. A stable ELBO does not prove that the approximation is adequate: mean-field factorization can understate posterior uncertainty when scores and loadings are strongly dependent. Different initializations can also lead to different local optima. For non-Gaussian observation models, the Gaussian coordinate formulas above do not apply directly; an appropriate expected-likelihood calculation or additional approximation is required.

Variational inference differs from expectation-maximization in its target: it optimizes a distribution over latent quantities, rather than using an E-step and M-step to optimize a point estimate of model parameters. Compared with posterior sampling methods such as Markov chain Monte Carlo, coordinate variational inference is often faster and easier to scale, but generally gives a less flexible representation of posterior dependence. These methods answer related but distinct computational needs; the appropriate comparison depends on the desired uncertainty representation and available computation.