Variational Inference for Single-Cell Topic Models

Why Approximate the Topic-Model Posterior?

In a single-cell topic model, a cell’s gene counts are explained through latent topic assignments: each observed count is associated with a topic, and each topic has a probability distribution over genes. The assignments and the cell’s topic proportions are uncertain, and their exact joint posterior can be difficult to compute. Mean-field variational inference replaces that posterior with a simpler distribution whose parameters can be updated efficiently.

The introductory material on latent Dirichlet allocation for cell gene counts defines the generative model, and the earlier material on the LDA joint distribution and conditional updates gives its joint probability and conditional structure. Here the focus is how a factorized approximation leads to an objective and practical coordinate updates. The approximation is an inference method, not a change to the assumed data-generating model.

Model assumptions: For one cell, let \(x_g\) be the observed count of gene \(g\), and let \(N=\sum_g x_g\) be the total count. There are \(K\) topics. The cell’s topic-proportion vector \(\theta\) has distribution \(\operatorname{Dir}(\alpha)\), where every \(\alpha_k>0\). Each topic \(k\) has fixed gene probabilities \(\beta_{kg}\geq 0\), with \(\sum_g\beta_{kg}=1\). For each of the \(N\) tokens, an assignment \(z_n\) is drawn from \(\operatorname{Categorical}(\theta)\), then its gene \(w_n\) is drawn from \(\operatorname{Categorical}(\beta_{z_n})\). Inference below conditions on the fixed topic probabilities and the observed counts.

A count \(x_g\) can be represented as \(x_g\) separate tokens with observed gene \(g\). This lets the inference equations be stated using token assignments. When identical genes are grouped, their assignment probabilities are shared, and the count \(x_g\) multiplies their contribution. The multinomial coefficient that relates the token likelihood to a count-vector likelihood is constant in the latent variables, so it does not affect the coordinate updates.

Mean-Field Approximation and Objective

The mean-field approximation assumes that the variational distribution separates the topic proportions from the token assignments and separates the assignments across tokens:

q(\theta,z_{1:N})=q(\theta)\prod_{n=1}^{N}q(z_n),\qquad q(\theta)=\operatorname{Dir}(\gamma),\qquad q(z_n=k)=\phi_{nk}.

Here \(\gamma_k>0\), and for each token \(n\), \(\phi_{nk}\geq 0\) and \(\sum_{k=1}^{K}\phi_{nk}=1\). This factorization makes the calculation manageable, but it excludes posterior dependence between \(\theta\) and the assignments. The evidence lower bound, or ELBO, is the objective optimized over these variational parameters:

\mathcal{L}(q) = \mathbb{E}_q[\log p(\theta,z_{1:N},w_{1:N}\mid\alpha,\beta)] -\mathbb{E}_q[\log q(\theta,z_{1:N})].

The first expectation rewards variational distributions that place mass on latent configurations compatible with the model and data; the second is the negative expected log variational density. More precisely, when the posterior is defined and the expectations are finite,

\log p(w_{1:N}\mid\alpha,\beta) =\mathcal{L}(q) +\operatorname{KL}\!\left(q(\theta,z_{1:N})\,\middle\|\,p(\theta,z_{1:N}\mid w_{1:N},\alpha,\beta)\right).

Because the Kullback–Leibler divergence is nonnegative, the ELBO is a lower bound on the log marginal probability of the observed tokens. Maximizing it minimizes the divergence from the exact posterior within the chosen variational family. It does not make the factorized approximation exact.

Definition: The digamma function is \(\psi(u)=\frac{d}{du}\log\Gamma(u)\), where \(\Gamma\) is the gamma function. Under \(q(\theta)=\operatorname{Dir}(\gamma)\), write \(\gamma_0=\sum_k\gamma_k\). Then \(\mathbb{E}_q[\log\theta_k]=\psi(\gamma_k)-\psi(\gamma_0)\).

For grouped counts, write \(\phi_{gk}\) for the assignment probability to topic \(k\) for each token of gene \(g\). The ELBO, omitting only a count-vector combinatorial constant, is

\begin{aligned} \mathcal{L}(\gamma,\phi) ={}&\log\Gamma(\alpha_0)-\sum_k\log\Gamma(\alpha_k) +\sum_k(\alpha_k-1)\bigl[\psi(\gamma_k)-\psi(\gamma_0)\bigr]\\ &+\log\Gamma(\gamma_0)-\sum_k\log\Gamma(\gamma_k) +\sum_k(\gamma_k-1)\bigl[\psi(\gamma_k)-\psi(\gamma_0)\bigr]\\ &+\sum_g x_g\sum_k\phi_{gk} \bigl[\psi(\gamma_k)-\psi(\gamma_0)+\log\beta_{kg}-\log\phi_{gk}\bigr]. \end{aligned}

The first line is the expected log prior for \(\theta\); the second is the entropy contribution \(-\mathbb{E}_q[\log q(\theta)]\); the final line combines the assignment prior, gene likelihood and assignment entropy. Terms with \(\phi_{gk}=0\) are interpreted by continuity, so \(\phi_{gk}\log\phi_{gk}=0\). A topic with \(\beta_{kg}=0\) cannot receive positive assignment probability for gene \(g\).

Coordinate Updates and Their Justification

Coordinate ascent optimizes one variational factor at a time while holding the others fixed. The general mean-field update for a latent variable \(u\) is \(q^*(u)\propto\exp\{\mathbb{E}_{q(\text{other variables})}[\log p(u,\text{other variables},w)]\}\), with the result normalized to integrate or sum to one. Applying this identity gives the two LDA updates.

Update assignments.
For a token of gene \(g\), the terms involving its topic \(k\) are \(\log\theta_k+\log\beta_{kg}\). Taking the expectation over \(q(\theta)\) gives an assignment probability proportional to \(\beta_{kg}\exp\{\psi(\gamma_k)-\psi(\gamma_0)\}\), normalized over topics.

Update topic proportions.
The prior contributes \((\alpha_k-1)\log\theta_k\), and each token contributes its expected assignment count to topic \(k\). The resulting factor is Dirichlet with parameter \(\alpha_k+\sum_g x_g\phi_{gk}\).

\phi_{gk} = \frac{\beta_{kg}\exp\{\psi(\gamma_k)-\psi(\gamma_0)\}} {\sum_{j=1}^{K}\beta_{jg}\exp\{\psi(\gamma_j)-\psi(\gamma_0)\}}, \qquad \gamma_k=\alpha_k+\sum_g x_g\phi_{gk}.

The assignment update applies only if its denominator is positive: the observed gene must have positive probability under at least one topic. If \(\beta_{kg}=0\) for every topic \(k\), then the gene has zero probability under the specified model. The observed data are impossible under that model, and the posterior for those data is undefined; assigning arbitrary normalized weights would not fix the problem.

These are conditional coordinate optima. In practice, initialize \(\gamma\) with positive entries, update \(\phi\) and \(\gamma\) in turn, and continue until the ELBO or the parameters change little. Exact coordinate updates do not decrease the ELBO, although they need not find its global maximum: the objective can have local optima.

Worked Example: Updating Assignments for Two Genes

Consider a hypothetical cell with one token each of genes A and B, two topics, and topic probabilities \(\beta_{1A}=0.8\), \(\beta_{1B}=0.2\), \(\beta_{2A}=0.3\), and \(\beta_{2B}=0.7\). These probabilities sum to one within each topic. Suppose the current variational parameter is \(\gamma=(3,1)\), so \(\gamma_0=4\). Both genes have positive probability under at least one topic, and in fact under both, so the assignment updates are defined.

For integer arguments, \(\psi(m)=H_{m-1}-c\), where \(H_r=\sum_{i=1}^{r}1/i\) and \(c\) is Euler’s constant. Therefore \(\psi(3)-\psi(4)=-1/3\) and \(\psi(1)-\psi(4)=-11/6\). The topic-specific exponential terms are approximately \(e^{-1/3}=0.7165\) and \(e^{-11/6}=0.1599\).

For gene A, the unnormalized weights are \(0.8e^{-1/3}\approx0.5732\) and \(0.3e^{-11/6}\approx0.0480\). Thus

\phi_{A1}=\frac{0.5732}{0.5732+0.0480}\approx0.9228, \qquad \phi_{A2}=1-\phi_{A1}\approx0.0772.

For gene B, the weights are \(0.2e^{-1/3}\approx0.1433\) and \(0.7e^{-11/6}\approx0.1119\). Normalizing gives \(\phi_{B1}\approx0.5615\) and \(\phi_{B2}\approx0.4385\). Each pair sums to one, as required.

With prior \(\alpha=(1,1)\), the subsequent proportion-factor update is \(\gamma_1=1+0.9228+0.5615=2.4843\) and \(\gamma_2=1+0.0772+0.4385=1.5157\). Their sum is \(4\), equal to \(\alpha_1+\alpha_2+N\), which checks the update’s arithmetic.

Worked Example: Aggregating Repeated Gene Counts

Now suppose a hypothetical cell has counts \(x_A=3\) and \(x_B=1\), and use the same current assignment probabilities calculated above. Let \(\alpha=(0.5,0.5)\). The expected assignment count for topic 1 is \(3(0.9228)+1(0.5615)=3.3299\). For topic 2 it is \(3(0.0772)+1(0.4385)=0.6701\).

The updated parameters are \(\gamma_1=0.5+3.3299=3.8299\) and \(\gamma_2=0.5+0.6701=1.1701\). Their total is \(5.0000\), which equals the prior total \(1.0\) plus the cell’s four tokens. The update uses expected, not integer, topic counts: uncertain assignments contribute fractions of a token to each topic.

Worked Example: A Gene With Zero Probability in One Topic

Suppose a hypothetical gene C has \(\beta_{1C}=0\) and \(\beta_{2C}=0.4\), with the same current \(\gamma=(3,1)\). The first unnormalized assignment weight is \(0\cdot e^{-1/3}=0\); the second is \(0.4e^{-11/6}\approx0.0640\). The normalized update is therefore \(\phi_{C1}=0\) and \(\phi_{C2}=1\). This is valid because the normalizing sum is positive.

If instead both topic probabilities for C were zero, the normalizing sum would be zero. No assignment distribution could explain this observed gene under the specified model. The correct response is to check the topic-gene probabilities and model support, not to divide by zero or replace the result with an arbitrary uniform distribution.

Diagnostics, Alternatives, and Limitations

Several checks catch common implementation and modeling failures:

Check normalization. Each assignment vector must have nonnegative entries summing to one, and every \(\gamma_k\) must remain positive. For each gene with positive count, verify that its topic probabilities have a positive sum.
Watch numerical scale. Exponentiating very small expected log proportions can underflow. Compute log weights, subtract their maximum, exponentiate, and normalize; this preserves the same probabilities while improving numerical stability.
Monitor the ELBO. Exact coordinate updates should not decrease it, apart from small numerical errors. A decrease can indicate an incorrect update, inconsistent normalization, or unstable arithmetic. A nearly unchanged ELBO indicates convergence of the chosen iteration, not correctness of the model or global optimality.
Check sensitivity to initialization. Different starting values can lead to different local optima. Comparing several initializations is useful when the inferred approximation matters.

A common mistake is to treat \(\phi_{gk}\) as the true posterior probability of a topic assignment without qualification. It is the probability under the optimized variational approximation, which imposes independence between assignments and topic proportions. Another mistake is to assume that a high ELBO proves a good biological interpretation: the ELBO is an optimization criterion for a specified model, not a test of biological validity.

Collapsed Gibbs sampling is a main alternative: it samples assignments while integrating out topic proportions, rather than replacing the posterior with a factorized distribution. It can represent posterior variation through samples, but requires iterative Monte Carlo computation and assessment of sampling behavior. Mean-field coordinate ascent is typically faster and deterministic for a fixed initialization, but its factorization can understate uncertainty and miss posterior dependence. Neither method repairs a misspecified topic model or data with zero probability under its assumptions.

Key takeaway: Mean-field LDA alternates normalized assignment probabilities, weighted by expected log topic proportions and gene probabilities, with Dirichlet parameters equal to prior parameters plus expected topic counts. Check model support, normalization and ELBO behavior, and interpret the result as an approximation to the posterior.

Check Your Understanding

Answer each question using the update definitions above. Solutions are included for checking.

What condition must hold for the assignment update of an observed gene to be normalized?
For a cell with counts \(x_A=2\), \(x_B=1\), prior \(\alpha=(1,1)\), and assignment probabilities \(\phi_{A}=(0.6,0.4)\), \(\phi_{B}=(0.3,0.7)\), calculate the updated \(\gamma\).
Why does a coordinate update use expected assignment counts rather than requiring each token to be assigned to exactly one topic in the calculation?
Does an ELBO that has stopped changing establish that the variational posterior is exact? Explain.

Solutions. (1) The sum of the unnormalized topic weights must be positive, equivalently the observed gene must have positive probability under at least one topic. (2) Topic 1 receives expected count \(2(0.6)+0.3=1.5\), and topic 2 receives \(2(0.4)+0.7=1.5\), so \(\gamma=(2.5,2.5)\). (3) The variational assignments are probability distributions, so their expected counts are the sufficient statistics in the coordinate optimum for \(q(\theta)\); a token can contribute fractionally across topics in expectation. (4) No. It indicates convergence of the optimization at that solution, but the factorized family may not represent the exact posterior, and the solution may be a local optimum.

Variational Inference for Single-Cell Topic Models

What You'll Learn