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