GenSBI is a powerful JAX-based library for Simulation-Based Inference (SBI) using state-of-the-art generative models, currently revolving around Optimal Transport Flow Matching and Diffusion Models.

It is designed for researchers and practitioners who need a flexible, high-performance toolkit to solve complex inference problems where the likelihood function is intractable.

• Modern SBI Algorithms: Implements cutting-edge techniques like Optimal Transport Conditional Flow Matching and Diffusion Models for robust and flexible posterior inference.
• Built on JAX and Flax NNX: Leverages the power of JAX for automatic differentiation, vectorization, and seamless execution on CPUs, GPUs, and TPUs.
• High-Level Recipes API: A simplified interface for common workflows, allowing you to train models and run inference with just a few lines of code.
• Powerful Transformer Models: Includes implementations of recent, high-performing models like Flux1, Flux1Join, and Simformer for handling complex, high-dimensional data.
• Modular and Extensible: A clean, well-structured codebase that is easy to understand, modify, and extend for your own research.

Examples for this library are available separately in the GenSBI-examples repository.

Some key examples include:

Unconditional Density Estimation:

• flow_matching_2d_unconditional.ipynb
  Demonstrates how to use flow matching in 2D for unconditional density estimation.
• diffusion_2d_unconditional.ipynb
  Demonstrates how to use diffusion models in 2D for unconditional density estimation.

Conditional Density Estimation:

• two_moons_flow_simformer.ipynb
  Uses the Simformer model for posterior density estimation on the two-moons benchmark.
• two_moons_flow_flux.ipynb
  Uses the Flux1 model for posterior density estimation on the two-moons benchmark.
• gaussian_linear_flow_flux1joint.ipynb
  Uses the Flux1Joint model for posterior density estimation on the Gaussian Linear benchmark.
• slcp_flow_simformer.ipynb
  Uses the Simformer model for posterior density estimation on the SLCP benchmark.

The following tasks are planned for future development:

• Implement OT flow matching techniques.
• Implement diffusion models (EDM and score matching).
• Implement Transformer-based models for conditional posterior estimation (Flux1 and Simformer).
• Unify the API for flow matching and diffusion models.
• Implement wrappers to make training of flow matching and diffusion models similar.
• Write tests for core functionalities.
• Consider implementing classifier free guidance for conditional models.
• Add more examples and benchmarks.
• Improve documentation and tutorials.
• Provide SOTA pre-trained models and checkpoints for some SBI benchmark cases

• bfloat16 support is currently limited and may lead to unexpected behavior.

If you use this library, please consider citing this work and the original methodology papers.

• Facebook Flow Matching library: https://github.com/facebookresearch/flow_matching
• Elucidating the Design Space of Diffusion-Based Generative Models: https://github.com/NVlabs/edm
• Simformer model: https://github.com/mackelab/simformer
• Flux1 model from BlackForest Lab: https://github.com/black-forest-labs/flux
• Simulation-Based Inference Benchmark: https://github.com/sbi-benchmark/sbibm