Streaming, Distributed Variational Inference for Bayesian Nonparametrics

Introduction

Bayesian nonparametric (BNP) stochastic processes are streaming priors – their unique feature is that they specify, in a probabilistic sense, that the complexity of a latent model should grow as the amount of observed data increases. This property captures common sense in many data analysis problems – for example, one would expect to encounter far more topics in a document corpus after reading $10^{6}$ documents than after reading $10$ – and becomes crucial in settings with unbounded, persistent streams of data. While their fixed, parametric cousins can be used to infer model complexity for datasets with known magnitude a priori , such priors are silent with respect to notions of model complexity growth in streaming data settings.

Bayesian nonparametrics are also naturally suited to parallelization of data processing, due to the exchangeability, and thus conditional independence, they often exhibit via de Finetti’s theorem. For example, labels from the Chinese Restaurant process are rendered i.i.d. by conditioning on the underlying Dirichlet process (DP) random measure, and feature assignments from the Indian Buffet process are rendered i.i.d. by conditioning on the underlying beta process (BP) random measure.

Given these properties, one might expect there to be a wealth of inference algorithms for BNPs that address the challenges associated with parallelization and streaming. However, previous work has only addressed these two settings in concert for parametric models , and only recently has each been addressed individually for BNPs. In the streaming setting, and developed streaming inference for DP mixture models using sequential variational approximation. Stochastic variational inference and related methods are often considered streaming algorithms, but their performance depends on the choice of a learning rate and on the dataset having known, fixed size a priori . Outside of variational approaches, which are the focus of the present paper, there exist exact parallelized MCMC methods for BNPs ; the tradeoff in using such methods is that they provide samples from the posterior rather than the distribution itself, and results regarding assessing convergence remain limited. Sequential particle filters for inference have also been developed , but these suffer issues with particle degeneracy and exponential forgetting.

The main challenge posed by the streaming, distributed setting for BNPs is the combinatorial problem of component identification. Most BNP models contain some notion of a countably infinite set of latent “components” (e.g. clusters in a DP mixture model), and do not impose an inherent ordering on the components. Thus, in order to combine information about the components from multiple processors, the correspondence between components must first be found. Brute force search is intractable even for moderately sized models – there are $K_{1}+K_{2}\choose K_{1}$ possible correspondences for two sets of components of sizes $K_{1}$ and $K_{2}$. Furthermore, there does not yet exist a method to evaluate the quality of a component correspondence for BNP models. This issue has been studied before in the MCMC literature, where it is known as the “label switching problem”, but past solution techniques are generally model-specific and restricted to use on very simple mixture models .

This paper presents a methodology for creating streaming, distributed inference algorithms for Bayesian nonparametric models. In the proposed framework (shown for a single node A in Figure 1), processing nodes receive a sequence of data minibatches, compute a variational posterior for each, and make asynchronous streaming updates to a central model using a mapping obtained from a component identification optimization. The key contributions of this work are as follows. First, we develop a minibatch posterior decomposition that motivates a learning-rate-free streaming, distributed framework suitable for Bayesian nonparametrics. Then, we derive the component identification optimization problem by maximizing the probability of a component matching. We show that the BNP prior regularizes model complexity in the optimization; an interesting side effect of this is that regardless of whether the minibatch variational inference scheme is truncated, the proposed algorithm is truncation-free. Finally, we provide an efficiently computable regularization bound for the Dirichlet process prior based on Jensen’s inequality Regularization bounds for other popular BNP priors may be found in the supplement.. The paper concludes with applications of the methodology to the DP mixture model, with experimental results demonstrating the scalability and performance of the method in practice.

Streaming, distributed Bayesian nonparametric inference

The proposed framework, motivated by a posterior decomposition that will be discussed in Section 2.1, involves a collection of processing nodes with asynchronous access to a central variational posterior approximation (shown for a single node in Figure 1). Data is provided to each processing node as a sequence of minibatches. When a processing node receives a minibatch of data, it obtains the central posterior (Figure 1(a)), and using it as a prior, computes a minibatch variational posterior approximation (Figure 1(b)). When minibatch inference is complete, the node then performs component identification between the minibatch posterior and the current central posterior, accounting for possible modifications made by other processing nodes (Figure 1(c)). Finally, it merges the minibatch posterior into the central variational posterior (Figure 1(d)).

In the following sections, we use the DP mixture as a guiding example for the technical development of the inference framework. However, it is emphasized that the material in this paper generalizes to many other BNP models, such as the hierarchical DP (HDP) topic model , BP latent feature model , and Pitman-Yor (PY) mixture (see the supplement for further details).

This decomposition suggests a simple streaming, distributed, asynchronous update rule for a processing node: first, obtain the current central posterior density $p(\theta,z_{o}|y_{o})$, and using it as a prior, compute the minibatch posterior $p(\theta,z_{m},z_{o}|y_{o},y_{m})$; and then update the central posterior density by using (1) with the current central posterior density $p(\theta,z_{i},z_{o}|y_{i},y_{o})$. However, there are two issues preventing the direct application of the decomposition rule (1):

Unknown component correspondence: Since it is generally intractable to find the minibatch posteriors $p(\theta,z_{m},z_{o}|y_{o},y_{m})$ exactly, approximate methods are required. Further, as (1) requires the multiplication of densities, sampling-based methods are difficult to use, suggesting a variational approach. Typical mean-field variational techniques introduce an artificial ordering of the parameters in the posterior, thereby breaking symmetry that is crucial to combining posteriors correctly using density multiplication . The use of (1) with mean-field variational approximations thus requires first solving a component identification problem.

Unknown model size: While previous posterior merging procedures required a 1-to-1 matching between the components of the minibatch posterior and central posterior , Bayesian nonparametric posteriors break this assumption. Indeed, the datasets $y_{o}$, $y_{i}$, and $y_{m}$ from the same nonparametric mixture model can be generated by the same, disjoint, or an overlapping set of cluster parameters. In other words, the global number of unique posterior components cannot be determined until the component identification problem is solved and the minibatch posterior is merged.

2 Variational component identification

Suppose we have the following mean-field exponential family prior and approximate variational posterior densities in the minibatch decomposition (1),

where $\zeta_{r}(\cdot)$, $r\in\{o,i,m\}$ are products of categorical distributions for the cluster labels $z_{r}$, and the goal is to use the posterior decomposition (1) to find the updated posterior approximation

As mentioned in the previous section, the artificial ordering of components causes the naïve application of (1) with variational approximations to fail, as disparate components from the approximate posteriors may be merged erroneously. This is demonstrated in Figure 3(a), which shows results from a synthetic experiment (described in Section 4) ignoring component identification. As the number of parallel threads increases, more matching mistakes are made, leading to decreasing model quality.

To address this, first note that there is no issue with the first $K_{o}$ components of $q_{m}$ and $q_{i}$; these can be merged directly since they each correspond to the $K_{o}$ components of $q_{o}$. Thus, the component identification problem reduces to finding the correspondence between the last $K^{\prime}_{m}=K_{m}-K_{o}$ components of the minibatch posterior and the last $K^{\prime}_{i}=K_{i}-K_{o}$ components of the intermediate posterior. For notational simplicity (and without loss of generality), fix the component ordering of the intermediate posterior $q_{i}$, and define $\sigma\!:\!\left[K_{m}\right]\to\left[K_{i}+K^{\prime}_{m}\right]$ to be the 1-to-1 mapping from minibatch posterior component $k$ to updated central posterior component $\sigma(k)$, where $\left[K\right]:=\{1,\dots,K\}$. The fact that the first $K_{o}$ components have no ordering ambiguity can be expressed as $\sigma(k)=k\,\,\forall k\in\left[K_{o}\right]$. Note that the maximum number of components after merging is $K_{i}+K^{\prime}_{m}$, since each of the last $K^{\prime}_{m}$ components in the minibatch posterior may correspond to new components in the intermediate posterior. After substituting the three variational approximations (2) into (1), the goal of the component identification optimization is to find the 1-to-1 mapping $\sigma^{\star}$ that yields the largest updated posterior normalizing constant, i.e. matches components with similar densities,

Despite its motivation from the Dirichlet process mixture, the component identification optimization (5) is not specific to this model. Indeed, the derivation did not rely on any properties specific to the Dirichlet process mixture; the optimization applies to any Bayesian nonparametric model with a set of “components” $\theta$, and a set of combinatorial “indicators” $z$. For example, the optimization applies to the hierarchical Dirichlet process topic model with topic word distributions $\theta$ and local-to-global topic correspondences $z$, and to the beta process latent feature model with features $\theta$ and binary assignment vectors $z$. The form of the objective in the component identification optimization (5) reflects this generality. In order to apply the proposed streaming, distributed method to a particular model, one simply needs a black-box variational inference algorithm that computes posteriors of the form (2), and a way to compute or bound the expectation in the objective of (5).

3 Updating the central posterior

To update the central posterior, the node first locks it and solves for $\sigma^{\star}$ via (5). Locking prevents other nodes from solving (5) or modifying the central posterior, but does not prevent other nodes from reading the central posterior, obtaining minibatches, or performing inference; the synthetic experiment in Section 4 shows that this does not incur a significant time penalty in practice. Then the processing node transmits $\sigma^{\star}$ and its minibatch variational posterior to the central processing node where the product decomposition (1) is used to find the updated central variational posterior $q$ in (3), with parameters

Finally, the node unlocks the central posterior, and the next processing node to receive a new minibatch will use the above $K$, $\zeta(z)$, and $\eta_{k}$ from the central node as their $K_{o}$, $\zeta_{o}(z_{o})$, and $\eta_{ok}$.

Application to the Dirichlet process mixture model

The expectation in the objective of (5) is typically intractable to compute in closed-form; therefore, a suitable lower bound may be used in its place. This section presents such a bound for the Dirichlet process, and discusses the application of the proposed inference framework to the Dirichlet process mixture model using the developed bound. Crucially, the lower bound decomposes such that the optimization (5) becomes a maximum-weight bipartite matching problem. Such problems are solvable in polynomial time by the Hungarian algorithm, leading to a tractable component identification step in the proposed streaming, distributed framework.

For the Dirichlet process with concentration parameter $\alpha>0$, $p(z_{i},z_{m},z_{o})$ is the Exchangeable Partition Probability Function (EPPF)

where $n_{k}$ is the amount of data assigned to cluster $k$, and $\mathcal{K}$ is the set of labels of nonempty clusters. Given that the variational distribution $\zeta_{r}(z_{r})$, $r\in\{i,m,o\}$ is a product of independent categorical distributions $\zeta_{r}(z_{r})=\prod_{j=1}^{N_{r}}\prod_{k=1}^{K_{r}}\pi_{rjk}^{\mathds{1}\left[z_{rj}=k\right]}$, Jensen’s inequality may be used to bound the regularization in (5) below (see the supplement for further details) by

where $C$ is a constant with respect to the component mapping $\sigma$, and

Note that the bound (12) allows incremental updates: after finding the optimal mapping $\sigma^{\star}$, the central update (10) can be augmented by updating the values of $s_{k}$ and $t_{k}$ on the central node to

As with $K$, $\eta_{k}$, and $\zeta$ from (10), after performing the regularization statistics update (14), a processing node that receives a new minibatch will use the above $s_{k}$ and $t_{k}$ as their $s_{ok}$ and $t_{ok}$, respectively.

2 Solving the component identification optimization

The optimization (5) can be rewritten in terms of $\mathbf{X}$ and $\mathbf{R}$ as

The first two constraints express the 1-to-1 property of $\sigma(\cdot)$. The constraint $\mathbf{X}_{kk}=1\forall k\in\left[K_{o}\right]$ fixes the upper $K_{o}\times K_{o}$ block of $\mathbf{X}$ to $\mathbf{I}$ (due to the fact that the first $K_{o}$ components are matched directly), and the off-diagonal blocks to $\mathbf{0}$. Denoting $\mathbf{X}^{\prime}$, $\mathbf{R}^{\prime}$ to be the lower right $\left(K^{\prime}_{i}+K^{\prime}_{m}\right)\times\left(K^{\prime}_{i}+K^{\prime}_{m}\right)$ blocks of $\mathbf{X}$, $\mathbf{R}$, the remaining optimization problem is a linear assignment problem on $\mathbf{X}^{\prime}$ with cost matrix $-\mathbf{R}^{\prime}$, which can be solved using the Hungarian algorithm For the experiments in this work, we used the implementation at github.com/hrldcpr/hungarian.. Note that if $K_{m}=K_{o}$ or $K_{i}=K_{o}$, this implies that no matching problem needs to be solved – the first $K_{o}$ components of the minibatch posterior are matched directly, and the last $K^{\prime}_{m}$ are set as new components. In practical implementation of the framework, new clusters are typically discovered at a diminishing rate as more data are observed, so the number of matching problems that are solved likewise tapers off. The final optimal component mapping $\sigma^{\star}$ is found by finding the nonzero elements of $\mathbf{X}^{\star}$:

Experiments

In this section, the proposed inference framework is evaluated on the DP Gaussian mixture with a normal-inverse-Wishart (NIW) prior. We compare the streaming, distributed procedure coupled with standard variational inference (SDA-DP) to five state-of-the-art inference algorithms: memoized online variational inference (moVB) , stochastic online variational inference (SVI) with learning rate $(t+10)^{-\frac{1}{2}}$, sequential variational approximation (SVA) with cluster creation threshold $10^{-1}$ and prune/merge threshold $10^{-3}$, subcluster splits MCMC (SC) , and batch variational inference (Batch) . Priors were set by hand and all methods were initialized randomly. Methods that use multiple passes through the data (e.g. moVB, SVI) were allowed to do so. moVB was allowed to make birth/death moves, while SVI/Batch had fixed truncations. All experiments were performed on a computer with 24 CPU cores and 12GiB of RAM.

Airplane Trajectories: This dataset consisted of $\sim$3,000,000 automatic dependent surveillance broadcast (ADS-B) messages collected from planes across the United States during the period 2013-03-22 01:30:00UTC to 2013-03-28 12:00:00UTC. The messages were connected based on plane call sign and time stamp, and erroneous trajectories were filtered based on reasonable spatial/temporal bounds, yielding 15,022 trajectories with 1,000 held out for testing. The latitude/longitude points in each trajectory were fit via linear regression, and the 3-dimensional parameter vectors were clustered. Data was split into minibatches of size 100, and SDA-DP used 16 parallel threads.

MNIST Digits : This dataset consisted of 70,000 $28\times 28$ images of hand-written digits, with 10,000 held out for testing. The images were reduced to 20 dimensions with PCA prior to clustering. Data was split into minibatches of size 500, and SDA-DP used 48 parallel threads.

SUN Images : This dataset consisted of 108,755 images from 397 scene categories, with 8,755 held out for testing. The images were reduced to 20 dimensions with PCA prior to clustering. Data was split into minibatches of size 500, and SDA-DP used 48 parallel threads.

Figure 4 shows the results from the experiments on the three real datasets. From a qualitative standpoint, SDA-DP discovers sensible clusters in the data, as demonstrated in Figures 4(a)–4(c). However, an important quantitative result is highlighted by Table 1(a): the larger a dataset is, the more the benefits of parallelism provided by SDA-DP become apparent. SDA-DP consistently provides a model quality that is competitive with the other algorithms, but requires orders of magnitude less computation time, corroborating similar findings on the synthetic dataset.

Conclusions

This paper presented a streaming, distributed, asynchronous inference algorithm for Bayesian nonparametric models, with a focus on the combinatorial problem of matching minibatch posterior components to central posterior components during asynchronous updates. The main contributions are a component identification optimization based on a minibatch posterior decomposition, a tractable bound on the objective for the Dirichlet process mixture, and experiments demonstrating the performance of the methodology on large-scale datasets. While the present work focused on the DP mixture as a guiding example, it is not limited to this model – exploring the application of the proposed methodology to other BNP models is a potential area for future research.

This work was supported by the Office of Naval Research under ONR MURI grant N000141110688.

References