Bayesian nonparametrics: Dirichlet process mixtures

Challenge C06 hands you 82 galaxy velocities and makes you pick the number of mixture components. This notebook is the sequel: put a prior on mixtures with an unbounded number of components and let the data decide how many get used. That is the promise of the Dirichlet process (DP) mixture. The fine print is the lesson: the data decide less than the promise suggests, the prior on the concentration parameter decides more, and "the number of clusters" turns out to be three different quantities. What the DP mixture is reliably good at - flexible density estimation, and flexible distributions inside bigger models - is where the notebook ends.

We work with the enzyme data of Richardson & Green's (1997) classic paper on mixtures with an unknown number of components: the activity of an enzyme that metabolises carcinogenic substances, measured in 245 unrelated people. The biological question is whether the population splits into "slow" and "fast" metabolisers (a genetic polymorphism).

#1 · The data

Richardson & Green's data page holds three samples in one text file (enzyme, lake acidity and the galaxy velocities of C06), each as a heading, a description, the sample size and the values. A few lines of parsing:

Two groups, but not two bell curves. About two thirds of the people sit in a tight, steep pile around 0.2 (the slow metabolisers). The rest spread from about 0.6 to 2.9 in a wide, right-skewed hump. Bechtel et al. fitted two skewed distributions. A mixture of normals will need more than two components to draw the same picture, and that is the first thing to keep in mind: a mixture component is a kernel, not a population.

#2 · The Dirichlet process in three pictures

A DP mixture says: there are infinitely many components, with weights $w_1, w_2, \dots$ that sum to one, drawn by stick-breaking. Take a stick of length 1, break off a fraction $v_1 \sim \text{Beta}(1, \alpha)$, that is $w_1$; break the fraction $v_2$ off what is left, that is $w_2 = v_2(1 - v_1)$; and so on:

$$w_k = v_k \prod_{\ell<k} (1 - v_\ell), \qquad v_k \sim \text{Beta}(1, \alpha).$$

Each component gets its own parameters $\theta_k$ drawn from a base distribution $G_0$. The single concentration $\alpha$ sets how fast the stick is used up: small $\alpha$ gives a few big weights, large $\alpha$ many small ones.

#What $\alpha$ says about the number of clusters, before any data

With $n$ observations, each assigned to a component with probability $w_k$, the number of occupied components $K_n$ is random, and its prior distribution depends on $\alpha$ and $n$ only: $E[K_n] = \sum_{i=1}^{n} \alpha / (\alpha + i - 1) \approx \alpha \log(1 + n/\alpha)$. So the DP does not leave the number of clusters open - $\alpha$ sets it, and it grows with $n$. The simplest way to see the prior is to simulate it. The Chinese-restaurant representation needs no weights at all: observation $i$ opens a new cluster with probability $\alpha / (\alpha + i - 1)$.

Three facts to take into the model. With $\alpha = 0.2$ the prior expects about two clusters, with $\alpha = 5$ about twenty - for the same 245 people. The popular $\alpha \sim \text{Gamma}(2, 1)$ is not "uninformative" about $K$: it spreads the prior over 2 to 21 occupied clusters with a median of 9. And the right panel shows the growth with $n$: a DP expects more clusters in a bigger sample, forever, which is a modelling assumption (new data keep revealing new kinds) and not an obvious one for a genetic polymorphism.

Pitman-Yor aside. The two-parameter Pitman-Yor process breaks the stick with $v_k \sim \text{Beta}(1 - d, \alpha + k d)$, $0 \le d < 1$. The discount $d$ makes later sticks longer than under the DP, so new clusters keep appearing at rate $n^d$ instead of $\log n$, and cluster sizes follow a power law - the right prior for words in a corpus or species in a sample, where there is no finite number of "types". For a handful of clusters in 245 measurements it only makes the problem below (too many small clusters) worse, so we stay with the DP.

#3 · A truncated DP mixture of normals

NUTS cannot sample infinitely many sticks, so we truncate at $T$ components (the last weight takes the remainder of the stick) and check afterwards that the truncation did not bite. The prior mass beyond $T$ is on average $(\alpha / (1+\alpha))^{T-1}$: with $\alpha \approx 2$ and $T = 20$ that is $4 \times 10^{-4}$.

pm.StickBreakingWeights(alpha=..., K=T-1) is the truncated stick-breaking distribution (note: K counts the breaks, the output has $K+1$ weights). pm.NormalMixture marginalises the discrete assignments, so the model is continuous and NUTS can sample it.

$$y_i \sim \sum_{k=1}^{T} w_k\, \text{Normal}(\mu_k, \sigma_k), \quad w \sim \text{SB}_T(\alpha),\quad \alpha \sim \text{Gamma}(2, 1),\quad \mu_k \sim \text{Normal}(\bar y, 2 s_y),\quad \sigma_k \sim \text{InvGamma}(2, 0.1).$$

The base distribution matters. $\sigma_k$ gets an inverse-gamma prior because it puts no mass near zero: a normal mixture's likelihood is unbounded when one component shrinks onto a single point, and a prior that allows $\sigma_k \to 0$ lets the sampler find those spikes. The prior median of $\sigma_k$ is about 0.06 and its upper 5% point 0.3: small enough to resolve the slow group, and several components can team up for the wide fast group.

#Diagnostics: the usual table is useless here

NUTS reports no divergences, $\alpha$ is fine, and the component means are a mess: $\hat R$ up to 1.36 and bulk ESS below 10 for some of them. This is not a sampling failure. The component labels are not identified - the same density can be written with the components in a different order - so "$\mu_3$" has no meaning of its own across chains. Stick-breaking tempers the problem (component 1 tends to be the biggest), it does not remove it. In C06 you fixed this with an ordering constraint; with 20 components, most of them empty, an ordering on the means helps little, and it is simpler to stop asking about labels.

We need diagnostics of quantities that do not care what the components are called: $\alpha$, the log density, and the fitted density $f(x) = \sum_k w_k \text{Normal}(x \mid \mu_k, \sigma_k)$ at a grid of points.

Every label-free quantity mixes: the density at five points has $\hat R \le 1.02$ and bulk ESS in the hundreds to thousands. The log density ($\hat R$ 1.02, ESS 181) and $\alpha$ (ESS 105) are the slowest, which is typical: $\alpha$ is informed only through the number of occupied components, and that number changes only when a component empties or fills.

#Is the truncation big enough?

If $T$ were too small, the last weight - which collects the whole rest of the stick - would be visibly positive. Here it is the posterior of $w_T$ and of the total weight in the last five components:

The last component carries on average $3 \times 10^{-4}$ of the mass, and even the 99th percentile of the last five together is 1.5%: 20 components are enough. The four chains agree on the sizes of the components (the sorted weights), although, as we will see, not necessarily on which component is which.

#4 · Label-invariant summaries

#The density, with bands

The quantity a DP mixture estimates best is the density itself.

#"How many clusters?" - three different answers

There are three quantities that could be called the number of clusters:

1. the number of components: $T = 20$ by construction (infinity for the untruncated DP);
2. the number of occupied clusters $K_{\text{occ}}$: the number of distinct labels when each person is allocated to a component. The assignments were marginalised out, so we draw them back: for posterior draw $s$, person $i$ goes to component $k$ with probability $\propto w_k^{(s)} \text{Normal}(y_i \mid \mu_k^{(s)}, \sigma_k^{(s)})$;
3. the number of components with a non-negligible weight, say $w_k > 0.02$.

The three counts disagree, and that is the point. $K_{\text{occ}}$ has a median of 7 and a long right tail: a DP routinely opens small clusters with one or two people in them, because under its prior a new cluster is always possible and a single outlying value is cheaply explained by one. Counting only components with more than 2% of the weight - or, nearly the same, with at least 5 people - gives 4 or 5. None of them is 2, the number of biological groups: the fast group is skewed, and normal kernels need two or three to draw it. Compared with the prior (grey), the data have ruled out 1-3 clusters and very many clusters, but the posterior of $K_{\text{occ}}$ is still wide. Miller & Harrison (2014) showed this is not a small-sample accident: the DP mixture posterior of the number of clusters is inconsistent - even with infinitely many observations from a finite mixture it does not concentrate on the true number. The DP is a tool for density estimation; counting its clusters is a separate, prior-driven question.

#Who belongs with whom: the co-clustering matrix

Labels switch; pairs do not. The posterior similarity matrix $P_{ij}$ = probability that persons $i$ and $j$ are in the same cluster is label-invariant, and the whole clustering posterior can be read from it. Because the data are one-dimensional and sorted, clusters show up as blocks on the diagonal.

Two dark off-diagonal blocks say that nobody below about 0.45 is ever clustered with anybody above 0.6: the slow/fast split is certain. Within the slow group the co-clustering is high but fades towards its upper edge (about 0.25-0.45), where people are sometimes put in a small extra component. The fast group is much less coherent: it is carved into two or three kernels in varying ways, so its pairwise probabilities sit between 0.3 and 0.8.

#A single "best" clustering

A report usually wants one partition. The Bayesian answer is the partition that minimises the posterior expected loss. Two losses are standard (Wade & Ghahramani 2018):

• Binder's loss counts the pairs on which two partitions disagree. Its expectation needs only the similarity matrix: $\sum_{i<j} |\mathbb{1}[c_i = c_j] - P_{ij}|$.
• The variation of information (VI) is an information-theoretic distance between partitions. Its expectation has a cheap lower bound that also uses only $P$: $\frac1N \sum_i \big[\log_2 \sum_j \mathbb{1}[c_j{=}c_i] + \log_2 \sum_j P_{ij}
  • 2 \log_2 \sum_j \mathbb{1}[c_j{=}c_i] P_{ij}\big]$.

Binder's loss is known to favour extra small clusters, VI to favour fewer, larger ones. Searching all partitions is impossible; the usual shortcut is to score the sampled partitions plus the cuts of a hierarchical clustering on $1 - P$, and keep the best.

Here the two losses tell different stories from the same posterior. The VI-optimal partition is exactly the biological answer: 151 slow metabolisers (activity up to 0.41) and 94 fast ones (from 0.47) - the split Bechtel et al. found with two skewed distributions. Binder's loss keeps five clusters, splitting both groups and isolating two people around 0.5. Neither is "the" answer; they are optimal for different questions (VI: a coarse, stable summary; Binder: penalise every mis-paired couple equally). Whichever you report, report the similarity matrix next to it.

#5 · How much does $\alpha$ decide?

Refit with $\alpha$ fixed at three values instead of learning it, and with the truncation cut to $T = 5$. Everything else is identical.

The density is robust, the cluster count is not. The four posterior median densities are nearly indistinguishable, but $K_{\text{occ}}$ runs from a median of 4 ($\alpha = 0.2$) to 14 ($\alpha = 5$): with 245 observations the prior on $\alpha$ still decides the count. Learning $\alpha$ does not escape this - it simply moves the choice to the hyperprior (Gamma(2, 1) gave 7).

Two truncation failures. With $\alpha = 5$ the stick is used up slowly: the last of 20 components carries 1.7% of the mass on average (up to 9%), so the truncation is biting and $T$ would have to go up (to about 50 for a tail below $10^{-4}$). With $T = 5$ it bites hard: $w_T$ averages 0.23, and the posterior of $K_{\text{occ}}$ is piled against the ceiling of 5. A $K_{\text{occ}}$ posterior that touches $T$ is the same alarm in another form. Notice that neither failure shows in the density, or in the sampler diagnostics: only the last-stick check catches it.

#Mixing across chains

The stick-breaking weights are not exchangeable: component 1 is a priori the biggest, component 2 the next, and so on. So relabelling two components changes the prior density, and the posterior has many unequal modes, one per plausible ordering. NUTS moves between them only rarely, so different chains can settle on different orderings - harmless for label-invariant summaries if each chain explores the same clustering distribution. Compare the chains on $K_{\text{occ}}$:

Chain 1 almost never visits $K_{\text{occ}} \le 3$, the others do in 35-50% of their draws, and the density at 0.15 has $\hat R = 1.15$ with a bulk ESS of 19. The chains are stuck in different orderings with different cluster structures, and with $\alpha = 0.2$ it is costly to open or close a component, so they do not trade. (In the fit with $\alpha$ learned, the same checks gave $\hat R \le 1.02$.) Remedies, from cheap to expensive: more chains and longer runs, judged by label-free $\hat R$; starting chains from different clusterings; or a sampler that updates assignments directly (Gibbs / split-merge moves for the conditional or marginal DP representation), which is what the specialised DP software does and what NUTS on the marginalised model cannot.

#6 · A sparse finite mixture

Malsiner-Walli, Frühwirth-Schnatter & Grün (2016) turn the question around: fit a finite mixture with deliberately too many components, $K = 10$, and a symmetric $\text{Dirichlet}(e_0, \dots, e_0)$ prior on the weights with a tiny $e_0$ (they recommend 0.01 or smaller). The prior then wants most components empty, and superfluous components are emptied rather than split off. Unlike the DP, the weights are exchangeable, and the prior number of clusters does not grow with $n$.

#The textbook setting, $e_0 = 0.01$, under NUTS

We try it twice: with PyMC's pm.Dirichlet, and with a reparameterisation that often helps with small Dirichlet parameters. A Dirichlet vector is a set of independent Gamma variables divided by their sum, $g_k \sim \text{Gamma}(e_0, 1)$, $w_k = g_k / \sum_j g_j$, so we can sample $\log g_k$ instead of the simplex: $K$ coordinates that are independent a priori, each with the smooth density $\propto \exp(e_0 x - e^x)$. PyMC derives the density of pt.log(Gamma) automatically. Everything then stays on the log scale ($\log w = \log g - \text{logsumexp}(\log g)$) and the mixture likelihood is written with logsumexp, because an empty component's weight can underflow to exactly zero.

Two different failures. With pm.Dirichlet, 70 divergences in three of four chains, step sizes around 0.01 and trees at or near the maximum depth of 10: the fitted density still agrees across chains ($\hat R$ 1.01), but the sampler is not exploring the prior region where most components are empty and it says so. The log-gamma version is worse: three chains without a single divergence ran with step sizes of $10^{-5}$ and saturated trees - frozen in place, $\hat R$ of the density 2.83, ESS 5 - and the fourth diverged on almost every draw. Zero divergences is not a clean bill of health; the step size and the label-free $\hat R$ gave it away.

The cause is the prior itself. $\text{Gamma}(0.01, 1)$ puts $\log g_k$ of an empty component on an exponential tail with scale $1/e_0 = 100$, while occupied components have $\log w_k$ determined to $\pm 0.1$: a thousandfold difference in scale that swaps between coordinates whenever a component empties or fills. No mass matrix fits both. A Gibbs sampler (what Malsiner-Walli et al. use) updates the allocations and never sees this geometry. For NUTS the practical fix is a less extreme $e_0$. What does that cost? The implied prior on $K_{\text{occ}}$, simulated the same way as for the DP:

With $e_0 = 0.05$ the prior still favours few clusters (median 3, and it will not grow with $n$), and NUTS is healthy: no divergences, step size 0.02-0.03, density $\hat R$ 1.00 and ESS above 900 - at the price of deep trees (several hundred gradient evaluations per draw). The posterior of $K_{\text{occ}}$ is tighter than the DP's, 4 or 5 with 75% probability and almost never above 7: the sparse prior empties superfluous components instead of keeping small ones alive, which is the sense in which it is better behaved for counting clusters. The VI-optimal partition is again the 151/94 slow/fast split; Binder's keeps four clusters.

#7 · Finite mixtures with $K$ chosen by LOO

The classical route: fit $K = 2, \dots, 5$ with ordered means (to kill label switching; starting values via pm.sample(initvals=...) because initval= breaks pm.compute_log_likelihood) and compare by PSIS-LOO. The DP and sparse mixtures go into the same comparison.

Two normals are clearly not enough (K = 2 is 8 elpd behind), three nearly so, and from four components on every model predicts about equally well: finite $K = 4$ and $K = 5$, the DP and the sparse mixture are within 3.2 elpd of each other, and ArviZ flags all these differences as below 4, too small to trust. So LOO says "at least four normal kernels" - a statement about how many normals it takes to draw a skewed density, not about how many kinds of people there are. (Note also $\hat R = 1.64$ for $K = 3$ even with ordered means: the chains disagree about how to spend three kernels on two skewed groups.) The DP and the sparse mixture reach the same predictive accuracy without choosing $K$ at all, and their co-clustering tells the biological story that no single finite fit does.

#8 · Where a DP mixture earns its keep: a flexible random-effects distribution

The most useful DP mixtures are not the ones whose clusters you interpret, but the ones that replace a Normal assumption inside a bigger model. Take the Mayo Clinic trial in primary biliary cholangitis: serum bilirubin measured repeatedly (1-16 visits) on 312 patients. On the log scale, each patient's trajectory is roughly a line,

$$\log \text{bili}_{ij} = a_i + b_i t_{ij} + \varepsilon_{ij}, \qquad t \text{ in years}.$$

The standard model puts $b_i \sim \text{Normal}(\mu_b, \sigma_b)$. But a liver disease has patients who are stable and patients who progress: the distribution of slopes may be skewed, and a Normal would then misjudge exactly the quantity a clinician cares about - how many patients progress fast. Replace the Normal by a DP mixture of normals, again with pm.NormalMixture, now as a prior on the 312 latent slopes.

Both fits are clean (no divergences; the 312 slopes have $\hat R \le 1.02$). The Normal model says the typical patient's bilirubin grows by about 18% a year ($\mu_b = 0.16$ on the log scale) with a patient-to-patient sd of 0.17. The DP mixture tells a different shape: a sharp peak of near-stable patients at slopes around 0.02, a shoulder around 0.15-0.2, and a long right tail of progressors reaching 0.5-0.7 per year. The Normal, forced to be symmetric, puts 18% of patients on a falling trajectory to cover the right tail; the DP puts 14% there and 18% (against 15%) in the "bilirubin doubles within two years" tail. The intervals of the two models overlap, so this is a shift of emphasis rather than a reversal - with one to sixteen visits per patient the data cannot resolve the shape of the slope distribution sharply, which the wide DP band shows honestly and the Normal's narrow band hides.

By LOO the DP is ahead by 6.7 elpd with a standard error of 7.4: not a decisive difference, and PSIS flags 89-113 observations with $\hat k > 0.7$, which is expected when every patient has their own two parameters and few visits. A leave-one-patient-out comparison (K-fold by patient with az.loo_kfold) would be the honest next step. The broader lesson stands: a DP mixture is a low-risk way to relax a Normal random-effects assumption, it costs a few seconds here, and it shows you where the Normal was wrong.

#9 · Take-aways

• A DP mixture's $\alpha$ is a prior on the number of clusters, and it grows with $n$. Simulate the implied prior of $K_{\text{occ}}$ before fitting; with 245 observations it still decided the posterior count (median 4 to 14).
• Truncate and check the last stick. Neither the density nor NUTS diagnostics warn when $T$ is too small.
• Component parameters are not identified. Diagnose label-free quantities (density on a grid, log density, $\alpha$), and summarise with the density, the co-clustering matrix and a loss-optimal partition - VI and Binder can disagree, so say which.
• "How many clusters" has at least three answers (components, occupied, non-negligible), and none of them is the number of populations when the kernel is wrong: skewed groups need several normals.
• A sparse finite mixture concentrates the count better, but a tiny $e_0$ creates a geometry NUTS cannot handle - with or without divergences. $e_0 = 0.05$ worked.
• The DP is at its best as a flexible distribution inside a bigger model: a random-effects distribution, a residual distribution, a prior you do not want to be Normal.

#Try it yourself

1. Lake acidity. rg_sample("Acidity data") gives Richardson & Green's second benchmark (155 lakes, log scale). Fit the DP and the sparse finite mixture. Their paper's posterior for $K$ peaks at 3-4; which of your three "numbers of clusters" agrees, and how much does it move between $\alpha \sim \text{Gamma}(2, 1)$ and $\text{Gamma}(1, 4)$?
2. Skewed kernels. Replace the normal kernel by a skew-normal (pm.SkewNormal, via pm.Mixture with a list of component distributions or a batched .dist) in the sparse finite mixture on the enzyme data. Does the posterior now put most of its mass on the two groups of Bechtel et al.? Compare the densities and the LOO.
3. Pitman-Yor mixture. Write the Pitman-Yor stick-breaking weights by hand ($v_k \sim \text{Beta}(1-d, \alpha + k d)$ with pm.Beta and a cumulative product) and fit it with $d = 0.25$. Predict first, from section 2, what happens to $K_{\text{occ}}$; then check the truncation, which is harder to satisfy for $d > 0$.