NUTS sampler - trying to fix the step size (step_scale vs step_size?)

Questions
1

Hi,

I am implementing an algorithm where I interleave the NUTS sampler within a larger program basically calling pm.sample() every 50/100 steps of my larger program…hence I cant afford to tune excessively each time I enter the sampling window.

I wanted to tune the sampler just once and extract the step size param so in subsequent sampling windows I can skip tuning (tune=0 and adapt_step_size=False) and just substitute the step size from the earlier run. I am writing it like below but I reckon I am missing the relationship between step_scale and step_size → the doc says it is something like (1/n**(1/4)) …I don’t quite understand it, like what is n?

When I run the below…the extracted step_size from trace.get_sampler_stats('step_size') doesnt bear any resemblance to the step_scale (but I think they mean different things?)

 step=pm.NUTS(step_scale=0.05, adapt_step_size=False)
 trace = pm.sample(100, tune=0, step=step, chains=1,discard_tuned_samples=False)
1 Like
2

I don’t know about step_size vs step_scale, but to be able to skip tuning you’ll need also “good” initial points and the mass matrix in addition to the step size, otherwise sampling will be very inefficient. I’ll try to add some references whenever I have some free time.

3

The best place to get an idea of setting up NUTS manually is to look at the init_nuts function

Something like below should work:

corr = np.asarray([[1., .65], [.65, 1]])
std = np.asarray([1., 5.])
cov = corr * np.outer(std, std)
chol_cov = np.linalg.cholesky(cov)

# Set up model
with pm.Model() as m:
    x = pm.MvNormal('x', mu=np.asarray([10., 0.]), chol=chol_cov, lower=True)
    y = pm.HalfNormal('y', 1.)

# Initial sample
with m:
    trace = pm.sample(return_inferencedata=False, compute_convergence_checks=False)

from pymc.step_methods.hmc import quadpotential

def fix_step_from_trace(trace, model):
    cov_est = pm.trace_cov(trace, vars=model.value_vars, model=model)
    potential = quadpotential.QuadPotentialFull(cov_est)
    # Use the largest step_size from previous tuning, 
    # alternatively we can compute the mean here
    step_size = np.max(trace.get_sampler_stats('step_size'))
    step_scale = step_size * (cov_est.shape[0] ** 0.25)
    step = pm.NUTS(vars=model.value_vars,
                   model=model,
                   potential=potential,
                   step_scale=step_scale,
                   adapt_step_size=False)
    step.tune = False
    return step

tuned_step = fix_step_from_trace(trace, m)

# New sample without tuning
with m:
    trace_new = pm.sample(initvals=trace[-1],
                          tune=0,
                          step=tuned_step, 
                          return_inferencedata=False,
                          compute_convergence_checks=False)
2 Likes