Speed of NUTS sampling

[narendramukherjee]

I am trying to detect changepoints on several trials of a process with categorical emissions. I find that NUTS sampling in really slow (something like 4s/it) - Metropolis is faster (of course) and samples about 200 samples/s, but I need to run it to about 500k samples to get reasonable Gelman-Rubin convergence values. I just wanted to check in and see if there was something incorrect (or open to improvement) in the way that I have set up my model, and if NUTS can be gotten to work better in this case.

Here's one example of what I am trying to do: 32 trials of a Categorical process with 17 emissions, each trial has two changepoints. There are 7 unique states of emissions in the set of trials - depending upon the type of trial, it changes from a start state (which is the same for all trials) to 1 (out of 2 possible) states at changepoint1, and then changes again to 1 (out of 4 possible) states at changepoint2. Each trial has 150 emissions in total:

with pm.Model() as model:
	# Dirichlet prior on the emission/spiking probabilities - 7 states (1 start state, 2 from changepoint1 to 
        changepoint2, 4 from changepoint2 to end of the trial)
	p = pm.Dirichlet('p', np.ones(num_emissions), shape = (7, num_emissions))

	# Uniform switch times
	# First changepoint
	t1 = pm.Uniform('t1', lower = 20, upper = 60, shape = num_trials)
	# Second changepoint
	t2 = pm.Uniform('t2', lower = t1 + 20, upper = 130, shape = num_trials)

	# Get the actual state numbers based on the switch times.
	states = []
	for i in range(num_trials):
		states1 = tt.switch(t1[i] >= np.arange(150), 0, set1[i])
		states2 = tt.switch(t2[i] >= np.arange(150), states1, set2[i])
		states.append(states2)

	# Define the log-likelihood function
	def logp(value):
		value = tt.cast(value, 'int32')
		loglik = 0
		for i in range(32):
			loglik += tt.sum(tt.log(p[states[i], value[i, :]]))
		return loglik

	# Categorical observations
	obs = pm.DensityDist('obs', logp, observed = {'value': data[0, :, :150]})

	# Inference button :D
	tr = pm.sample(500000, init = None, step = pm.Metropolis(), njobs = 2, start = {'t1': 25.0, 't2': 120.0})

Initializing with find_MAP() gives errors at times with this model, so I stuck to the start values I specified above as they seem to work well generally with Metropolis, although I have to sample in the range of 500-600k samples to get reasonable convergence (as I said above).

[aseyboldt]

On first glance it looks like the logp and gradient evaluations might just be very slow. You could try to replace the loops by broadcasting:

for i in range(num_trials):
    states1 = tt.switch(t1[i] >= np.arange(150), 0, set1[i])
    states2 = tt.switch(t2[i] >= np.arange(150), states1, set2[i])
    states.append(states2)

to something like this (untested, I don't have a working example):

states1 = tt.switch(t1 >= np.arange(150)[:, None], 0, set1)
states2 = tt.switch(t2 >= np.arange(150)[:, None], states1, set2)

And something similar in logp.
I'd probably try to rewrite the logp on paper somehow (including the stuff with arange), maybe you can find a nicer representation.
If that doesn't do it, you can have a look at the stats of the nuts sampler. If trace['tree_depth'] returns lots of numbers close to 10, then NUTS is looking at very long trajectories, which might slow it down even further.

[aseyboldt]

PS: closing this, as I don't think this is actually an issue with pymc3. But feel free to continue to write here :-)

[narendramukherjee]

Thanks for those suggestions, I figured out ways to broadcast the for loops with your suggestions, and here's the model I have:

with pm.Model() as model:
	# Dirichlet prior on the emission/spiking probabilities - 7 states (1 start state, 2 from changepoint1 to changepoint2, 4 from changepoint2 to end of the trial)
	p = pm.Dirichlet('p', np.ones(17), shape = (7, 17))

	# Uniform switch times
	# First changepoint
	t1 = pm.Uniform('t1', lower = 20, upper = 60, shape = 32)
	# Second changepoint
	t2 = pm.Uniform('t2', lower = t1 + 20, upper = 130, shape = 32)

	# Get the actual state numbers based on the switch times
	states1 = tt.switch(t1 >= np.repeat(np.arange(150)[:, None], 32, axis = 1), 0, set1[:])
	states2 = tt.switch(t2 >= np.repeat(np.arange(150)[:, None], 32, axis = 1), states1, set2[:])

	# Define the log-likelihood function
	def logp(value):
		value = tt.cast(value, 'int32')
		loglik = tt.sum(tt.log(p[states2.T, value]))
		return loglik

	# Categorical observations
	obs = pm.DensityDist('obs', logp, observed = {'value': data[:, :150]})

Metropolis sampling does speed up drastically with these two changes to about 500 samples/sec. NUTS sampling is still very slow though (about 2s/it) - I used the default initialization with ADVI. Funnily though, HamiltonianMC is much faster, doing about 80 samples/sec. I am surprised at the huge difference between NUTS and HMC.

As you had suggested last time, here are the sampler stats for NUTS sampling (I collected only about 80 samples in this run):

In [35]: tr.get_sampler_stats('tree_size')
Out[35]: 
array([   63.,    15.,   127.,  1023.,  1023.,  1023.,  1023.,  1023.,
        1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,
        1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,
        1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,
        1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,
        1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,
        1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,
        1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,
        1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,
        1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,
        1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.,  1023.])
In [36]: tr.get_sampler_stats('depth')
Out[36]: 
array([ 6,  4,  7, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
       10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
       10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
       10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
       10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
       10, 10, 10])

[narendramukherjee]

Finally, I am attaching datasets here so that you (and anyone else looking at this issue) can run the code I posted and look at the problem. Thanks so much!

data.zip

[aseyboldt]

That looks bad. NUTS does not try longer trajectories than 1023 by default (you could increase this by setting max_treedepth to something higher than 10). This means that it thinks the samples are still highly correlated. Usually, that indicates a problem with the parametrization or the model itself. You could try to use much more informative priors. Are the changepoints related somehow?
It's a bit late now, I'll have a closer look tomorrow.