GenSBI: Generative Methods for Simulation-Based Inference in JAX

Abstract: Flow and diffusion generative models have established themselves as widely adopted density estimators for simulation-based inference (SBI), extending naturally from neural posterior estimation to likelihood and joint density estimation. Their principled optimization objectives and freedom from architectural constraints have driven rapid adoption across the natural sciences. Yet the most widely used SBI libraries remain PyTorch-based, leaving researchers who develop their forward models and analysis pipelines in JAX without a native option. We present GenSBI, an open-source library that implements flow matching, score matching, and denoising diffusion entirely in JAX. The library offers three transformer-based architectures - SimFormer, Flux1, and a novel Flux1Joint that extends gate-modulated transformer blocks to joint density estimation - all interchangeable through a unified interface that decouples generative method, neural backbone, and inference mode. GenSBI provides an end-to-end workflow from training through posterior calibration (SBC, TARP, LC2ST) and supports custom architectures with domain-specific embedding networks. We validate the framework on standard SBI benchmarks, achieving near-ideal mean C2ST scores (0.50-0.56, where 0.50 is ideal) on SBIBM tasks with minimal per-task tuning and well-calibrated posterior coverage across all tested configurations. The code is publicly available at https://github.com/aurelio-amerio/GenSBI.

• The paper introduces GenSBI, a JAX-native library that unifies neural posterior, likelihood, and joint density estimation for efficient simulation-based Bayesian inference.
• It employs generative techniques—score matching, EDM diffusion, and flow matching—to achieve high-fidelity posterior recovery and robust calibration.
• Empirical benchmarks on tasks like gravitational wave estimation and strong lensing demonstrate competitive C2ST scores and reliable uncertainty quantification.

GenSBI: Generative Methods for Simulation-Based Inference in JAX

---

Overview and Motivation

GenSBI introduces a modular, JAX-native library for simulation-based inference (SBI) leveraging flow matching, score matching, and denoising diffusion generative models. The core challenge addressed is Bayesian inference in scientific domains where the likelihood function is not analytically tractable, but synthetic observations can be generated via complex forward simulators. Traditional approaches (ABC, surrogate likelihoods) suffer from limitations, especially in high dimensions. Neural SBI, based on expressive density estimators trained on simulator outputs, amortizes the inference cost and shifts the paradigm toward posterior, likelihood, or joint density learning. Existing SBI tools have been largely PyTorch-based, leaving a gap for JAX-users, especially those requiring tight integration with JAX's composable transformations, probabilistic programming (NumPyro), and hardware acceleration.

---

SBI Formulation and Neural Estimation Strategies

GenSBI formalizes inference as estimating the posterior $p(\theta|x_{\text{obs}})$ given an implicit likelihood $p(x|\theta)$ accessible only through simulation. It builds on neural density estimation paradigms:

• Neural Posterior Estimation (NPE): Direct conditional density modeling, yielding amortized inference.
• Neural Likelihood Estimation (NLE): Surrogate modeling of $p(x|\theta)$, decoupled from prior, enabling MCMC-based posteriors.
• Neural Ratio Estimation (NRE): Likelihood-to-evidence classification.
• Neural Joint Estimation (NJE): Joint modeling of $p(\theta, x)$, supporting arbitrary conditional queries.

GenSBI implements NPE, NLE (via exact log-prob evaluation backed by ODE integrators, albeit at higher computational cost than flows), and NJE, unifying these under a single pipeline interface.

Figure 1: Comparison of four neural density estimation strategies for simulation-based inference; GenSBI implements NPE, NLE, and NJE with flexible composition of architectures and generative methods.

---

Generative Model Design

The library incorporates three generative families optimized for density estimation in SBI:

• Score Matching / Diffusion Models: Continuous-time SDEs (VP/VE variants), trained by denoising score matching objectives, enabling both stochastic sampling and exact log-probability via probability flow ODE.
• EDM (Elucidated Diffusion Model): Deterministic ODE-based generation with preconditioned denoiser architectures and noise schedules for improved stability and sample quality.
• Flow Matching: Conditional optimal-transport-based ODEs trained with a regression objective, yielding straighter trajectories, empirically mass-covering posteriors, and efficient sampling.

  Figure 2: VE SDE unconditional density estimation; shows sample evolution under score-matched diffusion.

  

Figure 3: Individual sample trajectories: Score matching yields stochastic paths, EDM straightens them, flow matching produces nearly rectilinear transport.

---

Software Architecture and Modular Pipelines

GenSBI is structured for maximal composability:

• Three generative methods: Each encapsulated as a strategy class with path construction, loss, batch prep, sampler, and initial sample creation.
• Three pipeline types: Conditional, Joint, Unconditional, corresponding to inference over posteriors, full joint, or arbitrary densities.
• Three transformer architectures: SimFormer (joint), Flux1 (conditional), Flux1Joint (joint, combines SimFormer masking with Flux1 block expressiveness).

The adapter pattern ensures neural backbone agnostic pipeline operation, supporting end-to-end workflows from training to posterior calibration.

Figure 4: GenSBI's three-axis architecture: generative methods, pipelines, and model wrappers are fully modular.

Figure 5: Workflow: generative-method handles batch prep and loss; pipeline manages optimization and inference integration.

Figure 6: Model wrappers: translation of pipeline calls into native architecture interfaces for conditional, joint, or unconditional modes.

---

Neural Architectures

GenSBI leverages transformer-based architectures for high-capacity, scalable density approximation:

• SimFormer: Single-joint-sequence architecture with per-token value, ID, and condition embeddings. Flexible masking for arbitrary conditional queries post-training.
• Flux1: Dual-stream architecture, cross-attention and independent residuals, AdaLN-Zero modulation.
• Flux1Joint: Joint sequence, condition-masking, AdaLN-Zero modulation, attention-MLP computation. Superior expressivity for complex joint densities.

  Figure 7: SimFormer: embeddings with condition masking enable flexible conditioning at inference.

  

  Figure 8: Flux1: observation and conditioning streams, double and single blocks, AdaLN-Zero for time modulation.

  

  Figure 9: Flux1Joint: joint token sequence, AdaLN-Zero, merges SimFormer masking with Flux1 block design.

---

Benchmarking and Posterior Quality

GenSBI achieves strong numerical results on SBIBM and scientific benchmarks:

• C2ST scores (classifier two-sample test): Near-ideal values ($0.50$–$0.56$) across SBIBM tasks; matches or outperforms baselines such as SimFormer and OneFlowSBI.
• SLCP results: Flux1Joint achieves $0.534$ with $10^5$ training samples (superior to SimFormer and NPE).
• Well-calibrated posteriors: TARP and SBC diagnostics consistently confirm frequentist coverage, with ECP curves on the diagonal.

  Figure 10: Marginal posteriors for Two Moons task (Flux1Joint, flow matching), resolving bimodal structure.

  

  Figure 11: C2ST accuracy as function of simulation budget for Flux1 architecture; lower is better.

  

  Figure 12: C2ST accuracy for Flux1Joint; shows improvement over Flux1.

  

  Figure 13: C2ST comparison: GenSBI vs baselines, demonstrating competitive or superior performance.

---

Posterior Calibration

GenSBI includes integrated diagnostics:

• TARP, SBC, LC2ST, marginal coverage: Automated for each trained pipeline.
• Posterior coverage: TARP curves consistently fall on the diagonal across tasks, indicating reliable credible intervals.

  Figure 14: TARP expected coverage probability for Two Moons (Flux1Joint, flow matching). Diagonal alignment reflects calibration.

---

Structured Scientific Applications

GenSBI demonstrates practical efficacy on high-dimensional physical inference tasks:

• GW parameter estimation: Recovery of binary component masses from simulated LIGO strain, with convolutional encoder conditioning.
• Strong lensing: Posterior recovery for ring position and radius from synthetic $64 \times 64$ images, spatial embedding, patchification, and transformer inference.

Figure 15: Example gravitational wave application: strain time series and posterior marginals.

Figure 16: TARP curve for GW task; showing well-calibrated coverage.

Figure 17: Strong lensing image and posterior marginals for ring parameters.

Figure 18: TARP coverage for lensing task; establishes calibration with image-conditioned inference.

---

Methodological and Technical Implications

GenSBI brings several implications:

• Posterior Mass Coverage: Conditional flow matching achieves empirically mass-covering inference; diffusion offers formal guarantees via variational bounds (likelihood weighting).
• Transformer compatibility: Freed from bijective constraints, modern architectures (e.g., attention-based transformers) are deployable for density approximation in SBI.
• JAX ecosystem integration: Functional transformations (jit, vmap), differentiable simulation, probabilistic programming, and hardware acceleration are native.
• Calibration imperative: Built-in diagnostics for calibration address the practical necessity for reliable scientific inference.

---

Future Directions

Key areas for further development include:

• Normalizing flow backbones for efficient density evaluation in MCMC workflows.
• Scaling benchmarks to higher-dimensional parameter spaces with robust sample quality.
• Enhanced support for structured outputs (images, fields).
• Systematic tuning of EDM for low-budget, low-dimensional settings and inclusion of alternative loss weightings/noise schedules.
• Expanded neural architectures for bespoke applications in computational sciences.

---

Conclusion

GenSBI provides a robust, composable framework for neural SBI in the JAX ecosystem, offering interchangeable flow matching, score matching, and diffusion generative models, state-of-the-art transformer architectures, and rigorous calibration tools. Empirical validation on established benchmarks and real-world structured data tasks confirms its capability for high-fidelity posterior recovery and well-calibrated uncertainty quantification. Its modular design, ecosystem compatibility, and performance profile position it as a practical tool for scientific inference workflows, particularly for researchers relying on native JAX pipelines and demanding transformer-based density estimators.

---