Collapsed Variational Dirichlet Process Mixture Models

1 Introduction

Mixture modeling remains one of the most useful tools in statistics, machine learning and data mining for applications involving density estimation or clustering. One of the most prominent recent developments in this field is the application of nonparametric Bayesian techniques to mixture modeling, which allow for the automatic determination of an appropriate number of mixture components. Current inference algorithms for such models are mostly based on Gibbs sampling, which suffer from a number of drawbacks. Most importantly, Gibbs sampling is not efficient enough to scale up to the large scale problems we face in modern-day data mining. Secondly, sampling requires careful monitoring of the convergence of the Markov chain, both to decide on the number of samples to be ignored for burn-in and to decide how many samples are needed to reduce the variance in the estimates. These considerations have lead researchers to develop deterministic alternatives which trade off variance with bias and are easily monitored in terms of their convergence. Moreover, they can be orders of magnitude faster than sampling, especially when special data structures such as KD trees are used to cache certain sufficient statistics [Moore, 1998; Verbeek et al., 2003; Kurihara et al., 2006].

∗This material is based in part upon work supported by the National Science Foundation under Grant Number D/HS-0535278. YWT thanks the Lee Kuan Yew Endowment Fund for funding.

[Blei and Jordan, 2005] recently applied the framework of variational Bayesian (VB) inference to Dirichlet process (DP) mixture models and demonstrated significant computational gains. Their model was formulated entirely in the truncated stick-breaking representation. The choice of this representation has both advantages and disadvantages. For instance, it is very easy to generalize beyond the DP prior and use much more flexible priors in this representation. On the flip side, the model is formulated in the space of explicit, nonexchangeable cluster labels (instead of partitions). In other words, randomly permuting the labels changes the probability of the data. This then requires samplers to mix over cluster labels to avoid bias [Porteous et al., 2006].

In this paper we propose and study alternative approaches to VB inference in DP mixture models beyond that proposed in [Blei and Jordan, 2005]. There are three distinct contributions in this paper: in proposing an improved VB algorithm based on integrating out mixture weights, in comparing the stick-breaking representation against the finite symmetric Dirichlet approximation to the DP, and in the maintaining optimal ordering of cluster labels in the stick-breaking VB algorithms. These lead to a total of six different algorithms, including the one proposed in [Blei and Jordan, 2005]. We experimentally evaluate these six algorithms and compare against Gibbs sampling.

In Section 2.1 we explore both the truncated stick-breaking approximation and the finite symmetric Dirichlet prior as finite dimensional approximations to the DP. As opposed to the truncated stick-breaking approximation, the finite symmetric Dirichlet model is exchangeable over cluster labels. Theoretically this has important consequences, for example a Gibbs sampler is not required to mix over cluster labels if we are computing averages over quantities invariant to cluster label permutations (as is typically the case).

In Section 2.2 we explore the idea of integrating out the mixture weights π, hence collapsing the model to a lower dimensional space. This idea has been shown to work well for LDA models [Teh et al., 2006] where strong dependencies exist between model parameters and assignment variables. Such dependencies exist between mixture weights and assignment variables in our mixture model context as well, thus collapsing could also be important here. This intuition is reflected in the observation that the variational bound on the log evidence is guaranteed to improve.

In Section 3 we derive the VB update equations corresponding to the approximations in Section 2. We also consider optimally reordering cluster labels in the stick-breaking VB algorithms. As mentioned, the ordering of the cluster labels is important for models formulated in the stick-breaking representation. In the paper [Blei and Jordan, 2005] this issue was ignored. Here we also study the effect of cluster reordering on relevant performance measures such as the predictive log evidence.

The above considerations lead us to six VB inference methods, which we evaluate in Section 4. The methods are: 1) the truncated stick-breaking representation with standard VB (TSB), 2) the truncated stick-breaking representation with collapsed VB (CTSB), 3) the finite symmetric Dirichlet representation with standard VB (FSD), 4) the finite symmetric Dirichlet presentation with collapsed VB (CFSD), and 5) and 6) being TSB and CTSB with optimal reordering (O-TSB and O-CTSB respectively).

2 Four Approximations to the DP

We describe four approximations to the DP in this section. These four approximations are obtained by a combination of truncated stick-breaking/finite symmetric Dirichlet approximations and whether the mixture weights are marginalized out or not. Based on these approximations we describe the six VB inference algorithms in the next section.

The most natural representation of DPs is using the Chinese restaurant process, which is formulated in the space of partitions. Partitions are groupings of the data independent of cluster labels, where each data-point is assigned to exactly 1 group. This space of partitions turns out to be problematic for VB inference, where we wish to use fully factorized variational distributions on the assignment variables, Q(z) = ∏n q(zn). Since the assignments z1 = 1, z2 = 1, z3 = 2 represent the same partition (1, 2)(3) as z1 = 3, z2 = 3, z3 = 2, there are intricate dependencies between the assignment variables and it does not make sense to use the factorization above. We can circumvent this by using finite dimensional approximations for the DP, which are formulated in the space of cluster labels (not partitions) and which are known to closely approximate the DP prior as the number of explicitly maintained clusters grows [Ishwaran and James, 2001; Ishwaran and Zarepour, 2002]. These finite approximations are what will we discuss next.

2.1 TSB and FSDApproximations

In the first approximation we use the stick-breaking representation for the DP [Ishwaran and James, 2001] and truncate it after T terms,

where B(v; 1, α) is a beta density for variable v with parameters 1 and α, and one can verify that ∑Ti=1 πi = 1. Incorporating this into a joint probability over data items X = {xn}, n = 1, ...,N, cluster assignments z = {zn}, n = 1, ...,N, stick-breaking weights v = {vi}, i = 1, ..., T and cluster parameters η = {ηi}, i = 1, ..., T we find

where π(v) are the mixture weights as defined in (3). In this representation the cluster labels are not interchangeable, i.e. changing labels will change the probability value in (5). Note also that as T → ∞ the approximation becomes exact.

A second approach to approximate the DP is by assuming a finite (but large) number of clusters, K, and using a symmetric Dirichlet prior D on π [Ishwaran and Zarepour, 2002],

This results in the joint model,

The essential difference with the stick-breaking representation is that the cluster labels remain interchangeable under this representation, i.e. changing cluster labels does not change the probability [Porteous et al., 2006]. The limit K → ∞ is somewhat tricky because in the transition K → ∞ we switch to the space of partitions, where states that result from cluster relabelings are mapped to the same partition. For example, both z1 = 1, z2 = 1, z3 = 2 and z1 = 3, z2 = 3, z3 = 2 are mapped to the same partition (1, 2)(3).

In figure 1 we show the prior average cluster sizes under the truncated stick-breaking (TSB) representation (left) and under the finite symmetric Dirichlet (FSD) prior (middle) for two values of the truncation level and number of clusters respectively. From this figure it is apparent that the cluster labels in the TSB prior are not interchangeable (the probabilities are ordered in decreasing size), while they are interchangeable for the FSD prior. As we increase T and K these priors approximate the DP prior with increasing accuracy.

One should note however, that they live in different spaces. The DP itself is most naturally defined in the space of partitions, while both TSB and FSD are defined in the space over cluster labels. However, TSB and FSD also live in different spaces! More precisely, one can transform a sample from the FSD prior into the stick-breaking representation by performing a size-biased permutation of the mixture weights π (i.e. after every sample from D(π) we sample an ordering according to π without replacement). As it turns out, for finite K this does not exactly recover the left hand figure in 1, but rather samples from a prior very closely related to it shown in the right pane of figure 1. This prior is given by a stick-breaking construction as in eqn.(3) with stick-lengths sampled from,

Conversely, we can obtain samples from the FSD prior by applying a random, uniformly distributed permutation on the cluster weights obtained from eqn.(8). Although these two stick-breaking constructions are slightly different, for large enough K, T they are very similar and we do not expect any difference in terms of performance between the two.

2.2 Marginalizing out the Mixture Weights

The variational Bayesian approximations discussed in the next section assume a factorized form for the posterior distribution. This means that we assume that parameters are independent of assignment variables. This is clearly a very bad assumption because changes in π will have a considerable impact on z. Ideally, we would integrate out all the parameters, but this is too computationally expensive. There is however a middle ground: we can marginalize out π from both methods without computational penalty if we make another approximation which will be discussed in section 3.3. For both TSB and FSD representations the joint collapsed model over X, z, η is given by,

with different distributions over cluster labels p(z) in both cases. For the TSB representation we have,

with

and N≥i = Ni + N>i. For FSD we find instead,

3 Variational Bayesian Inference

The variational Bayesian inference algorithm [Attias, 2000; Ghahramani and Beal, 2000] lower bounds the log marginal likelihood by assuming that parameters and hidden variables are independent. The lower bound is given by,

where θ is either {η, v}, {η,π} or {η} in the various DP approximations discussed in the previous section. Approximate inference is then achieved by alternating optimization of this bound over Q(z) and Q(θ). In the following we will spell out the details of VB inference for the proposed four methods. For the TSB prior we use,

where q(v) is not used in the TSB model with v marginalized out. For the FSD prior we use,

As well, q(π) is left out for the collapsed version.

3.1 Bounds on the Evidence

Given the variational posteriors we can construct bounds on the log marginal likelihood by inserting Q into eqn.(13). After some algebra we find the following general form,

where the "extra term" depends on the particular method. For the TSB prior we have,

On the other hand for the FSD prior we find,

For both collapsed versions these expressions are replaced by,

3.2 VB Update Equations

Given these bounds it is now not hard to derive update equations for the various methods. Due to space constraints we will refer to the papers [Blei and Jordan, 2005; Ghahramani and Beal, 2000; Penny, 2001; Yu et al., 2005] for more details on the update equations for the un-collapsed methods and focus on the novel collapsed update equations.

Below we will provide the general form of the update equations where we do not assume anything about the particular form of the prior p(ηi). The equations become particularly simple when we choose this prior in the conjugate exponential family. Explicit update equations for q(ηi) can be found in the papers [Ghahramani and Beal, 2000; Blei and Jordan, 2005; Penny, 2001; Yu et al., 2005].

For q(ηi) we find the same update for both methods,

while for q(zn) we find the update

where the conditional p(zn|z¬n) is different for the FSD and TSB priors. For the TSB prior we use (10), giving the conditional

where N ¬n i = Ni - I(zn = i), N ¬n >i = N>i - I(zn > i) are the corresponding counts with zn removed. In contrast, for the FSD prior we have,

3.3 Gaussian Approximation

The expectation required to compute the update (21) seems intractable due to the exponentially large space of all assignments for z. It can in fact be computed in polynomial time using convolutions, however this solution still tended to be too slow to be practical. A much more efficient approximate solution is to observe that both random variables Ni and N>i are sums over Bernoulli variables: Ni = ∑n I(zn = i) and N>i = ∑n I(zn > i). Using the central limit theorem these sums are expected to be closely approximated by Gaussian distributions with means and variances given by,

To apply this approximation to the computation of the average in (21), we use the following second order Taylor expansion,

This approximation has been observed to work extremely well in practice, even for small values of m.

3.4 Optimal Cluster Label Reordering

As discussed in section 2.1 the stick-breaking prior assumes a certain ordering of the clusters (more precisely, a size-biased ordering). Since a permutation of the cluster labels changes the probability of the data, we should choose the optimal permutation resulting in the highest probability for the data. The optimal relabelling of the clusters is given by the one that orders the cluster sizes in decreasing order (this is true since the average prior cluster sizes are also ordered). In our experiments we assess the effect of reordering by introducing algorithms O-TSB and O-CTSB which always maintain this optimal labelling of the clusters. Note that optimal ordering was not maintained in [Blei and Jordan, 2005].

4 Experiments

In the following experiments we compared the six algorithms discussed in the main text in terms of their log-probability on held out test data. The probability for a test point, xt, is then given by,

where the expectation E[p(zt|ztrain)]q(ztrain) is computed using the techniques introduced in section 3.3. All experiments were conducted using Gaussian mixtures with vague priors on the parameters.

In the first experiment we generated synthetic data from a mixture of 10 Gaussians in 16 dimensions with a separation coefficient 1 c = 2. We studied the accuracy of each algorithm as a function of the number of data cases and the truncation level of the approximation. In figures 2 and 3 we show the results as we vary N (keeping T and K fixed at 30) while in figures 4 and 5 we plot the results as we vary T and K (keeping N fixed at 200). We plot both the absolute value of the log probability of test data and the value relative to a Gibbs sampler (GS). We 50 iterations for burn-in, and run another 200 iterations for inference. Error bars are computed on the relative values in order to subtract variance caused by the different splits (i.e. we measure variance on paired experiments).

Figure 5: Relative average log probability per data-point for test data as a function of T (for TSB methods) or K (for FSD methods).

1 Following [Dasgupta, 1999], a Gaussian mixture is c-separated if for each pair (i, j) of components we have ||mi - mj||2 ≥ c2D max(λmax i , λmax j ) , where λmax denotes the maximum eigenvalue of their covariance.

Figure 6: 15 most populated clusters found by the various algorithms in descending order of E[Ni]. Algorithms were trained on a random subset of 10,000 images from MNIST and dimensionality reduced to 50 dimensions using PCA. Log probability of 10,000 test images are given by, L = -574.05 ± 0.52 (TSB), L = -574.03 ± 0.53 (O-TSB), L = -573.90 ± 0.54 (CTSB), L = -573.89 ± 0.54 (O-CTSB), L = -574.06 ± 0.50 (FSD), and L = -573.89 ± 0.51 (CFSD). Standard error over differences relative to O-CTSB are given by: dL = -0.17 ± 0.13 (TSB), dL = -0.14 ± 0.11 (O-TSB), dL = -0.01 ± 0.05 (CTSB), dL = -0.17 ± 0.13 (FSD), and dL = -0.00 ± 0.10 (CFSD).

Results were averaged over 30 independently sampled training/testing datasets, where the number of test instances was always fixed at 1000. 2

In the second experiment we have run the algorithms on subsets of MNIST. Images of size 28 × 28 were dimensionality reduced to 50 PCA dimensions as a preprocessing step. We trained all algorithms on 30 splits of the data, each split containing 5000 data-cases for training and 10,000 data-cases for testing. Truncation levels were set to 80 for all algorithms. Unfortunately, the dataset was too large to obtain results with Gibbs sampling. All algorithms typically find between 32 and 36 clusters. Results are shown in figure 4.

4.1 Discussion

In this paper we explored six different ways to perform variational Bayesian inference in DP mixture models. Besides an empirical study of these algorithms our contribution has been to introduce a new family of collapsed variational algorithms where the mixture weights are marginalized out. To make these algorithms efficient, we used the central limit theorem to approximate the required averages.

We can draw three conclusions from our study. Firstly, there is very little difference between variational Bayesian inference in the reordered stick-breaking representation and the finite mixture model with symmetric Dirichlet priors. Secondly, label reordering is important for the stick-breaking representation. Thirdly, variational approximations are much more efficient computationally than Gibbs sampling, with almost no loss in accuracy.

2 For N=200 and T or K=40, TSB, O-TSB, CTSB, O-CTSB, FSD, CFSD and GS took 0.48, 0.70, 0.84, 1.22, 0.67, 1.16 and 1,019 seconds on average, respectively. Note that the computational complexities of the variational algorithms are the same.

We are currently working towards models where the parameters η are marginalized out as well. We expect this to have a more significant impact on test accuracy than the current setup which only marginalizes over π, especially when clusters are overlapping. Unfortunately, it seems this will come at the cost of increased computation.

Collapsed variational inference has also been applied to LDA models [Teh et al., 2006], where preliminary results indicate significant performance improvement. We are currently also exploring collapsed variational inference for hierarchical DP models [Teh et al., 2004].