Skip to content

About

A prototype compiled language for Bayesian inference: write the model as mathematics, and the compiler derives gradients, lays out memory and integrates out what it recognises before NUTS runs.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Latest commit

 

History

69 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Martin: Tony the eagle carrying a crate stencilled gross 37,753, net 253 parameters

Martin is a prototype language for Bayesian inference and small numerical programs. You write the model as mathematics, with its intent stated in the types: shapes, positivity, symmetric positive definiteness, which quantities are data and which are unknown. The compiler uses those facts to reject wrong programs, to derive gradients, to lay out memory, and to decide what work the inference does not need at all.

Martin has its own compiler, mintc, written in Rust: a parser, type checker, compile-time differentiation, layout and fusion passes, and an LLVM IR emitter. LLVM turns that IR into machine code; nothing interprets a Martin program while it runs. A small runtime library, written in C, supplies file I/O, printing, a Cholesky solve and the NUTS sampler that calls the compiled model. (The tools still carry the earlier name: the compiler is mintc, programs end in .mint, and the runtime's variables start with MINT_.)

Martin is named after the author's father, and Tony the eagle after him too: why Martin?

It is a research prototype: one machine (Ryzen 9 5900X, Linux), a handful of models, and every number below comes from files in bench/.

Speed

Same sampler, different languages. Martin, hand-written Rust and Stan's C++ (through BridgeStan) were run under the same NUTS implementation, with the same warmup, settings, seeds and thread budget; what differs is the code computing the log density and its gradient (and, for Martin, the parameter order inside the sampler: see below). All of them agree on the gradient to about 3e-14 (rules and full results). These runs were made at commit de58040; the runtime has changed since (streaming summaries, the Kalman collapse, the low-rank merge), and a quick rerun at the current commit gave similar ratios.

Time per gradient, µs:

model Martin Rust, max effort Rust with Martin's tricks Stan
time series, 3,171 parameters 3.56 4.10 3.73 51.0
time series, 37,901 parameters 44.4 52.9 46.6 654
same, 3 threads 17.0 single-threaded 18.4 single-threaded
logistic regression (n = 5000, p = 20) 23.6 38.0 42.7
eight schools 0.025 0.025 0.50

Whole runs, 4 chains, 1000 warmup + 1000 draws (medians over 5 seeds, 3 for the large model):

model Martin Rust, max effort Rust with Martin's tricks Stan
time series, 3,171 parameters 3.32 s 3.76 s 3.50 s 28.3 s
time series, 37,901 parameters 62.4 s 98.6 s 65.2 s 737 and 741 s
logistic regression 0.38 s 0.59 s 0.64 s
eight schools 5.5 ms 5.7 ms 16.6 ms

What to read into this:

  • On the time series, Martin's lead comes from its memory layout. Its compiler stores the random-walk matrix in blocks of 8 series, column by column, and the runtime converts the sampler's vector to that order (mint_set_layout). The Rust baselines kept the user's row-major order; the runtime offers them the same hook, but the harness did not use it. With the layout switched off, Martin is slower than both Rust versions (7.35 µs against 4.10 and 3.73 µs, and 96.2 µs against 52.9 and 46.6 µs). So the result is that Martin's compiler finds this layout from 18 lines, not that Rust cannot have it.
  • "Rust, max effort" is hand-written with AVX2 intrinsics and glibc's vector math. "Rust with Martin's tricks" (baselines/dynpois_par.rs, 827 lines) also has Martin's table-driven exp and splits its gradient across threads; against it Martin's lead is 5 to 8%.
  • On the large model every implementation gets 3 sampler threads per chain, but the max-effort Rust and Stan compute their gradient on one of them, so the like-for-like whole-run comparison is 62.4 s against 65.2 s for the threaded Rust, not against 98.6 s. Stan's large runs were not interleaved with the others. On eight schools, Stan's column is mostly BridgeStan's per-call cost.
  • The logistic and Newton Rust baselines were written early and not re-tuned while Martin's compiler gained the techniques behind these numbers (its own vector exp and log1p, the tiled Gram kernel, row fusion). On the time series, giving the Rust Martin's tricks cut the gap from 15 to 19% to 5 to 8%; a similar effort on these two was not made.
  • Newton's method for logistic regression (200,000 rows, 50 features, 10 iterations) takes 0.114 s against 0.209 s for the max-effort Rust (merged build; the Rust figure was measured once and frozen).
  • The Martin programs are 15 to 18 lines; the max-effort Rust is 92 to 105 lines plus shared SIMD helpers, and 592 lines for the time series. The eight-schools Rust is plain Rust.
  • A pass of the small whole runs taken while other jobs loaded the machine reversed the time-series ranking (max-effort Rust 7.28 s, Martin 9.98 s; bench/same_sampler/results/whole_small_loaded.json).

Every framework on the hardest model

The large dynamic Poisson panel (250 series, 150 time points, 37,901 parameters), written as it stands in each framework with no reformulation, and with each framework's own defaults apart from the options listed: 4 chains of 1000 warmup and 1000 draws (Martin's opt-in fast warmup uses 200 of them; rustmc keeps 125 draws thinned by 8 because of its storage cap), every run pinned to the same 12 cores, one run at a time. A run counts only if it finished within 20 minutes, compilation included, and mixed: R-hat at most 1.01 and bulk and tail ESS at least 400 over 501 quantities, from one ArviZ script for all. Every implementation's log density and gradient matched the specification to 2e-13 or better, except rustmc, which exposes no log density and was checked by its posterior means only (rules, settings and caveats; every run).

Seed 1, on a quiet machine; the mixing column counts all three seeds where three were run:

implementation lines of code total time peak memory mixed? verdict
Martin, opt-in sampler options 18 42 s 0.68 GiB 3 of 3 seeds ranked
Martin, default settings 18 82 s 0.31 GiB 3 of 3 seeds ranked
Rust gradient under Martin's sampler 982 84 s 0.31 GiB 2 of 3 seeds ranked on 2 seeds
Rust gradient with nuts-rs, diagonal 1,166 118 s 0.53 GiB 0 of 3 seeds did not mix
Rust gradient with nuts-rs, low rank 1,166 828 s 6.25 GiB no (1 seed) did not mix
rustmc (elliptical slice sampling) 12 111 s 0.20 GiB 0 of 3 seeds did not mix
PyMC + nutpie, numba 24 264 s 5.11 GiB 0 of 3 seeds did not mix
PyMC + nutpie, JAX 24 596 s 5.39 GiB 0 of 3 seeds did not mix
nutpie with the Stan model, diagonal 44 1,139 s 5.01 GiB no (seeds 2 and 3 timed out) did not mix
NumPyro 25 1,488 s 2.75 GiB yes (1 seed) over 20 minutes
nutpie with the Stan model, low rank 44 1,878 s 5.45 GiB no (1 seed) over 20 minutes
Stan, reduce_sum, 3 threads per chain 57 2,399 s 1.59 GiB yes (1 seed) over 20 minutes
Stan 48 4,076 s 1.86 GiB no (1 seed; R-hat 1.017, bulk ESS 363) over 20 minutes

What to read into this:

  • This is the out-of-the-box test: the model as written, nothing reformulated. Most engines failed by not mixing within 1000 tuning iterations, not by running slowly. A statistician who reparameterised the model or integrated the walk out would change these results; that comparison, with the expert's extra code and time measured, has not been run yet.
  • The sampler matters more than the language here. The same tuned Rust gradient mixed in 2 of 3 seeds under Martin's sampler and in 0 of 3 under nuts-rs. Against that Rust under the same sampler, Martin took 82 s against 84 s on seed 1; the later seeds ran under different outside loads, so they do not separate the two. Martin mixing in 3 of 3 seeds against the Rust's 2 of 3 is not a finding either: Martin's default seed 1 passed at R-hat 1.0097 against the 1.01 bar, the Rust's failed at 1.017, and one more seed could change either verdict. Martin's sampler is a port of Stan's NUTS with its own warmup and threading; the opt-in options (a low-rank metric, and 200 warmup iterations instead of 1000) are runtime algorithms, not compiler output. The fast warmup starts every chain near the same optimised point, which makes split R-hat a weaker test of mixing for those runs.
  • The 20-minute limit was chosen after the first results had been seen; real times are shown for every run. Seeds 2 and 3 ran while other work loaded the machine, from 14% to 76% slower than seed 1 depending on the configuration. Some seed-1 runs were redone after fixes to the harness (Martin's default run after a post-processing fix that had inflated its memory, 87.5 s first time; the Rust under Martin's sampler, 88.0 s first time; nuts-rs diagonal, 214 s first time); the first records are kept in bench/shootout/results/probes/. Configurations over 20 minutes on seed 1 were mostly run once. This is a model Martin had already been tuned on.
  • Lines of code count the model plus loading the data and running it; the Rust counts include a self-test and instrumentation. Memory is each framework's default trace handling (nutpie keeps the whole trace in memory; Martin streams draws to a file). Martin uses 3 threads per chain; Stan without reduce_sum, nutpie and rustmc use one; JAX can use all 12 cores.
  • Not tried: longer warmup for the nuts-rs samplers, nutpie's draw_diag adaptation, and longer runs.

Work the compiler removes

A latent random walk integrated out (details). In examples/random_walk_panel.mint, a panel of Gaussian random walks observed with Gaussian noise and unknown scales, mintc finds the walk, integrates it out exactly with a Kalman filter per series, and NUTS samples the 23 (G = 20) or 253 (G = 250) remaining parameters instead of 3,023 or 37,753; the walk is drawn back afterwards. Over three sampler seeds, on one simulated data set per size, the lowest effective sample size of the remaining parameters per gradient was 100 to 146 times that of full NUTS on the non-centred form. (The centred form as written did not converge in some runs, so its larger ratios are not quoted.) Each collapsed gradient costs 1.8 to 3.6 times as much, and drawing the walk back afterwards costs further gradients' worth per kept draw, which these ratios leave out; the recorded runs were also made before that step was rewritten. No hand-collapsed baseline (such as Stan's gaussian_dlm_obs) was run. The log density matches a dense Gaussian computation to about 1e-15 on small panels, and the posterior matches full NUTS within Monte Carlo error on small panels (see kalman.md for how the threshold was set). Only a local-level walk with Gaussian observations is recognised so far.

Sampler options (details). Two opt-in changes to the runtime's NUTS cut the gradients needed per effective draw: MINT_WARMUP=fast (an L-BFGS starting point, a shorter warmup, chains pooling their adaptation) and MINT_METRIC=lowrank (Stan's diagonal metric plus up to 24 directions estimated from warmup gradients). On the small time-series model, over six seeds, the low-rank metric raised effective draws per second from 546 to 1,232 (median; data), and the fast warmup gave 1.35 to 1.90 times as many effective draws per gradient over eight seeds (data). Both were measured before the streaming summaries, whose ESS estimate can read low for slowly mixing parameters, so reruns are not directly comparable. They are not the defaults yet: four models is too few to rule out posteriors where they do worse, and the fast warmup starts every chain near the same point, which weakens split R-hat as a convergence check. They are runtime algorithms, not compiler output: a Rust gradient linked to the same runtime would get them too (not measured).

Two examples

Newton's method for L2-regularised logistic regression. The checker proves H is symmetric positive definite, so solve compiles to a Cholesky solve; diag(...) is never built as an n-by-n matrix, and only one triangle of the Gram product is computed.

fn fit(X: Matrix[n, p], y: Vector[n], lambda: Positive) -> Vector[p] {
    let mut w = zeros(p)
    repeat 10 {
        let mu = sigmoid(X * w)
        let g  = X' * (mu - y) + lambda * w
        let H  = X' * diag(mu .* (1 - mu)) * X + lambda * I(p)
        w = w - solve(H, g)
    }
    w
}

Replace H with X' * X and the program no longer compiles:

error: solve needs a matrix the compiler can prove is SPD, but it is only known to be PSD, which allows it to be singular
help: SPD comes from I(p), from PSD + (Positive * I(p)), from sums and positive multiples of SPD matrices, or from assume_spd(A), which is checked at run time

Bayesian logistic regression, sampled with NUTS. The gradient is derived by the compiler:

model Logistic {
    data X: Matrix[n, p]
    data y: Vector[n]

    param alpha: Real
    param beta: Vector[p]

    alpha ~ Normal(0, 2.5)
    beta  ~ Normal(0, 1)
    y     ~ BernoulliLogit(alpha + X * beta)
}

fn main() {
    let X: Matrix[n, p] = read("data/logit_X.f64")
    let y: Vector[n]    = read("data/logit_y.f64")
    let post = sample(Logistic(X, y), draws = 1000, warmup = 1000, chains = 4, seed = 7)
    print(post)
}

The straightforward Rust equivalent of the model block is a hand-written logp function with the gradient worked out on paper: priors and their derivatives, the stable softplus, the residuals y - sigmoid(eta), and the X' r accumulation. See baselines/logistic_bayes.rs. Declare a scale parameter Real instead of Positive and Martin refuses it:

error: the scale of Normal must be Positive, but `sigma` is Real
help: declare the parameter as `param sigma: Positive`; Martin then samples log(sigma) and adds the Jacobian itself

All examples are in examples/: logistic_newton.mint, logistic_bayes.mint, linear_bayes.mint, eight_schools.mint, dynamic_poisson.mint and random_walk_panel.mint. Programs that must fail to compile are in examples/errors/.

Quick start

Requirements: Rust (cargo), clang with its OpenMP runtime (libomp), and glibc's vector math library (libmvec). It has been tested only on Linux x86-64 with rustc 1.89, clang 22.1 and glibc 2.44; other versions are untested.

./bench/setup.sh            # builds mintc and writes the example datasets to data/
./compiler/target/release/mintc build examples/logistic_bayes.mint -o build/logistic_bayes
./build/logistic_bayes

mintc check FILE type-checks only; mintc emit FILE writes the LLVM IR; mintc explain FILE reports what the compiler found and did (below).

The full benchmark (it builds everything, generates data and runs about 6 minutes; the max-effort baselines need nightly Rust) and the test suite:

python3 bench/bench.py 7    # writes bench/results.json
./tests/run.sh              # compile errors, runtime checks, gradients, posteriors
(cd compiler && cargo test --release)

Environment variables for compiled programs (threads, draws, the sampler's metric and warmup) and the compiler switches that turn each optimisation off are listed in docs/options.md.

What the compiler did: mintc explain

mintc explain FILE [flags] compiles the program as build does, with the same flags (a test checks that explain -o writes the same IR as emit), and prints what the compiler decided: per model, the parameters and their transforms, how each ~ statement was compiled (including which latent random walk the Kalman collapse integrated out, with which filter, and what NUTS samples instead, or why a candidate was refused), the data layouts and the data that may be read narrow; per function, the row fusion groups, the kernels and the type facts behind solve. The decisions and their reasons are recorded by the code generators where they make them; the descriptions of what a chosen kernel does are fixed text kept next to its code. Sizes that depend on the data are symbolic. An excerpt for examples/dynamic_poisson.mint (... marks omitted lines):

model DynamicPoisson
  ...
  parameters: NUTS samples G x T + G + T + 1 unconstrained values
  ...
    line 20: y ~ PoissonLog(beta + state)
      with state = cumsum(shared + innov, T)
      ...
      fused scan kernel over Matrix[G, T]: 1 running sum along T, each carried in a register along its row
      vector code: 4 lanes (<4 x double>), one row per lane, in groups of 8 rows (2 vectors); ...
      ...
      absorbed: line 17 (innov ~ Normal(0, 0.08)), in the reverse loop
      owned gradient: innov; ...

and for the Newton example, why solve may use Cholesky:

      solve(H, g): Cholesky solve (mint_chol_solve), allowed because H is proved SPD:
        H = X' * diag(mu .* (1 - mu)) * X + lambda * I(p) is SPD: PSD + SPD
          X' * diag(mu .* (1 - mu)) * X is PSD: a Gram product with weights >= 0 (the weights are Prob)
            mu .* (1 - mu) is Prob: Prob .* Prob
            ...
          lambda * I(p) is SPD: Positive * SPD

Choices that depend on the data values, such as narrow copies, are reported as the checks that sample() makes when it starts.

The language in one page

Types.

  • Scalars: Real, Positive, Prob (the open interval 0 to 1), Int.
  • Vectors: Vector[n], and vectors whose entries lie in a domain: Positive[n], Prob[n].
  • Matrices: Matrix[m, n], PSD[n], SPD[n].

Dimension names such as n are bound by function parameters or by an annotated read, and are checked everywhere after. A vector combined with a matrix is matched to the dimension with the same name: a Vector[G] plus a Matrix[G, T] repeats across T. If both dimensions have the same name, that is a compile error, because it is ambiguous.

Expressions.

  • + - * / and ^. The elementwise operators are .* and ./, and A' is the transpose.
  • * means the mathematical product. Between a scalar and anything it scales; A * v is a matrix-vector product and v' * w is a dot product. Between two vectors it is an error, because it is ambiguous.
  • Functions: exp, log, log1p, sqrt, sigmoid and abs apply elementwise; sum, dot and norm reduce.
  • cumsum(v) and cumsum(M, T) are running sums along the last dimension.
  • Constructors: zeros(p), ones(p), I(p), diag(w) (inside a product only), and literals like [1, 2, 3].
  • solve(H, g), which needs a proved-SPD H.
  • assume_spd(A), which checks symmetry and positive definiteness at run time.

Statements.

  • let x = ..., let mut x = ..., x = ...
  • repeat N { ... }
  • print(...)
  • let X: Matrix[n, p] = read("file.f64")
  • clock()

Models.

  • data and param declarations. Parameters are Real, Positive, Vector[n], Positive[n] or Matrix[m, n]. Positive parameters are sampled on the log scale, with the Jacobian added for you.
  • let definitions and x ~ Distribution(...) statements, with Normal, BernoulliLogit, PoissonLog and Exponential. The two sides of a ~ can be scalars, vectors or matrices.
  • sample(Model(data...), draws =, warmup =, chains =, seed =) runs NUTS with Stan's warmup and returns a posterior; print summarises it (mean, sd, quantiles, ESS, split R-hat).

Data files (.f64) hold two little-endian u64 values (rows, cols), followed by row-major little-endian f64.

Documents

About

A prototype compiled language for Bayesian inference: write the model as mathematics, and the compiler derives gradients, lays out memory and integrates out what it recognises before NUTS runs.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages