arXiv:2605.27499cs.LGastro-ph.CO2026-05

JAX用户可用的生成式模拟推断工具,支持流、扩散等模型。

GenSBI: Generative Methods for Simulation-Based Inference in JAX

论文配图:GenSBI: Generative Methods for Simulation-Based Inference in JAX
图 1 · 摘自论文原文
  • 基于JAX实现流匹配、得分匹配和去噪扩散模型。
  • 在SBIBM任务中平均C2ST得分0.50-0.56,接近理想值0.50。
  • 提供统一接口,适合用JAX构建模型的研究者快速上手。

流和扩散生成模型已成为模拟推断(SBI)中广泛使用的密度估计器,自然扩展至神经后验估计、似然和联合密度估计。其合理的优化目标与架构无关性推动了在自然科学领域的快速应用。然而,当前主流SBI库仍基于PyTorch,使使用JAX开发前向模型和分析流程的研究者缺乏原生支持。本文提出GenSBI,一个完全在JAX中实现的开源库,支持流匹配、得分匹配和去噪扩散模型。该库提供三种基于Transformer的架构——SimFormer、Flux1,以及一种新提出的Flux1Joint,可扩展门控Transformer块以支持联合密度估计。所有组件通过统一接口互换,解耦生成方法、神经主干和推断模式。GenSBI提供从训练到后验校准(SBC、TARP、LC2ST)的全流程支持,并兼容自定义架构与领域特定嵌入网络。在标准SBI基准测试中,其在SBIBM任务上达到0.50–0.56的近理想平均C2ST得分(理想值为0.50),且所有配置下后验覆盖率均良好校准。代码已公开于https://github.com/aurelio-amerio/GenSBI。

原文摘要 · Abstract (English)

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.

生成模型模拟推断JAX密度估计

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。