Deriving the LDA Joint Distribution and Conditional Updates

From Cell Counts to a Joint Model

Latent Dirichlet allocation (LDA) represents each cell as a mixture of topics and each topic as a distribution over genes. Its joint distribution makes the model's conditional independence assumptions explicit. It also provides the starting point for deriving updates: the conditional distribution of a topic assignment, given the other assignments and the observed genes, and the conditional distributions of the topic proportions and gene distributions.

Let \(D\) be the number of cells, \(V\) the number of genes, and \(K\) the number of topics. Expand the count for each cell-gene pair into individual gene tokens: \(N_d=\sum_{v=1}^V x_{dv}\) tokens for cell \(d\), where \(x_{dv}\) is the observed count of gene \(v\). Write \(w_{dn}\in\{1,\ldots,V\}\) for the gene index of token \(n\) in cell \(d\), and \(z_{dn}\in\{1,\ldots,K\}\) for its latent topic assignment. Thus, \(x_{dv}\) is the number of tokens in cell \(d\) for which \(w_{dn}=v\).

For each cell, \(\theta_d=(\theta_{d1},\ldots,\theta_{dK})\) is its topic-proportion vector. For each topic, \(\phi_k=(\phi_{k1},\ldots,\phi_{kV})\) is its gene-probability vector. Both are probability vectors. LDA assigns \(\theta_d\) a Dirichlet prior with positive parameter vector \(\alpha\), and \(\phi_k\) a Dirichlet prior with positive parameter vector \(\eta\). The prior vectors are fixed hyperparameters.

This specification assumes that, conditional on the cell's topic proportions and the topic gene distributions, tokens are generated independently. The cell token totals \(N_d\) are treated as fixed; this model describes gene identities among those tokens, not the mechanism that generated the total number of counts. It also assumes that a token's topic determines its gene probabilities and that the stated Dirichlet priors are independent. These are modeling assumptions, not consequences of the count data.

The Joint Distribution and Count Representation

For a probability vector \(q\), the categorical probability of category \(j\) is \(q_j\). Using this fact and the independence structure above, the joint density of the parameters and the token-level latent and observed variables is

The final product has two factors for every token: \(\theta_{d,z_{dn}}\), the probability of selecting its topic in cell \(d\), and \(\phi_{z_{dn},w_{dn}}\), the probability of its observed gene under that topic. This joint is for labeled token positions. If the token order is discarded, probabilities for count matrices must account for the number of token sequences that produce the same counts.

Define \(n_{dk}\) as the number of tokens in cell \(d\) assigned to topic \(k\), and \(m_{kv}\) as the number of tokens of gene \(v\) assigned to topic \(k\). Let \(m_k=\sum_v m_{kv}\). Given \(\theta\) and \(\phi\), the token factors can equivalently be collected as \(\prod_{d,k}\theta_{dk}^{n_{dk}}\prod_{k,v}\phi_{kv}^{m_{kv}}\). The counts \(n_{dk}\) and \(m_{kv}\) summarize the assignments and genes for these likelihood factors; they do not eliminate the need to specify how observed counts map to token assignments.

Conditional Updates from Dirichlet Conjugacy

For positive parameters \(a_1,\ldots,a_J\), a Dirichlet distribution has density proportional to \(\prod_j q_j^{a_j-1}\) on probability vectors \(q\). Multiplying this density by categorical likelihood factors adds their category counts to the corresponding exponents. Consequently, the conditional distributions of the LDA probability vectors are Dirichlet distributions with prior parameters plus assigned-token counts.

For example, the likelihood contribution involving \(\theta_d\) is \(\prod_k\theta_{dk}^{n_{dk}}\). Multiplication by its prior density changes the exponent of \(\theta_{dk}\) from \(\alpha_k-1\) to \(\alpha_k+n_{dk}-1\), which is the exponent in the stated Dirichlet conditional. For \(\phi_k\), the likelihood contribution is \(\prod_v\phi_{kv}^{m_{kv}}\); the same exponent argument gives its conditional. The two updates follow separately because the priors factor and each likelihood contribution involves its own probability vector.

These are conditional distributions, not fixed estimates. Their means are also useful summaries: for example, \(E[\theta_{dk}\mid z,\alpha]=(\alpha_k+n_{dk})/(\sum_j\alpha_j+N_d)\). A posterior mean is not itself a draw from the conditional distribution.

Deriving the Collapsed Topic-Assignment Update

One can update a token's topic after integrating out \(\theta\) and \(\phi\), rather than conditioning on particular values for those vectors. This is called a collapsed conditional update. It uses all assignments and observed genes except the token currently being updated. Superscript \(-i\) below means that the token at position \(i=(d,n)\) is excluded from the relevant counts. Let its observed gene be \(v=w_{dn}\).

The Dirichlet-multinomial integral gives the assignment-and-gene probability after integrating out the probability vectors. For fixed labeled tokens, the terms depending on assignments are

where \(A=\sum_k\alpha_k\), \(E=\sum_v\eta_v\), and \(\Gamma\) is the gamma function, which satisfies \(\Gamma(a+1)=a\Gamma(a)\). To obtain the conditional for token \(i\), compare this expression under the candidate assignment \(z_i=k\) with the expression excluding the token. In the cell factor, the only changing term is \(\Gamma(\alpha_k+n_{dk}^{-i}+1)/\Gamma(\alpha_k+n_{dk}^{-i})=\alpha_k+n_{dk}^{-i}\). In the topic-gene factor, the corresponding numerator ratio is \(\eta_v+m_{kv}^{-i}\), while the topic-total denominator contributes \(m_k^{-i}+E\). The cell denominator \(N_d^{-i}+A\) is the same for every candidate topic, so it cancels on normalization. Thus, the conditional probabilities are proportional to

Normalize these nonnegative weights by dividing each by their sum over \(k=1,\ldots,K\). This derivation requires positive Dirichlet parameters and the token-level LDA model specified above. The update is conditional on all other assignments and the observed genes; it is not an unconditional probability that a gene belongs to a topic.

Assumptions, Diagnostics, and Alternatives

The updates are exact for this LDA probability model, but exact algebra does not make the model assumptions automatically appropriate for every count matrix. Check that token counts are nonnegative integers and that the definition of \(N_d\) agrees with the included genes. The calculation also depends on the analyzed vocabulary and on using the same gene indexing in the observed counts and topic distributions. Zero counts are permitted; positive Dirichlet parameters keep the conditional weights well-defined even when a topic has not yet been assigned a particular gene.

A common implementation failure is to forget to remove the current token before computing its conditional. If the token is included, the resulting weight uses information from the assignment being sampled and is not the conditional derived above. Another warning sign is a probability vector that fails to sum to 1 after normalization, or a purported posterior Dirichlet mean whose components do not sum to 1. These usually indicate inconsistent count totals, indexing errors, or an incorrect denominator. The number \(m_k^{-i}\) must count all genes assigned to topic \(k\), while \(m_{kv}^{-i}\) counts only gene \(v\) within that topic.

The collapsed update integrates out the probability vectors and samples assignments using their dependence induced by the Dirichlet priors. An uncollapsed update instead retains \(\theta\) and \(\phi\): it samples an assignment using a weight proportional to \(\theta_{dk}\phi_{kv}\), then draws the probability vectors from their Dirichlet conditionals. Both approaches target the same LDA posterior when implemented correctly, but their intermediate calculations differ. Directly enumerating all assignment configurations can provide an exact check for very small toy data, but the number of configurations grows exponentially with the number of tokens and is not a practical general strategy.

Finally, LDA's probability model does not by itself establish that a topic is a distinct biological program. Topics are latent statistical components, and their names or biological interpretations require evidence beyond the assignment equations. Likewise, this token model conditions on cell totals and does not explain variation in total RNA counts. Its joint distribution and conditionals describe what follows if the specified model is used; they do not show that its assumptions fit a particular experiment.