[论文解读] JAXNS: a high-performance nested sampling package based on JAX

Joshua G. Albert|arXiv (Cornell University)|Dec 30, 2020

Gaussian Processes and Bayesian Inference参考文献 11被引用 4

一句话总结

JAXNS 是一个基于 JAX 和 XLA 原语构建的高性能嵌套采样包，其证据计算速度比 PolyChord、MultiNEST 和 dynesty 快 2–3 个数量级，同时保持了精度。它利用静态内存计算和即时（JIT）编译技术，加速了各类贝叶斯推断问题中的似然评估，使其实现了在机器学习和数据科学中的实际应用。

ABSTRACT

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}.

研究动机与目标

• 解决时间受限和高吞吐量数据科学应用中嵌套采样存在的计算瓶颈问题。
• 为复杂、多模态且非基于梯度的后验分布提供高效、开箱即用的贝叶斯证据计算能力。
• 将嵌套采样集成到现代机器学习框架 JAX 中，而此前该框架尚无原生实现。
• 克服现有嵌套采样工具的性能限制，使其适用于实时或大规模推断任务。
• 在未来工作中将嵌套采样扩展至离散先验，以及可扩展、内存高效的采样策略。

提出的方法

• 通过 JAX 的 XLA 编译原语，将嵌套采样重新表述为静态内存算法，实现在加速器上的高性能执行。
• 利用 JAX 的 JIT 编译和自动微分实现似然评估，支持快速、向量化计算。
• 利用 XLA 的低级优化技术，最小化内存开销并最大化计算吞吐量。
• 设计一种静态内存版本的水库采样，以在不存储所有中间点的情况下保持代表性样本。
• 确保与连续和离散随机变量的兼容性，与 Skilling (2004) 中原始定义一致。
• 构建一个模块化、开源的包（JAXNS），提供 Pythonic 接口，实现与概率编程工作流的无缝集成。

实验结果

研究问题

• RQ1 嵌套采样能否被充分加速，从而在数据科学中实现对实时和大规模贝叶斯推断的实际应用？
• RQ2 与传统实现相比，JAX 的 XLA 编译和静态内存设计在多大程度上提升了嵌套采样的性能？
• RQ3 嵌套采样能否被高效地集成到现代机器学习后端（如 JAX）中，从而实现在神经网络中的基于证据的模型选择？
• RQ4 JAXNS 在实现相比现有工具显著提速的同时，如何保持贝叶斯证据估计的准确性？
• RQ5 静态内存嵌套采样的局限性是什么？在线重采样策略在多大程度上可以缓解这些局限？

主要发现

• 在顺序评估模式下，JAXNS 的证据计算速度比 PolyChord、MultiNEST 和 dynesty 快 2–3 个数量级。
• 性能提升归因于 XLA 编译的静态内存计算，该计算优化了加速器上的内存访问和计算效率。
• JAXNS 在证据计算方面保持了与现有工具相同的精度，确保了贝叶斯模型比较的可靠性。
• 该软件包已公开发布于 https://github.com/joshuaalbert/jaxns，便于研究社区广泛采用。
• JAXNS 实现了此前无法实现的应用，例如强化学习中的基于证据的训练，以及具有多模态后验分布的贝叶斯神经网络。
• 未来扩展将包括对离散先验的支持，以及一种静态内存水库采样变体，以在不牺牲统计保真度的前提下提升内存效率。