Inference for LDA: Collapsed Gibbs Sampling

Having established the Bayesian formulation of Latent Dirichlet Allocation (LDA) in the previous section, we now turn to the challenge of performing posterior inference. Our goal is to compute the posterior distribution of the latent variables, primarily the topic assignments $z$, given the observed words $w$ and the hyperparameters $\alpha$ and $\beta$. That is, we want $p(z, \theta, \phi | w, \alpha, \beta)$.

While standard Gibbs sampling could be applied by iteratively sampling each latent variable ($z$, $\theta$, $\phi$) from its conditional distribution, the high dimensionality of the continuous parameters $\theta$ (document-topic distributions) and $\phi$ (topic-word distributions) makes this inefficient. Furthermore, we are often most interested in the topic assignments $z$ themselves, or the expected values of $\theta$ and $\phi$.

Collapsed Gibbs sampling offers an elegant and often effective alternative for LDA. The core idea is to integrate out, or "collapse," the continuous parameters $\theta$ and $\phi$ analytically, leveraging the conjugacy between the Dirichlet priors and the Multinomial likelihoods inherent in the LDA model. This leaves us with a Gibbs sampler that only needs to iterate through the discrete topic assignments $z_{d,n}$ for each word $n$ in each document $d$.

The Collapsed Gibbs Sampling Update Rule

The heart of the collapsed Gibbs sampler for LDA is the conditional probability of assigning a specific word token $(d,n)$ (the $n$-th word in document $d$) to a particular topic $k$, given all other topic assignments $z_{\neg(d,n)}$, the observed words $w$, and the hyperparameters $\alpha, \beta$. We denote the vocabulary word corresponding to $w_{d,n}$ as $v$.

Using Bayes' theorem and exploiting the conditional independencies and Dirichlet-Multinomial conjugacy, we can derive this conditional probability. We integrate out $\theta_d$ (the topic proportions for document $d$) and $\phi_k$ (the word distribution for topic $k$):

$$p(z_{d,n}=k | z_{\neg(d,n)}, w, \alpha, \beta) \propto p(w_{d,n}=v | z_{d,n}=k, z_{\neg(d,n)}, \beta) \times p(z_{d,n}=k | z_{d, \neg n}, \alpha)$$

Let's break down the two terms on the right-hand side:

1. Document-Topic Term $p(z_{d,n}=k | z_{d, \neg n}, \alpha)$: This term reflects how likely topic $k$ is in document $d$, considering the assignments of other words in the same document ($z_{d, \neg n}$) and the document-topic prior $\alpha$. Integrating out $\theta_d$ yields a probability proportional to the count of words already assigned to topic $k$ in document $d$ (excluding the current word $n$), plus the prior parameter $\alpha_k$. Let $N_{d,k}^{\neg n}$ be the count of words in document $d$ (excluding word $n$) assigned to topic $k$. Assuming a symmetric prior $\alpha$ for simplicity (i.e., $\alpha_k = \alpha$ for all $k$):

   $$p(z_{d,n}=k | z_{d, \neg n}, \alpha) \propto N_{d,k}^{\neg n} + \alpha$$

   This term favors assigning the word to topics already prevalent in the document.

2. Topic-Word Term $p(w_{d,n}=v | z_{d,n}=k, z_{\neg(d,n)}, \beta)$: This term reflects how likely the specific word $v$ is under topic $k$, considering all other word assignments across the corpus ($z_{\neg(d,n)}$) and the topic-word prior $\beta$. Integrating out $\phi$ yields a probability proportional to the count of times word $v$ has been assigned to topic $k$ elsewhere in the corpus, plus the prior parameter $\beta_v$. Let $N_{k,v}^{\neg (d,n)}$ be the count of word $v$ assigned to topic $k$ across all documents, excluding the current instance $(d,n)$. Let $N_k^{\neg (d,n)}$ be the total count of words assigned to topic $k$, excluding the current instance. Assuming a symmetric prior $\beta$ (i.e., $\beta_v = \beta$ for all $v$) and $V$ is the vocabulary size:

   $$p(w_{d,n}=v | z_{d,n}=k, z_{\neg(d,n)}, \beta) \propto \frac{N_{k,v}^{\neg (d,n)} + \beta}{N_k^{\neg (d,n)} + V\beta}$$

   This term favors assigning the word to topics that frequently generate this specific word type $v$.

Combining these, the full conditional probability for sampling the topic assignment $z_{d,n}$ is:

$$p(z_{d,n}=k | z_{\neg(d,n)}, w_{d,n}=v, \alpha, \beta) \propto (N_{d,k}^{\neg n} + \alpha) \times \frac{N_{k,v}^{\neg (d,n)} + \beta}{N_k^{\neg (d,n)} + V\beta}$$

This formula provides the unnormalized probability for assigning the current word token $w_{d,n}$ to each topic $k$. We compute this value for all $K$ topics and then normalize them to form a valid probability distribution from which we sample the new topic assignment.

Dependencies in the Collapsed Gibbs sampling update for topic assignment $z_{d,n}=k$. The probability depends on counts related to the document ($N_{d,k}$) and counts related to the topic and word type ($N_{k,v}$, $N_k$), modulated by the priors ($\alpha$, $\beta$).

The Collapsed Gibbs Sampling Algorithm for LDA

The algorithm proceeds as follows:

1. Initialization:

   • Choose the number of topics $K$.
   • Set hyperparameters $\alpha$ and $\beta$ (often symmetric, e.g., $\alpha = 50/K$, $\beta = 0.01$).
   • For each word token $w_{d,n}$ in the corpus, randomly assign a topic $z_{d,n} \in \{1, ..., K\}$.
   • Initialize count matrices based on these random assignments:
     • $N_{d,k}$: number of words in document $d$ assigned to topic $k$.
     • $N_{k,v}$: number of times word $v$ is assigned to topic $k$ across all documents.
     • $N_k$: total number of words assigned to topic $k$ across all documents ($N_k = \sum_v N_{k,v}$).

2. Iteration (MCMC Sampling):

   • Repeat for a desired number of iterations (including burn-in and sampling phases):
     • For each document $d = 1 \dots M$:
       • For each word position $n = 1 \dots N_d$ (where $N_d$ is the length of document $d$):
         • Let $k_{old} = z_{d,n}$ be the current topic assignment and $v = w_{d,n}$ be the word type.
         • Decrement counts: Decrease the counts associated with the current assignment $(d, n, k_{old}, v)$: $N_{d,k_{old}} \leftarrow N_{d,k_{old}} - 1$ $N_{k_{old},v} \leftarrow N_{k_{old},v} - 1$ $N_{k_{old}} \leftarrow N_{k_{old}} - 1$ (These counts now represent $N^{\neg(d,n)}$ as required by the formula).
         • Calculate sampling probabilities: For each topic $k' = 1 \dots K$, compute the unnormalized probability $p'_{k'}$ using the formula: $$p'_{k'} = (N_{d,k'}^{\neg n} + \alpha) \times \frac{N_{k',v}^{\neg (d,n)} + \beta}{N_{k'}^{\neg (d,n)} + V\beta}$$ (Note: the counts used here are the decremented ones).
         • Normalize: Create a normalized probability distribution $P = [p_1, ..., p_K]$ where $p_{k'} = p'_{k'} / \sum_{j=1}^K p'_j$.
         • Sample new topic: Draw a new topic $k_{new}$ for word $(d,n)$ from the multinomial distribution $P$. $k_{new} \sim \text{Multinomial}(1, P)$
         • Update assignment: Set $z_{d,n} = k_{new}$.
         • Increment counts: Increase the counts associated with the new assignment $(d, n, k_{new}, v)$: $N_{d,k_{new}} \leftarrow N_{d,k_{new}} + 1$ $N_{k_{new},v} \leftarrow N_{k_{new},v} + 1$ $N_{k_{new}} \leftarrow N_{k_{new}} + 1$

3. Output:

   • Discard samples from the initial "burn-in" phase.
   • Use the counts from one or more samples after burn-in to estimate the model parameters.

Estimating Model Parameters ($\theta$ and $\phi$)

Although $\theta$ and $\phi$ were integrated out during sampling, we can estimate their posterior expectations using the final counts from the sampler (typically from the last iteration after sufficient burn-in, or averaged over several post-burn-in samples).

The expected document-topic distribution for document $d$ is:

$$\hat{\theta}_{d,k} = \frac{N_{d,k} + \alpha_k}{\sum_{k'=1}^K (N_{d,k'} + \alpha_{k'})}$$

The expected topic-word distribution for topic $k$ is:

$$\hat{\phi}_{k,v} = \frac{N_{k,v} + \beta_v}{\sum_{v'=1}^V (N_{k,v'} + \beta_{v'})}$$

These estimated distributions $\hat{\theta}$ and $\hat{\phi}$ represent the learned topic mixtures for each document and the word probabilities defining each topic, respectively.

Discussion

Collapsed Gibbs sampling is a widely used inference technique for LDA due to its relative simplicity and the fact that it uses the model's conjugacy properties effectively. By avoiding direct sampling of the continuous parameters, it can sometimes explore the posterior distribution of topic assignments $z$ more efficiently than a standard Gibbs sampler.

However, it is still an MCMC method. Convergence needs to be assessed (e.g., by monitoring the log-likelihood of the data or topic coherence metrics), and a suitable burn-in period is required. The sampler processes words sequentially, which can lead to slow mixing, especially on large datasets or with a high number of topics. Correlations between assignments can mean that many iterations are needed to obtain independent samples from the posterior.

Despite these potential limitations, Collapsed Gibbs sampling provides a solid baseline and a clear way to perform Bayesian inference for LDA. It contrasts with optimization-based approaches like Variational Bayes, which we will explore next, trading sampling variability for potentially faster convergence via deterministic approximation.