PyMC for nonparametric clustering: Dirichlet process to estimate Gaussian mixture's parameters fails to cluster

Question

Problem setup

One of the first toy problems I wanted to apply PyMC to is nonparametric clustering: given some data, model it as a Gaussian mixture, and learn the number of clusters and each cluster's mean and covariance. Most of what I know about this method comes from video lectures by Michael Jordan and Yee Whye Teh, circa 2007 (before sparsity became the rage), and the last couple of days' reading Dr Fonnesbeck's and E. Chen's tutorials [fn1], [fn2]. But the problem is well-studied and has some reliable implementations [fn3].

In this toy problem, I generate ten draws from a one-dimensional Gaussian $\mathcal{N}(\mu=0, \sigma=1)$ and forty draws from $\mathcal{N}(\mu=4, \sigma=2)$. As you can see below, I didn't shuffle the draws, to make it easy to tell which samples came from which mixture component.

I model each data sample $y_i \sim \mathcal{N}(\mu_{z_i}, \sigma_{z_i})$, for $i=1,...,50$ and where $z_i$ indicates the cluster for this $i$th data point: $z_i \in [1,...,N_{DP}]$. $N_{DP}$ here is the length of the truncated Dirichlet process used: for me, $N_{DP}=50$.

Expanding the Dirichlet process infrastructure, each $z_i$ cluster ID is a draw from a categorical random variable, whose probability mass function is given by the stick-breaking construct: $z_i \sim Categorical(p)$ with $p \sim Stick(\alpha)$ for a concentration parameter $\alpha$. Stick-breaking constructs the $N_{DP}$-long vector $p$, which must sum to 1, by first obtaining $N_{DP}$ i.i.d. Beta-distributed draws that depend on $\alpha$, see [fn1]. And since I'd like to the data to inform my ignorance of $\alpha$, I follow [fn1] and assume $\alpha \sim Uniform(0.3, 100)$.

This specifies how each data sample's cluster ID is generated. Each of the $N_{DP}$ clusters has an associated mean and standard deviation, $\mu_{z_i}$ and $\sigma_{z_i}$. Then, $\mu_{z_i} \sim \mathcal{N}(\mu=0, \sigma=50)$ and $\sigma_{z_i} \sim Uniform(0, 100)$.

(I was previously following [fn1] unthinkingly and placing a hyperprior on $\mu_{z_i}$, that is, $\mu_{z_i} \sim \mathcal{N}(\mu_0, \sigma_0)$ with $\mu_0$ itself a draw from a fixed-parameter normal distribution, and $\sigma_0$ from a uniform. But per https://stats.stackexchange.com/a/71932/31187, my data doesn't support this kind of hierarchical hyperprior.)

In summary, my model is:

$y_i \sim \mathcal{N}(\mu_{z_i}, \sigma_{z_i})$ where $i$ runs from 1 to 50 (the number of data samples).

$z_i \sim Categorical(p)$ and can take on values between 0 and $N_{DP}-1=49$; $p \sim Stick(\alpha)$, a $N_{DP}$-long vector; and $\alpha \sim Uniform(0.3, 100)$, a scalar. (I now slightly regret making the number of data samples equal to the truncated length of the Dirichlet prior, but I hope it's clear.)

$\mu_{z_i} \sim \mathcal{N}(\mu=0, \sigma=50)$ and $\sigma_{z_i} \sim Uniform(0, 100)$. There's $N_{DP}$ of these means and standard deviations (one for each of the $N_{DP}$ possible clusters.)

Here's the graphical model: the names are variable names, see the code section below.

Problem statement

Despite several tweaks and failed fixes, the parameters learned are not at all similar to the true values that generated the data.

Currently, I'm initializing most of the random variables to fixed values. The mean and standard deviation variables are initialized to their expected values (i.e., 0 for the normal ones, the middle of their support for the uniform ones). I initialize all $z_i$ cluster IDs to 0. And I initialize the concentration parameter $\alpha=5$.

With such initializations, 100'000 MCMC iterations simply can't find a second cluster. The first element of $p$ is close to 1, and nearly all draws of $\mu_{z_i}$ for all data samples $i$ are the same, around 3.5. I show every 100th draw here for the first twenty data samples, i.e., $\mu_{z_i}$ for $i=1,...,20$:

Recalling that the first ten data samples were from one mode and the rest were from the other, the above result clearly fails to capture that.

If I allow random initialization of the cluster IDs, then I obtain more than one cluster but the cluster means all wander around the same 3.5 level:

This suggests to me that it's the usual problem with MCMC, that it can't reach another mode of the posterior from the one it's at: recall that these different results happen after just changing the initialization of the cluster IDs $z_i$, not their priors or anything else.

Am I making any modeling mistakes? Similar question: https://stackoverflow.com/q/19114790/500207 wants to use a Dirichlet distribution and fit a 3-element Gaussian mixture and is running into somewhat similar problems. Should I consider setting up a fully conjugate model and using Gibbs sampling for this kind of clustering? (I implemented a Gibbs sampler for the parametric Dirichlet distribution case, except using a fixed concentration $\alpha$, back in the day and it worked well, so expect PyMC to be able to solve at least that problem handily.)

Appendix: code

import pymc
import numpy as np

### Data generation

# Means and standard deviations of the Gaussian mixture model. The inference
# engine doesn't know these.
means = [0, 4.0]
stdevs = [1, 2.0]

# Rather than randomizing between the mixands, just specify how many
# to draw from each. This makes it really easy to know which draws
# came from which mixands (the first N1 from the first, the rest from
# the secon). The inference engine doesn't know about N1 and N2, only Ndata
N1 = 10
N2 = 40
Ndata = N1+N2

# Seed both the data generator RNG  as well as the global seed (for PyMC)
RNGseed = 123
np.random.seed(RNGseed)

def generate_data(draws_per_mixand):
    """Draw samples from a two-element Gaussian mixture reproducibly.

    Input sequence indicates the number of draws from each mixand. Resulting
    draws are concantenated together.

    """
    RNG = np.random.RandomState(RNGseed)
    values = np.hstack([RNG.normal(means[i], stdevs[i], ndraws)
                        for (i,ndraws) in enumerate(draws_per_mixand)])
    return values

observed_data = generate_data([N1, N2])


### PyMC model setup, step 1: the Dirichlet process and stick-breaking

# Truncation level of the Dirichlet process
Ndp = 50

# "alpha", or the concentration of the stick-breaking construction. There exists
# some interplay between choice of Ndp and concentration: a high concentration
# value implies many clusters, in turn implying low values for the leading
# elements of the probability mass function built by stick-breaking. Since we
# enforce the resulting PMF to sum to one, the probability of the last cluster
# might be then be set artificially high. This may interfere with the Dirichlet
# process' clustering ability.
#
# An example: if Ndp===4, and concentration high enough, stick-breaking might
# yield p===[.1, .1, .1, .7], which isn't desireable. You want to initialize
# concentration so that the last element of the PMF is less than or not much
# more than the a few of the previous ones. So you'd want to initialize at a
# smaller concentration to get something more like, say, p===[.35, .3, .25, .1].
#
# A thought: maybe we can avoid this interdependency by, rather than setting the
# final value of the PMF vector, scale the entire PMF vector to sum to 1? FIXME,
# TODO.
concinit = 5.0
conclo = 0.3
conchi = 100.0
concentration = pymc.Uniform('concentration', lower=conclo, upper=conchi,
                             value=concinit)

# The stick-breaking construction: requires Ndp beta draws dependent on the
# concentration, before the probability mass function is actually constructed.
betas = pymc.Beta('betas', alpha=1, beta=concentration, size=Ndp)

@pymc.deterministic
def pmf(betas=betas):
    "Construct a probability mass function for the truncated Dirichlet process"
    # prod = lambda x: np.exp(np.sum(np.log(x))) # Slow but more accurate(?)
    prod = np.prod
    value = map(lambda (i,u): u * prod(1.0 - betas[:i]), enumerate(betas))
    value[-1] = 1.0 - sum(value[:-1]) # force value to sum to 1
    return value

# The cluster assignments: each data point's estimated cluster ID.
# Remove idinit to allow clusterid to be randomly initialized:
idinit = np.zeros(Ndata, dtype=np.int64)
clusterid = pymc.Categorical('clusterid', p=pmf, size=Ndata, value=idinit)

### PyMC model setup, step 2: clusters' means and stdevs

# An individual data sample is drawn from a Gaussian, whose mean and stdev is
# what we're seeking.

# Hyperprior on clusters' means
mu0_mean = 0.0
mu0_std = 50.0
mu0_prec = 1.0/mu0_std**2
mu0_init = np.zeros(Ndp)
clustermean = pymc.Normal('clustermean', mu=mu0_mean, tau=mu0_prec,
                          size=Ndp, value=mu0_init)

# The cluster's stdev
clustersig_lo = 0.0
clustersig_hi = 100.0
clustersig_init = 50*np.ones(Ndp) # Again, don't really care?
clustersig = pymc.Uniform('clustersig', lower=clustersig_lo,
                          upper=clustersig_hi, size=Ndp, value=clustersig_init)
clusterprec = clustersig ** -2

### PyMC model setup, step 3: data

# So now we have means and stdevs for each of the Ndp clusters. We also have a
# probability mass function over all clusters, and a cluster ID indicating which
# cluster a particular data sample belongs to.

@pymc.deterministic
def data_cluster_mean(clusterid=clusterid, clustermean=clustermean):
    "Converts Ndata cluster IDs and Ndp cluster means to Ndata means."
    return clustermean[clusterid]

@pymc.deterministic
def data_cluster_prec(clusterid=clusterid, clusterprec=clusterprec):
    "Converts Ndata cluster IDs and Ndp cluster precs to Ndata precs."
    return clusterprec[clusterid]

data = pymc.Normal('data', mu=data_cluster_mean, tau=data_cluster_prec,
                   observed=True, value=observed_data)

References

1. fn1: http://nbviewer.ipython.org/urls/raw.github.com/fonnesbeck/Bios366/master/notebooks/Section5_2-Dirichlet-Processes.ipynb
2. fn2: http://blog.echen.me/2012/03/20/infinite-mixture-models-with-nonparametric-bayes-and-the-dirichlet-process/
3. fn3: http://scikit-learn.org/stable/auto_examples/mixture/plot_gmm.html#example-mixture-plot-gmm-py

Answer 1

I'm not sure if anyone is looking at this question any more but I put your question in to rjags to test Tom's Gibbs sampling suggestion while incorporating insight from Guy about the flat prior for standard deviation.

This toy problem might be difficult because 10 and even 40 data points are not enough to estimate variance without an informative prior. The current prior σzi∼Uniform(0,100) is not informative. This might explain why nearly all draws of μzi are the expected mean of the two distributions. If it does not alter your question too much I will use 100 and 400 data points respectively.

I also did not use the stick breaking process directly in my code. The wikipedia page for the dirichlet process made me think p ~ Dir(a/k) would be ok.

Finally it is only a semi-parametric implementation since it still takes a number of clusters k. I don't know how to make an infinite mixture model in rjags.

library("rjags")

set1 <- rnorm(100, 0, 1)
set2 <- rnorm(400, 4, 1)
data <- c(set1, set2)

plot(data, type='l', col='blue', lwd=3,
     main='gaussian mixture model data',
     xlab='data sample #', ylab='data value')
points(data, col='blue')

cpd.model.str <- 'model {
  a ~ dunif(0.3, 100)
  for (i in 1:k){
    alpha[i] <- a/k
    mu[i] ~ dnorm(0.0, 0.001)
    sigma[i] ~ dunif(0, 100)
  }
  p[1:k] ~ ddirich(alpha[1:k])
  for (i in 1:n){
    z[i] ~ dcat(p)
    y[i] ~ dnorm(mu[z[i]], pow(sigma[z[i]], -2))
  }
}' 


cpd.model <- jags.model(textConnection(cpd.model.str),
                        data=list(y=data,
                                  n=length(data),
                                  k=5))
update(cpd.model, 1000)
chain <- coda.samples(model = cpd.model, n.iter = 1000,
                      variable.names = c('p', 'mu', 'sigma'))
rchain <- as.matrix(chain)
apply(rchain, 2, mean)

Answer 2

The poor mixing that you are seeing is most likely because of the way that PyMC draws samples. As explained in section 5.8.1 of the PyMC documentation, all elements of an array variable are updated together. In your case, that means it will try to update the entire clustermean array in one step, and similarly for clusterid. PyMC doesn't do Gibbs sampling; it does Metropolis where the proposal is chosen by some simple heuristics. This makes it unlikely to propose a good value for an entire array.