Subclassing existing sampling algorithms / kernels?

[mhinne]

Hi all,

I've been toying around with Blackjax for a bit and I'm loving the experience. Boy are things fast :-)

Now that I've built up some experience with the existing samplers, I'm working on creating some extensions of these samplers for my specific needs. For instance, I've used the suggestions in #275 to construct a Gibbs kernel that combines elliptical slice sampling and RMH for a Gaussian process + its hyperparameters. This works very well, but now that I want to move on from ad-hoc fixes and a bit more clean design, I run into two (closely related) issues:

1. I wanted to subclass the ESS sampler to exploit some independencies in my specific likelihood. However, I cannot import from blackjax.util. I can of course copy-paste the generate_gaussian_noise and linear_map functions from the source code, but this seems clunky for an otherwise straightforward procedure.
2. I'm trying to use the Gibbs sampler from Metropolis-within-Gibbs #275 within SMC, to faster explore the whole space of the Gaussian process hyperparameters (cf. Svensson et al., 2015, except that I am using a non-conjugate likelihood), but this requires that I replace my ad-hoc Gibbs step function (as suggested in Metropolis-within-Gibbs #275) with a proper Gibbs MCMCSamplingAlgorithm object. However, I cannot import this type from blackjax.base.

Am I missing an obvious way of importing these objects?

Thanks for your input :-)

[rlouf]

Hey,

1. Could you please paste here the import that generates an error? generate_gaussian_noise is imported in a couple places in the lirbary using from blackjax.util import generate_gaussian_noise.
2. You can create a new SMC kernel by calling from blackjax.smc.base import kernel (here in the code). This constructor requires two functions:
   • a mcmc_kernel_factory (you can see what this corresponds to here), which is just a function that returns a sampling step function given a logprob function.
   • a mcmc_kernel_generator, which given a position in the sample space returns the state that your sampling step function takes as an input. This is important for HMC, for instance: SMC takes positions as inputs, but HMC needs extra information.

This means that you don't have to subclass MCMCSamplingAlgorithm at all, and can just work with pure step functions like the Gibbs function you constructed! You are touching on some minor clunkiness with the SMC sampler that I'm trying to solve here #279, and the lack of documentation 😬 but you should totally be able to do it. This is a perfectly reasonable use case, thus it shouldn't be too difficult to do; if it turns out to be we'll need to refactor the internals a little.

If that answer is not enough, please share a simple example and we'll walk through it together.

[rlouf]

Awesome! Is that Metropolis-within-Gibbs-within-SMC then? 😁

[mhinne]

Even better! It is (Elliptical Slice Sampling and Metropolis)-within-Gibbs-within-SMC :-)

[Hesterhuijsdens]

Hi @rlouf ,

I am working together with @mhinne to combine the SMC approach with our Gibbs kernel. As mentioned above, we thought we could just create an instance of MCMCSamplingAlgorithm (or, when using the new Blackjax version, an init function and a step function) for our Gibbs kernel, and feed this into the tempered_smc sampler.

However, we now noticed that, within the tempered SMC script, the tempered_logposterior_fn is defined and given to the (Gibbs) kernel. However, since we want to use Elliptical Slice Sampling within the Gibbs kernel, we would require a tempered log likelihood function instead.

We have now adapted our local version of the tempered script so that it feeds a tempered_loglikelihood_fn into the Gibbs kernel, but I was wondering if there is a nicer way to achieve this, without having to adapt the Blackjax scripts? 🙂

[maxhinne]

Would it be possible to have the SMC temperature parameter (lmbda) be part of the mcmc_init_fn or mcmc_step_fn, in lines 128 and 131 in https://github.com/blackjax-devs/blackjax/blob/main/blackjax/smc/tempered.py? That way, any custom kernel (that doesn't take the tempered_logposterior_fn as input, like our Gibbs kernel) can use the temperature internally :-)

[rlouf]

That's a fair point, please open an issue that describes your problem.

[AdrienCorenflos]

By the way, have you considered using this instead of ESS? It's been proven to work better on a number of applications. We have an implementation in BlackJAX. See here for example
https://blackjax-devs.github.io/blackjax/examples/LogisticRegressionWithLatentGaussianSampler.html

[mhinne]

Hmm I had not, mostly out of ignorance. It seems this approach can not only outperform ESS for the latent Gaussians themselves, but also for their hyperparameters which is a big challenge in our project. So I'll definitely give this a look, thanks for the tip! :-)

[AdrienCorenflos]

Check their "z-" sampler.

[mhinne]

That looks spot-on indeed! Not to sound too lazy, but this is not part of BlackJAX (yet), right, only the marginal sampler? :)