NUTS: the No-U-Turn Sampler (self-tuning trajectory length)

[davidkhjo]

NUTS — the last item on the roadmap. HMC that tunes both of its knobs away.

What

• Multinomial NUTS (Hoffman & Gelman 2014): each draw doubles its trajectory until the path makes a whole-span U-turn, so there is no leapfrog_steps to set; and a dual-averaging warmup tunes the step size to target_accept (default 0.8), so there is no step-size sweep. Unbiased — targets exactly p ∝ exp(-E).
• Batched via lockstep-with-freeze: all B chains build their trees to a shared depth each doubling, but a chain that has stopped (U-turn or divergence) is frozen — every state write is gated by a per-chain active mask (torch.where). This yields, per chain, exactly the draw an independent single-chain NUTS would produce (no cross-chain coupling; a stopped chain's state is never rewritten). Wasted leapfrogs on finished chains are the only cost.
• Numerically careful: log-space multinomial weights (logaddexp) for progressive sampling; non-finite / ΔH > max_delta_h treated as divergence (weight 0, chain frozen) so a NaN can't leak across chains; trajectory math stays off-graph and the proposal is detached. Reuses HMC's leapfrog, _energy_grad, and AdaptiveMALA's dual-averaging recursion.
• Exposes last_tree_depth and a divergences count for diagnostics.

Validation (closed-form / distributional)

• Recovers a standard and a correlated Gaussian's covariance; tunes to 0.8 acceptance with zero divergences on the Gaussian.
• Enters Neal's funnel (finite, spreads — with the honest caveat that an identity metric under-samples the neck).
• Terminates by U-turn: on an easy Gaussian the tree depth stays small ({1,2,3}), never bottoming out at max_depth.

Notes

• Example examples/nuts_sampling.py (funnel + per-draw tree-depth histogram). torch-only; public API preserved; accumulates under the unreleased 0.16.0.
• This completes the roadmap end to end — no deferred features remain.

🤖 Generated with Claude Code