I really enjoyed the pedagogy of Edwin Chen’s introduction to infinite mixture models, but I was a little disappointed that it does not go as far as presenting the details of the Dirichlet process Gaussian mixture model (DPGMM), as he uses sklearn’s variational Bayes DPGMM implementation.

For this reason, I will try and give here sufficient information to implement a DPGMM with collapsed Gibbs sampling. This is not an in-depth evaluation of which conjugate priors to use, nor an analysis of the parameters and hyper-parameters (that should have their own priors! ;)).

Prerequisites

On Dirichlet processes, Chinese Restaurant processes, Indian Buffet processes, there is the excellent blog post by Edwin Chen. Another excellent introduction to Dirichlet processes is provided by Frigyik, Kapila and Gupta.

If you lack some knowledge about clustering or density estimation (unsupervised learning), you can read Chapters 20 (p. 284) to (at least) 22 of MacKay’s ITILA, that you can find as a free ebook; or chapter 9 of Bishop’s PRML. As a refresher, the Wikipedia article on mixture models, and the sklearn documentation on GMM are more efficient.

DPGMM: the model

Let’s say we have $N$ observations and $K$ clusters, $i \in [1\dots N]$ is the indice for the observations, while $k \in [1\dots K]$ is the indice for the clusters. With $z_i$ the cluster assignment of observation $x_i$, and $\theta_k$ the parameter of mixture $k$:

$$P(x_{1:N}) = \prod_{i=1}^N \sum_{k=1}^K P(x_i|\theta_{z_i}) P(z_i=k)$$

So, the generative story of a DPGMM is as follows:

$\pi \sim Stick(\alpha)$ (mixing rates)
$z_i \sim \pi$ (cluster assignments)
$\theta_k \sim H(\lambda)$ (parameters)
$x_i \sim F(\theta_{z_i})$ (values)

Fitting the data

Notation:

$$\begin{align} & z_{-i} = \{z_j | j \neq i\}\\ & x = x_{1:N}\\ & k^*\mathrm{\ is\ a\ new\ cluster}\\ & x_{-i,c} = \{x_j | z_j = c, j \neq i\}\\ & N_{k,-i} = card(x_{-i,c}) \end{align}$$

Let’s decompose the probability that the observation $i$ belongs to cluster $k$ into its two independent factors:

$$P(z_i = k | z_{-i}, x, \alpha, \lambda) \propto P(z_i = k | z_{-i}, \alpha)P(x_i | x_{-i}, z_i=k, z_{-i}, \lambda)$$

Then:

$$\boxed{P(z_i = k | z_{-i}, \alpha) = \left\{ \begin{array}{l} \frac{N_{k,-i}}{\alpha + N - 1} \mathrm{\ if\ }k\ \mathrm{has\ been\ seen\ before}\\ \frac{\alpha}{\alpha + N - 1} \mathrm{\ if\ }k\ \mathrm{is\ a\ new\ cluster}\\ \end{array} \right.}$$ $$P(x_i | x_{-i}, z_i=k, z_{-i}, \lambda) = P(x_i | x_{-i,k}, \lambda) = \frac{P(x_i, x_{-i, k}, \lambda)}{P(x_{-i,k}|\lambda)}$$ $$\boxed{P(x_i, x_{-i, k}, \lambda) = \int P(x_i | \theta_k)\left[\prod_{j \neq i, z_j = k}P(x_j | \theta_k)\right] H(\theta_k | \lambda) d\theta_k}$$

is the marginal likelihood of all the data assigned to cluster $k$, including $i$.

If $z_i = k^*$ (new cluster) then:

$$P(x_i | x_{-i}, z_i=k^*, z_{-i}, \lambda) = \boxed{P(x_i | \lambda) = \int P(x_i | \theta)H(\theta | \lambda) d\theta}$$

Conjugate priors

Now we should choose $H$ for it to be conjugate to $F$ and have easy to compute parameters posterior. As we want $F$ to be multivariate normal: we can look on Wikipedia’s page of conjugate priors under multivariate normal with unknown $\mu$ and $\Sigma$ to see that $H$ should be normal-inverse-Wishart with prior parameters:

• $\mu_0$ initial mean guess [In my code further, I set it to the mean of whole the dataset.]
• $\kappa_0$ mean fraction (smoothing parameter) [A common value is 1. I set it to 0.]
• $\nu_0$ degrees of freedom [I set it to the number of dimensions.]
• $\Psi_0$ pairwise deviation product (matrix) [I set it to $10 \times I_d$ ($I_d$ is the $d\times d$ identity matrix). Indentity matrix makes this prior Gaussian circular, the $10$ factor should be dependant on the dataset, for instance on the mean distance between points.]

This gives us MAP estimates on parameters, for one of the clusters:

$$\begin{align} & \mu_n = \frac{\kappa_0 \mu_0 + n\tilde{x}}{\kappa_0 +n} = \mu\\ & \kappa_n = \kappa_0 + n\\ & \nu_n = \nu_0 + n \\ & \Psi_n = \Psi_0 + C + \frac{\kappa_0 n}{\kappa_0 + n}(\tilde{x} - \mu_0)(\tilde{x} - \mu_0)^T\\ & \Sigma = \frac{\kappa_n + 1}{\kappa_n * (\nu_n - d + 1)}\Psi_n \end{align}$$

with $\tilde{x}$ the sample mean and $C=\sum_{i=1}^n (x_i-\tilde{x})(x_i-\tilde{x})^T$.

Set $\kappa_{0} = 0$ to have no effect of the prior on the posterior mean. This reduces to MLE estimates if:

$$\kappa_{0} = 0, \nu_{0} = d, \|\Psi_{0}\| = 0$$

So now we can compute the posterior predictive for cluster $k$ evaluated at $x_i$

$$P(x_i | x_{-i}, z_i=k, z_{-i}, \lambda) \propto \mathcal{N}(\mu_{k,-i}, \Sigma_{k,-i})$$

Collapsed Gibbs sampling

Here is the pseudo-code of collapsed Gibbs sampling adapted from algorithm 3 of Neal’s seminal paper:

while (not converged on mus and sigmas):
    for each i = 1 : N in random order do:
        remove x[i]'s sufficient statistics from old cluster z[i]
        if any cluster is empty, remove it and decrease K
        for each k = 1 : K do
            compute P_k(x[i]) = P(x[i] | x[-i]=k)
            N[k,-i] = dim(x[-i]=k)
            compute P(z[i]=k | z[-i], Data) = N[k,-i] / (alpha + N - 1)
        compute P*(x[i]) = P(x[i] | lambda)
        compute P(z[i]=* | z[-i], Data) = alpha / (alpha + N - 1)
        normalize P(z[i] | ...)
        sample z[i] from P(z[i] | ...)
        add x[i]'s sufficient statistics to new cluster z[i]
        (possibly increase K)

Results

Here is the result of our implementation of collapsed Gibbs sampling DPGMM compared to scikit-learn’s implementation of variational Bayes DPGMM:

Code

Here is my quick-and-dirty code implementing this version of Gibbs sampling for DPGMM. You may want to comment out scikit-learn (that I used for the comparison above) if you do not have it installed.