Is my NUTS normal?

An actual serious question:

For me the optimal step size found so far (step_size_bar) during tuning steps always follow the same trajectory for all models I’ve seen/tried. Namely that is dropping straight to zero and then raising slowly up to a (hopefully) stable point (say around 0.8 for example).

Those with lots of experience; is this normal?

I tried setting nuts_kwargs={'step_scale':1.0} to help the step size converge to 0.8 (which seemed to be best for my model, but again it just drops straight to 0 at start of tuning anyway so doesn’t improve the time to become optimal.

Thoughts?

If there is no warning usually this is fine. However it is true that it looks a bit weird that it first drop to a very small value, seems maybe there is room for improvement? cc @aseyboldt

I agree that there might be room for improvement in the step size adaptation, but I think this particular phenomenon isn’t as bad as it might seem at first: step_size_bar is the step size it would take if tuning were to be disabled in the next step, not the step size it actually using (that is trace.step_size). You can have a look at the section 3.4 of the nuts paper for some more info (https://arxiv.org/pdf/1111.4246.pdf)
If you want to play with this, you could try to change gamma, t0 and k in the nuts_kwargs of pm.sample.

1 Like