flowMC:用于Jax中概率推断的归一化流增强采样包
flowMC: Normalizing-flow enhanced sampling package for probabilistic inference in Jax
摘要:
flowMC是一个Python库,用于加速马尔可夫链蒙特卡罗(MCMC),利用深度生成建模。它建立在机器学习库JAX和Flax之上。在其核心,flowMC使用一个本地采样器和一个可学习的全局采样器串联,以有效地采样后验分布。当局部采样器的多个链在目标参数空间中的感兴趣区域上生成样本时,该包使用这些样本来训练归一化流模型,然后使用它来提出跨参数空间的全局跳跃。flowMC采样器可以处理非平凡的几何形状,例如多峰分布和具有局部相关性的分布。flowMC的主要特性总结如下:* 由于flowMC构建在JAX之上,因此它通过自动微分(如MALA和Hamilton Monte Carlo(HMC))支持基于梯度的采样器。* flowMC使用最先进的标准化流模型,如二次样条来支持其全局采样器。这些模型在相对较短的训练时间内捕获重要特征方面非常有效。* 本机支持使用GPU和TPU等加速器。该代码还支持使用具有SIMD并行性的多个加速器。* 默认情况下,使用即时(JIT)编译来进一步加快采样过程。* 我们为希望通过默认参数使用flowMC的用户提供了一个简单的黑盒界面,同时提供了一个广泛的指南,解释了在调整采样器参数时的权衡。所有上述功能的紧密集成使flowMC成为一个高性能但易于使用的统计推断包。
英文摘要:
flowMC is a Python library for accelerated Markov Chain Monte Carlo (MCMC) leveraging deep generative modeling. It is built on top of the machine learning libraries JAX and Flax. At its core, flowMC uses a local sampler and a learnable global sampler in tandem to efficiently sample posterior distributions. While multiple chains of the local sampler generate samples over the region of interest in the target parameter space, the package uses these samples to train a normalizing flow model, then uses it to propose global jumps across the parameter space. The flowMC sampler can handle non-trivial geometry, such as multimodal distributions and distributions with local correlations. The key features of flowMC are summarized in the following list: * Since flowMC is built on top of JAX, it supports gradient-based samplers through automatic differentiation such as MALA and Hamiltonian Monte Carlo (HMC). * flowMC uses state-of-the-art normalizing flow models such as Rational-Quadratic Splines to power its global sampler. These models are very efficient in capturing important features within a relatively short training time. * Use of accelerators such as GPUs and TPUs are natively supported. The code also supports the use of multiple accelerators with SIMD parallelism. * By default, Just-in-time (JIT) compilations are used to further speed up the sampling process. * We provide a simple black box interface for the users who want to use flowMC by its default parameters, yet provide at the same time an extensive guide explaining trade-offs while tuning the sampler parameters. The tight integration of all the above features makes flowMC a highly performant yet simple- to-use package for statistical inference.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。