(this post is by Charles)

Last week, BayesComp 2023 took place in Levi, Finland. The conference covered a broad range of topics in Bayesian computation, with many high quality sessions, talks, and posters. Here’s a link to the talk abstracts. I presented two posters at the event. The first poster was on assessing the convergence of MCMC in the many-short-chains regime. I already blogged about this research (link): here’s the poster and the corresponding preprint.

The second poster was also on the topic of running many chains in parallel but in the context of models based on ordinary differential equations (ODEs). This was the outcome of a project led by Stanislas du Ché, during his summer internship at Columbia University. We examined several pharmacometrics models, with likelihoods parameterized by the solution to an ODE. Having to solve an ODE inside a Bayesian model is challenging because the behavior of the ODE can change as the Markov chains journey across the parameter space. An ODE which is easy-to-solve at some point can be incredibly difficult somewhere else. In the past, we analyzed this issue in the illustrative planetary motion example (Gelman et al (2020), Section 11). This is the type of problem where we need to be careful about how we initialize our Markov chains and we should not rely on Stan’s defaults. Indeed, these defaults can start you in regions where your ODE is nearly impossible to solve and completely kill your computation! A popular heuristic is to draw the initial point from the prior distribution. On a related note, we need to construct priors carefully to exclude patently absurd parameter values and (hopefully) parameter values prone to frustrate our ODE solvers.

Even then—and especially if our priors are weakly informative—our Markov chains will likely journey through challenging regions. A common manifestation of this problem is that some chains lag behind because their random trajectories take them through areas that frustrate the ODE solver. Stanislas observed that this problem becomes more acute when we run many chains. Indeed, as we increase the number of chains, the probability that at least some of the chains get “stuck” increases. Then, even when running chains in parallel, the efficiency of MCMC as measured by effective sample size per second (ESS/s) eventually goes down as we add more chains because we are waiting for the slowest chain to finish!

Ok. Well, we don’t want to be punished for throwing more computation at our problem. What if we instead waited for the fastest chains to finish? This is what Stanislas studied by proposing a strategy where we stop the analysis after a certain ESS is achieved, even if some chains are still warming up. An important question is what bias does dropping chains introduce? One concern is that the fastest chains are biased because they fail to explore a region of the parameter space which contains a non-negligible amount of probability mass and where the ODE happens to be more difficult to solve. Stanislas tried to address this problem using stacking (Yao et al 2018), a strategy designed to correct for biased Markov chains. But stacking still assumes all the chains somehow “cover” the region where the probability mass concentrates and, when properly weighted, produce unbiased Monte Carlo estimators.

We may also wonder about the behavior of the slow chains. If the slow chains are close to stationarity, then by excluding them we are throwing away samples which would reduce the variance of our Monte Carlo estimators, however, it’s not worth waiting for these chains to finish if we’ve already achieved the wanted precision. What is more, as Andrew Gelman pointed out to me, slow chains can often be biased, for example if they get stuck in a pathological region during the warmup and never escape this region—as was the case in the planetary motion example. But we can’t expect this to always be the case.

In summary, I like the idea of waiting only for the fastest chains and I think understanding how to do this in a robust manner remains an open question. This work posed the problem and took steps in the right direction. There was a lot of traffic at the poster and I was pleased to see many people at the conference working on ODE-based models.

(This post is by Charles, not Andrew)

I’ve been a Stan developer for a bit more than 5 years now (!!) and this summer, I interned at Google Research with the team behind TensorFlow Probability (TFP). I wanted to study a new probabilistic programming language and gain some perspective on how Stan does things. I’ll compare the two languages but rather than focus on performance, I want to talk about design and user experience. Which is fun because this aspect gets a bit lost in academic writing.

As an applied statistician, I love the Stan user experience because the focus is on modeling. Stan’s default inference is general-purpose and can be used on a very wide range of models. Other than tweaking a few tuning parameters, I usually rely on Stan’s default inference engine, which is a dynamic Hamiltonian Monte Carlo sampler. This means I can spend more time worrying about the model and whether or not it gives me the insights I seek. Being able to use “sampling” statements in the model block, i.e.

y ~ normal(theta, sigma)

is really nice because this is how I write models on paper. No graphs, no names chosen from a menu of options, no formula. For me, nothing beats writing down the generative process. Pedagogically this is a little bit confusing because any statement in Stan’s model block is not an actual sampling statement, since nothing gets sampled. Rather it’s a call to increment the log of the joint density or, in Stan’s jargon, the “target”. So maybe we should only teach the equivalent statement

target += normal_lpdf(y | theta, sigma)?

I’m not sure. Currently I teach both. I also like Stan’s coding blocks. Since a Bayesian model is specified by a joint distribution over observed variables and parameters, declaring variables in a data and parameter block further resonates with my statistical formulation of the model. There’s also another advantage: the blocks each have a different computation cost which is helpful to think about when writing computationally intensive models (see for example the discussion in this paper).

TFP acts more as a Python package than its own language. Within the same script, we’ll process the data, specify our model, run our inference, and do all the post-processing. This doesn’t have the pedagogical appeal of building blocks, nor its rigidity. Based on the tutorials I worked through, there are no pseudo-sampling statements, which means less clarity in terms of what model we’re writing down but also less confusion in terms of what computation takes place. Matt Hoffman tells me TFP has something akin to Stan’s sampling but it has less of a DSL (domain specific language) feel to it. When writing the model in TFP we can use the many functions available in Python with one super important caveat: the code needs to be differentiable, if we use gradient-based algorithms (which we probably do). With that in mind, TFP is designed to be on top of two automatic differentiation libraries: TensorFlow (surprise, surprise) and Jax. Stan has its own autodiff library so *broadly speaking* everything you write in Stan is safely differentiable. I was advised to use Jax which is pretty awesome — the documentation itself is a fascinating read. One compelling feature is that Jax supports many NumPy functions. And NumPy is great because of all the mathematical operations it offers and because it seamlessly parallelizes these operations across CPU cores. Jax takes NumPy a step further by parallelizing operations across accelerators, such as GPUs, and making a lot of these functions differentiable and therefore amiable to gradient-based inference. For those who are familiar with NumPy, this is great. For me there was a bit of a learning curve.

It’s worth pointing out that there is quite a bit of overlap between the functions Jax and Stan support but there are differences which reflect the motivating problems behind each project. Over the years, Stan’s support for ODEs has become very good and quite broad. Jax offers more specific ODE support, which works well for certain problems (e.g. Neural ODEs), less so in other cases. Note that there are pending pull requests extending Jax’s ODE support. But let’s focus on design here. Stan’s ODEs report error messages when they fail: for instance the numerical integrator reaches the maximum number of steps. Jax on the other hand reports… nothing. This is because apparently exceptions don’t exist on XLA (accelerated linear algebra, see this discussion on GitHub). Not having error messages caused me a lot of debugging grief. Overall I did find Stan did a better job issuing error and warning messages. Another example is divergent transitions: Stan automatically warns you when divergences occur. In TFP, you have to dig them out, which means you need to already know what a divergent transition is, and that may not be the case for many users.

Ok, this is where things get fun. TFP asks you to specify your inference algorithm. If you call tfp.mcmc you need to define the transition kernel of your Markov chain and usually, you should also write down an adaptive scheme. This takes a few lines of code, even if you use TFP’s built-in features. You then probably want to specify a “trace” function which records various properties of your sampler across iterations, such as whether a divergent transition occurs. This is more effort than relying on Stan’s very good defaults. But it also makes it easier to try out new transition kernels and moreover new inference algorithms! How easy is it to do this in Stan? It’s hard and this is not what Stan was designed for. That said, the underlying C++ code is open-source, well-written and well documented, and so it can be hacked. People do it. I’ve done it. And this leads to very interesting forks of Stan.

The question is: as an applied Statistician, do I resent not being able to specify my own inference algorithm in Stan? This short answer is occasionally. I’m biased because many of my scientific collaborations begin with people not being able to fit their model in Stan. Often we do eventually manage to do it by reparameterizing the model, baking more information into the priors, rewriting the ODE integrator, etc. Working out a new inference algorithm and efficiently implementing it is a bigger leap. My research does focus on the problems that frustrate our inference and for which we want to develop better methods. But this takes time. In the meantime, I believe there is a strong case to be made for an algorithm that has been stress-tested throughout the years, theoretically well-studied, and which empirically has worked for many problems.

On the other hand, as a researcher in the field of statistical algorithms, I do enjoy the modularity TFP provides and having the ability to easily try out new methods. TFP also has “experimental” features which are not fully supported but can nevertheless be tried out. I pitched the idea of including features explicitly labeled as “experimental” in Stan, as a middle-ground between making available some of the tech we’re developing and being upfront about what remains ongoing research. This gets into questions of what Stan’s mission is and how much we trust users to take heed of our warning messages. Long story short, for the time being, experimental features won’t be part of Stan’s official releases (except ADVI, which does issue a warning message).

Comparing TFP and Stan reminds me a little bit of the historical conflict between Microsoft and Apple, that is a modular approach versus end-to-end control designed to create a seamless user experience, at least on the algorithmic front. In the past I’ve approached this question by thinking about specific problems, the performance benefits of using specialized versus general techniques, and whether it makes sense to have a universal inference algorithm, or at least a universal default. But I think a big component of software development is design and user experience. How do “everyday scientists”, if you’re willing to accept this term, interact with the software? What do they want to focus on? What balance between guidance and freedom works best for a broad audience? Now when it comes to model specification, both Stan and TFP champion modularity and giving the users a large control over their model, which we could contrast to other packages out there. This too is a design choice and one we could ponder about.