Sampling with MCLMC

[yomori]

Hi,
Thanks for the nice package, and also for the implementation of MCLMC.
I'm currently trying to sample a very large problem using the mclmc sampler that has two different types of parameters, the initial condition volume ~O(10^6) and two other parameters which are just scalars. My original code was written for numpyro and now I'm trying to call BlackJax's MCLMC sampler following the instructions https://num.pyro.ai/en/latest/tutorials/other_samplers.html
and here https://blackjax-devs.github.io/sampling-book/algorithms/mclmc.html

I have a few very basic questions and I'm wondering if anyone can give me some hints:

1. Suppose I want to start my chains from the "truth". After I do blackjax.mclmc_find_L_and_step_size, and look at blackjax_state_after_tuning, it looks very different from the what was fed in as initial_state ( to the point that it looks like its picking a random point from the prior). This is presumably because in the process of determining L and step_size, it made wrong steps (with inappropriate choices of L and step_size) that lead to divergence. Is there any way to avoid this from happening or do I have to manually search for those tuning parameters that give reasonable posteriors (it seems like a lot of fine tuning)?

2. Suppose I do find the correct L and step_size from trial and error. How critical is it to use blackjax_state_after_tuning for the initial_state argument in blackjax.util.run_inference_algorithm? Can I just use the predetermined values for L and step_size and feed in the actual initial state instead?

3. When I look at the chains, the way they move is very different from what I'm familiar with (such as NUTS) -- they look like they are traveling along some trajectory instead of jumping around. Is this expected? or is this just an consequence of bad choice of L and step_size?

Appreciate any help!

[junpenglao]

pinging our MCLMC ninja @reubenharry

[reubenharry]

Hi,

Great to hear you're trying MCLMC! I'm travelling this weekend, so can give a proper answer after, but for now just some quick pointers:

1. So in theory you could manually tune L and eps, but this is generally not practical. The goal is that the tuning algorithm should work well in most cases, so I think the next step would be for me to help you debug the tuning.
2. Not at all critical to use blackjax_state_after_tuning. It's just the final state from the chain used in tuning. If tuning works, it should just be a draw from the target distribution after burn in.
3. I think the chains should look fairly similar to NUTS in general. This suggests to me a problem with the choice of L.

Have you tried running your problem on blackjax's NUTS implementation? This would be a good way to establish whether there's a specific problem with MCLMC here, or if the move from numpyro to blackjax incurred some other issue.

After that, I can suggest what to look at next, or if you can share your code, I can also try to work out what's going wrong in the tuning. In our experience, if NUTS works, MCLMC will too, but it's always possible this isn't the case, in which case, this would be an interesting problem for us to study.

[yomori]

Hi @reubenharry

Thanks for your reply!
Ok so given your 3 points, I'm highly suspicious of the connecting piece between Numpyro and Blackjax so I'll take a look at that again, and perhaps try it with a simpler model first.

Have you tried running your problem on blackjax's NUTS implementation? This would be a good way to establish whether there's a specific problem with MCLMC here, or if the move from numpyro to blackjax incurred some other issue.

This is a very good suggestion. I will try this out.

After that, I can suggest what to look at next, or if you can share your code, I can also try to work out what's going wrong in the tuning.

Thanks that would be amazing (but I'll make sure to exhaust all the possibilities on my end first before wasting your time)!

In our experience, if NUTS works, MCLMC will too, but it's always possible this isn't the case, in which case, this would be an interesting problem for us to study.

Ok thats also good to know too. I don't think the problem is particularly unique other than the fact that the parameter space is extremely large, so I'm hopeful that this will work eventually.