GenSBI: JAX Library for Simulation-Based Inference

• GenSBI is an open-source, JAX-native library for simulation-based inference that implements continuous-time generative modeling techniques.
• It combines flow matching, score matching, and EDM diffusion with transformer-based architectures for joint density estimation and efficient posterior calibration.
• Its modular design supports various inference modes, including Neural Posterior, Likelihood, and Joint Estimation, facilitating rigorous calibration and benchmark evaluations.

GenSBI is an open-source library for simulation-based inference implemented entirely in JAX, built around flow matching, score matching, and denoising diffusion, with three transformer-based architectures—SimFormer, Flux1, and Flux1Joint—and an end-to-end workflow from training through posterior calibration with SBC, TARP, and LC2ST (Amerio, 26 May 2026). In the surrounding literature, the label also overlaps with broader formulations of “generalized SBI” and with Generator Based Inference, where simulation is generalized to any generator, including data-driven and hybrid generators (Cheng et al., 30 May 2025). The term should therefore be read in context: in its explicit 2026 usage it denotes a JAX-native software framework for neural SBI, while adjacent papers use closely related language for wider classes of generator-based inference.

1. Definition and terminological scope

In (Amerio, 26 May 2026), GenSBI addresses Bayesian inference in the standard likelihood-free setting: given parameters $\theta$, a stochastic simulator $x \sim \mathrm{Sim}(\theta)$, and an observation $x_{\text{obs}}$, the objective is the posterior

$$p(\theta \mid x_{\text{obs}}) = \frac{p(x_{\text{obs}} \mid \theta)p(\theta)}{p(x_{\text{obs}})}\,,$$

even though $p(x\mid\theta)$ cannot be evaluated (Amerio, 26 May 2026). Neural SBI approaches this problem by generating simulation pairs $\{(\theta^{(i)},x^{(i)})\}$ and training neural density estimators to approximate posteriors, likelihoods, ratios, or joints (Amerio, 26 May 2026).

The same term also sits near a broader conceptual move in high-energy physics. “Generator Based Inference (GBI)” explicitly treats Simulation-Based Inference as a special case of a more general framework in which the generator need not be a physics simulator: it can be a physics simulator, a data-driven generative model, or a hybrid (Cheng et al., 30 May 2025). That paper states that this direction is exactly what many people loosely call “generalized SBI” or “GenSBI” (Cheng et al., 30 May 2025).

A separate terminological ambiguity arises from genomics, where “SBI” denotes “Single Base Indicator,” a $4\times 4$ matrix of base-level matches and mismatches in NGS preprocessing (Chierici et al., 2011). This usage is unrelated to simulation-based inference and should not be conflated with GenSBI in the SBI literature.

2. Statistical setting and supported inference modes

GenSBI operates in the standard SBI regime in which one can sample

$$\theta \sim p(\theta), \qquad x \sim p(x\mid\theta),$$

but cannot evaluate $p(x\mid\theta)$ in closed form (Amerio, 26 May 2026). Within this regime, the library implements three inference modes.

First, it supports Neural Posterior Estimation (NPE), modeling $q_\phi(\theta\mid x)$ with conditional flow or diffusion models and amortizing inference across observations (Amerio, 26 May 2026). Second, it supports Neural Likelihood Estimation (NLE) in principle through log_prob from the continuous change-of-variables formula for a probability flow ODE, although the paper emphasizes that this is too slow for MCMC inner loops and is therefore supported only in a limited sense (Amerio, 26 May 2026). Third, it supports Neural Joint Estimation (NJE), in which a model learns the joint $x \sim \mathrm{Sim}(\theta)$0 and uses masking so that any subset of entries of $x \sim \mathrm{Sim}(\theta)$1 can be clamped and the rest sampled (Amerio, 26 May 2026). This allows posterior, likelihood, and evidence queries from a single model (Amerio, 26 May 2026).

GenSBI does not implement NRE; for that, the paper points to sbi and swyft (Amerio, 26 May 2026). This omission is substantive rather than cosmetic, because the library’s emphasis is on continuous-time generative modeling—flow matching, score-based diffusion, and EDM diffusion—rather than ratio-estimation pipelines.

The joint formulation is central to the software design. For joint estimation, GenSBI uses a binary mask $x \sim \mathrm{Sim}(\theta)$2, where $x \sim \mathrm{Sim}(\theta)$3 denotes observed or conditioned dimensions and $x \sim \mathrm{Sim}(\theta)$4 denotes dimensions to be generated (Amerio, 26 May 2026). Dynamics act only on unobserved dimensions, and the loss is computed only on unobserved entries (Amerio, 26 May 2026). This masked formulation generalizes SimFormer’s diffusion-based training to flow matching and underlies the library’s support for posterior, likelihood, and marginal queries (Amerio, 26 May 2026).

3. Generative methods and mathematical formulation

GenSBI’s methodological core is a set of interchangeable continuous-time generative models (Amerio, 26 May 2026).

Score-based models are implemented through forward SDEs of the form

$x \sim \mathrm{Sim}(\theta)$5

with VP and VE variants, together with denoising score matching for training and reverse-SDE or probability-flow ODE solvers for sampling (Amerio, 26 May 2026). The denoising score matching loss is

$x \sim \mathrm{Sim}(\theta)$6

and with likelihood weighting, $x \sim \mathrm{Sim}(\theta)$7, the loss is an upper bound on $x \sim \mathrm{Sim}(\theta)$8 (Amerio, 26 May 2026). The reverse SDE and the probability flow ODE supply stochastic and deterministic samplers, respectively (Amerio, 26 May 2026).

EDM diffusion is implemented through the generalized PF-ODE parameterization and a preconditioned denoiser

$x \sim \mathrm{Sim}(\theta)$9

with training loss

$x_{\text{obs}}$0

and sampling by a Heun ODE solver in $x_{\text{obs}}$1-space, optionally with churn controlled by $x_{\text{obs}}$2 (Amerio, 26 May 2026). The paper characterizes this as a denoising-diffusion implementation specialized to the EDM parameterization (Amerio, 26 May 2026).

Flow matching is implemented through deterministic ODE transport

$x_{\text{obs}}$3

trained with Conditional Flow Matching (CFM) rather than direct regression to the intractable marginal velocity field (Amerio, 26 May 2026). GenSBI uses affine paths with the CondOT schedule

$x_{\text{obs}}$4

so trajectories become nearly straight between $x_{\text{obs}}$5 and $x_{\text{obs}}$6, improving solver efficiency (Amerio, 26 May 2026). The resulting CFM loss is

$x_{\text{obs}}$7

The library provides deterministic ODE solvers and stochastic SDE variants, including FMODESolver, ZeroEndsSolver, and NonSingularSolver (Amerio, 26 May 2026).

Within the broader “GenSBI” landscape, these choices are distinct from adversarial methods such as GATSBI, which learns implicit conditional posteriors $x_{\text{obs}}$8 via a conditional GAN objective and targets the posterior through adversarial matching of $x_{\text{obs}}$9 pairs (Ramesh et al., 2022). That contrast separates explicit or semi-explicit continuous-time density modeling from implicit posterior generation.

4. Architectures and software organization

GenSBI couples these generative methods to three transformer-based architectures (Amerio, 26 May 2026).

SimFormer is a single-stream transformer on the joint vector $$p(\theta \mid x_{\text{obs}}) = \frac{p(x_{\text{obs}} \mid \theta)p(\theta)}{p(x_{\text{obs}})}\,,$$0. Each scalar component becomes a token with a value embedding, an ID embedding, and a condition embedding, and training uses random conditioning masks so that the same model can be used for posterior, likelihood, joint, and marginal sampling (Amerio, 26 May 2026). GenSBI adapts SimFormer so the architecture can be used with flow matching as well as diffusion (Amerio, 26 May 2026).

Flux1 is a conditional two-stream architecture derived from the Flux.1 image model. One stream processes noisy $$p(\theta \mid x_{\text{obs}}) = \frac{p(x_{\text{obs}} \mid \theta)p(\theta)}{p(x_{\text{obs}})}\,,$$1, the other processes the conditioning observation $$p(\theta \mid x_{\text{obs}}) = \frac{p(x_{\text{obs}} \mid \theta)p(\theta)}{p(x_{\text{obs}})}\,,$$2, with double-stream blocks followed by single-stream blocks (Amerio, 26 May 2026). Time conditioning uses adaLN-Zero, and the architecture supports learned absolute IDs, 1D sinusoidal or RoPE embeddings, and 2D sinusoidal or RoPE embeddings, with patchification for 2D inputs (Amerio, 26 May 2026). Flux1 is conditional only and does not implement joint masking (Amerio, 26 May 2026).

Flux1Joint is introduced as a new architecture combining SimFormer’s joint and masked design with Flux1’s single-stream blocks, adaLN-Zero modulation, and parallel attention-plus-MLP structure (Amerio, 26 May 2026). It therefore supports joint density estimation while retaining a more expressive block design aligned with modern diffusion and flow architectures (Amerio, 26 May 2026).

A major contribution of the library is its modular decomposition into three axes: generative method, inference pipeline, and neural backbone (Amerio, 26 May 2026). The abstract GenerativeMethod interface is instantiated by ScoreMatchingMethod, DiffusionEDMMethod, and FlowMatchingMethod; inference modes are handled by ConditionalPipeline, JointPipeline, and UnconditionalPipeline; and model wrappers reconcile different backbone call signatures through ConditionalWrapper, JointWrapper, and UnconditionalWrapper (Amerio, 26 May 2026). The training loop uses AdamW with cosine schedule and linear warmup, EMA of parameters, checkpointing via Orbax, and early stopping based on validation loss ratio (Amerio, 26 May 2026).

The software stack is entirely JAX-native. The paper specifies Flax nnx for neural networks, diffrax for ODE and SDE integration, NumPyro distributions for priors and log-probabilities, and Orbax for checkpointing, with jit, vmap, and accelerator support used throughout (Amerio, 26 May 2026). The code is publicly available at https://github.com/aurelio-amerio/GenSBI (Amerio, 26 May 2026).

5. Calibration, benchmarks, and application domains

GenSBI includes a calibration and evaluation suite consisting of SBC, TARP, LC2ST, and marginal coverage (Amerio, 26 May 2026). SBC computes rank statistics under the prior predictive and checks them against uniformity; TARP measures expected coverage probability as a function of nominal credible level and reports Jeffreys confidence intervals; LC2ST trains a local classifier to distinguish simulator joint samples from approximate posterior samples at a specific $$p(\theta \mid x_{\text{obs}}) = \frac{p(x_{\text{obs}} \mid \theta)p(\theta)}{p(x_{\text{obs}})}\,,$$3; and marginal coverage compares nominal and empirical coverage in each parameter dimension (Amerio, 26 May 2026). Together, these diagnostics assess both global and local posterior calibration (Amerio, 26 May 2026).

On SBIBM, the paper reports near-ideal C2ST performance, where $$p(\theta \mid x_{\text{obs}}) = \frac{p(x_{\text{obs}} \mid \theta)p(\theta)}{p(x_{\text{obs}})}\,,$$4 is ideal (Amerio, 26 May 2026).

Task	C2ST (mean ± std)
Gaussian Linear	$$p(\theta \mid x_{\text{obs}}) = \frac{p(x_{\text{obs}} \mid \theta)p(\theta)}{p(x_{\text{obs}})}\,,$$5
Gaussian Mixture	$$p(\theta \mid x_{\text{obs}}) = \frac{p(x_{\text{obs}} \mid \theta)p(\theta)}{p(x_{\text{obs}})}\,,$$6
Two Moons	$$p(\theta \mid x_{\text{obs}}) = \frac{p(x_{\text{obs}} \mid \theta)p(\theta)}{p(x_{\text{obs}})}\,,$$7
SLCP	$$p(\theta \mid x_{\text{obs}}) = \frac{p(x_{\text{obs}} \mid \theta)p(\theta)}{p(x_{\text{obs}})}\,,$$8
Bernoulli GLM	$$p(\theta \mid x_{\text{obs}}) = \frac{p(x_{\text{obs}} \mid \theta)p(\theta)}{p(x_{\text{obs}})}\,,$$9

These “ceiling performance” results are reported for Flux1Joint with flow matching and $p(x\mid\theta)$0 simulations (Amerio, 26 May 2026). At lower budgets, the paper states that for simpler tasks all methods converge to approximately $p(x\mid\theta)$1–$p(x\mid\theta)$2 by $p(x\mid\theta)$3–$p(x\mid\theta)$4 simulations, while for SLCP and Bernoulli GLM, flow matching and score matching converge faster than EDM and ultimately achieve better scores (Amerio, 26 May 2026). It also reports that Flux1Joint consistently matches or improves over Flux1, especially on SLCP, indicating that joint modelling helps (Amerio, 26 May 2026).

The library is also demonstrated on domain-specific tasks without reference posteriors. In a gravitational-wave example, Flux1 with flow matching and a CNN encoder is applied to $p(x\mid\theta)$5 detector time series and produces posteriors over $p(x\mid\theta)$6 that enclose the true parameters, with TARP expected-coverage curves close to diagonal (Amerio, 26 May 2026). In strong lensing, a 2D CNN encoder with 2D RoPE is used on $p(x\mid\theta)$7 images of simple rings, again yielding close-to-diagonal calibration diagnostics (Amerio, 26 May 2026).

The paper also reports practical runtimes. On Two Moons, training on a Tesla V100 for 50k steps with batch size $p(x\mid\theta)$8 runs at approximately $p(x\mid\theta)$9 iterations per second for Flux1 and approximately $\{(\theta^{(i)},x^{(i)})\}$0 iterations per second for Flux1Joint in a smaller configuration; sampling $\{(\theta^{(i)},x^{(i)})\}$1k posterior draws takes $\{(\theta^{(i)},x^{(i)})\}$2–$\{(\theta^{(i)},x^{(i)})\}$3 seconds for flow matching with $\{(\theta^{(i)},x^{(i)})\}$4 ODE steps, $\{(\theta^{(i)},x^{(i)})\}$5–$\{(\theta^{(i)},x^{(i)})\}$6 seconds for score matching with $\{(\theta^{(i)},x^{(i)})\}$7 SDE steps, and $\{(\theta^{(i)},x^{(i)})\}$8–$\{(\theta^{(i)},x^{(i)})\}$9 seconds for EDM with $4\times 4$0 steps (Amerio, 26 May 2026). The models fit on consumer GPUs such as an RTX 4070 with 12GB at batch size $4\times 4$1 (Amerio, 26 May 2026).

6. Broader landscape, limitations, and future directions

GenSBI sits in a broader literature on generative methods for SBI. GATSBI provides an adversarial alternative in which a conditional GAN learns an implicit posterior generator $4\times 4$2, supports implicit priors, and is fully amortized across observations (Ramesh et al., 2022). Its reported strengths are high-dimensional posterior spaces and implicit priors, while its limitations include GAN training instability, sensitivity to hyperparameters, and the fact that it yields only an implicit posterior sampler rather than an explicit density (Ramesh et al., 2022).

A different generalization appears in Generator Based Inference, where SBI is made explicit as a special case of a larger framework in which the “simulation” component is generalized to “generator” (Cheng et al., 30 May 2025). In that formulation, background can be provided by a data-driven generator trained in sidebands, signal can be simulation-defined or non-parametric, and anomaly detection outputs are converted into full inference with parameters and confidence intervals (Cheng et al., 30 May 2025). This suggests a broader sense of “GenSBI” as inference over composite generative models that mix simulation and data-driven components.

The GenSBI library in (Amerio, 26 May 2026) is nevertheless specific in its current scope. The paper lists several limitations: no NRE; no normalizing flows; log_prob for FM and SM is available but NLE remains impractical because each evaluation requires ODE integration; benchmarks cover up to $4\times 4$3, so very high-dimensional posteriors are not yet empirically studied; EDM underperforms at small budgets on some tasks, including poor C2ST on SLCP at $4\times 4$4 simulations; architectures treat $4\times 4$5 as a vector and do not explicitly handle structured outputs; and there is no SPMD model sharding yet, although JAX and Flax support it (Amerio, 26 May 2026).

The future directions named in the paper are correspondingly concrete: adding normalizing flows for efficient NLE and model comparison, broader high-dimensional benchmarks, more specialized architectures for structured outputs, better EDM tuning for low-dimensional SBI problems, and closer integration with BlackJAX and NumPyro pipelines (Amerio, 26 May 2026). Within the broader literature, a plausible implication is that “GenSBI” is evolving along two coupled axes: as a software stack for modern generative SBI in JAX, and as a wider research program in which simulators, learned generators, and hybrid generative mechanisms are all treated as valid substrates for inference (Cheng et al., 30 May 2025, Amerio, 26 May 2026).