I’ve been independently building stanwasm, an implementation of a subset of
the Stan language that parses, compiles and samples models entirely in the
browser. It is an independent implementation rather than a port: it does not
use stanc3 or Stan Math, and it is not affiliated with stan-dev.
I recently came across the threads on
stanli and
running stanr / brms client-side via webR,
and they made me realise there are now several quite different approaches to
making Stan usable on the web — approaches that reuse as much of the existing
stack as possible, and, at the other end, ones that give that up for a smaller
self-contained runtime. stanwasm sits at the second end. I’d like to describe
what that costs and what it buys, and get feedback on whether I have drawn the
line in a useful place.
Try it: gallery (nothing to install)
· from Python, on JupyterLite
· repository (Apache-2.0, stanwasm on npm)
const model = new StanModel(stanCode, JSON.stringify(data));
const draws = model.sample(model.randomInit(42n), 1000, 1000, 42n);
// starting point, warmup, sampling draws, seed
What it is, in numbers
- Sampling:
nuts-rsby Adrian
Seyboldt and the PyMC developers, compiled to wasm. - Frontend and runtime: the parser, evaluator, autodiff tape and code
generator are written from scratch in Rust in this project. - Model specialisation: at model load time the computation is recorded onto
an autodiff tape. That tape can be walked directly, or rewritten into a
model-specific wasm module holding the forward and reverse passes. - Coverage: 141 of the 147 posteriors in posteriordb load, evaluate a
gradient, and compile. - Gradients: across sixteen models, log densities and gradients agree with
CmdStan 2.38.0 to 3.3e-13 at worst, thirteen of them to 1e-14 or better. - Draws: the sampler is independent code, so agreement on the gradient does
not imply the same draws. Across the 45 posteriordb posteriors that ship
reference draws, 42 have posterior means within 0.2 sd of the reference
mean, and none beyond 0.5 sd. That is a check on the mean only — I have not
compared
R-hat, ESS or divergences yet. - Size: 664 KB, 253 KB gzipped, shipped inside the npm package. Nothing is
fetched per model. - Cold start: blank page to first posterior draw, 0.4 s over 50 Mbps and
1.6 s over 1.6 Mbps. - Gradient execution: the model-specific wasm path is 1.5x to 12x faster
per gradient than walking stanwasm’s own recorded tape. That is an internal
comparison between my two execution paths, not a benchmark against Stan or
anything else.
The cold-start number is the one I actually optimise for. If inference is
embedded in a blog post or a teaching page, a reader deciding whether to try it
cares about time to first draw rather than throughput after a long setup.
Why a separate implementation
The conventional Stan path is
model.stan → stanc3 → model.hpp → a C++ compiler → an executable, and Stan
Math’s reverse mode works through that: instantiate the templates with a
reverse-mode type and the compiled program contains the operations that build
and run the derivative.
What I wanted was slightly different — after seeing a particular model in the
browser, turn its recorded forward and reverse computation directly into
another wasm module. The difference that makes that possible is that stanwasm
keeps differentiation explicitly as data: the tape in
crates/stanwasm-autodiff/src/lib.rs
is a list of Op::Add, Op::Mul, Op::Log and so on, which the code
generator can walk and emit instructions for.
So the question was never “can Stan Math run in wasm” — clearly it can. It was:
if I don’t want to ship a C++ toolchain into the browser, but I do want to
specialise machine code for a model that only arrives after the page has
loaded, what representation belongs between Stan source and the sampler?
stanwasm’s answer is to make the autodiff graph an explicit IR and lower that
IR to wasm.
There is more than one reasonable answer to that. A static graph over a
precompiled kernel set is one; emitting a new function for the specific model,
which is what stanwasm does, is another. They trade compatibility, artifact
size and specialisation differently, and I do not think my own project is the
one that gets to say which trade is better. That is genuinely one of the things
I would like feedback on.
I would also be glad if anything here turns out to be useful elsewhere. Nothing
about the tape-to-wasm step is specific to my subset, and I would rather see
the idea tested against a fuller implementation than kept to this one.
A small experiment on binary size
I also wanted to know whether Stan Math itself is what makes a browser artifact
large, so I measured rather than assumed — building it with emcc 6.0.9 against
stan-dev/math main.
| Stan Math built for one specific model | wasm | gzip |
|---|---|---|
rev.hpp + normal_lpdf |
0.10 MB | 0.05 MB |
math.hpp (the umbrella header) + normal_lpdf + student_t_lpdf |
0.11 MB | 0.05 MB |
The second row is what surprised me. Even with #include <stan/math.hpp>
making the whole library available, the compiler only instantiates the
templates the program actually calls, so for one known model Stan Math’s
contribution is about 0.1 MB. (I ran the result under node to check the value:
lp = -9.993098.)
That suggests the distinction is not “Stan Math is large” but something more
specific:
a model known at build time
→ instantiate only what it calls → a small model-specific binary
an arbitrary model loaded later
→ the precompiled runtime must carry the
whole supported operation surface → a larger fixed runtime
an arbitrary model, specialised in the browser
→ a small generic runtime, plus a
model-specific module emitted in the browser
stanwasm takes the third row: avoid the fixed cost by owning a much smaller
mathematical surface and generating code once the model is known. The price is
that I now own that surface.
The trade
Gained. 253 KB gzipped; 1.6 s to the first draw on a throttled connection;
no server and no per-model download; Chromium, Firefox and WebKit, so iOS too;
and a model-specialised execution path alongside tape replay.
Given up. The language is a subset of Stan, not the full language. The
sampler is independent
code, so the same seed will not reproduce Stan’s draws. Fixes in Stan Math do
not flow in. Special functions (lgamma, digamma, erf, the incomplete
beta) and each distribution’s edge cases had to be written and tested here.
And the one I am least sure about: dynamic computation shapes sit badly with
a record-once architecture. ODE integrators are the clearest case — an
adaptive solver does parameter-dependent work, so the recorded graph would
freeze at whatever the tracing point chose. stanwasm refuses rather than
freezing, but I have not built a dynamic fallback for that path.
What I’ve built with it
pystanwasm — Stan models from
Python on Pyodide, two independent wasm modules side by side in one tab, no
server and nothing to install:
JupyterLite /
marimo.
onboard-bayes — on-device
Bayesian calibration of phone motion sensors (gyroscope and accelerometer bias,
walking stride, heading) inside a WebView.
This 2019 thread
on compiling Stan to wasm mentioned online learning of a magnetic field for
indoor navigation as one motivation; this is an attempt at that, with the model
in Stan syntax and the sensor data never leaving the device. That shape —
Bayesian inference as a local capability of a web application rather than
something behind an API — is what originally motivated the project.
What I’d like feedback on
The subset is mostly an accident of working through posteriordb rather than
something I designed, so I have no strong view yet on where the boundary
belongs. Three things I’d genuinely like opinions on:
- Are there browser-embedded uses where a small runtime and a fast cold start
matter enough to justify an independent implementation at all? - Which unsupported feature would immediately stop you embedding a model you
actually care about? - Is emitting code per model an interesting direction for browser Stan at all,
or does a reusable precompiled graph already capture most of what it would
buy, at far better compatibility?
And concretely: if you have a Stan model you’d want to put in a page and
stanwasm doesn’t load it, that is the most useful thing you could send me.
Thanks for reading — criticism of the architecture is as welcome as ideas for
use cases.