Barker proposal

[jecampagne]

Hello,
Does the Barker proposal in BlackJAX follows the same implementation as the Numpyro BarkerMC?
How the warm up optimisation is done: ie. is there any mass-matrix as well as step size optimisation?
Thanks.

[AdrienCorenflos]

Well, I'd argue their implementation follows ours based on timelines :D, it looks similar but I can't (don't have time) to comment on the specifics.

As a design rule, we do not implement warm up/adaptation as part of kernels, but as outer utilities, so the answer is neither. Looking at their code, they use HMC adaptation, so you can also do this here with very little trouble:

I'm rather sure that if you change nuts for barker below, it will just work

warmup = blackjax.window_adaptation(blackjax.nuts, joint_logdensity)

# we use 4 chains for sampling
n_chains = 4
rng_key, init_key, warmup_key = jax.random.split(rng_key, 3)
init_keys = jax.random.split(init_key, n_chains)
init_params = jax.vmap(init_param_fn)(init_keys)

@jax.vmap
def call_warmup(seed, param):
    (initial_states, tuned_params), _ = warmup.run(seed, param, 1000)
    return initial_states, tuned_params

warmup_keys = jax.random.split(warmup_key, n_chains)
initial_states, tuned_params = jax.jit(call_warmup)(warmup_keys, init_params)

(example is from here https://blackjax-devs.github.io/sampling-book/models/change_of_variable_hmc.html)

This question makes me think that Zanella and friends have derived results on calibration which we may want to have a look at eventually https://academic.oup.com/biomet/article/110/3/579/6764577

[ismael-mendoza]

Sorry for the delay @AdrienCorenflos , I believe it might be helpful to include the mass matrix in the kernel as at the moment I'm not sure how I would adapt the Barker kernel with the HMC adaptation otherwise? I'm pretty new to blackjax so I'll defer to what you both think is best

[AdrienCorenflos]

Ok, feel free to send a PR and I'll review (or someone else will) it then. I've got my PhD defence next week, so I can't commit to coding the change myself, but it should not be too much and is anyway a fairly good first crack at the code base.

[AdrienCorenflos]

#725

[ismael-mendoza]

Thanks @AdrienCorenflos I can try to take a stab at it in the coming days

[ismael-mendoza]

@AdrienCorenflos I opend a draft PR in #730 , it's still work in progress but most of the code to follow Appendix G of the Barker paper is there. I'll post additional questions there for the new implementation. Hope your defense goes well!

[jecampagne]

Hello @AdrienCorenflos and @junpenglao
I get in touch with Samuel Livingstone and he points me a possible improvement of the Barker proposal :
"
It’s also worth mentioning that there is a variation on the Barker algorithm called ‘bi-modal Barker’, which is introduced in Section 4.1 of this paper and described in lines 26-38 of the code here. We find that on average the bi-modal version improves ESSs by a factor of 2 over the regular version, so it may be worth trying in your example.
"

[AdrienCorenflos]

Yes it's what I was mentioning by

This question makes me think that Zanella and friends have derived results on calibration which we may want to have a look at eventually https://academic.oup.com/biomet/article/110/3/579/6764577

I think this is a bit more work than just the pre-conditioning so I'd rather order the two changes by 1. putting the mass matrix in, 2. allowing the user to change the noise distribution

But on the principle, yes, I'd be happy if we added this to the library.

[jecampagne]

Ho, sorry I do not realize that Samuel was mentioning the same paper as the on you mentioned too.

[AdrienCorenflos]

No worries, that's fine, I agree that it's a good addition, just second in the list of priorities (if anything because of simplicity of use and implementation).