Online Learning for Latent Dirichlet Allocation

Matthew D. Hoffman
Department of Computer Science
Princeton University
Princeton, NJ
mdhoffma@cs.princeton.edu

David M. Blei
Department of Computer Science
Princeton University
Princeton, NJ
blei@cs.princeton.edu

Francis Bach
INRIA—Ecole Normale Supérieure
Paris, France
francis.bach@ens.fr

Abstract

We develop an online variational Bayes (VB) algorithm for Latent Dirichlet Allocation (LDA). Online LDA is based on online stochastic optimization with a natural gradient step, which we show converges to a local optimum of the VB objective function. It can handily analyze massive document collections, including those arriving in a stream. We study the performance of online LDA in several ways, including by fitting a 100-topic topic model to 3.3M articles from Wikipedia in a single pass. We demonstrate that online LDA finds topic models as good or better than those found with batch VB, and in a fraction of the time.

1 Introduction

Hierarchical Bayesian modeling has become a mainstay in machine learning and applied statistics. Bayesian models provide a natural way to encode assumptions about observed data, and analysis proceeds by examining the posterior distribution of model parameters and latent variables conditioned on a set of observations. For example, research in probabilistic topic modeling—the application we will focus on in this paper—revolves around fitting complex hierarchical Bayesian models to large collections of documents. In a topic model, the posterior distribution reveals latent semantic structure that can be used for many applications.

For topic models and many other Bayesian models of interest, however, the posterior is intractable to compute and researchers must appeal to approximate posterior inference. Modern approximate posterior inference algorithms fall in two categories—sampling approaches and optimization approaches. Sampling approaches are usually based on Markov Chain Monte Carlo (MCMC) sampling, where a Markov chain is defined whose stationary distribution is the posterior of interest. Optimization approaches are usually based on variational inference, which is called variational Bayes (VB) when used in a Bayesian hierarchical model. Whereas MCMC methods seek to generate independent samples from the posterior, VB optimizes a simplified parametric distribution to be close in Kullback-Leibler divergence to the posterior. Although the choice of approximate posterior introduces bias, VB is empirically shown to be faster than and as accurate as MCMC, which makes it an attractive option when applying Bayesian models to large datasets [1, 2, 3].

Nonetheless, large scale data analysis with VB can be computationally difficult. Standard “batch” VB algorithms iterate between analyzing each observation and updating dataset-wide variational parameters. The per-iteration cost of batch algorithms can quickly become impractical for very large datasets. In topic modeling applications, this issue is particularly relevant—topic modeling promises

to summarize the latent structure of massive document collections that cannot be annotated by hand. A central research problem for topic modeling is to efficiently fit models to larger corpora [4, 5].

To this end, we develop an online variational Bayes algorithm for latent Dirichlet allocation (LDA), one of the simplest topic models and one on which many others are based. Our algorithm is based on online stochastic optimization, which has been shown to produce good parameter estimates dramatically faster than batch algorithms on large datasets [6]. Online LDA handily analyzes massive collections of documents and, moreover, online LDA need not locally store or collect the documents—each can arrive in a stream and be discarded after one look.

In the subsequent sections, we derive online LDA and show that it converges to a stationary point of the variational objective function. We study the performance of online LDA in several ways, including by fitting a topic model to 3.3M articles from Wikipedia without looking at the same article twice. We show that online LDA finds topic models as good as or better than those found with batch VB, and in a fraction of the time (see figure 1). Online variational Bayes is a practical new method for estimating the posterior of complex hierarchical Bayesian models.

2 Online variational Bayes for latent Dirichlet allocation

Latent Dirichlet Allocation (LDA) [7] is a Bayesian probabilistic model of text documents. It assumes a collection of $K$ “topics.” Each topic defines a multinomial distribution over the vocabulary and is assumed to have been drawn from a Dirichlet, $\beta_k \sim \text{Dirichlet}(\eta)$. Given the topics, LDA assumes the following generative process for each document $d$. First, draw a distribution over topics $\theta_d \sim \text{Dirichlet}(\alpha)$. Then, for each word $i$ in the document, draw a topic index $z_{di} \in {1, \ldots, K}$ from the topic weights $z_{di} \sim \theta_d$ and draw the observed word $w_{di}$ from the selected topic, $w_{di} \sim \beta_{z_{di}}$. For simplicity, we assume symmetric priors on $\theta$ and $\beta$, but this assumption is easy to relax [8].

Note that if we sum over the topic assignments $z$, then we get $p(w_{di} | \theta_d, \beta) = \sum_k \theta_{dk} \beta_{kw}$. This leads to the “multinomial PCA” interpretation of LDA; we can think of LDA as a probabilistic factorization of the matrix of word counts $n$ (where $n_{dw}$ is the number of times word $w$ appears in document $d$) into a matrix of topic weights $\theta$ and a dictionary of topics $\beta$ [9]. Our work can thus...

be seen as an extension of online matrix factorization techniques that optimize squared error [10] to more general probabilistic formulations.

We can analyze a corpus of documents with LDA by examining the posterior distribution of the topics $\beta$, topic proportions $\theta$, and topic assignments $z$ conditioned on the documents. This reveals latent structure in the collection that can be used for prediction or data exploration. This posterior cannot be computed directly [7], and is usually approximated using Markov Chain Monte Carlo (MCMC) methods or variational inference. Both classes of methods are effective, but both present significant computational challenges in the face of massive data sets. Developing scalable approximate inference methods for topic models is an active area of research [3, 4, 5, 11].

To this end, we develop online variational inference for LDA, an approximate posterior inference algorithm that can analyze massive collections of documents. We first review the traditional variational Bayes algorithm for LDA and its objective function, then present our online method, and show that it converges to a stationary point of the same objective function.

2.1 Batch variational Bayes for LDA

In Variational Bayesian inference (VB) the true posterior is approximated by a simpler distribution $q(z, \theta, \beta)$, which is indexed by a set of free parameters [12, 13]. These parameters are optimized to maximize the Evidence Lower BOund (ELBO):

$$\log p(w | \alpha, \eta) \geq L(w, \phi, \gamma, \lambda) \triangleq E_q[\log p(w, z, \theta, \beta | \alpha, \eta)] - E_q[\log q(z, \theta, \beta)].$$

Maximizing the ELBO is equivalent to minimizing the KL divergence between $q(z, \theta, \beta)$ and the posterior $p(z, \theta, \beta | w, \alpha, \eta)$. Following [7], we choose a fully factorized distribution $q$ of the form

$$q(z_{di} = k) = \phi_{dwik}; \quad q(\theta_d) = \text{Dirichlet}(\theta_d; \gamma_d); \quad q(\beta_k) = \text{Dirichlet}(\beta_k; \lambda_k),$$

The posterior over the per-word topic assignments $z$ is parameterized by $\phi$, the posterior over the per-document topic weights $\theta$ is parameterized by $\theta$, and the posterior over the topics $\beta$ is parameterized by $\lambda$. As a shorthand, we refer to $\lambda$ as “the topics.” Equation 1 factorizes to

$$L(w, \phi, \gamma, \lambda) = \sum_d \left{ E_q[\log p(w_d | \theta_d, z_d, \beta)] + E_q[\log p(z_d | \theta_d)] - E_q[\log q(z_d)] + E_q[\log p(\theta_d | \alpha)] - E_q[\log q(\theta_d)] + (E_q[\log p(\beta | \eta)] - E_q[\log q(\beta)]) / D \right}.$$

Notice we have brought the per-corpus terms into the summation over documents, and divided them by the number of documents $D$. This step will help us to derive an online inference algorithm.

We now expand the expectations above to be functions of the variational parameters. This reveals that the variational objective relies only on $n_{dw}$, the number of times word $w$ appears in document $d$. When using VB—as opposed to MCMC—documents can be summarized by their word counts,

$$L = \sum_d \sum_w n_{dw} \sum_k \phi_{dwk} (E_q[\log \theta_{dk}] + E_q[\log \beta_{kw}] - \log \phi_{dwk}) - \log \Gamma(\sum_k \gamma_{dk}) + \sum_k (\alpha - \gamma_{dk}) E_q[\log \theta_{dk}] + \log \Gamma(\gamma_{dk}) + (\sum_k - \log \Gamma(\sum_w \lambda_{kw}) + \sum_w (\eta - \lambda_{kw}) E_q[\log \beta_{kw}] + \log \Gamma(\lambda_{kw})) / D + \log \Gamma(K \alpha) - K \log \Gamma(\alpha) + (\log \Gamma(W \eta) - W \log \Gamma(\eta)) / D$$

$$\triangleq \sum_d \ell(n_d, \phi_d, \gamma_d, \lambda),$$

where $W$ is the size of the vocabulary and $D$ is the number of documents. $\ell(n_d, \phi_d, \gamma_d, \lambda)$ denotes the contribution of document $d$ to the ELBO.

$L$ can be optimized using coordinate ascent over the variational parameters $\phi, \gamma, \lambda$ [7]:

$$\phi_{dwk} \propto \exp{E_q[\log \theta_{dk}] + E_q[\log \beta_{kw}]}; \quad \gamma_{dk} = \alpha + \sum_w n_{dw} \phi_{dwk}; \quad \lambda_{kw} = \eta + \sum_d n_{dw} \phi_{dwk}.$$

The expectations under $q$ of $\log \theta$ and $\log \beta$ are

$$E_q[\log \theta_{dk}] = \Psi(\gamma_{dk}) - \Psi(\sum_{i=1}^K \gamma_{di}); \quad E_q[\log \beta_{kw}] = \Psi(\lambda_{kw}) - \Psi(\sum_{i=1}^W \lambda_{ki}),$$

where $\Psi$ denotes the digamma function (the first derivative of the logarithm of the gamma function).

The updates in equation 5 are guaranteed to converge to a stationary point of the ELBO. By analogy to the Expectation-Maximization (EM) algorithm [14], we can partition these updates into an “E” step—iteratively updating $\gamma$ and $\phi$ until convergence, holding $\lambda$ fixed—and an “M” step—updating $\lambda$ given $\phi$. In practice, this algorithm converges to a better solution if we reinitialize $\gamma$ and $\phi$ before each E step. Algorithm 1 outlines batch VB for LDA.

Algorithm 1 Batch variational Bayes for LDA

Initialize $\lambda$ randomly.

while relative improvement in $L(w, \phi, \gamma, \lambda) > 0.00001$ do

E step: for $d = 1$ to $D$ do Initialize $\gamma_{dk} = 1$. (The constant 1 is arbitrary.) repeat Set $\phi_{dwk} \propto \exp{E_q[\log \theta_{dk}] + E_q[\log \beta_{kw}]}$ Set $\gamma_{dk} = \alpha + \sum_w \phi_{dwk} n_{dw}$ until $\frac{1}{K} \sum_k |\text{change in } \gamma_{dk}| < 0.00001$ end for

M step: Set $\lambda_{kw} = \eta + \sum_a n_{dw} \phi_{dwk}$ end while

2.2 Online variational inference for LDA

Algorithm 1 has constant memory requirements and empirically converges faster than batch collapsed Gibbs sampling [3]. However, it still requires a full pass through the entire corpus each iteration. It can therefore be slow to apply to very large datasets, and is not naturally suited to settings where new data is constantly arriving. We propose an online variational inference algorithm for fitting $\lambda$, the parameters to the variational posterior over the topic distributions $\beta$. Our algorithm is nearly as simple as the batch VB algorithm, but converges much faster for large datasets.

A good setting of the topics $\lambda$ is one for which the ELBO $L$ is as high as possible after fitting the per-document variational parameters $\gamma$ and $\phi$ with the E step defined in algorithm 1. Let $\gamma(n_d, \lambda)$ and $\phi(n_d, \lambda)$ be the values of $\gamma_d$ and $\phi_d$ produced by the E step. Our goal is to set $\lambda$ to maximize

$$L(n, \lambda) \triangleq \sum_d \ell(n_d, \gamma(n_d, \lambda), \phi(n_d, \lambda), \lambda),$$

where $\ell(n_d, \gamma_d, \phi_d, \lambda)$ is the $d$th document’s contribution to the variational bound in equation 4. This is analogous to the goal of least-squares matrix factorization, although the ELBO for LDA is less convenient to work with than a simple squared loss function such as the one in [10].

Online VB for LDA (“online LDA”) is described in algorithm 2. As the $t$th vector of word counts $n_t$ is observed, we perform an E step to find locally optimal values of $\gamma_t$ and $\phi_t$, holding $\lambda$ fixed. We then compute $\tilde{\lambda}$