ClusterCluster: Parallel Markov Chain Monte Carlo for Dirichlet Process Mixtures

Introduction

Bayesian nonparametric models are a remarkable class of stochastic objects that enable one to define infinite dimensional random variables that have tractable finite dimensional projections. This projective property often makes it possible to construct probabilistic models which can automatically balance simplicity and complexity in the posterior distribution. The Gaussian process (see, e.g., Adler and Taylor (2007); Rasmussen and Williams (2006)), Dirichlet process (Ferguson, 1973, 1974) and Indian buffet process (Griffiths and Ghahramani, 2006; Ghahramani et al., 2007) are the most common building blocks for Bayesian nonparametric models, and they have found uses in a wide variety of domains: natural language models (Teh et al., 2006), computer vision (Sudderth et al., 2006), activity modeling (Fox et al., 2009a), among many others.

Most commonly, Bayesian nonparametric models use the infinite dimensional construction to place priors on the latent parameters of the model, such as in Dirichlet process mixtures (Escobar and West, 1995; Rasmussen, 2000), Gaussian Cox processes (Møller et al., 1998; Adams et al., 2009a), and latent feature models (Fox et al., 2009b). This approach to priors for latent structure is appealing as the evidence for, e.g., a particular number of components in a mixture, is often weak and we wish to be maximally flexible in our specification of the model. Unfortunately, the use of Bayesian nonparametric priors for latent structure often yields models whose posterior distribution cannot be directly manipulated; indeed a density is often unavailable. In practice, it is therefore common to perform approximate inference using Markov chain Monte Carlo (MCMC), in which posterior computations are performed via Monte Carlo estimates from samples. These samples are obtained via a Markov chain that leaves the posterior distribution invariant. Remarkably, MCMC moves can be simulated on practical finite computers for many Bayesian nonparametric models, despite being infinite-dimensional, e.g., Rasmussen (2000); Neal (2000); Walker (2007); Papaspiliopoulos and Roberts (2008); Adams et al. (2009b). This property arises when finite data sets recruit only a finite projection of the underlying infinite object. Most practical Bayesian nonparametric models of interest are designed with this requirement in mind.

Markov chain Monte Carlo, however, brings with it frustrations. Chief among these is the perception that MCMC is computationally expensive and not scalable. This is conflated with the observation that the Markovian nature of such inference techniques necessarily require the computations to be sequential. In this paper, we challenge both of these conventional wisdoms for one of the most important classes of Bayesian nonparametric model, the Dirichlet process mixture (DPM). We take a novel approach to this problem that exploits invariance properties of the Dirichlet process to reparameterize the random measure in such a way that conditional independence is introduced between sets of atoms. These induced independencies, which are themselves inferred as part of the MCMC procedure, enable transition operators on different parts of the posterior to be simulated in parallel on different hardware, with minimal communication. Unlike previous parallelizing schemes such as Asuncion et al. (2008), our approach does not alter the prior or require an approximating target distribution. We find that this parallelism results in real-world gains as measured by several different metrics against wall-clock time.

The Dirichlet Process

The Dirichlet process defines a distribution over probability measures in terms of a base probability measure $H$ on a sample space $\Theta$ and a concentration parameter ${\alpha>0}$. In its most general form, a Dirichlet process is characterized by the property that any finite measurable partition of a $\Theta$ leads to a finite Dirichlet distribution over the associated probability measure. That is, if $\Theta$ is partitioned into ${A_{1},A_{2},\cdots,A_{M}}$, then the probability measure ${G(A_{1}),G(A_{2}),\cdots,G(A_{M})}$ has a finite Dirichlet distribution with parameters ${\alpha H(A_{1}),\alpha H(A_{2}),\cdots,\alpha H(A_{M})}$. As an alternative to this somewhat abstract definition, DP probability measures can also be constructed from a stick breaking process:

where it is clear from this construction that the $G$ are discrete with probability one. To achieve continuous density functions, the DP is often used as a part of an infinite mixture model:

where $F_{\theta}$ is a parametric family of component distributions and the base measure $H$ is now interpreted as a prior on this family. This Dirichlet process mixture model (DPM) is frequently used for model-based clustering, in which data belonging to a single $F_{\theta}$ are considered to form a group. The Dirichlet process allows for this model to possess an unbounded number of such clusters. This view leads to a related object called the Chinese restaurant process in which the $\pi_{j}$ are integrated out and one considers the infinitely exchangeable distribution over groupings alone.

Nesting Partitions in the Dirichlet Process

Then draw $K$ random distributions from $K$ independent Dirichlet processes with base measure $H$ and concentration parameters $\alpha\mu_{k}$. These are then mixed together with the $\gamma_{k}$:

Marginalizing out the sticks of each local Dirichlet process results naturally in a Chinese restaurant process with concentration parameter $\alpha\mu_{k}$. Interestingly, we can also integrate out the $\gamma_{k}$ to construct a two-stage Chinese restaurant variant. Each “customer” first chooses one of the $K$ “restaurants” according to its popularity:

This corresponds to the predictive distribution of the Dirichlet-multinomial over superclusters. In the second stage, the customer chooses a table $z_{n}$ at their chosen restaurant $k$ according to its popularity among other customers at that restaurant:

Markov Transition Operators for Parallelized Inference

We introduce some notation for convenience when referring to counts:

We use these to examine the prior over the $z_{n}$ and $s_{j}$:

Note that the terms cancel out so that the result is a marginal Chinese restaurant process multiplied by a multinomial over how the components are distributed over the superclusters.

This operation must be centralized, but is lightweight. Each supercluster communicates its number of clusters $J_{k}$ and these are used to sample from the conditional (assuming prior $p(\alpha)$):

This can be done with slice sampling or adaptive rejection sampling.

Updating base measure hyperparameters:

It is often the case in Bayesian hierarchical models that there are parameters governing the base measure $H$. These are typically hyperparameters that determine the priors on cluster-specific parameters $\theta$ and constrain the behavior of $F_{\theta}$. Updates to these parameters are performed in the reduce step, based on sufficient statistics transmitted from the map step.

These are model-specific updates that can be done in parallel, as each $\theta_{j}$ is only asked to explain data that belong to supercluster $s_{j}$.

This is typically the most expensive MCMC update for Dirichlet process mixtures: the hypothesis over cluster assignments must be modified for each datum. However, if only local components are considered, then this update can be parallelized. Moreover, as the reparameterization induces $K$ conditonally independent Dirichlet processes, standard DPM techniques, such as Neal (2000), Walker (2007), or Papaspiliopoulos and Roberts (2008) can be used per supercluster without modification. Data cannot move to components on different machines (in different superclusters), but can instantiate previously-unseen clusters within its local superclusters in the standard way. Note that the $\mu_{k}$ scaling causes these new components to be instantiated with the correct probability.

As data can only move between clusters that are local to their machine, i.e., within the same supercluster, it is necessary to move data between machines. One efficient strategy for this is to move entire clusters, along with their associated data to new superclusters. This is a centralized update, but it only requires communicating a set of data indices and one set of component parameters. The update itself is straightforward: Gibbs sampling according to the Dirichlet-multinomial, given the other assignments. In particular, we note that since $\theta_{j}$ moves with the cluster, the likelihood does not participate in the computation of the transition operator. We define $J_{k\backslash j}$ to be the number of extant clusters in supercluster $k$, ignoring cluster $j$. The conditional posterior update is then

Distributed implementation

The conditional independencies inherent in the Markov chain transition operators we have defined correspond naturally with an efficient, distributed implementation in terms of Map-Reduce (Dean and Ghemawat, 2008). Figure 3 describes the workflow. These operators, implemented as mappers and reducers, act on a distributed representation of the latent state, that is also based on the independencies in our auxiliary variable representation. Intuitively, each mapper performs MCMC updates on an independent clustering problem (within the supercluster it corresponds to), assuming fixed hyperparameters. The reduce step collects the latent state together and updates hyperparameters, while the shuffle step broadcasts the new hyperparameters and shuffles clusters amongst the superclusters.

Our system software implementation, described in Figure 5, is based on Python implementations, with modest optimization (in Cython) for the most compute-intensive inner loops. We use the Hadoop open source framework for Map-Reduce, and perform experiments using on-demand compute clusters from Amazon’s Elastic Compute Cloud. Typical configurations for our experiments involved 10-50 machines, each with 2-4 cores, stressing gains due to parallelism were possible despite significant inter-machine communication overhead.

We used a uniform prior over the superclusters, i.e., ${\mu_{k}=1/K}$. For initialization, we perform a small calibration run (on 1-10% of the data) using a serial implementation of MCMC inference, and use this to choose the initial concentration parameter $\alpha$. We then assign data to superclusters uniformly at random, and initialize the clustering via a draw from the prior using the local Chinese restaurant process. This is sufficient to roughly estimate (within an order of magnitude) the correct number of clusters, which supports efficient density estimation from the distributed implementation.

There is considerable room for further optimization. First, if we were pushing for the largest achievable scale, we would use a C++ implementation of the map, reduce and shuffle operations. Back-of-the-envelope suggestions suggest performance gains of 100x should be feasible with only minimal memory hierarchy optimization. Second, we would focus on use of a small number of many-core machines. Third, we would use a distributed framework such as HaLoophttps://code.google.com/p/haloop/ or MPI, which would permit simultaneous, asynchronous computation and communication. Fourth, it is known that tuning various parameters that control Hadoop can result in significant performance enhancements (Herodotou and Babu, 2011). For datasets larger than 100GB (roughly 50x larger than the datasets considered in this paper), we would want to distribute not just the latent state, but also the data itself, perhaps using the Hadoop File System (Shvachko, 2010). These advanced distributed implementation techniques are beyond the scope of this paper.

Williamson et al. (2013), concurrently with and independently of our previous preliminary work (Lovell et al., 2012), investigated a related parallel MCMC method based on the same auxiliary variable scheme. They focus on multi-core but single-machine implementations and on applications to admixtures based on the hierarchical Dirichlet process (HDP) (Teh et al., 2006). The compatibility of our transition operators with a Map-Reduce implementation enables us to analyze datasets with 100x more dimensions than those from Williamson, even in the presence of significant inter-machine communication overhead. We also rely purely on MCMC throughout, based on initialization from our prior, avoiding the need for heuristic initialization based on k-means. Combined, our approaches suggest many models based on the DP may admit principled parallel schemes and scale to significantly larger problems than are typical in the literature.

Intuitively, this scheme resembles running ”restricted Gibbs” scans over subsets of the clusters, then shuffling clusters amongst subsets, which one might expect to yield slower convergence to equilibrium than full Gibbs scans. Our auxiliary variable representation shows this can be interpreted in terms of exact transition operations for a DPM.

Experiments

We explored the accuracy of our prototype distributed implementation on several synthetic and real-world datasets. Our synthetic data was drawn from a balanced finite mixture model. Each mixture component $\theta_{j}$ was parameterized by a set of coin weights drawn from a $\text{Beta}(\beta_{d},\beta_{d})$ distribution, where $\{\beta_{d}\}$ is a set of cluster component hyperparameters, one per dimension $d$ of the data. The binary data were Bernoulli draws based on the weight parameters $\{\theta_{j}\}$ of their respective clusters. Our implementation collapsed out the coin weights and updated each $\beta_{d}$ during the reduce step using a Griddy Gibbs (Ritter and Tanner, 1992) kernel.

Figure 5 shows results supporting the accuracy of our inference scheme as a density estimator given high-dimensional datasets with large numbers of clusters. We see reliable convergence to predictive probabilities close to the true entropy of the generating mixture.

Figure 6 shows the convergence behavior of our sampler: predictive densities (and joint probabilities) typically asymptote quickly, while latent structure estimates (such as the number of clusters, or the concentration parameter) converge far more slowly. As the Dirichlet process is known to not result in consistent estimates of number of clusters for finite mixture datasets (Miller and Harrison, 2013) (which have no support under the DP prior), it is perhaps not surprising that predictive likelihood converges more quickly than estimates of number of clusters – especially given our auxiliary variable representation, which may encourage representations with multiple predictively equivalent clusters. It would be interesting to characterize the regimes where these dynamics occur, and to determine whether they are also present for approximate parallelization schemes or variational algorithms based on our auxiliary variable representation. Figure 9 shows a typical pattern of efficiencies for parallel computation: for a given problem size, speed increases until some saturation point is reached, after which additional compute nodes slow down the computation. In future work we will explore means of separating the components of this tradeoff due to communication costs, initialization, and any components due to convergence slowdown.

Finally, in Figures 9 and 10, we show a representative run on a 1MM vector subset of the Tiny Images (Torralba, 2008) dataset, where we use the Dirichlet process mixture to perform vector quantization. The input data are binary features obtained by running a randomized approximation to PCA on 100,000 rows and thresholding the top 256 principal components into binary values at their component-wise median. After one day of elapsed time and 32 CPU-days of computation, the sampler is still making significant progress compressing the data, and has converged to the vicinity of 3000 clusters. Serial MCMC (not shown) is intractable on this problem.

Conclusion

We have introduced an auxiliary variable representation of the Dirichlet process and applied it to mixture models, where it yields superclusters that cluster the clusters. We have shown how this representation enables an exact parallelization of the standard MCMC schemes for inference, where the DP learns how to parallelize itself, despite the lack of conditional independencies in the traditional form of the model. We have also shown that this representation naturally meshes with the Map-Reduce approach to distributed computation on a large compute cluster, and developed a prototype distributed implementation atop Amazon’s Elastic Compute Cloud, tested on over 50 machines and 100 cores. We have explored its performance on synthetic and real-world density estimation problems, including runs on over 1MM row, 256-dimensional data sets.

These results point to a potential path forward for “big data” applications of Bayesian statistics for models that otherwise lack apparent conditional independencies. We suspect searching for auxiliary variable representations that induce independencies may lead to new ways to scale up a range of nonparametric Bayesian models, and may perhaps also lead to further correspondences with established distributed computing paradigms for MCMC inference. We hope our results present a useful step in this direction.

This work was partially supported by DARPA XDATA contract FA8750-12-C-0315. RPA was funded in part by DARPA Young Faculty Award N66001-12-1-4219. VKM was partially supported by gifts from Google and Analog Devices.

References