idioms for mixture models and inference convergence/speed?

[yarden]

I'm having trouble fitting simple Gaussian mixture models in PyMC3. On only 500 data points, default inference (NUTS with advi) takes 3-4 minutes. Metropolis inference doesn't seem to work well on this problem. It's unclear if this is a failure (or bug) of inference, or if it's an issue in the way I've defined the model. The code is here and the core definition of the mixture uses NormalMixture (as was suggested by @twiecki here):

   with pm.Model() as basic_model:
        weights = pm.Dirichlet("weights",
                               np.ones_like(np.ones(num_mixtures)))
        # uniform prior on mean (mu)
        mu = pm.Uniform("mu", mu_min, mu_max, shape=(num_mixtures,))
        # uniform prior on standard deviation (sd)
        sd = pm.Uniform("sd", sd_min, sd_max, shape=(num_mixtures,))
        obs = pm.NormalMixture("obs", weights, mu, sd=sd,#tau=tau,
                               observed=data)

Part of the confusion for me in debugging this is that there are several very different ways to define (finite) mixture models in PyMC3 in the documentation and other resources, and it's not clear which one is the "preferred" one -- or at least, which inference approach should be used with each model writing idiom. The different ways of writing simple mixtures I spotted with PyMC3 were:

1. Using explicit categorical variables for assignments (with pm.Categorical), which appears to make the simple Metropolis() step method fail (as described here in detail). (It seems that with categorical variables for assignments in the model, one needs either Gibbs sampling or categorical variable specific steps to make sampling-based inference work?)
2. Using DensityDist and custom log-likelihood functions -- here, ADVI inference is used.
3. Using categorical pm.Potential functions (to solve label-switching problem, it seems).
4. UsingNormalMixture (or Mixture class), without explicitly modeling categorical assignments, as suggested by @twiecki here.

All these approaches encode essentially the same Gaussian mixture model. It'd be helpful to have some guidelines on the "canonical" solution here, because these models are so frequently used and familiar, and some suggestions about the right inference method to use with each. For instance, why is the default (NUTS + advi) so slow given the way I've written the mixture model?

Thanks very much.

Best, Yarden

[twiecki]

That's indeed not explained consistently. Should improve those docs.

4. is the best solution. Does the trace / convergence look OK?

[aseyboldt]

About the long run time: Most of the time is spent in advi, the number of iterations is way larger than it should be (set n_init to maybe 30000 for this problem). But even then, nuts is quite slow and the effective sample size isn't great either. Part of the problem is the mass matrix. The code uses a diagonal mass matrix, which is appropriate for high dimensional problems, but for a 5 dimensional problem with covariances in the posterior I'd rather use a full mass matrix. Something like this should work better:

with basic_model:
    trace1 = pm.sample(1000, init='advi', n_init=50000, tune=500)
    samples = np.array([basic_model.dict_to_array(p) for p in trace1[500:]])
    cov = np.cov(samples.T)
    step = pm.NUTS(scaling=cov, is_cov=True)
    start = [trace1[-i - 1] for i in [0, 100, 200, 300]]
    trace2 = pm.sample(2000, step=step, init=None, start=start, njobs=4, tune=1000)

On my laptop that takes about 1,5min (or 1min on master). Still not great if we compare it to stan, which takes about a minute to compile and then fits the model in 8s

data {
    int<lower=1> K;
    int<lower=1> N;
    real y[N];
}
parameters {
    simplex[K] theta;
    real mu[K];
    real<lower=0> sigma[K];
}
model {
    real ps[K];
    sigma ~ cauchy(0, 2.5);
    mu ~ normal(0, 10);
    
    for (n in 1:N) {
        for (k in 1:K) {
            ps[k] = log(theta[k]) + normal_lpdf(y[n] | mu[k], sigma[k]);
        }
        target += log_sum_exp(ps);
    }
}

I guess reasons for the difference might be:

• stan is using a diagonal mass matrix (I couldn't get it to converge with a full mass matrix) and pymc3 uses a full mass matrix.
• The mass matrix adaptation works differently: stan does its adaptation in the warmup, pymc3 in advi and a separate chain.
• pymc3 spends time building the computational graph. That's part of the compilation phase for stan
• I guess there is quite a bit of python overhead for small models like this. (eg In the grad evaluations Split takes up about 35% of the runtime).

The raw sampling speed after adaptation is about 800 samples/s for stan and 200 samples/s for pymc3 (running only one chain at a time). Gradient evaluations take 120μs for stan and 300μs for pymc3.

[twiecki]

@aseyboldt Thanks for looking into this. Split is being coded in C: Theano/Theano#5499 If you want, I'd be quite curious if that speeds things up.

I also have a PR that tries to do nuts init more similar to Stan: #1738

[aseyboldt]

Sorry, there was a typo in my stan code. I wrote log(K) instead of log(k).

The timings of the gradient evaluations didn't change much, but the posterior is more complicated now. Stan takes something like 12s for most chains and from time to time the sampler seems to get stuck in a really bad local optimum. I did not observe this with pymc3. I guess using advi to initialize the chains has the advantage (or disadvantage depending on how you want to look at it) that you can force the sampler into a particular region of the posterior.

[fonnesbeck]

I've been meaning to submit a quick PR that drastically reduces the default n_init for ADVI initialization. It seems like there is very little benefit to running ADVI close to the posterior mode, given what we know about its relationship to the typical set, right? Has anyone played with this at all?

[twiecki]

Really what we should do is have a dynamic measure of convergence and stop dynamically.

[fonnesbeck]

We don't require convergence when using ADVI for initialization, do we? All we want is a point that is in or near the typical set, then we let NUTS (or whatever step method) take over and sample.

[twiecki]

I think the cov matrix is more critical and would feel more comfortable if we had convergence there.

[yarden]

Thanks for all these comments. I was able to narrow down a case where there appears to be a convergence problem using the recommended method of NormalMixture. The script is: test_pymc_mixture.py. It runs the same Metropolis-based inference twice on identical data, which gives these very different results (the first set of result appear to be close to the correct posterior mean; the second set of results are very off). It might be a helpful test case.

I'd be interested in any suggestions you might have on the best way to solve this inference problem in PyMC? I'm not sure ADVI would solve this issue but regardless, having a shorter default initialization sounds like a good idea as @fonnesbeck said.

[aseyboldt]

This is pretty much expected behaviour. Your posterior has multiple modes (both chains describe "reasonable" explanations for the dataset) and if they are separated by a large enough energy barrier, nuts won't jump between them. If you initialize several chains from the same mode, you can more or less force all of them to end up in the same mode – that's what I wrote earlier –, but nuts and metropolis won't give you an overview about which modes exist.
How to sample from multimodal posteriors is an open question, but tempering should help. I've seen a few papers about using tempering in hamiltonian mcmc, but I don't know of any general purpose implementation. In a low dimensional problem like this you might be able to get away with using some kind of metropolis sampling (maybe using http://dan.iel.fm/emcee/current/user/pt/, which you should be able to feed a pymc3 model).

[yarden]

I'm not sure the posterior here is multimodal; or at least, I don't think these two modes are equal. It looks to me like the second solution should have a substantially poorer fit to the data - though I'm not sure how to access the log scores in PyMC3? In the first solution, the green Gaussian explains the bulk of the data and the red Gaussian accounts for the "bump" on the right.

[twiecki]

Mixture models are almost always multi-modal. You could add potentials to break the symmetry (e.g. http://pymc-devs.github.io/pymc3/notebooks/gaussian_mixture_model.html).

[yarden]

By symmetry you mean multimodes induced by label switching? If so it doesn't look like the two modes PyMC finds in my example are symmetric. I tried to correct for label switching when processing the posterior samples in my code, but I am not sure if this is right (lines 69-74 in code).

[twiecki]

True, it seems like you legitimately have two modes. At least your plot above strongly suggests so (you need to add the two distributions to see how well it matches the data). These are challenging for sure. You could try SVGD which was recently merged, or, as @aseyboldt suggested emcee (#1689). I'm closing this but feel free to ask more questions.