论文标题

JAXNS:基于JAX的高性能嵌套采样包

JAXNS: a high-performance nested sampling package based on JAX

论文作者

Albert, Joshua G.

论文摘要

自从John Skilling于2004年首次亮相以来,嵌套采样已证明是科学家的宝贵工具,为复杂的后验分布提供了假设证据计算和参数推断,尤其是在天文学领域。由于其计算复杂性和长期运行的性质,在过去,嵌套的采样保留用于离线型贝叶斯推断,留下了诸如变异推理和MCMC之类的工具,用于在线型,时间约束,贝叶斯计算。这些工具不容易处理复杂的多模式后代,离散的随机变量以及缺乏梯度的后代,也无法实现贝叶斯证据的实际计算。因此,对于可以在计算时间内缩小差距的高性能嵌套采样包的开口仍然是一个开口,并让嵌套采样成为数据科学工具箱中的常见位置。我们提出了基于JAX的嵌套采样(JAXNS),这是一种使用JAX编写的高性能嵌套采样包,并表明它比当前可用的Polychord,Multinest和Dynesty的嵌套采样实现的速度快几个数量级,同时维持了相同的证据计算准确性。 JAXNS软件包可公开可在\ url {https://github.com/joshuaalbert/jaxns}上获得。

Since its debut by John Skilling in 2004, nested sampling has proven a valuable tool to the scientist, providing hypothesis evidence calculations and parameter inference for complicated posterior distributions, particularly in the field of astronomy. Due to its computational complexity and long-running nature, in the past, nested sampling has been reserved for offline-type Bayesian inference, leaving tools such as variational inference and MCMC for online-type, time-constrained, Bayesian computations. These tools do not easily handle complicated multi-modal posteriors, discrete random variables, and posteriors lacking gradients, nor do they enable practical calculations of the Bayesian evidence. An opening thus remains for a high-performance out-of-the-box nested sampling package that can close the gap in computational time, and let nested sampling become common place in the data science toolbox. We present JAX-based nested sampling (JAXNS), a high-performance nested sampling package written in XLA-primitives using JAX, and show that it is several orders of magnitude faster than the currently available nested sampling implementations of PolyChord, MultiNEST, and dynesty, while maintaining the same accuracy of evidence calculation. The JAXNS package is publically available at \url{https://github.com/joshuaalbert/jaxns}.

扫码加入交流群

加入微信交流群

微信交流群二维码

扫码加入学术交流群,获取更多资源