Inference for LDA: Variational Bayes

While Collapsed Gibbs Sampling provides a way to perform posterior inference in Latent Dirichlet Allocation (LDA), its iterative nature, processing one token assignment at a time, can be computationally intensive for large corpora. Variational Bayes (VB) offers a deterministic alternative that often scales more effectively. Instead of sampling, VB frames inference as an optimization problem.

The Variational Bayes Approach for LDA

The core challenge in LDA inference is computing the posterior distribution $p(\theta, z, \beta | w)$, where $\theta$ represents the document-topic distributions, $z$ the topic assignments for each word, $\beta$ the topic-word distributions, and $w$ the observed words. This posterior is intractable due to the complex dependencies between the variables.

Variational Bayes tackles this by introducing a simpler, approximate distribution, $q(\theta, z, \beta)$, chosen from a family of distributions $\mathcal{Q}$. The goal is to find the distribution $q \in \mathcal{Q}$ that is "closest" to the true posterior $p$. Closeness is typically measured using the Kullback-Leibler (KL) divergence, $KL(q || p)$. Minimizing this KL divergence is equivalent to maximizing a lower bound on the log marginal likelihood of the data, known as the Evidence Lower Bound (ELBO):

$$\mathcal{L}(q) = \mathbb{E}_q[\log p(\theta, z, \beta, w)] - \mathbb{E}_q[\log q(\theta, z, \beta)] \le \log p(w)$$

For LDA, a common choice for the variational family $\mathcal{Q}$ is the mean-field family, which assumes the latent variables are mutually independent in the approximate posterior:

$$q(\theta, z, \beta | \gamma, \phi, \lambda) = \prod_{d=1}^M q(\theta_d | \gamma_d) \prod_{k=1}^K q(\beta_k | \lambda_k) \prod_{n=1}^{N_d} q(z_{dn} | \phi_{dn})$$

Here:

• $q(\theta_d | \gamma_d)$ is a Dirichlet distribution approximating the posterior over the topic mixture for document $d$, parameterized by the variational Dirichlet parameter vector $\gamma_d = (\gamma_{d1}, ..., \gamma_{dK})$.
• $q(\beta_k | \lambda_k)$ is a Dirichlet distribution approximating the posterior over the word distribution for topic $k$, parameterized by the variational Dirichlet parameter vector $\lambda_k = (\lambda_{k1}, ..., \lambda_{kV})$, where $V$ is the vocabulary size.
• $q(z_{dn} | \phi_{dn})$ is a Multinomial distribution approximating the posterior over the topic assignment for the $n$-th word in document $d$, parameterized by the variational Multinomial parameter vector $\phi_{dn} = (\phi_{dn1}, ..., \phi_{dnK})$, where $\sum_k \phi_{dnk} = 1$.

The parameters $\gamma$, $\phi$, and $\lambda$ are the variational parameters we need to optimize to maximize the ELBO.

The mean-field approximation assumes independence between latent variables ($\theta$, $\beta$, $z$) in the variational distribution $q$, simplifying the intractable dependencies present in the true posterior $p$.

Coordinate Ascent Variational Inference (CAVI) for LDA

The ELBO can be maximized using an iterative coordinate ascent approach. We cycle through the variational parameters ($\phi$, $\gamma$, $\lambda$), updating each one while holding the others fixed. The update equations are derived by taking the derivative of the ELBO with respect to each parameter, setting it to zero, and solving. The resulting updates often have an intuitive form, resembling the updates in the Expectation-Maximization (EM) algorithm.

1. Update $\phi_{dnk}$: This update estimates the probability that word $w_{dn}$ belongs to topic $k$. It depends on the current estimates of the expected log probabilities of the topic proportions for document $d$ (related to $\gamma_d$) and the expected log probabilities of word $w_{dn}$ under topic $k$ (related to $\lambda_k$).

   $$\phi_{dnk} \propto \exp(\mathbb{E}_{q(\theta_d|\gamma_d)}[\log \theta_{dk}] + \mathbb{E}_{q(\beta_k|\lambda_k)}[\log \beta_{k, w_{dn}}])$$

   The expectations involve the digamma function $\Psi$, as $\mathbb{E}[\log \theta_{dk}] = \Psi(\gamma_{dk}) - \Psi(\sum_{j=1}^K \gamma_{dj})$. After calculating the right-hand side for all $k$, $\phi_{dn}$ is normalized so that $\sum_k \phi_{dnk} = 1$.

2. Update $\gamma_{dk}$: This update aggregates the responsibilities ($\phi_{dnk}$) assigned to topic $k$ by all words within document $d$, adding the original Dirichlet prior parameter $\alpha_k$. It represents the expected number of words in document $d$ assigned to topic $k$.

   $$\gamma_{dk} = \alpha_k + \sum_{n=1}^{N_d} \phi_{dnk}$$

3. Update $\lambda_{kv}$: This update aggregates the responsibilities ($\phi_{dnk}$) assigned to topic $k$ for all instances of word $v$ (where $w_{dn}=v$) across all documents, adding the original Dirichlet prior parameter $\eta_v$. It represents the expected number of times word $v$ was generated by topic $k$.

   $$\lambda_{kv} = \eta_v + \sum_{d=1}^M \sum_{n: w_{dn}=v} \phi_{dnk}$$

These updates are iterated until the ELBO converges, meaning the variational parameters stabilize.

Interpreting the Results

Once converged, the optimized variational parameters provide approximations to the posterior distributions:

• $\gamma^*$: Parameterizes the approximate posterior Dirichlet distribution $q(\theta_d | \gamma^*_d)$ for the topic mixture of each document $d$. The expected topic proportions are $\mathbb{E}[\theta_{dk}] = \gamma^*_{dk} / \sum_{j=1}^K \gamma^*_{dj}$.
• $\lambda^*$: Parameterizes the approximate posterior Dirichlet distribution $q(\beta_k | \lambda^*_k)$ for the word distribution of each topic $k$. The expected word probabilities are $\mathbb{E}[\beta_{kv}] = \lambda^*_{kv} / \sum_{v'=1}^V \lambda^*_{kv'}$.
• $\phi^*$: Represents the optimized posterior probability $q(z_{dn}=k | \phi^*_{dn})$ that word $w_{dn}$ belongs to topic $k$.

VB vs. Collapsed Gibbs Sampling

Feature	Variational Bayes (VB)	Collapsed Gibbs Sampling
Approach	Optimization (Maximize ELBO)	Sampling
Nature	Deterministic	Stochastic
Speed	Generally faster per iteration	Can be slower per iteration
Scalability	Better for large datasets	Can struggle with large data
Convergence	Converges to a local optimum of ELBO	Converges to true posterior (in limit)
Accuracy	Provides an approximation (biased)	Asymptotically exact
Implementation	Update equations can be complex	Simpler sampling
Parallelism	More amenable to certain parallelization	Can be parallelized (e.g., per document)

VB often converges much faster than Gibbs sampling, especially on large datasets. However, because it optimizes a lower bound based on an approximate distributional family (the mean-field assumption), it might converge to a solution that is a poorer approximation of the true posterior compared to a well-run Gibbs sampler. The mean-field assumption tends to underestimate the variance of the posterior distribution.

Scaling Further: Stochastic Variational Inference

For truly massive datasets, even standard batch VB can be too slow as it requires processing the entire corpus in each update step for $\lambda$. Stochastic Variational Inference (SVI) addresses this by using stochastic gradient ascent on the ELBO. It processes mini-batches of documents at each step, allowing for online learning and significant speedups, making VB applicable to web-scale datasets.

In summary, Variational Bayes provides a powerful and scalable inference framework for LDA. By framing inference as optimization and using a factorized approximation, it allows us to efficiently estimate topic models from large text collections, complementing the sampling-based approaches like Collapsed Gibbs Sampling.