Variational Inference for Bayesian Integrative Clustering

Why use a variational approximation?

A Bayesian integrative mixture represents each sample by a shared latent cluster label while allowing each omics modality to have its own cluster-specific profile. The posterior distribution over labels and parameters can be difficult to explore when many samples, features, or clusters are involved. Gibbs sampling and split-merge moves target that posterior by simulation, but may require many iterations and can be computationally demanding. Variational inference instead searches for a tractable distribution that approximates the posterior and improves a lower bound on the log marginal likelihood.

The trade-off is important: variational inference is often faster and produces convenient approximate distributions, but its factorization can omit posterior dependence and its optimization can settle at a local optimum. The resulting assignment probabilities are therefore probabilities under the fitted approximation, not exact posterior probabilities.

Model and mean-field family

Consider \(N\) samples, \(K\) clusters, and \(M\) modalities. For sample \(i\), modality \(m\), and feature \(j\), let \(x_{imj}\) be a binary observation and \(z_i\) the shared cluster label, represented by \(z_{ik}=1\) when sample \(i\) belongs to cluster \(k\) and zero otherwise. Suppose \(x_{imj}\), conditional on \(z_i=k\) and its profile parameter \(\theta_{kmj}\), follows a Bernoulli distribution with success probability \(\theta_{kmj}\). Let the cluster proportions be \(\pi=(\pi_1,\ldots,\pi_K)\).

For this derivation, assume the observations are conditionally independent across samples and features and across modalities given the shared labels and the modality-specific profiles. The priors are \(\pi\sim\operatorname{Dirichlet}(\alpha_1,\ldots,\alpha_K)\), \(z_i\mid\pi\sim\operatorname{Categorical}(\pi)\), and independently \(\theta_{kmj}\sim\operatorname{Beta}(a_{mj},b_{mj})\). Missing feature values are omitted from the likelihood and sufficient statistics; this treatment assumes the observation pattern does not itself carry modeled information about the labels or profiles.

This factorization is an approximation, not a conditional-independence assertion about the exact posterior. In particular, the exact posterior generally couples labels to mixture proportions and profiles. The variational family suppresses those dependencies in order to make optimization manageable.

The evidence lower bound and coordinate updates

Let \(X\) denote all observed values and \(Z\) all labels. The evidence lower bound (ELBO) is

The identity \(\log p(X)=\mathcal{L}(q)+\operatorname{KL}\{q(Z,\pi,\theta)\,\|\,p(Z,\pi,\theta\mid X)\}\) shows that the ELBO is a lower bound, since Kullback–Leibler divergence is nonnegative. Maximizing it reduces this divergence within the chosen family, but does not ensure that the best member of the family is close to the exact posterior.

For any one factor \(q_j\), with all other factors fixed, the coordinate-optimal update is proportional to the exponential of the expected log joint density, where the expectation is over all variables except the one in that factor. To see this, write the terms involving \(q_j\) as \(\int q_j(u)h(u)\,du-\int q_j(u)\log q_j(u)\,du\), plus a constant. Normalizing \(\exp\{h(u)\}\) gives the maximum: the difference from that value is a negative Kullback–Leibler divergence. Thus the coordinate update improves or leaves unchanged the ELBO.

For the model above, these updates preserve conjugate forms. Define \(\psi\) as the digamma function, the derivative of the log gamma function. Then

For each profile parameter, using only observed values of that feature,

The assignment update is a normalized exponential of the expected log mixture weight plus the expected log likelihood for all observed features and modalities:

This update combines evidence across modalities because a sample has one shared label. It does not imply that each modality has the same profile: the parameters \(\theta_{kmj}\) remain modality- and feature-specific.

Approximate assignment uncertainty

The factor \(q(z_i)\) is a categorical distribution with probabilities \(r_{i1},\ldots,r_{iK}\). Its entropy, measured in nats, is \(-\sum_k r_{ik}\log r_{ik}\), with \(0\log0\) interpreted as zero. Entropy is small when one approximate assignment probability dominates and large when probability is spread across clusters. It describes uncertainty within the fitted variational factor, not all uncertainty in the exact posterior.

Optimization, diagnostics, and alternatives

A practical coordinate-ascent algorithm alternates updates of the label factors with the conjugate parameter factors. Initialization matters because the ELBO is generally not jointly concave in all variational parameters. Each exact coordinate update should not decrease the ELBO, apart from numerical error, but the final solution may be only a local optimum. Multiple initializations are therefore useful; compare their final ELBO values and inspect whether the resulting fitted profiles and allocations are substantively plausible.

Numerically, use a log-sum-exp calculation to normalize assignment scores: subtract the largest score before exponentiating, then divide by the sum. Otherwise, very negative expected log scores can underflow to zero. A decreasing ELBO can indicate an implementation error, an inexact update, or numerical instability. A rising ELBO alone is not a model check: it confirms progress in the chosen optimization criterion, not that the assumed likelihoods or conditional independences adequately describe the data.

Compared with Gibbs sampling or split-merge sampling, mean-field variational inference avoids producing a long chain of posterior draws and can be substantially faster for large data sets. Sampling can represent posterior dependence more faithfully in principle, but computational mixing and multimodality can be serious concerns. Collapsed sampling integrates out some parameters before updating labels, whereas the variational updates above retain variational factors for prevalence and profiles. These approaches answer the same posterior-inference problem with different computational approximations; none removes the need to examine model adequacy.

The main limitation of this mean-field family is its imposed independence. It may understate uncertainty when labels and profiles are strongly coupled, or when several parameters must vary together to explain the data. A sharply concentrated \(q(z_i)\) can therefore reflect the approximation as well as strong evidence. Further, a model with many features can accumulate large likelihood contributions, so a modality with more features may exert more influence; this tutorial’s updates do not correct a poorly chosen likelihood or an unsuitable feature scale. Evaluate the model and the modality contributions rather than interpreting a high ELBO or confident-looking responsibilities as proof of a correct clustering.

Common pitfall

• Wrong reasoning: “The variational responsibility is 0.99, so the exact posterior probability that this sample belongs to this cluster is 0.99.”
• Correction: The value 0.99 is \(q(z_i=k)\) under a particular optimized factorized approximation. The omitted posterior dependencies and local optimization can make it overconfident. Report it as an approximate assignment probability and assess sensitivity to initialization and the model assumptions.