Blackjax MCLMC support

[reubenharry]

MCLMC is an algorithm intended as a statistically more efficient alternative to NUTS HMC. I'm interested in adding the option to PyMC to use it via the Blackjax implementation, in the spirit of pm.sample(nuts_sampler="blackjax"), so maybe something like pm.sample(mclmc_sampler="blackjax").

It comes with either an unadjusted version (no MH, controlled asymptotic bias, faster) or an MH adjusted version, both of which are faster than NUTS for a range of benchmarks of interest.

Basically my question is: would PyMC be open to a PR along these lines? For public facing visibility, we'd like to put the algorithm into PyMC rather than use PyMC indirectly (i.e. extract a density from a probabilistic program in PyMC), and I thought this might be the best way.

[ricardoV94]

Is bayeux a better tool to allow combining pymc with blackjax samplers? CC @ColCarroll

sample is already a very overloaded function, so I wouldn't want to add a new mclmc_sampler kwarg to it.

[ricardoV94]

Needs someone to have a go at it. I don't have the complete mental map to anticipate edge cases. First guess is that it may be a more scalable solution.

An external sampler class and a monolithic external sampler subclass that tell pymc stuff will be sampled externally and must sample the whole model. The external nuts samplers should be implemented like this then as well

[twiecki]

@reubenharry This is probably a bit more than what you bargained for I assume?

[twiecki]

Claude Code at a go at the interface for nutpie in #7719

[reubenharry]

Thanks for this discussion! Not necessarily more than I bargained for, but might take me some time to understand how to approach (perhaps with Claude's help, as you suggest).

[ricardoV94]

We're happy to help, this is something that would be nice to have so its easier to plug alternative samplers in the future

[reubenharry]

Just to give a sense of the performance benefits, here's our paper, with a comparison between NUTS and Metropolis Adjusted Microcanonical Sampling (MAMS) on some standard test problems, in table 1: https://arxiv.org/pdf/2503.01707.

MAMS is nowhere near as battle-tested as NUTS, but it seems to give a 2x or greater speedup in most cases we tried, is easier to parallelize, and we're interested in getting input from pymc users with difficult problems (hence this discussion).

[ricardoV94]

Ricardo wanted

Ricardo wants* :D It's just too good not to

[reubenharry]

Ah yes, I had forgotten that. But yes, it would be nice to try obtaining the nutpie mass matrix or normalizing flow for a given problem, and then subsequently run MAMS instead of NUTS. Our trajectory length tuning is described in section 5.3 of the paper linked above. (It could actually also be used for HMC, but we haven't tried that).

[aseyboldt]

nutpie.sample(model, store_mass_matrix=True).sample_stats.mass_matrix_inv :-)

[ColCarroll]

[aseyboldt]

Sorry Colin 😅

[reubenharry]

Just following up on this, as my group is still interested in getting MCLMC into PyMC. Re. the above, I agree that bayeux would work, but since our main motivation is discoverability by general PyMC users, our main focus is on seeing if we can get it into PyMC directly.

From the discussion above, it sounds like a class for external samplers is needed (ala https://github.com/pymc-devs/pymc/pull/7719/files) - does that sound right? I've been putting off approaching this, but will need to sooner or later

[twiecki]

Thanks! Yeah, I don't know if you'd be able to take that PR over but that I think would help a lot.

[ricardoV94]

@reubenharry I'll try to do something about it this week (70% confidence). Keep an eye on that PR to see how it goes

[reubenharry]

thanks, I'll keep an eye on it!

[ricardoV94]

#7880

[nataziel]

Since this is a bit of a new/experimental thing you could be a little cheeky and just handle it within the sample_blackjax_nuts function. Around here you might be able to implement whatever logic is necessary.

Obviously long term it's not ideal but I for one would definitely give it a go on a real world problem if all I had to change was a kwarg like {nuts_kwargs: {adapdation_kwargs: {algorithm: mclmc}}}