Skip to content

v0.6.1: sequential_monte_carlo + hamiltonian_monte_carlo + nested linear_combination

Choose a tag to compare

@xuda-ye-math xuda-ye-math released this 01 Jun 15:39

A backward-compatible follow-up to v0.6.0 (no breaking changes — hmc keeps working via an alias).

New

sequential_monte_carlo — flow-free annealed-Langevin SMC. Transports particles from exp(-beta_source * source) to exp(-beta_target * target) (same space, no flow) through a ladder of linearly-interpolated bridge potentials, alternating reweight + multinomial resample + Langevin on each bridge u_k:

samples, ess = sequential_monte_carlo(x0, source, target,
                                      beta_source=1.0, beta_target=1.0,
                                      ladder=20, step=5e-3, iters=30, chunk=1)

The bridge u_k = (1 - k/M)*beta_source*source + (k/M)*beta_target*target is built once with linear_combination and retuned per rung via set_coeffs. It returns (samples, ess), where ess is a list[float] of the per-rung effective sample size (pre-resample) — a ladder-spacing diagnostic. This differs from annealed_importance_sampling_F/_G, which use a trained flow as the proposal and rejuvenate only in the target. The per-rung reweighting uses the compiled .eval fast path when enable_eval() is set.

Changed

hmchamiltonian_monte_carlo (clearer canonical name); hmc is kept as an alias, so existing code is unaffected. The optimization = lbfgs / rejuvenation = langevin aliases now sit directly after their function definitions.

Nested linear_combination is supported. The previous "do NOT nest" caveat is removed: a linear_combination may be a child of another — the combined .grad / .eval closures compose recursively, enable_grad() / enable_eval() cascade to all leaves, and the fresh-allocation buffer-aliasing safety holds at every level (correct under both mode='default' and mode='reduce-overhead').

Tests & docs

  • New tests/_verify_nested.py exercises nested combinations end-to-end (depths 2–3, enable cascade, set_coeffs propagation, tensor coeffs, default==reduce-overhead, runtime gating, and MALA/HMC reproducing an analytic Gaussian) — focused on the compiled grad/eval fast paths.
  • README: a warning that the compiled fast paths are not a free lunch — at high d with large conditioners (e.g. d=256, hidden_features=(256,256)) compiling the NSF for_ladj/inv_ladj maps is extremely slow to compile, with no effective fix (drop enable_* and use the identical plain maps) — plus a "Why is torch.compile so slow?" FAQ entry.

Verification

tests/_verify_potential.py, _verify_flow.py, _verify_utils.py, and _verify_nested.py all pass on Linux + NVIDIA GPU.