Investigate: ADR-0059 rejects simulator-path HMC on an unmeasured cost model ("~1x a solve per leapfrog") — measure it against the bngsim JAX bridge, and settle it with numbers either way

[wshlavacek]

This is an investigation, not a feature request. It asks for a measurement, and proposes recording the result in ADR-0059 either way so the question is settled with numbers instead of being re-argued.

What ADR-0059 decided

ADR-0059 ships job_type = hmc as a blackjax-NUTS reference sampler restricted to the analytical / BYO-expression model, and puts simulator-path HMC in "Out (rejected, not deferred)":

Simulator-path HMC (HMC on a BNGL/SBML posterior via bngsim CVODES forward sensitivities). Rejected: ~1×-a-solve gradient per leapfrog × many leapfrogs per draw is dominated by running the gradient-free samplers for the same budget; it would not help anyone. The bngsim JAX differentiable_solve bridge is real but is not adopted here.

That is a good decision written honestly, and it names its own load-bearing assumption: a cost model. "~1× a solve per leapfrog × many leapfrogs per draw, dominated by am/dream at equal wall-clock" is an arithmetic claim, and as far as I can tell it was reasoned rather than measured.

Why it is worth re-measuring now

bngsim/python/bngsim/_jax_bridge.py registers CVODE as a JAX custom_jvp primitive dispatching to CVODES forward sensitivities, and describes itself as:

CVODE-quality numerics (0.1ms per solve, 4 orders of magnitude faster than Diffrax) — Exact sensitivities (CVODES forward method, same integration direction, no adjoint brittleness) — JAX composability (automatic chain rule, jax.grad, jax.value_and_grad, jax.jacfwd)

0.1 ms per solve is a different regime from the one the rejection implicitly assumes. Separately, PyBNF's own gradient path (#386) and the EFIM assembly (#481) landed after the ADR, so a simulator-path log-density gradient now has two possible routes, not zero.

But the cost model may be wrong in the other direction too, and that is the real point of this issue.

The specific thing I want measured

ADR-0059 says the gradient costs ~1× a solve. For forward sensitivities that is only true for a handful of parameters. Forward sensitivity analysis integrates an additional n_p sensitivity systems alongside the state, so a full gradient costs roughly O(k) solves, not O(1). On this corpus that is not a rounding error:

slug	k	implied gradient cost
Perelson_Science1996	3	~3× a solve
Boehm_JProteomeRes2014	9	~9×
Borghans_BiophysChem1997	23	~23×
Weber_BMC2015	36	~36×
Giordano_Nature2020	50	~50×

So the honest expectation is a crossover: the ADR's rejection is probably too pessimistic at small k (where a 0.1 ms solve makes NUTS cheap and its per-draw efficiency is enormous) and too optimistic at large k (where the true cost is 50× a solve per leapfrog and the rejection understates it). Finding where that crossover sits is the deliverable.

A related technical question to resolve first

The bridge is registered with @jax.custom_jvp — a forward-mode rule — while the docstring advertises jax.grad and jax.value_and_grad (reverse mode). A custom_jvp primitive is not generally transposable, so either JAX is falling back to something equivalent to jacfwd (i.e. k forward passes per gradient, consistent with the O(k) estimate above), or there is a path I am not seeing. Establish which, before benchmarking anything — it determines the entire cost model and may make the rest of the investigation short.

Requested measurements

1. Gradient cost vs. k. Wall-clock for one log-density gradient through differentiable_solve, on Perelson (k=3), Boehm (k=9), Borghans (k=23), Weber (k=36), as a multiple of one plain solve. Confirms or refutes the O(k) model and locates the crossover.
2. ESS per wall-clock second — the only metric that decides this. NUTS-on-simulator vs. am vs. dream vs. p_dream, matched wall-clock, on at least a small and a mid-size slug. Use the existing rank-normalized split-R̂ / bulk-tail ESS diagnostics so the comparison is apples to apples.
3. Does it sample the geometries the gradient-free samplers fail on? ADR-0059's own framing is that HMC is the yardstick. Borghans is the corpus's last unsolved slug and every gradient-free optimizer converges to the same no-dynamics attractor (ADR-0108). Whether NUTS explores that posterior differently is a qualitative result that could matter more than throughput.
4. Stiffness / divergence behaviour. A stiff biological ODE plus a leapfrog integrator is exactly where divergences appear. The existing divergence-count signal should be reported — a fast sampler with 40% divergences has not solved anything.

Kill criterion, and what to do either way

If ESS/sec does not beat dream at any k tested: close this, and record the measured numbers in ADR-0059 as a supersession note. The point of the investigation is to convert a reasoned rejection into a measured one so it does not get re-litigated a third time.

If it wins below some k: that argues for a narrowly-gated hmc on simulators — available when the parameter count and the model's stiffness are inside the measured envelope, refused with a pointed message outside it, in the style of the existing gradient-path capability gates (GradientNotSupported). Not a new default, and not a general inference engine — ADR-0059's second rejection ("HMC as a general/production inference engine") is a separate question this issue does not reopen.

Notes for whoever picks this up

• hmc already bypasses the dask → execute() → score-column loop and builds its logdensity_fn in process, precisely because "the gradient cannot survive the per-pset dask round-trip." That architectural pattern is the one a simulator path would need, and it already exists — so this is not a from-scratch build.
• The priors already have logpdf_jax across all 16 families, and the unconstraining bijector path (ADR-0059 item 5) is model-agnostic. Both carry over unchanged.
• If the O(k) cost model is confirmed, that is independently an argument for adjoint sensitivities (CVODES ASA, already vendored in bngsim's SUNDIALS), whose cost is independent of k. bngsim explicitly chose forward and called adjoint brittle, so that is its own conversation — but this investigation is where the evidence for it would come from.