Adaptive SMC as in pyMC?

[jucor]

Hi folks!

This question is probably best suited for @ciguaran who worked on the SMC implementation and the Tuning SMC notebook tuning the inner-SMC MCMC kernel, handling ), which handles Design Choice (c) of Section 2.1.3 of Buchholz, Chopin, Jacob (2020) for the case of an IRMH.

1. Do you have any plan to merge this implementation of the tuning of MCMC kernel into the blackjax lib itself please? To avoid copy/pasting/modifying from the specific case in the notebook? Easier to use ;-) I realise this might be blackjax.smc.inner_kernel_tuning, so would you have an example of usage please?

2. I Have you implemented anywhere Design Choice (b), i.e. adapting the number of MCMC steps based on Algorithm 3, by any chance? I believe blackjax.smc.adaptive_tempered implements Algorithm 2 for Design Choice (a): the temperature schedule, based on ESS.

Reason I'm asking: I'm moving a codebase from PyMC to Blackjax, to benefit from jax's vmap (each likelihood evaluation requires solving an ODE: costly, but vmap with diffrax makes it embarassingly parallelizable on GPU). PyMC's SMC sampler, while regrettably not GPU-parallelizable (at least that I know of), has all those adaptations already baked in, which makes it quite easy to use by applied statisticians who are not SMC specialists.

It would be terrific to have those in blackjax, please :)

Thanks!

[ciguaran]

Hei @jucor nice to meet you and thanks for your interest in all these!.
Let me try to answer/give context.

Hi folks!

This question is probably best suited for @ciguaran who worked on the SMC implementation and the Tuning SMC notebook tuning the inner-SMC MCMC kernel, handling ), which handles Design Choice (c) of Section 2.1.3 of Buchholz, Chopin, Jacob (2020) for the case of an IRMH.

1. Do you have any plan to **merge this implementation of the tuning of MCMC kernel into the blackjax lib itself** please? To avoid copy/pasting/modifying from the specific case in the notebook? Easier to use ;-)  I realise this might be [`blackjax.smc.inner_kernel_tuning`](https://github.com/blackjax-devs/blackjax/blob/main/blackjax/smc/inner_kernel_tuning.py), so would you have an example of usage please?

Apologies on this one. I don't super understand the question. The notebook does make use of inner_kernel_tuning and some primitives that can be found inside https://github.com/blackjax-devs/blackjax/tree/main/blackjax/smc/tuning there to compute means, stds, covariance matrixes, etc from particles. Regarding the inference loop, I think the core contributors decision has always been not to include them in the library. LMK is this is what you meant or I am missing some part of the question. See Note at the bottom regarding pre-tuning. Also LMK if you had in mind a different parameter tuning procedure.

2. I **Have you implemented anywhere Design Choice (b), i.e. adapting the number of MCMC steps based on Algorithm 3, by any chance?**   I believe [`blackjax.smc.adaptive_tempered`](https://github.com/blackjax-devs/blackjax/blob/main/blackjax/smc/adaptive_tempered.py) implements Algorithm 2 for  Design Choice (a): the temperature schedule, based on ESS.

Regarding the number of MCMC steps. Based on convos with @nchopin, Waste-Free seemed like a better option. #721. Although is not the only reason, WF is better suited for Jax since the number of steps is the same on all chains.

Indeed, blackjax.smc.adaptive_tempered implements the temperature schedule.

Reason I'm asking: I'm moving a codebase from PyMC to Blackjax, to benefit from jax's vmap (each likelihood evaluation requires solving an ODE: costly, but vmap with diffrax makes it embarassingly parallelizable on GPU). PyMC's SMC sampler, while regrettably not GPU-parallelizable (at least that I know of), has all those adaptations already baked in, which makes it quite easy to use by applied statisticians who are not SMC specialists.

Right, pymc already has things ready out of the box. If using Blackjax you would need to assemble the pieces. You might be interested in the code https://github.com/pymc-devs/pymc-experimental/blob/main/pymc_experimental/inference/smc/sampling.py which I implemented some time ago. The only caveat is that there's a current issue with pytensor blocking merging this fix PR: pymc-devs/pymc-experimental#374 . Second caveat is that the parameter tuning would need improvement (see below)

It is in my near future roadmap to implement the pre-tunning approach considered in the paper you cited. Right now you can only tweak parameters based on the particles or the sampler state. With the pre-tuning algorithm, an MCMC step would be executed, performance of parameters would be measured, and we would probabilistically select a population of parameters suited for the distribution at step T.

With pre-tunning, my understanding is that all the pieces would be in place for the Blackjax implementation to be competitive with state of the art when the goal is sampling from a bayesian posterior distribution.

It would be terrific to have those in blackjax, please :)

Thanks!

Hope it helps!

[jucor]

Thanks a lot @ciguaran for the detailed answer, super helpful!

1. Great, I was not aware of the conscious decision not to include the core loops. This makes sense, so I'll live without. This is indeed what gave me the impression that the notebook's features had not been included in the library -- I had overlooked that the notebook already made use of inner_kernel_tuning. My bad, thanks a lot for clarifying!

2. I was not familiar with waste-free. I indeed had spotted the varying number of MCMC steps per particle would be sub-optimal for vmapping over all particles, but that's anyway already so in my application: each likelihood evaluation requires a call to an ODE sovler which has a different adaptive step-size for each particle, so I already have all the particles having to wait until the slowest one is finished. Worth it anyway, compared to the sequential approach. But since Waste Free SMC available for adaptive tempered and tempered SMC.  #721 is there, that's great, I'll be delighted to use that!

3. On PyMC, having a way to call a vmapped blackjax SMC sampler from pymc would be super convenient. However, I have already re-coded my likelihood in pure JAX/numpy, because calling Diffrax (the ODE solver) within pyMC involved yet another layer of wrappers that seemed at risk to introduce too much complexity. On a meta level, the bug with pytensor that you mention reminds me of why I'm uneasy about using a lib that depends on a "smaller" tensor library like pytensor than, say, pytorch or jax. I was already surprised to see the fork of pytensor from aesara (itself a fork of the abandoned theano), fragmenting a developer base that is itself on the smaller side: it makes me ... uneasy. But that's just my personal feeling :)

4. As to pre-tuning, not sure if it would help much in my case: I suspect that the geometry of my problem is possibly quite nasty (highly correlated parameters of ODEs), possibly with high correlations, and possibly curvature so the pairwise correlation changes sign depending on the location in parameter-space. So not sure how well pre-tuning would help. I'm also keen on staying derivative-free (thus no HMC) for now, due to bad experience with NUTS which diverged easily due to accumulated error in the calculation of the derivative via the extended ODE. It might perform better with grad on diffrax, as this time the whole ODE solver is differentiable, but if I can get derivative-free working well enough I'll be happy.

Thanks a lot for your super helpful answer, I'll now get to assemble the main inference loops. (If you know where I could read about the rationale for not including them in the library, I'd love to know!)

Thanks!