NUTS Algorithm
The No-U-Turn Sampler (NUTS) is the primary MCMC algorithm used in bgms for continuous parameter sampling. It extends Hamiltonian Monte Carlo by adaptively selecting the trajectory length, eliminating the need to tune the number of leapfrog steps.
The implementation is in src/mcmc/algorithms/nuts.h and nuts.cpp, based on the NUTS algorithm described in Hoffman & Gelman (2014) with the generalized U-turn criterion from Betancourt (2017).
Leapfrog integration
Each elementary step of HMC is a leapfrog integrator that approximately preserves the Hamiltonian \(H(q, p) = -\log \pi(q) + \frac{1}{2} p^\top M^{-1} p\):
- Half-step momentum: \(p \leftarrow p + \frac{\epsilon}{2} \nabla \log \pi(q)\)
- Full-step position: \(q \leftarrow q + \epsilon\, M^{-1} p\)
- Half-step momentum: \(p \leftarrow p + \frac{\epsilon}{2} \nabla \log \pi(q)\)
The implementation uses a diagonal mass matrix \(M\), stored as the vector of inverse diagonal elements (inv_mass_diag). The integrator lives in leapfrog.h/leapfrog.cpp; the kinetic-energy and initial-step-size helpers live in hamiltonian_utils.h/hamiltonian_utils.cpp.
Memoizer
The Memoizer class in leapfrog.h is a single-entry cache for \((\log \pi(q), \nabla\log \pi(q))\) evaluations.
Why a single-entry cache? NUTS tree building calls the gradient repeatedly as it extends the trajectory. Each leapfrog step produces a new unique \(q\), so a multi-entry hash map would almost never hit — and hashing an arma::vec element-by-element is expensive. A single-entry cache exploits the fact that at tree-node boundaries the same \(q\) appears as both the end of one subtree and the start of the next. The cache uses memcmp for equality, which is a single comparison per step.
Joint evaluation. Models often share computation between \(\log \pi\) and \(\nabla\log \pi\) (e.g., normalization constants, sufficient statistics). The Memoizer accepts a joint function (q) → (logp, grad) and caches both values from a single call. Two access methods retrieve the cached values — cached_log_post(q) and cached_grad(q) — and both trigger a joint evaluation if the cache misses.
Constraints via parameterization
Edge selection and sparse graphs need no special integrator. Models with continuous variables encode excluded edges through the constrained Cholesky parameterization: the sampler’s position vector contains only the free coordinates (dimension \(p + |E|\)), and every position corresponds to a positive-definite precision matrix with exact zeros at excluded edges. NUTS therefore always runs the standard unconstrained leapfrog integrator, for every model type.
Binary tree construction
NUTS builds a balanced binary tree by doubling the trajectory in a randomly chosen direction (forward or backward) at each depth level. The tree grows until one of these stopping conditions is met:
- U-turn — The trajectory starts bending back. Detected by checking whether the momentum sum across the subtree is no longer moving away from the endpoints (the generalized U-turn criterion from Betancourt, 2017).
- Divergence — The Hamiltonian error exceeds a threshold (default: \(\Delta H > 1000\)), indicating the integrator has entered a region of high curvature.
- Maximum depth — The tree reaches
max_tree_depth(default 10, corresponding to \(2^{10} = 1024\) leapfrog steps).
BuildTreeResult
Each subtree returns a BuildTreeResult struct containing:
| Field | Purpose |
|---|---|
theta_min, r_min |
Leftmost (backward) position and momentum |
theta_plus, r_plus |
Rightmost (forward) position and momentum |
theta_prime, r_prime |
Proposed sample from this subtree, with its momentum (for energy diagnostics) |
logp_prime |
Log-posterior at theta_prime (avoids a re-evaluation in nuts_step) |
rho |
Running momentum sum (for U-turn check) |
p_sharp_beg, p_sharp_end |
Mass-preconditioned momentum \(M^{-1}p\) at subtree endpoints (used in the generalized U-turn criterion) |
p_beg, p_end |
Raw momentum at subtree endpoints (used in cross-subtree U-turn checks at merge time) |
log_sum_weight |
\(\log \sum_i \exp(H_0 - h_i)\) across the subtree (multinomial candidate weight) |
s_prime |
Continuation flag (0 = stop) |
alpha, n_leapfrog |
Acceptance-probability accumulator and the number of leapfrog steps contributing to it |
divergent |
Whether a divergence was detected |
Proposal selection
bgms follows Stan’s base_nuts.hpp and uses multinomial candidate weighting throughout. Each leaf \(i\) in a subtree carries weight \(\log w_i = H_0 - h_i\), the Hamiltonian offset from the initial state; weights are accumulated in log_sum_weight with log-sum-exp. Within a subtree (build_tree) sibling subtrees are combined symmetrically: accept the new subtree with probability \(\exp(w_{\text{new}} - \mathrm{LSE}(w_{\text{old}},\, w_{\text{new}}))\). At the top level (nuts_step) the new subtree is accepted with probability \(\min(1,\, W_{\text{subtree}} / W_{\text{trajectory}})\), where the capital \(W\) are the summed weights on the linear scale, recovered from the log-scale accumulators by exponentiation. This rule is what Stan calls biased progressive sampling: biased because the new subtree competes with its full weight rather than its proportional share, which on average prefers later, farther-out subtrees over staying near the starting point. Divergence is detected directly by \(h - H_0 > 1000\); no slice variable is used.
Divergence detection
A divergence occurs when a single leapfrog step produces a Hamiltonian error \(\Delta H = H' - H_0 > 1000\). Divergences indicate that the step size is too large relative to the local curvature of the posterior. A divergent leaf contributes no weight to candidate selection (log_sum_weight \(= -\infty\)), but its alpha is still summed into the trajectory-level acceptance diagnostic. The number of divergent transitions is recorded in StepResult::NUTSDiagnostics and surfaced to the user via fit$nuts_diag.
A non-finite Hamiltonian (NaN or infinity from a broken state) is guarded at the leaf: it is treated as \(h = +\infty\), which lands in the same divergence branch — the leaf gets zero weight and the trajectory stops — rather than propagating NaN through the tree.
Mass matrix: full versus active layout
Under edge selection the active parameter set changes whenever an edge is added or removed, so a mass matrix stored per active coordinate would scramble its entries at every graph move. The adaptation controller therefore estimates the diagonal mass matrix on the full (zero-padded) theta layout, where every possible coordinate keeps a fixed slot. At each step, the sampler gathers the entries for the currently active coordinates (get_active_inv_mass()) and hands that vector to the integrator.
When the controller emits a new mass matrix (at a stage-2 window boundary; see Warmup Schedule), the sampler re-runs the initial step-size heuristic under the new metric and restarts dual-averaging from the result. At the stage-3c boundary — when edge selection activates — dual-averaging is restarted as well, so the step size can retune quickly to the changed geometry.
Warmup
NUTS adapts the step size \(\epsilon\) and the diagonal mass matrix \(M\) during a multi-stage warmup period before sampling begins. The three core stages (1, 2, 3a) follow Stan’s warmup scheme for NUTS. Stages 3b and 3c are bgms extensions that accommodate the spike-and-slab edge selection machinery.
The full warmup schedule — stage boundaries, budget allocation, adaptation algorithms, and warning system — is documented in Warmup Schedule.
NUTS diagnostics
Each NUTS step records four diagnostics:
- Tree depth — Number of doubling steps (0 to
max_tree_depth). Consistently hitting the maximum indicates the step size may be too small. - Divergent — Whether a divergence was detected in any leapfrog step during tree building.
- Energy — The Hamiltonian \(H = -\log \pi(q) + \frac{1}{2} p^\top M^{-1} p\). Used to compute E-BFMI (energy Bayesian fraction of missing information).
- Mean acceptance probability — The Metropolis acceptance probability averaged over every leapfrog step in the trajectory, paralleling Stan’s
accept_stat__. Surfaced asfit$nuts_diag$accept_prob(chains × iterations) with a per-chain mean in the diagnostic summary; also drives dual-averaging step-size adaptation.
These diagnostics are collected per chain and accessible via fit$nuts_diag in R (built by summarize_nuts_diagnostics() in R/diagnostics_nuts.R).