probabilistic-programming

How it is built

Architecture overview

A probabilistic programming system lets a statistician write down a model ("these observations come from this distribution with these unknown parameters and these priors") and get back the posterior distribution of the parameters given the data. The user writes the model; the system supplies the inference. Stan and PyMC are the standard examples, with NumPyro, Pyro and Turing.jl on the same pattern.

The one idea is to reduce every model to a single function: the log of the unnormalised posterior density over an unconstrained real vector. Once the model is that function and its gradient, one set of general algorithms applies to every model. Parameters with constraints (a variance must be positive, a probability must lie in [0, 1], a covariance matrix must be positive definite) are mapped to the whole real line by smooth transforms, with a Jacobian term added to keep the density correct. Reverse-mode automatic differentiation then produces the gradient at a cost close to that of evaluating the function. The inference engines consume only these two things: Hamiltonian Monte Carlo with the No-U-Turn sampler (the default in Stan and PyMC), variational inference when a fast approximation is enough, optimisation for point estimates, and sequential Monte Carlo when the data arrive in time order or the model has no usable gradient.

Around the engines sit the parts that decide whether the answer can be believed: convergence diagnostics across several chains, careful floating-point arithmetic in the log density, and reproducible random streams.

Parts:

• Model compiler and parameter space — turns the model into a log density on unconstrained space — Constrained parameter transforms (no page yet), Cholesky and LU factorisation
• Gradient engine — derivatives of the log density — Reverse-mode autodiff
• Numerics — keeps the density finite and the sums accurate — Floating-point error
• Gradient-based sampler — draws correlated samples that explore the posterior — Hamiltonian Monte Carlo and NUTS (no page yet), Markov chain Monte Carlo (no page yet), Verlet leapfrog
• Approximate inference — fits a simple distribution to the posterior — Variational inference (no page yet), AdamW, L-BFGS
• Sequential inference — particles that follow a changing target — Sequential Monte Carlo (no page yet), Monte Carlo
• Random numbers — reproducible parallel streams — PCG and xoshiro generators
• Diagnostics — decide whether to trust the draws — MCMC convergence diagnostics (no page yet)

 model code (Stan / PyMC / NumPyro ...)
      │ compile or trace
      ▼
 log p(θ_constrained | data)  ──► transforms to unconstrained y, add log|Jacobian|
      │                                   │
      │ value                             │ gradient ∇ log p(y)   (reverse-mode autodiff)
      ▼                                   ▼
 ┌───────────────┬──────────────────┬──────────────────┬───────────────────┐
 │ NUTS / HMC    │ variational (VI) │ L-BFGS (mode)    │ SMC (particles)   │
 │ leapfrog +    │ ELBO + SGD       │ MAP / start point│ reweight/resample │
 │ adaptation    │ (Adam-type)      │                  │ (no gradient ok)  │
 └──────┬────────┴────────┬─────────┴──────────────────┴─────────┬─────────┘
        ▼                 ▼                                      ▼
   chains of draws   fitted q(y)                         weighted particles
        │ several chains, seeded PRNG streams
        ▼
 diagnostics: R-hat, effective sample size, divergences ─► summaries, predictions

How the parts fit together

From model to log density

A model is a program that assigns a log probability to each combination of parameters and data. The system either compiles it (Stan translates the model to C++) or traces it by running it (PyMC and Pyro build a graph or trace from Python). Either way, the output is a function log_p(y) over an unconstrained vector y. The contract with every engine below is exactly this: evaluate the density and its gradient at a point, in a vector of real numbers, with no knowledge of the model's structure.

The step that makes this possible is Constrained parameter transforms (no page yet). A positive scale σ becomes y = log σ; a probability becomes logit p; simplexes and correlation matrices use dedicated transforms built from stick-breaking or Cholesky factors. Each transform changes the density by the absolute determinant of its Jacobian, so log_p(y) includes log|J|; forgetting that term samples from the wrong distribution while looking fine. Stan documents its transforms and the Jacobian adjustment in its reference manual. Dense multivariate normals, Gaussian processes and correlation structures need a matrix factorisation inside the density; Cholesky and LU factorisation factors an SPD matrix once in O(n³) and solves each right-hand side in O(n²), so these models are bounded by that cubic cost in the matrix size.

Gradients

Reverse-mode autodiff records the forward computation of log_p on a tape and sweeps it backwards once to produce the gradient with respect to all parameters. For a scalar output and d inputs this costs a small constant multiple of one evaluation, instead of d evaluations as with finite differences; that is why HMC remains usable with thousands of parameters. The price is memory for the tape. The numerical engine around it (Stan Math in Stan; the tensor library underneath PyMC, NumPyro and Pyro) provides the operations with derivatives defined. A user-written function that autodiff cannot see through, such as a call into foreign code, breaks the gradient, which is the most common reason models cannot use HMC.

Numerical care in the density

Log densities are sums of many terms and are evaluated millions of times, so rounding error in them is a systematic risk. Floating-point error explains the two habits that matter: sum in the log domain, with the log-sum-exp trick for mixtures, and use compensated or pairwise summation for long sums. Gradient-based sampling is also sensitive to a density that is nearly flat or has extreme curvature, because the leapfrog trajectory then amplifies tiny errors; models are reparameterised (non-centred forms for hierarchical models) as often as the code is optimised.

Hamiltonian Monte Carlo and NUTS

Plain Markov chain Monte Carlo (no page yet), such as random-walk Metropolis or Gibbs sampling, moves slowly in high dimensions because it proposes blindly. Hamiltonian Monte Carlo and NUTS (no page yet) uses the gradient: it treats -log_p as a potential energy, gives the position a random momentum, simulates the motion with the Verlet leapfrog integrator (half-step on momentum, full step on position, half-step on momentum), and accepts or rejects the end point with a Metropolis correction. Leapfrog is chosen because it is time-reversible and volume-preserving, which is what makes the accept probability correct, and its energy error stays bounded over long trajectories. Two parameters matter, the step size and the number of steps; Hoffman and Gelman's NUTS removes the second by building the trajectory recursively until it starts to double back on itself (source: the NUTS paper). Warm-up adapts the step size and a mass matrix, the metric that rescales the parameters, so that the geometry is closer to a unit sphere. Betancourt's conceptual introduction explains why this behaves well in high dimensions and what a divergent transition means: the integrator has run into a region of extreme curvature that it cannot follow, and the sampler is probably missing part of the posterior.

Variational inference and optimisation

When sampling is too slow, Variational inference (no page yet) replaces the posterior with a simple family, typically independent Gaussians in the unconstrained space ("mean-field"), and maximises the evidence lower bound. Kucukelbir et al.'s automatic differentiation variational inference (ADVI) is the template: transform to unconstrained space, use the reparameterisation trick so that a Monte Carlo estimate of the gradient is available, and run stochastic gradient ascent (source: the ADVI paper). The step rule is typically an Adam-family optimiser such as AdamW, and the expectation is a plain Monte Carlo average. The trade-off is honest: VI is fast and scales to large data, but mean-field Gaussians usually underestimate posterior variance and can miss correlations and multiple modes. For a point estimate, L-BFGS finds the posterior mode (MAP) from the same gradient, and it is also used to find a good starting point.

Sequential Monte Carlo

Sequential Monte Carlo (no page yet) runs a population of weighted particles through a sequence of distributions, with importance weights and resampling at each step. It fits where HMC does not: streaming data and state-space models, discrete latent variables, and models without gradients. It also gives an estimate of the marginal likelihood as a by-product. Resampling draws from a discrete distribution many times, where an alias table (alias-method-sampling, no page yet) gives constant-time draws, and importance-sampling (no page yet) supplies the reweighting rule. The cost is particle count: accuracy improves with the number of particles, and particle degeneracy (almost all weight on one particle) is the failure to watch.

Random numbers and diagnostics

Each chain needs its own reproducible random stream; PCG and xoshiro generators support seeding and independent parallel streams, so four chains can run on four cores and a result can be repeated exactly from a seed. Draws are correlated, so the number of samples overstates the information. MCMC convergence diagnostics (no page yet) turn several chains into a verdict: R-hat compares between-chain and within-chain variance and flags chains that disagree, and the effective sample size estimates how many independent draws the correlated sample is worth. Vehtari et al. show that the classic R-hat can miss problems with heavy tails or unequal variances, and propose a rank-normalised version with separate bulk and tail effective sample sizes (source: the rank-normalisation paper). Divergences from HMC are a third signal, and they are specific to this class of sampler.

Data and control flow

One Stan-style fit (HMC/NUTS):

1. The model file is parsed and translated to code with the log density and a gradient via autodiff; the data are read and validated.
2. The system draws initial values for the unconstrained parameters, with a different random stream per chain.
3. Warm-up: each iteration builds a NUTS tree with leapfrog steps (each step is one gradient evaluation, which dominates run time); step size and metric are adapted from the early draws.
4. Sampling: adaptation is frozen; each iteration produces a draw, stored in the constrained space (the transform is inverted).
5. After all chains finish, diagnostics are computed. If R-hat, effective sample size or divergence counts are bad, the user reparameterises the model or lengthens the run, and goes to step 1.
6. Draws are summarised: posterior means, intervals, predictions for new data.

One variational fit:

1. The model is transformed to the same unconstrained log density.
2. Initialise the Gaussian means and scales. Each iteration draws a few standard-normal samples, maps them to parameters, evaluates the log density and its gradient, and updates the variational parameters with an adaptive step rule.
3. Stop when the ELBO stops improving; draw samples from the fitted Gaussian for summaries.
4. If accuracy matters, check against a short HMC run on the same model.

Scaling

What limits you first, and the lever to pull:

1. Cost of one gradient. The model's likelihood over N data points. Levers: vectorised likelihoods, sufficient statistics, GPU/TPU backends (NumPyro and PyMC through JAX), subsampling with variational inference.
2. Number of leapfrog steps per draw. Posterior geometry decides this: correlated or funnel-shaped posteriors need many steps. Levers: reparameterisation, a dense mass matrix when parameters are correlated (at O(d²) memory), a better prior.
3. Matrix operations. Gaussian processes and dense covariances cost O(n³). Levers: sparse or low-rank structure, Kronecker products, inducing points.
4. Chains and cores. Chains are independent and run in parallel. Levers: more chains for better diagnostics, but each still needs its own warm-up.
5. Beyond HMC. At very large N or d, switch to variational inference or stochastic-gradient MCMC and accept an approximation; for streaming data use SMC.

Failure modes

Real implementations

The alternative: for a model with a closed form (conjugate priors) or a small generalised linear model, a direct formula or maximum likelihood is faster and enough. A probabilistic programming system pays off when the model has hierarchy, latent structure or custom likelihoods and the uncertainty itself is the product.

Sources

• M. Hoffman, A. Gelman, The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo, JMLR 15, 2014. paper
• A. Kucukelbir, D. Tran, R. Ranganath, A. Gelman, D. Blei, Automatic Differentiation Variational Inference, JMLR 18, 2017. paper
• B. Carpenter et al., Stan: A Probabilistic Programming Language, Journal of Statistical Software 76(1), 2017. doi
• M. Betancourt, A Conceptual Introduction to Hamiltonian Monte Carlo, 2017. arXiv
• R. Neal, MCMC using Hamiltonian dynamics, in the Handbook of Markov Chain Monte Carlo, 2011. arXiv
• A. Vehtari, A. Gelman, D. Simpson, B. Carpenter, P.-C. Bürkner, Rank-normalization, folding, and localization: An improved R-hat for assessing convergence of MCMC, 2019. arXiv
• A. Gelman, D. Rubin, Inference from Iterative Simulation Using Multiple Sequences, Statistical Science 7(4), 1992. doi
• Stan reference manual. docs · PyMC. site