This is Bob.

JAX on my mind

I’ve been thinking a lot about JAX lately. JAX is appealing to a computer scientist like me, due to its beautifully compositional architecture for coding autodiff on GPUs. Had JAX existed when we started coding Stan in 2011, we would’ve used that rather than rolling our own autodiff system. Apparently Theano did exist at that time, but we didn’t hear about it until long after releasing the first version of Stan.

Why JAX?

Originally, my interest was sparked by Matt Hoffman’s work on Cheese and Meads samplers (too lazy to look up their capitalization pattern, so going Andrew style) that use massively parallel HMC steps on GPU (Matt was driven to abandon NUTS because its recursive structure is anathema to GPU acceleration). It continues to be kindled by statements from people like Elizaveta Semenova, who announced during the recent StanCon at Oxford that she had to give up Stan and moved to NumPyro because she couldn’t code neural nets easily in Stan and Stan doesn’t scale on the GPU (we have some operations that can be sent to GPU, but we can’t keep the whole eval in kernel). The full blown fire is due to Justin Domke and Abhinav Agrawal’s work on normalizing flows—they have a repo, vistan, that implements their methods, including a greatly improved version of autodiff variational inference (ADVI) and a version of real non-volume preserving (realNVP) normalizing flows. Gilad Turok, who’s a research analyst here and is applying to grad school for next year, is almost done with a better engineered version we can submit to Blackjax.

We haven’t gotten far, but so far we haven’t found a density yet that the realNVP flows, followed by importance resampling, didn’t fit nearly perfectly. For example, it could fit a 1000-parameter hierarchical IRT 2PL without any identifying strategy other than priors, which is a model that Stan cannot fit well at all due to the product of funnels structure of the posterior geometry. Justin and Abhinav’s ADVI and realNVP variants rely on JAX’s ability to massively parallelize the log density and gradient calculations in order to stabilize the stochastic gradients used to calculate KL divergence up to a constant (i.e., the ELBO).

Comparing Stan and JAX performance

At the recent StanCon in Oxford, Simon Maskell presented a result I’ve been meaning to generate for a year or two, which is a comparison of JAX and Stan on CPU. Here’s Simon (or at least his arm!) presenting the result:

Working out the algebra problem Simon’s slide title presents (JG = 20 * JC, JG= 5 * SC, where JG is JAX on GPU, JC is JAX on CPU, and SC is Stan on CPU), we see that JAX on GPU is 5 times faster than Stan on CPU, which is in turn 4 times faster than JAX on CPU. For the evaluation, Simon was using a high-end consumer-grade (RTX) NVIDIA GPU—we will re-run Simon’s results on our state-of-the-art GPUs over the next couple of months and report back. Tensorflow did a similar evaluation with similar results a few years ago, but I wasn’t able to find it.

CPU to GPU over next decade?

This is actually kind of a bummer in that it means we couldn’t really move to JAX as a back end for most of our users because it’s not very performant on CPUs. My guess is that the median Stan user in terms of computation has a 5-year old Window’s notebook with no GPU. Part of my interest in things like JAX is my feeling that in 10 years, this is no longer going to be the case, and almost everyone trying to fit a statistical model will be working on a cluster.