Skip to content

Add a harness that reproduces the sbibm benchmark with neuralsbi - #35

Open
pedroliman wants to merge 2 commits into
mainfrom
claude/neural-sbi-benchmarks-tz3zx7
Open

Add a harness that reproduces the sbibm benchmark with neuralsbi#35
pedroliman wants to merge 2 commits into
mainfrom
claude/neural-sbi-benchmarks-tz3zx7

Conversation

@pedroliman

Copy link
Copy Markdown
Owner

The Level-3 claim so far has been checked against a Python sbi run we
do ourselves, which answers "do the two implementations agree" but not
"are we where the field is". Lueckmann et al. (2021) published per-run
C2ST numbers for ten tasks, three simulation budgets and ten
observations, along with the reference posteriors they scored against.
Those are a better target: fixed, public, and not something we can
accidentally tune against.

dev/benchmarks/ scores neuralsbi against them directly. It reimplements
all ten sbibm tasks in R, vectorized over parameter rows so a 100k-draw
budget is tractable, and reproduces the paper's NPE and NLE settings from
the config files in sbi-benchmark/results. Observations, true parameters
and reference posteriors are read out of an sbibm checkout rather than
re-derived. 03_report.R prints our C2ST beside the paper's per cell with
a MATCH/BETTER/WORSE verdict.

Three pieces needed care:

  • Three tasks depend on constants sbibm froze as pickled torch tensors
    (the GLM stimulus and design matrix, the SLCP distractor mixture).
    They came out of numpy's RNG and cannot be regenerated in R, so
    R/pt_io.R reads them straight out of the legacy pickle format. No
    Python is needed anywhere in the harness.
  • The paper's C2ST is a scikit-learn MLP, not a logistic regression, so
    neuralsbi::c2st() is not comparable to it. R/c2st.R reimplements the
    sklearn recipe: architecture, Glorot init, Adam with sklearn's bias
    correction, and its convergence rule.
  • SIR and Lotka-Volterra need an ODE solver. A vectorized fixed-step RK4
    keeps integration error far below the observation noise and avoids
    making deSolve a benchmark dependency.

smoke_test.R is the guard, 58 checks in a few minutes. The two that
matter most: it recomputes sbibm's published GLM summary statistics from
sbibm's published raw spike trains (agrees to 4e-6, float32 rounding),
and for every task it checks each shipped observation against our
simulator at the parameters that generated it, which is what caught the
species-major flattening in Lotka-Volterra.

Only a spot check has been run so far (two_moons, NPE, 10^3, three
observations: 0.739 against the paper's 0.677). The grid itself is
compute, not code.

The Level-3 claim so far has been checked against a Python `sbi` run we
do ourselves, which answers "do the two implementations agree" but not
"are we where the field is". Lueckmann et al. (2021) published per-run
C2ST numbers for ten tasks, three simulation budgets and ten
observations, along with the reference posteriors they scored against.
Those are a better target: fixed, public, and not something we can
accidentally tune against.

dev/benchmarks/ scores neuralsbi against them directly. It reimplements
all ten sbibm tasks in R, vectorized over parameter rows so a 100k-draw
budget is tractable, and reproduces the paper's NPE and NLE settings from
the config files in sbi-benchmark/results. Observations, true parameters
and reference posteriors are read out of an sbibm checkout rather than
re-derived. 03_report.R prints our C2ST beside the paper's per cell with
a MATCH/BETTER/WORSE verdict.

Three pieces needed care:

* Three tasks depend on constants sbibm froze as pickled torch tensors
  (the GLM stimulus and design matrix, the SLCP distractor mixture).
  They came out of numpy's RNG and cannot be regenerated in R, so
  R/pt_io.R reads them straight out of the legacy pickle format. No
  Python is needed anywhere in the harness.
* The paper's C2ST is a scikit-learn MLP, not a logistic regression, so
  `neuralsbi::c2st()` is not comparable to it. R/c2st.R reimplements the
  sklearn recipe: architecture, Glorot init, Adam with sklearn's bias
  correction, and its convergence rule.
* SIR and Lotka-Volterra need an ODE solver. A vectorized fixed-step RK4
  keeps integration error far below the observation noise and avoids
  making deSolve a benchmark dependency.

smoke_test.R is the guard, 58 checks in a few minutes. The two that
matter most: it recomputes sbibm's published GLM summary statistics from
sbibm's published raw spike trains (agrees to 4e-6, float32 rounding),
and for every task it checks each shipped observation against our
simulator at the parameters that generated it, which is what caught the
species-major flattening in Lotka-Volterra.

Only a spot check has been run so far (two_moons, NPE, 10^3, three
observations: 0.739 against the paper's 0.677). The grid itself is
compute, not code.
sbibm integrates these two one parameter set at a time with Julia's
DifferentialEquations, an adaptive solver. The hand-rolled vectorized RK4
they had here was fast, but it was also a fixed-step approximation
standing where the reference implementation puts a real solver, and it
was ours to get wrong. deSolve::ode() (lsoda: adaptive step, automatic
stiff/non-stiff switching) is the closest thing R has to what sbibm
calls, and it removes the one place in the harness where a task differed
from sbibm by more than a library boundary.

Cost is 1.8 ms per SIR solve and 2.7 ms per Lotka-Volterra solve, so a
100k-simulation budget is a few minutes rather than seconds. That is
affordable next to training, and NEURALSBI_BENCH_CORES fans the solves
out across cores when it is not. Only the deterministic solves are
parallelised and the noise model stays in the parent process, so the
draws for a given seed are bit-identical whatever the core count, which
the scratch check confirmed for both tasks.

Failure handling now mirrors sbibm's: a solve that errors, stops short of
the last requested time, or returns anything non-finite becomes a NaN
row, and those rows get dropped before training. The smoke test grew a
check on that rate, since a solver quietly failing on prior tails would
shrink the effective simulation budget without saying so. It reports 0%
of 300 prior draws unsolvable on both tasks.

Both tasks still agree with sbibm's shipped observations (mean rank 0.49
and 0.50 over the ten observations), and sir NPE at 10^3 simulations now
runs end to end at C2ST 0.993 against the paper's 0.972 for the same
three observations.
@pedroliman
pedroliman force-pushed the claude/neural-sbi-benchmarks-tz3zx7 branch from 5e78d09 to dfa8f36 Compare August 1, 2026 12:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant