This post is from Bob
Nutpie laps NUTS
The short story is that the Nutpie sampler is usually twice as fast as Stan’s sampler for Stan models and twice as fast as PyMC’s samplers for PyMC models. The speed improvement is due entirely to better mass-matrix adaptation. Everything else is the same—dual averaging for step size adaptation and the biased-progressive, multinomial no-U-turn sampler for Hamiltonian Monte Carlo. Note that the speed is not coming from Python or Rust; we have a C++ implementation in the works we plan to interface with Python and then R.
The backstory
Adrian Seyboldt, a member of PyMC Labs, designed the sampler and released the open-source implementation linked above. He built it so he could try out Rust! Adrian’s also behind Stan’s game-changing new sum-to-zero parameterization along with Sean Pinkney. He’s a physicist by training and has amazing geometric intuition.
I tried Nutpie and it didn’t disappoint. So I wanted to understand the algorithm. I couldn’t understand the Rust code and Adrian hadn’t written any pseudocode or algorithm description. Eliot Carlson, with a recently minted MA in stats from Columbia, showed up at my office on the recommendation of Philip Greengard to look for a computational stats project. I asked him what he thought about Rust and like any great grad student, he wasn’t daunted by the fact that he knew nothing about Rust and was a statistician by training. With serious digging and some help from Adrian, Eliot extracted pseudocode descriptions and we were able to formulate the algorithm in what I hope is an understandable way. Eliot now works full time at PyMC Labs.
The paper
The arXiv paper just came out yesterday.
- Adrian Seyboldt, Eliot L. Carlson, Bob Carpenter. 2026. Preconditioning Hamiltonian Monte Carlo by minimizing Fisher Divergence. arXiv 2603.18845.
We’d love to hear any comments, corrections, or suggestions people might have.
My role in this was just helping to formulate the whole thing as a stats paper in a language that computational statisticians could understand. Adrian and Eliot worked out the proofs, from which I learned a lot.
Why is Nutpie better?
There are two basic reasons why Nutpie’s diagonal mass-matrix adaptation is better than Stan’s:
1. Better initialization: Use the gradients at the initial point to estimate mass matrix and make the whole sampler scale invariant.
2. Two sources of information: Use the position and the gradient of the log density at the position as two sources of information with which to estimate a mass matrix
The motivation for (1) is that the outer product of the scores is a one-sample Monte Carlo estimate of the expected negative Hessian of the target log density. The one-sample estimate is regularized by taking a geometric average with an identity matrix—using a few more draws here may be more stable.
Will wonders never cease, but it turns out to be relatively simple to take the geometric mean as the midpoint of a geodesic in the affine-invariant Riemannian manifold (AIRM) of positive-definite matrices. It works out to regular geometric mean in the diagonal case, making it simple to implement.
The motivation for (2) is that in a multivariate normal distribution with covariance Sigma, the covariance of the draws is Sigma and the covariance of the scores (gradients of log density of the draws) is Sigma inverse. These are then geometrically averaged to produce an estimate. This can converge very fast—many fewer iterations than the dimensionality of the problem.
As with Stan, the default in Nutpie is to use a diagonal mass matrix. Unlike Stan, Nutpie also lets you use a low-rank plus diagonal approximation (as in L-BFGS optimization), which can go a long way to decorrelating hierarchical models. Using a rank K approximation, what would otherwise be N^3 operations like solves or matrix products become K^2 * N, which is manageable for K of order 10 or less.
Mathematical motivation
The geometric average of the covariance of the draws and the inverse covariance of the scores minimizes the Fisher divergence from a standard normal distribution to the preconditioned target density. Fisher divergence is like Kullback-Leibler divergence, but instead measures the expected distance between gradients rather than the difference between log densities. This turns out in practice to lead to a better diagonal preconditioner. The paper provides proofs and an appendix with geometric motivations.
What we do not prove and what I do not believe anyone understands is how to select an optimal static preconditioner when the target log density has varying curvature, in any form: diagonal, low rank plus diagonal, or dense. I think what we really need is dynamic adaptation that isn’t cubic like Riemannian and doesn’t require implicit integrators.
Is there a plan to implement something similar in Stan?
Nutpie works with Stan models
Kudos to putting this together in a paper!
I’ve had some good luck swapping over heavier long-running models to nutpie for performance gains. Feels great to essentially get that for free.
Is there any desire to integrate this sampler into Stan? Between this and the various other samplers in the works (WALNUTS, etc.?) I’m curious what the roadmap looks like.
I’m working with Steve Bronder and Brian Ward to implement all of this in a memory-efficient way in multi-threaded C++, then provide Nutpie-like wrappers in Python and R. The current version is here:
GitHub: flatironinstitute/walnuts
The plan is roughly:
1. (implemented) Use WALNUTS’s per-leapfrog-step step-size adaptation with biased-progressive, multinomial NUTS.
2. (implemented) Use a continuously adapting version of Nutpie’s adaptation (initially diagonal only) rather than the blocked version implemented in Nutpie and described in the paper. The idea is to have a schedule with which to exponentially discount the past so that it achieves convergence for long runs.
3. (implemented) Replace dual averaging with Adam. I
4. (prototyped) Implement automatic stopping for adaptation and for sampling using Welford-accumulators and lock-free synchronization (triple buffer, and simple atomic, respectively). I have a draft of a paper writing up the algorithms for (2) and (3).
5. (to do, but modular) Swap out the standard HMC leapfrog integrator with the isokenitic (aka micro canonical) sampler. Tore Kleppe has been evaluating it and replicating the results from the papers by Jakob Robnik, Uroš Seljak, and Reuben Cohn-Gordon (e.g., Metropolis Adjusted Microcanonical Hamiltonian Monte Carlo).
6. (easy, but lots of design/coding work) Write wrapper in Python a la TinyStan.
7. (easy, but lots of coding work) Write wrapper in R following the same interface.
I think we should be able to get this out relatively soon. Hopefully well before StanCon. I submitted a StanCon talk abstract on (1)–(7).
Having said all this, I think the immediate-term future is going to be JAX for scalability, so the plan’s to port as much of the algorithm as we can to JAX. At the same time, Brian Ward, Mitzi Morris, Matthijs Vákár, Andrew Gelman, and I are writing a paper on how to code Stan-like models in JAX without using a graphical-modeling language like PyMC or NumPyro. There’s a draft with working title JAX à la Stan, if you’d like to see what we’re up to or comment. We hope to get a complete draft out soon with some more serious examples like time series and discussions of loops and conditionals (i.e., scan vs. select).
Very exciting stuff here, Bob, thanks for sharing.
I’m very interested in seeing how the JAX work progresses; the draft paper was interesting. Recently, I’ve found myself fitting some rather large IRT 2PL models on the order of ~200-300k parameters and ~500k observations. Stan fits it well, albeit slowly. Nutpie also fits great and a bit faster! I tried to go to NumPyro on GPU but couldn’t get anything stable enough to fit properly (likely because I was using a consumer GPU that only had fast F32 performance and F64 was slower than CPU). Being able to push the scalability up will certainly open up Bayesian modeling up to a larger class of applied problems, which is exciting.
I’ve never been too keen on the graphical modeling approach, so something like `densejax` would be really attractive in the future. I’d miss the aesthetics of Stan proper (although it seems feasible that Stan -> JAX is a possible transpilation target?), but these seems to a promising direction.
That sounds like one big model! I’m also really curious about microcanonical HMC, I think that should help with high dimensional posteriors like this? It’s also on my todo list for nutpie.
You can also use nutpie with a JAX density function by the way, and run the gradient evaluation on the gpu. It comes with more overhead than I would like because JAX still doesn’t have a good ffi interface, but for a model that big the overhead might not matter all that much?
I think we still have a lot of room for improvement on the GPU front. To really make good use of GPUs we will probably have to change the algorithms, but even if we stick with what we currently have, we could do much better than current nutpie or numpyro on a GPU.
If different chains run on different CUDA streams, we can keep the GPU busy even while one chain has to wait for the results of a u-turn check. With a couple of simple example models I got a big speedup that way. But currently we don’t have a clean way to get low level gradient functions for models unless you write them by hand…
You don’t have to convince me that Stan’s cleaner—I’m the one who designed the language! When you define a standalone language, the only limit is how well you can write parsers and code generators. When I was a grad student and faculty member, I spent a lot of my time working on programming language semantics—I even wrote a book on programming languages for NLP about the same time as BDA 1 came out). We have since picked up a couple more experts in programming language theory and practice, Matthias Vákár and Brian Ward, who are both co-authors on the JAX paper.
The paper will be more interesting after we add a discussion of conditionals/loops and autoregressive and spatial models. We’re going to try to write the densejax interface so that the models are as clean as possible.
We could transpile some of Stan to JAX, but JAX doesn’t support the same level of control flow as Stan because Stan uses dynamic autodiff (like PyTorch) rather than statically compiled traces like JAX. The other issue is that Python users are highly resistant to learning a new language, especially one with different C-like structure, types, and indexing from 1 (the indexing makes it awkward to use CmdStanPy if you need to index things like data offsets. The other advantage of JAX is that it’s integrated into Python cleanly. And it’s much easier to install! One of our motivations is that scientists tend to learn Python more than R these days, even in fields like biology, epidemiology, and ecology that were once dominated by R. We were really struck by Elizaveta Seminova’s invited talk at the last StanCon (Oxford) where she began by saying she’s no longer a Stan user because of the speed boost from JAX. Mitzi Morris just moved from Andrew’s group to the CDC and they have also largely abandoned Stan in favor of NumPyro.
The graphical modeling approach does have some advantages, which we will try to emphasize cleanly in the paper. Namely automated posterior predictive inferences, simulation-based calibration, etc. The problem, as we will point out, is that it winds up not being efficient for both density evaluation and posterior predictive inference in things like time-series models.
The usual problem with IRT 2PL is the identifiability because of the additive (ability – difficulty) and multiplicative (discrimatinion * ability) non-identifiability. One of our first grants for Stan was from the Institute for Educational Sciences (IES) to fit IRT models. Sophia Rabe-Hesketh helped Daniel Lee and I understand how to identify them so that we could use real demos in the grant application. What they really wanted was a Stan Stata interface, which in a very lucky coincidence for us, Robert Grant contributed to Stan soon after we got the grant! If you’d like to share your model either publicly or with me ([email protected]), I’d be happy to review it for efficiency.
P.S. One of our first stress tests for Stan was a large-ish baseball batting model where you had batters (offense) and pitchers (defense) rather than students and questions. We fit all of the historical baseball data, which involved about 100K parameters and about 10M data points (every at bat recorded in Baseball Reference circa 2012). I was disappointed to see that it didn’t show abilities getting better over time—I’d still like to see a model that could capture how much better athletes are now than they were 50 years ago.
Adirian—Nawaf Bou-Rabee and Tore Kleppe have replicated the efficiency of isokinetic sampling (rebranded as “microcanonical” by physicists Jakob Robnik et al.). There’s also the interesting option there of using unadjusted HMC. Jakob, Uroš Seljak, and Reuben Cohn-Gordon have convinced me it’s a good idea in the sense that it’s just another way to trade some bias for a reduction in variance. I think the right metric is how fast we can get to the same standard errors as we get with 100 (or pick another number) independent draws (i.e., estimates with standard errors of 1/10th of the posterior standard deviation). I’m writing this up now in a paper on the WALNUTS implementation, which is going to include micro-canonical sampling and run through all of the implementation steps.
Maybe I’m just dense, but when I read “JAX a la Stan”, it really looks a lot more cumbersome and less straightforward to write and fit a model in JAX compared to Stan. Maybe it’s just because I am used to it, but Stan seems much cleaner and more elegant from the standpoint of writing down your model and then translating to code than JAX does. Also, I like that Stan is a separate language and not part of a scripting language. “The downside is that this kind of static typing is unfamiliar to most of Stan’s intended users.” – I’m wondering who are “most” of Stan’s intended users ? I would have thought statisticians and scientists, but now I am wondering if the Stan team is thinking more ML crowd.
Jd:
I agree. I too find Stan much more intuitive to use than Jax. But Jax can be much faster than Stan for some models. Also, the flexibility of Jax allows us to encapsulate hierarchical models in a way that can’t currently be done in Stan.
One goal of the Jax a la Stan paper is to help users like me (and you, perhaps) who would like to get the speed of Jax and the ability to encapsulate that’s in Jax while retaining the benefits of the Stan approach to writing models.
Another goal is to help people who are already using Jax to get some of the statistical benefits of the Stan approach to modeling.
Hi, jd: As I said in the last comment, I really like Stan too—I designed it to be as clean as possible, which is a lot easier in a standalone language. That also gives it portability of a sort across R and Python, which was one of the motivations Facebook had building Prophet. The main drawback people find is that it’s another language to learn and interface. I thought that would be trivial, but this was my own dumb thinking coming from a CS/programming language theory background. That also means we don’t have anything like decent debugging tools for Stan, no IDEs, etc., all of which come for free with an embedded language like JAX (or higher-level wrappers like NumPyro).
The intended users we targeted in designing Stan are applied statisticians who need a tool for real modeling work either for papers or for production (as opposed to, say, as a pedagogical tool in classes about Bayesian statistics).
We’ve gotten an incredible (to me) amount of pushback on just how hard our users find the static typing. Andrew, for example, complains about it incessantly and finds it nearly impossible to convert between arrays and vectors/matrices. I took the design pattern from Eigen, the C++ matrix library, but it’s been very confusing for users and I now consider it the biggest design error I made with Stan. We should have followed the ndarray pattern from NumPy. I believe the problem the applied statisticians have is that neither R nor Python have static typing and these users don’t first learn languages like C/C++, Haskell/OCaml, or Rust or any other statically compiled language.
I also believe Python will eventually dominate R, not because it’s perfect, but because it’s much better behaved than R as a programming language and package manager and it’s what many applied scientists are learning now even in fields like epidemiology and ecology that were once almost entirely dominated by R. This isn’t even considering those who switch for efficiency.
For anyone interested in trying this out on their Stan models from R, I just finished a 1.0 version of my R package nutpieR: https://github.com/andytimm/nutpieR
It’s a pretty minimal set of R bindings to expose nuts-rs. I haven’t yet ported all of the functionality, especially anything related to their normalizing flows work.
Thanks to r-universe, I’ve been able to expose precompiled binaries for most platforms, which should make this relatively easy to install – no Rust toolchain needed. You’ll still need the usual C++ tools for Stan, since models are compiled at runtime via BridgeStan.
Nice! The interface looks very clean in R. And it’s nice you can plug the output into the posterior package. I also appreciate the MIT license. This should be enough to get Andrew to try it.
My only suggestion from the top-level README (which is great!) is to read the Stan program in from a file. When using a string for a Stan model you don’t get any syntax highlighting, auto-formatting, etc., and depending on what you choose for quotes, you either can’t use the transpose operator (‘) or the quote for prints (“) without escaping. It’s also easy to reference line numbers of errors in a standalone file. And it makes the Stan program usable with other R programs and even in Python or Julia.
ah thanks for the suggestion- referencing a .stan already works.
I think the readme could use a tweak to emphasize that better though, will do that today.
Thanks for posting — I have an R-using friend who has wanted to try out nutpie
That’s great! I was hoping someone would do that eventually, I really don’t know R well, so I don’t think I’m the right person to do this…