[논문 리뷰] JAXNS: a high-performance nested sampling package based on JAX

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

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)를 구축하여 확률 프로그래밍 워크플로우에 원활하게 통합하기 위해.

실험 결과

연구 질문

• RQ1 내재 샘플링이 데이터 과학 분야에서 실시간 및 대규모 베이지안 추론에 실용적이기 위해 충분히 가속화될 수 있는가?
• RQ2 JAX의 XLA 컴파일링과 정적 메모리 설계가 기존 구현 대비 내재 샘플링의 성능을 얼마나 향상시킬 수 있는가?
• RQ3 내재 샘플링이 JAX와 같은 현대 머신러닝 백엔드에 효율적으로 통합될 수 있는가? 이를 통해 신경망에서 증거 기반 모델 선택이 가능해지는가?
• RQ4 JAXNS는 기존 도구들보다 뚜렷한 성능 향상을 이룰 수 있었지만, 베이지안 증거 추정의 정확도를 어떻게 유지하는가?
• RQ5 정적 메모리 내재 샘플링의 한계는 무엇이며, 온라인 리샘플링 전략은 이를 어떻게 완화할 수 있는가?

주요 결과

• JAXNS는 순차적 평가 모드에서 PolyChord, MultiNEST 및 dynesty보다 2~3개 정도의 주기 빠른 증거 계산 성능을 달성한다.
• 성능 향상은 XLA 컴파일된 정적 메모리 계산 덕분이며, 이는 가속기에서 메모리 접근과 계산을 최적화하기 때문이다.
• JAXNS는 기존 도구들과 동일한 정확도를 유지하여 신뢰할 수 있는 베이지안 모델 비교를 보장한다.
• 패키지는 https://github.com/joshuaalbert/jaxns에서 공개되어 있어 연구 커뮤니티에서 널리 채택될 수 있도록 한다.
• JAXNS는 이전에는 불가능했던 응용, 예를 들어 강화학습에서의 증거 기반 학습 및 다중 모달 사후 분포를 가진 베이지안 신경망의 학습을 가능하게 한다.
• 향후 확장에서는 이산 사전 분포 지원과 정적 메모리 리저보아 샘플링 변형을 통해 메모리 효율성을 향상시키되 통계적 정밀도를 손상시키지 않도록 할 계획이다.