(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.
Charles:
Great post. I also wanted to emphasize the value in having multiple competing probabilistic programming languages. We can develop Nuts or ADVI for Stan and then other people implement these in other languages. They make other developments that we then can include in Stan.
Also, I wouldn’t want your second-to-last paragraph mislead people into thinking that Stan isn’t improving. We continue to add new features and capabilities, and it seems likely that Pathfinder will be in Stan soon, also maybe stacking and nested Laplace.
Absolutely. There is a reason we have releases on a regular basis! There are a ton of new features being added to Stan and it’s fun to take a peek at the Pull Requests on the GitHub page (https://github.com/stan-dev).
I can’t push this enough. Take a look at Julia, in particular Turing. A huge motivation for Julia is the solution of the “two language problem” (or the more general N language problem).
For example, Stan, it’s a great self-contained modeling language… But to do anything much with it, you’ll need something like R or Python to do the data manipulation and visualize/analyze the resulting sample. And of course Stan is written in C++, so there’s a couple different languages involved at the least. With TFP you’ve got Python and C for Numpy and then TFP itself, and maybe some C++ or something.
If you want to extend or enhance or bug fix stuff, you’re going to have to wade into multiple languages. With Julia, most of the Julia ecosystem is written in Julia. Julia compiles to fast machine code that is comparable to C or C++ and Julia’s semantics make it easier for the compiler to optimize code to be fast. As long as you read the performance tips in the docs : https://docs.julialang.org/en/v1/manual/performance-tips/
Plus, Julia uses Julia itself as a macro language. So you can write code that writes code. This is how domain specific languages are handled. So Turing.jl is a Julia library that uses the @model macro to take in valid Julia syntax but translate it into Probabilistic Programming Semantics (that is things like a ~ Normal(b,c) mean basically the same thing within a @model block as they do in Stan even though this isn’t meaningful in base Julia).
Essentially All of Julia is available to you in the Turing DSL, because the Turing DSL is just a macro that writes Julia code to implement a function that calculates a logpdf so that the sample statement can run NUTS on the resulting logpdf.
The best and fastest way to solve ODEs in either Python or R is to call Julia’s DifferentialEquations.jl (In python diffeqpy, and in R diffeqr). Seriously, the best way to solve differential equations across all languages available today is using Julia, the ecosystem over there is HUGE because it’s attracted a lot of interest from that community. They handle neural-odes, differential algebraic equations, delay differential equations, stochastic differential equations, and there’s a mathematical modeling language for describing systems of equations again using the Julia macro language.
It’s hard to understand exactly how amazing the Julia language is unless you’re either a Common Lisp hacker from way back, or you have experienced it for at least about 6 months.
I’ll also point out that it’s relatively trivial to call from Julia into either R or Python (or C or C++ or Fortran). So if there’s something well developed in those other languages that handles a specific task you need, you don’t have to leave that behind.
Oh and I should have said, if you want to run Stan programs from Julia, Stan.jl will make that easy.
I really hate working iteratively with compiled languages. I actually think it’s great to have one interpreted language (in my case python) calling my high performance kernels (in this case Stan).
FWIW, there is an interpreter for Julia (https://julialang.org/blog/2019/03/debuggers/#julia interpreter) which should allow for a more Python-like experience in the iterative development.
Julia is an interactive language just like Python or R, even though compiled. If you’re talking about the “time to first plot” issue, this is largely easy to solve by compiling a sysimage with all the major packages you use frequently. I posted example scripts showing how to do this here:
https://discourse.julialang.org/t/a-julia-dataanalysis-sysimage-from-packagecompiler-its-so-easy-you-should-do-it-too/68127
Thanks both of you, I’ll have a look!
I keep thinking that I should look at Julia, but sometimes it reminds me of self-driving cars — it’s always in just a few more years it will be ready/everyone will be using it. It sounds like a great language, but by now I wonder if it’s always going to be a relatively niche one no matter its advantages over competitors.
It probably will be niche because it really appeals to people with a certain background, which is a combination of computational sophistication and mathematical/scientific computing expertise. But being niche doesn’t mean it’s not the right tool for the job.
I do think it has a few rough edges still, but the community and ecosystem has exploded rapidly since I first started using it in 2019. Today I can get all the stuff I’d normally do in R done with WAY less annoyance in Julia. Plus there are a ton of things I’d never do in R (without say calling out to Stan and/or a variety of R packages written in C/C++/Fortran)
If you learn languages because they’re marketable for a job… then yeah, learn Python. If you learn a language because you want a single language for doing all aspects of scientific computing, including metaprogramming and DSL creation… then learn Julia.
Thanks for shimming in Daniel. When I started looking for internships last year, I was bent on either using Python or Julia (being otherwise a regular R user). In another universe, this blog post would’ve been written about Turing and Stan. But for the time being, I know very little about Julia.
I takes time and effort to learn a new language but knowing more than one PPL certainly puts things in perspective. As much as there are “redundancies” — and per Andrew’s comment, synergies — between softwares, I ultimately find that the designs of these languages reflect different motivations. Understanding those can help us work out how these software might best be used in practice. But this decision is heavily perturbed by what we happen to be already familiar with.
Charles: What’s up with these “neural odes”? Can we do them in Stan? What are they even exactly, can you recommend a good primer?
Based on https://proceedings.neurips.cc/paper/2018/file/69386f6bb1dfed68692a24c8686939b9-Paper.pdf
My impression is that a “neural ODE” is where instead of thinking of a neural network as a set of discrete layers, with inputs to the first layer feeding through the layer into the second layer and into the third layer etc, think of it as a family of models in which the number of layers is a parameter. As the number of layers goes to infinity, the limiting case is basically a differential equation. So the flow of inputs to outputs behaves as the solution to a differential equation. But that’s about as far as I understand it.
Hm, my understanding is very limited but I don’t think this is accurate. I *think* the layers are used to approximate the vector field driving an ode, which then gets integrated using a regular ode solver. But I may be mistaken!
I’ll have a look at your link as well.
I’d need to think a bit before providing a useful summary of what Neural ODEs do. As I recall, the paper is very readable.
The reason I brought up Neural ODEs is because they involve large systems of ODEs with many parameters. This in turn motivates using a scalable differentiation method to propagate derivatives. The solution here is the adjoint method. While scaling very well, the adjoint method incurs quite a bit of overhead cost, and it is not the best solution for problems that involve smaller ODEs, such as the ones that arise in Pharmacometrics. Last time I checked, Jax also lacked a stiff ODE solver, although there was pending Pull Request addressing this.
Stan has both non-stiff and stiff ODE solvers, which support forward and adjoint differentiation. It took a lot of effort to get there. My colleagues on the Stan Dev team deserve a ton of credit for building this.
(p.s: I understand Julia has very extensive support for differential equations, as Daniel mentioned, and also their differentiation)
Thanks Charles!
One thing that I really appreciate about Stan is that its role as a standalone language helps bridge gaps between the various “language camps”. I’m a longtime Python user and I’ve been able to interact with a lot of R users through the common ground of Stan. My experience has been that learning a DSL is less burdensome than trying to pick up a new language and ecosystem like R or Julia. So picking up Stan for me has opened up a community (of more stats focused individuals) without requiring me to step too far out of my Python comfort zone.
None of this to say that communities surrounding Python-focused PPLs like TFP or PyMC are worse, just different groups of people with different focuses.
Thanks for writing this! I’ve never used TFP, so I’m glad to learn how they do things…
One thing I’ve been thinking about in this context is that it might be sensible to introduce some sort of plugin structure to Stan where each inference algorithm would be a plugin. If we allowed people to develop and install their own plugins, it would be much easier to have “an experimental plugin” as a middle ground between “experimental, not ready for use, available via a fork” and “thoroughly validated, included in default Stan”. I think the biggest benefit would be that it would reduce pressure on putting new stuff (e.g. pathfinder) into the main releases of Stan and thus allow for longer and better testing in a sort of “public beta” versions of algorithms.
But I know I spoke about this on a couple occasions already and nobody seemed to like it very much, so I’ll stop posting about it for a while :-)
Hi Martin,
These are all very good points!
A different context where the idea of a plug-in came up was when we were developing Torsten (an extension of Stan for Pharmacometrics). We introduced super helpful routines which were too specific to make it into the general Stan language. Torsten exists as a fork, but it would be nice if it could be installed on top of Stan’s official release.
Plugins on the level of Stan routines sound useful, bit a bit more challenging to implement – a plugin for a new inference engine for CmdStan could IMHO be as simple as adding a C++ include path and providing a single entry point that gets called for ” “, taking the model class as a template parameter and all the additional command line parameters as strings as an additional parameter to the plugin to handle. So I think it actually could be quite easy to implement.
I’m not sure on what type of probability models you are making, but for modeling the distribution of the label in a regression my lab found mixture models to be very flexible.
https://www.biorxiv.org/content/10.1101/2022.04.22.489230v2