Numpyro: Probabilistic programming with NumPy powered by Jax

I read through the second edition, and would recommend as well [1] (has numpyro, pyro, pymc code). He also has a great lecture series on youtube [2]

[1] https://xcelab.net/rm/

[2] https://www.youtube.com/watch?v=FdnMWdICdRs&list=PLDcUM9US4X...

it's not that complicated that you need to read a whole book just to get a rough idea: it's just a cutesy way to specify a "plate model"[1] and then run inference using that model.

[1]: https://en.wikipedia.org/wiki/Plate_notation

There are examples from the book done in PyMC:

https://github.com/pymc-devs/pymc-resources/tree/main/Rethin...

or, for first edition

https://github.com/pymc-devs/pymc-resources/tree/main/Rethin...

MCMC is AFAIU overly prohibitive for neural networks. If you are interested in incorporating uncertainty awareness and improving calibration of your neural net classifiers in a somewhat scalable manner, I think linearized Laplace is a good place to look. I'd suggest you look at `laplace-torch` [1], it's an easy way of doing that.

[1]: https://aleximmer.github.io/Laplace/

I would say that, at least for me, PyMC’s main advantage was in DX. I just found model construction much more straightforward and better aligned with how I wanted to assemble the model.

This textbook/walkthrough is great: https://a.co/d/9dWXDTK

I tired both a while back, but nothing too big or serious. One thing that numpyro benefits from is JAX's speed, so it might be faster for larger models. Though PyTensor, which is the backend for PyMC can apparently also generate JAX code, so the difference might not be drastic. The PyMC API also seemed to me easier to get started with for those learning Bayesian stats.

One thing I remember that I disliked about PyMC was the PyTensor API, it feels too much like Theano/TensorFlow. I much prefer using JAX for writing custom models.

You'll lose a lot of the PyMC convenience functions with Numpyro but gain a lot of control and flexibility over your model specification. If you're doing variational inference, Numpyro is the way to go.

You can use the Numpyro NUTS sampler in PyMC with pm.sample(nuts_sampler="numpyro") and it will significantly speed up sampling. It is less stable in my experience.

This is maybe not the place, but we did some apples to apples comparisons between PyMC, Dynesty, and the Julia Turing.jl package.

A little to my surprise, despite being a Julia fan, Turing really outperformed both the Python solutions. I think JAX should be competitive in raw speed, so it might come down to the maturity of the samplers we used.

I find the PyMC situation confusing. But PyMC should be able to use JAX in the backend: https://www.pymc.io/projects/examples/en/latest/samplers/fas...

PyMC can use NumPyro as a backend. PyMC's syntax and primitives for declaring models are much nicer than (Num)Pyro's, as is the developer experience overall. But those come at the cost of having to deal with PyTensor (a fork of a fork of Theano), which is quite bad IMO, instead of just working with Numpy or PyTorch.

There are a class of probabilistic models that have exactly this property -- exponential family models. While pedagogic examples of such models tend to be very simple, they need not be. A huge class of graphical models fall in this class and can be very flexible. The underlying statistical model of a Kalman filter is in this class.

These models have what are called sufficient statistics, that can be computed on data, s = f(D) where s is the sufficient statistics, D is the past data. The clincher is that there is a very helpful group theoretic property:

s = f(D ∪ d) = g( s', d) where s' = f(D). D is past data, d is new data, D ∪ d is the full complete data.

This is very useful because you don't have to carry the old data D around.

This machinery is particularly useful when

(i) s is in some sense smaller than D, for example, when s is in some small finite dimension.

(ii) the functions f and g are easy to compute

(iii) the relation between s and the parameters, or equivalently, the weights ⊝ of the model is easy to compute.

Even when models do not possess this property, as long models are differentiable one can do a local approximate update using the gradient of the parameters with respect to the data.

⊝_new = ⊝_old + ∇M * d.

(∇M being the gradient of the parameters with respect to data, also called score)

With exponential family models updates can be exact, rather than approximate.

This machinery applies both to Bayesian as well as more classical statistical models.

There are nuances also, where you can drop some of the effects of old data under the presumption that they no longer represent the changed model.

The algorithmic family that in general seems called for is "sequential monte carlo" or "particle filtering". While often this is presented for models and problems where the latent dimension is small and the problem is about how something evolves over time, the "sequential" part of it can just be about the order in which data are received (I think this traces to "A Sequential Particle Filter Method for Static Models", N Chopin, 2002).

On the PPL side, I think SMC has often been a secondary target, but there has been good work on getting good, efficient inference be more turn-key. I think this has often focused on being able to automatically provide better/adaptive proposal distributions. For example, "SMCP3: Sequential Monte Carlo with Probabilistic Program Proposals" by Lew et al 2023 (Stuart Russell is among the contributors).

One way I've seen this done in practice is to construct an offline model that produces an initial set of posterior samples, then construct a second, online model that takes posterior samples and new observations as input and constructs a new posterior. This probably wouldn't make sense computationally in a high-frequency streaming context, but (micro)batching works fine.

I've seen lots of other approaches proposed for this over the years, here's a recent Stan forum thread with some links: https://discourse.mc-stan.org/t/updating-model-based-on-new-...

Update the model with the new data and recompute it (hopefully the unchanged parts are cached). Most probabilistic programming models are not black boxes like neural nets.

> What is probabilistic programming actually useful for?

You can think of a probabilistic programming language as a set of building blocks for building statistical models. In the olden days, people used very simple frequentist models based on standard reference distributions like the normal, Student's t, chi2, etc. The models were simple because the computational capabilities were limited.

In modern days, thanks to widespread compute and the inference algorithms, you can "fit" a much wider class of models, so researchers now tend to build bespoke models adapted for each particular application they are interested in. Probabilistic programming language are used to build those "custom" models.

Yeah interpretability is the main reason that I know. If you need to make decisions based on statistics of the data, I think it's a lot easier with explicit statistical models.

MCMC and Bayesian models.

https://en.wikipedia.org/wiki/Probabilistic_programming

They're trying to drum up interest in Jax. They do this time to time, but yeah, good for them doing what they like.