Dear Numpyro experts
I’m using NUTS sampler and I’m wandering if one can explain me how some variables of the progress bar are computed. Here is a current example
sample: 45% | | 4725/10500 [17:42:54<99:27:39, 62.00s/it, 255 steps of size 6.87e-06. acc. prob=1.0]
For the NUTS parameters I only change max_tree_depth=8 and I ask for 500 warmup and 10,000 samples. At the beginning (warmup + a thousand of samples) the expected overall time was arround 10 hours with few steps of the same kind of step size. So, I wander why currently the NUTS algo uses 255 steps, so the maximum of steps (ie. 2^max_tree_depth-1) ???
May be you will invite me to perform some tests or trry some optimisations.
Thanks
I think you got it right. What is your concern here? If you want to use a small number of steps you can reduce max tree depth. You can also specify different tree depths for warmup and sample phases: e.g. by providing a pair (5, 10) to max tree depth, warmup phase will have maximum numsteps is 31, sample phase will have maximum numsteps is 1023.
Well, I was wandering if it is a normal behaviour of NUTS algo to start (even after warm-up) with low number of steps and then to progressively increase the number of steps up the maximum??? my concern is that it is very long… 100h for 10,000 samples, but I guess even using JAX autodiff my use-case is not an easy-one.
It is hard to tell - it seems to be normal during warmup phase. If your model is large, you can use approximated methods like SVI.