Overthinking Bayesian Mixture Models with PyMC

Much has changed in the ten years since I published Density Estimation with Dirichlet Process Mixtures using PyMC3. To name a few new developments:

• PyMC is now on version 6, it was on version 3.
• ArviZ and PreliZ did not exist.
• Nutpie did not exist.
• I understood the math, probability, and statistics in question less well than I do now.

This post revisits the first example in the former post, waiting times between eruptions of the Old Faithful in Yellowstone National Park, from a slightly different perspective.

In this post, we will use the Old Faithful data to explore (normal) mixture models and computational considerations in their Bayesian estimation. We will begin with a simple finite mixture model, explore its identifiability, then progress to nonparametric Dirichlet process mixture models as in the prior post.

First we make the necessary Python imports and do some light configuration.

Next we load the data from GitHub.

Exploratory data analysis¶

First we visualize the distribution of the (standardized) wait times between eruptions as a histogram.

Visually, it seems plausible that this data is from a combination (mixture) of two normal distributions, one centered just below $-1$ and one between $0.5$ and $1$.

Adding a Gaussian kernel density estimator to the histogram seems to confirm this. Our goal in this post is to show several methods to model this density using Bayesian statistics, and discuss some of the computational properties of those models.

Mixture models¶

A crash course¶

Mixture models are useful when we suspect, as with the Old Faithful data, that the data generating process is a combination of simpler component distributions. Formally, assume there are $K$ component distributions, each with probability density function $f_k(\cdot\ |\ \Theta_k)$ for $k = 1, 2, \ldots, K$. We assume the probability that an observation is generated from the $k$-th component of the mixture is $w_k$ for $k = 1, 2, \ldots, K$.

The simplest way to sample from such a mixture model is

1. sample $\hat{k} \sim \text{Cat}(w_1, w_2, \ldots, w_K)$, then
2. sample $x \sim f_{\hat{k}}(\cdot\ |\ \Theta_{\hat{k}}).$

Below we illustrate this sampling scheme for a three-component normal mixture.

In addition to a histogram of the samples, this plot shows the component and combined mixture probability density functions.

Bayesian mixture models¶

In the rest of this post, we will explore Bayesian inference for mixture models. To make these models Bayesian, we need to specify priors on the mixture weights $w_k$ and the component distribution parameters $\Theta_k$.

We start with the simplest case where the target number of normal components, $K$, is known. In this case of Old Faithful, we strongly suspect $K = 2$.

We now turn to implementing a Bayesian normal mixture model in PyMC. While we could translate the sampling scheme above fairly straightforwardly, we would not be able to take advantage of Nutpie's strong sampling performance, which does not work for discrete latent variables. Instead, we use one of the applied Bayesian statistician's go-to tricks, marginalization of discrete latent parameters. Marginalization removes the need to sample the discrete values $\hat{k}$ explicitly, instead combining the prior and component likelihoods and only then sampling.

First, we use a Dirichlet prior on the mixture weights,

$$w_k \sim \text{Dir}(1, 1).$$

Note that for these hyperparameters, this distribution is uniform on the simplex.

Next we place priors on the component parameters

$$ \begin{align} \mu_1, \mu_2 & \sim N(0, 2.5^2) \\ \sigma_1, \sigma_2 & \sim \text{Half-}N(2.5^2). \end{align} $$

Finally, we add the marginalized normal mixture likelihood. PyMC's NormalMixture class handles the marginalization for us transparently.

With this model defined, we are ready to sample from its posterior distribution using Nutpie.

We examine the Gelman-Rubin ($\hat{R}$) statistics to see if there were any potential sampling issues.

We see that the $\hat{R}$ statistics for $\mu$ and $w$ are significantly larger than one, which indicates potential sampling issues. To diagnose these, we visualize the joint distribution of $\mu_0$ and $\mu_1$.

We see that the joint distribution is bimodal, and those modes appear to be reflections of each other around the line $\mu_1 = \mu_0$. This symmetry strongly suggests we are experiencing label switching, a phenomenon in many discrete latent parameter models that arises because their likelihood is invariant under permutations of the latent discrete parametrs (and the corresponding permutations of the component parameters).

To validate this conclusion, we plot the mean traces for the first two chains.

We see that indeed the actual values of the means are the same in both chains (different colored plots), but that what is labelled $\mu_0$ in the first chain is labelled $\mu_1$ in the second, and vice versa.

Ordering constraints¶

The standard method for breaking this permutation symmetry that makes the model non-identifiable is to introduce an ordering constraint on one of the component distribution parameters. Here we build almost the same model, but apply PyMC's ordered transformation to ensure $\mu_0 < \mu_1$.

We now sample from this ordered-mean model and examine its $\hat{R}$ statistics.

We see $\hat{R}$ values quite close to one, indicating much better sampling. The pair and trace plots also confirm the lack of label switching for this model.

Now that we are satisfied with our sampling, we visualize the posterior distribution of the mixture probability density.

We see that the posterior expected value is a reasonable fit to the data. Unlike with standard kernel density estimation (like we used for exploratory data analysis), the Bayesian approach naturally quantifies the uncertainty in this distribution.

Nonparametric Bayesian mixture models¶

Outside of having to introduce an ordering constraint on the mean to identify the model, the above analysis was fairly straightforward. It did, however, rely on the fact that we had a good idea of the number of mixture components a priori. In the spirit of Bayesian statistics, it would be ideal to place a prior on the number of mixture components and then perform inference on the number of these components. This desire puts us firmly in the domain of Bayesian nonparametric statistics, which has two primary tools, Gaussian processes and Dirichlet processes. Dirichlet processes are more applicable to our situation.

Dirichlet processes mixtures¶

The previous post gives a much deeper exposition of Diriclet processes in general (a fascinating topic, to be sure). Since this post is focused on mixture models, we skip that exposition and focus on the stick-breaking process formulation most relevant to these models.

To facilitate Bayesian inference on the number of comoponents, a Dirichlet process mixture assumes that there are may be infinitely many, with weights $w_1, w_2, \ldots $. Given a concentraction parameter $\alpha \geq 0$, the stick-breaking process constructs a distribution on these weights as follows.

1. Sample $\beta_1, \beta_2, \ldots \sim \text{Beta}(1, \alpha)$.
2. Let $w_k = \beta_k \cdot \prod_{\ell = 1}^{k - 1} (1 - \beta_{\ell})$.

To get a sense of how the concentration parameter, $\alpha$, controls the distribution of the weights, note that

$$\mathbb{E}(\beta_k) = \frac{1}{1 + \alpha},$$

so

$$\mathbb{E}(w_k) = \frac{1}{1 + \alpha} \cdot \left(\frac{\alpha}{1 + \alpha}\right)^{k - 1} = \frac{1}{\alpha} \cdot \left(\frac{\alpha}{1 + \alpha}\right)^k.$$

Below we visualize the expected values of the first ten component weights for various values of $\alpha$.

We see that smaller values of $\alpha$ favor more weight in early components and a quickly decaying tail, whereas higher values favor more uniformly distributed and slowly decaying early components.

In practice, we cannot simulate infinitely many weights, so we turn to truncated Dirichlet process mixtures. In a truncated Dirichlet process mixture, we set the maximum possible number of components, $K_{\max}$, to be well above the likely true number of components in the data generating process. The truncated stick breaking process is defined as follows.

1. Sample $\beta_1, \beta_2, \ldots, \beta_{K_{\max} - 1} \sim \text{Beta}(1, \alpha)$.
2. Let

$$\tilde{w}_k = \begin{cases} \beta_k \cdot \displaystyle\prod_{\ell = 1}^{k - 1} (1 - \beta_{\ell}) & \text{if } k = 1, 2, \ldots, K_{\max} - 1 \\ 1 - \displaystyle\sum_{k = 1}^{K_{\max} - 1} \tilde{w}_k & \text{if } k = K_{\max} \end{cases}.$$

When truncating, we use the final component's weight to ensure the sum-to-one constraint is satisfied. For those readers interested in more details on truncated stick-breaking processes, Ishwaran and James's Gibbs Sampling Methods for Stick-Breaking Priors is an excellent reference.

The above expression for $\mathbb{E}(w_k)$ allows us to derive a heuristic for choosing $K_{\max}$. Since

$$1 = \sum_{k = 1}^{\infty} w_k = \sum_{k = 1}^{K_{\max} - 1} \tilde{w}_k + \sum_{k = K_{\max}}^{\infty} w_k,$$

we have that

$$ \begin{align*} \mathbb{E}(\tilde{w}_{K_{\max}}) & = 1 - \sum_{k = 1}^{K_{\max} - 1} \mathbb{E}(\tilde{w}_k) \\ & = \sum_{k = K_{\max}}^{\infty} \mathbb{E}(w_k) \\ & = \frac{1}{\alpha} \cdot \sum_{k = K_{\max}}^{\infty} \left(\frac{\alpha}{1 + \alpha}\right)^k \\ & = \frac{1}{\alpha} \cdot \left(\frac{\alpha}{1 + \alpha}\right)^{K_{\max}} \cdot \sum_{k = 0}^{\infty} \left(\frac{\alpha}{1 + \alpha}\right)^k \\ & = \left(\frac{\alpha}{1 + \alpha}\right)^{K_{\max} - 1}. \end{align*} $$

Therefore, if we want to ensure the expected weight truncated into the final component is less than $w^*$, we get

$$(K_{\max} - 1) \cdot (\log \alpha - \log(1 + \alpha))) < \log w^*,$$

which gives

$$K_{\max} > 1 + \frac{\log w^*}{\log \alpha - \log(1 + \alpha))}.$$

To choose $K_{\max}$ for our Old Faithful model, we must first place a prior on $\alpha$. The original paper that introduced Dirichlet process mixutres, Antoniak's Mixtures of Dirichlet Processes with Applications to Bayesian Nonparametric Problems, shows that the expected number of active components in a Dirichlet process mixture is

$$\mathbb{E}(K) \approx \mathbb{E}(\alpha) \cdot \log n,$$

where $n$ is the number of data points observed. For the Old Faithful data, this means our prior on $\alpha$ should have expected value

We use PreliZ's maximum entropy functionality to choose a lognormal prior that has the desired mean and 99% of its mass in the interval $[0, 5]$, which contains the hypothesized number of components.

Using this prior, we can calculate $K_{\max}$ using $\alpha = 5$ to be safe even if we have underestimated the number of required components.

We now specify this prior on $\alpha$. We use PyMC's StickBreakingWeights distribution for the prior on $w$.

The rest of the model is analagous to the finite ordered model above.

We now sample from this model.

There are some divergences, which often indicate model geometry issues. The the $\hat{R}$ statistics also exceed 1.01, confirming that we should investigate these samples closely.

Interestingly, even in the unordered finite mixture model, the $\hat{R}$ for $\sigma$ was not too large, so it makes sense to start exploring there.

Here we show the posterior distributions for the scale parameter of six mixture components that we expect to be largely unused. Below, we plot the scale parameters from the finite ordered model for comparison.

Immediately we see that the Dirichlet process mixture posteriors are much more concentrated near zero in many cases. This is a function of the prior geometry of the half-normal prior, as shown below.

We see here that the half-normal prior assigns the highest probability density to zero. In the case of the finite mixture model, enough samples correspond to each component for the data to overwhelm the prior and move the posterior density well away from zero. By design in the Dirichlet process mixture model, many components will have no corresponding samples, causing instablity around the hard zero-boundary.

To remedy this issue, we will place a lognormal prior on $\sigma_k$.

Note that this prior does not assign any density to zero. In fact, a log-transformed version of it will be unconstrained with a normal density, which can be sampled from very efficiency.

We now reimplement the ordered Dirichlet process mixture model using this prior and sample from it.

There are some divergences, although fewer than with the previous model. The the $\hat{R}$ statistics also exceed 1.01, though are lower than with the previous model as well.

Due to the elevated $\hat{R}$ values, we look again for sampling issues before turning to inference. Examining the pairplot for the two component means, we see that many samples are concentrated on the border line $\mu_0 = \mu_1$. While the samples seem to want to cross this boundary, the ordering constraint on the means prohibits that behavior.

Notably, we do not see this behavior in the finite ordered model, so it is worth investigating further.

To investigate further, we plot examine pair plots of the first two component means and weights.

The plot of $w_0$ against $\mu_0$ is interesting, so we examine it more closely.

We see that many of the samples for the first component's weight are near zero and that when that weight is near zero, the mean can vary freely, as it does not contribute anything to the model. We confirm this behavior by looking at the plots for the a single trace.

What we see here is two prior constraints conflicting with one another. The ordered transformation on the means ensure that $\mu_0$ must be the smallest mean, but a small expected value of $\alpha$ encourages most of the mass from the stick-breaking prior on $w$ to concentrate in the first component.

However, it is clear that the second mixture component of the (standardized) waiting times, with the larger mean, should have the highest mixture weight.

Therefore we must relax one of these two constraints. Since we chose the prior on $\alpha$, and therefore $w$, to be informative based on the hypothesized number of mixture components, we remove the ordering constraint from the means.

We sample from this model.

There are no divergences, is an improvement. Examining the $\hat{R}$ statistics, we see that they are high, which is expected due to label switching, as we have removed the ordering constraint from the means.

The following trace plot confirms that we are experiencing label switching as expected.

The quantity we really care about, however, is the mixture density, so we calculate that and examine its $\hat{R}$.

This plot shows how the mixture density's $\hat{R}$ varies over the grid on which we have evaluated it. We see that $\hat{R} < 1.01$ in regions that are covered by the data and only increase when we venture intro extrapotation, so we are satisfied with these results.

Now that we are satisfied with the quality of these samples, we visualize the posterior mixture likelihood.

We see excellent agreement with the finite ordered mixture model's posterior mixture density, without having had to specify the number of mixture components.

We also see that the most likely number of posterior active components is two, but that there are a decent number of samples with a few more.

It is a well-known property of Dirichlet process mixtures that they are not consistent estimators of the number of true components and can be prone to overfitting. See Miller and Harrison's Inconsistency of Pitman-Yor Process Mixtures for the Number of Components for a detail discussion of this phenomenon.

In this post we have explored five different Bayesian mixture models on a toy dataset, starting with a simple finite mixture model and working our way up to a nonparametric (theoretically) infinite Dirichlet process mixture model. Along the way we saw what is necessary and achievable in terms of identification, sampling performance, and estimator consistency.

I hope you have enjoyed reading this post as much as I have enjoyed writing it!

This post is available as a Jupyter notebook here.

AI usage transparency: I have used various AI models to refine, simplify, and validate some of the mathematics and modeling in this post. I have also used AI to check spelling and grammar for correctness and quality. However, AI has not directly created, changed, or edited any of the content in this post. It started with my writing and all AI suggestions were reviewed and incorporated manually by me.