Chainsail, a web service for sampling multimodal distributions: opinions and beta testers wanted!

Sharing
1

Hi PyMC community,

Colleagues and I over at Tweag created a web service called Chainsail (https://chainsail.io) that can drastically improve sampling of multimodal distributions, which occur often in models with unidentifiable parameters or when you have ambiguous data. Chainsail has flexible support for models and probability distributions defined in PyMC, Stan, or hand-written Python.
The secret sauce in Chainsail is an autotuning Replica Exchange algorithm that uses cloud computing to scale dynamically beyond the computing resources available on single machines. If you want to learn more about Replica Exchange, here’s a shameless plug: I wrote a blog post about it.

We’re currently looking for probabilistic programming practitioners who might be interested in beta-testing it. If you have multimodal distributions to sample, shoot us an email at support@chainsail.io to get your email address authorized! Don’t hesitate either to tell us a bit about your sampling problem or ask us for a demo - we would be happy to chat and show you around. Currently, Chainsail is deployed on Tweag’s premises and is available to beta testers for free (within reasonable computing time limits).

Future Chainsail development depends mostly on beta tester feedback, but faster Stan support, a better HMC implementation (possibly using BlackJAX) and more choices for the tempering schemes (for example, applying a temperature only to the likelihood) would be among the next things to work on.
Chainsail is currently closed-source, but it is highly likely that we will eventually make at least parts of the service, if not all of it, open-source.
If you’d like to learn more about Chainsail, here’s a couple of additional resources:

  • we shot a ~15 minutes walkthrough video that demonstrates how to use Chainsail,
  • we wrote an announcement blog post that presents Chainsail and the kind of problems it solves,
  • we have a repository that provides more detailed documentation and example probability distributions. It also hosts the chainsail-helpers Python package that provides probability distribution interfaces for PyMC and Stan and a helper script to process the downloaded sampling results: GitHub - tweag/chainsail-resources: Examples, documentation and other additional resources related to Chainsail,
  • and, in a couple of days, we will publish another blog post in which we use Chainsail to analyze a soft k-means model / fitting a Gaussian mixture with unequal weights (in one dimension for easy visualization) and show how using Replica Exchange via Chainsail makes a difference.

Chainsail is in an early stage and currently has a couple major limitations:

  • it uses only a very basic, untuned HMC implementation (no NUTS, fixed and preset number of integration steps, simple heuristic to tune the timestep),
  • only the most important parameters can currently be set by the user,
  • it’s in its infancy, so expect glitches and rough edges :slight_smile:

If you have any questions about Replica Exchange, what the Chainsail service can and cannot do and if you’d like to test it, please don’t hesitate to let me know in this thread or email us: support@chainsail.io. Looking forward to hearing your questions, opinions and ideas!

5 Likes
2

As promised above, we now published another blog post in which we perform a detailed analysis of (a variation of) the soft k-means model aka fitting a a Gaussian mixture both using a single Markov chain and with multiple, connected chains via Chainsail :slightly_smiling_face:

1 Like
3

Thanks for sharing! Replica Exchange HMC is highly underappreciated and it is lovely to see a nice product build on top of it.
Not sure how much you can share, but I am curious about the where the sampling actually happen (ie how the RE-HMC is implemented). Since both PyMC and Stan model are accepted, I am assuming the implementation is written in something easy to compiled to run on multiple device?

4

I can share everything I want, I believe, so don’t hold back with questions :slightly_smiling_face:

The RE-HMC implementation is largely based on a code for Replica Exchange I wrote a long time ago for my scientific research. It’s essentially pure Python + MPI (Message Passing Interface) for inter-process communication. The interface of this simple: it expects just a Python object modeling a distribution with a log_prob method and a sampler object that has a sample method, which performs a single MCMC move. The sampler object holds the current state of the Markov chain internally.
This is exactly the interface we use in Chainsail - we wrote a HMC sampler that conforms to this interface and the probability distribution is provided by the user in a Python module.

Re compilation: I’m not sure I understood correctly, but if you’re asking whether there is compilation happening elsewhere than in the user-defined probability, the answer is no. Essentially, if you provide a PyMC model, it runs model.compile_logp and model.compile_dlogp and exposes the result as a the log_prob method of the probability distribution object. For Stan, the probability distribution object talks to a httpstan instance, which compiles the model and evaluates the log-probability and its gradient. You can find both implementations here.

As for Replica Exchange being underappreciated - I fully agree and would love to understand better why this is the case. Do you have any insights you could share? A naive explanation could be that models with multimodal posteriors are just not as common in problems the probabilistic programming community solves as they are in computational physics / biomolecular simulation, which is my background.

2 Likes