We are happy to announce the official release of Walnutpie version 0.0.1.
Walnutpie is an MCMC sampler for continuously differentiable densities coded in Python, accepting models coded in Stan, PyMC, NumPyro, JAX, and plain old Python.
Walnutpie is not an official Stan project
I thought this was worth saying up front. It may eventually migrate to Stan, but for now, we followed the Nutpie approach of building a standalone sampling package that worked with a range of packages for defining models.
pip installable
Walnutpie is on PyPI, so it’s pip installable. The documentation includes information on getting started, running models, and posterior analysis.
Stan through C++
Walnutpie runs Stan models through C++ using BridgeStan, so there is no Python dispatch overhead for Stan sampling. The basic architecture of the API is based on Adrian Seyboldt’s sampler Nutpie. We are working on doing that for other packages like NumPy and JAX to the extent that we can.
GitHub source
Development discussions and source code are managed through GitHub.
Features of Walnutpie
We are almost ready to release the paper on arXiv with all of the gory pseudocode details all of the algorithms used for Walnutpie. It will explain the following points in detail.
-
Walnuts: The underlying Hamiltonian Monte Carlo sampler is Walnuts. Walnuts uses Nuts for choosing the number of steps per iteration. It further allows step sizes within the Hamiltonian dynamics simulation to be lowered when necessary to preserve simulation accuracy. This helps with robustness and with accuracy in multi-scale distributions (i.e., ones where the curvature as represented by the Hessian varies around the posterior). With a high tolerance threshold, Walnuts reverts to Nuts’s behavior.
-
Mass-matrix warmup: The mass-matrix warmup strategy is an online form of Nutpie (links to: the paper and software). Nutpie minimizes Fisher divergence by estimating the inverse mass matrix as the midpoint (in the appropriate manifold) between an estimate based on the variance of the draws and the covariance of the scores (gradients of the log density). The target is better than Nuts’s variance of draws in both convergence speed and sampling efficiency. Walnutpie only supports diagonal mass matrices (Nuts supports dense matrices and Nutpie supports low-rank plus diagonal and even more general normalizing flow approaches). The approach is online in the sense that it is not blocked like warmup in Nuts or Nutpie—it updates every iteration by exponentially discounting the past to mimic Stan’s exponentially increasing history sizes. We also borrow Nutpie’s mass matrix initialization based on a regularized outer product of gradients at the initial point.
-
Step-size warmup: The step size adaptation strategy has not changed, but the underlying stochastic gradient descent algorithm is Adam rather than dual averaging. We found Adam to be faster to converge and much more stable. Matt Hoffman included a hack in the original Nuts approach to stabilize dual averaging, but even with that it is not as stable as Adam.
-
Concurrency and automatic stopping: The underlying sampler is multi-threaded (using C++11 threads) with shared data. On top of the multi-threading, we have layered a convergence monitor in a separate thread that communicates with the chains through lock-free, latest-only, single-producer/single-consumer (SPSC) buffers (specifically, a triple buffer). The monitor automatically stops warmup when the mass matrices and step sizes have converged within tolerance to their cross-chain averages. The monitor automatically stops sampling when a target (traditional, non-split, non-ranked) \widehat{R} threshold is satisfied for the unnormalized log density, which typically converges more slowly than any of the individual parameters. It can also be configured to run for a fixed number of warmup and/or sampling iterations. The link between the original \widehat{R} and effective sample size makes this essentially an unscaled ESS target.
-
Ragged chain summaries: Asynchronous concurrent execution of chains with automatic stopping produces chains of different lengths. Because ArviZ does not accept ragged chain input of this kind, we have included posterior analysis tools for means, variances/standard deviations, quantiles, traditional R-hat, effective sample size, and Monte Carlo standard error that work with ragged chains.
-
C++20: Walutpie is implemented in C++20. As a programming language type fanatic (I’ve written two books with “type” in the title!), I don’t know how I survived without C++ concepts before C++20.
-
ctypes FFI:The foreign function interface in Python uses ctypes rather than a higher-level interface, which sidesteps the requirement of ABI compatibility of C++ binaries.
Developers
We would also like to welcome new developers who may want to get involved. There are already a stack of improvements we’d like to make, which we have enumerated on the GitHub issues.
Sources of algorithms
The Walnuts algorithm was a joint effort among Nawaf Bou-Rabee, Sifan Liu, Tore Kleppe, and Milo Marsden. Nutpie was developed by Adrian Seyboldt. Nuts, in the form used currently in Stan, was originally developed by Matt Hoffman and Andrew Gelman, then improved with multinomial sampling and mass matrix adaptation by Michael Betancourt. We haven’t yet added cross-chain adaptation as developed by Ben Bales, but the pieces are all in place to do so.
Brian Ward and I wrote all of the version 0.0.1 code with design advice and code review from Steve Bronder. Claude (the LLM) helped with code review and testing, but we didn’t use it to write the actual code (not out of principle, but because Brian and I both prefer the control of doing things manually).
Feedback
We would very much appreciate any feedback people have, including feedback on the documentation and ease of use, the source code, and performance.
We are happy to get feedback through issues on GitHub, through replies to this post, or through mail to one of the developers.
