Add a harness that reproduces the sbibm benchmark with neuralsbi - #35
Open
pedroliman wants to merge 2 commits into
Open
Add a harness that reproduces the sbibm benchmark with neuralsbi#35pedroliman wants to merge 2 commits into
pedroliman wants to merge 2 commits into
Conversation
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
force-pushed
the
claude/neural-sbi-benchmarks-tz3zx7
branch
from
August 1, 2026 12:48
5e78d09 to
dfa8f36
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The Level-3 claim so far has been checked against a Python
sbirun wedo 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:
(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.
neuralsbi::c2st()is not comparable to it. R/c2st.R reimplements thesklearn recipe: architecture, Glorot init, Adam with sklearn's bias
correction, and its convergence rule.
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.