[논문 리뷰] BlackJAX: Composable Bayesian inference in JAX

Alberto Cabezas, Adrien Corenflos외 20명|arXiv (Cornell University)|2024. 02. 16.

Bayesian Methods and Mixture Models인용 수 6

한 줄 요약

BlackJAX는 MCMC, SMC, SG-MCMC 및 VI에 대한 JAX 기반의 기능적이고 구성 가능한 베이지안 추론 빌딩 블록 세트를 제공하여, 사용자가 저수준 구성 요소로부터 새로운 알고리즘을 구성할 수 있도록 합니다. 이는 CPU, GPU 및 TPU를 대상으로 하며 비정규화된 로그-밀도 타깃과의 통합을 지원합니다.

ABSTRACT

BlackJAX is a library implementing sampling and variational inference algorithms commonly used in Bayesian computation. It is designed for ease of use, speed, and modularity by taking a functional approach to the algorithms' implementation. BlackJAX is written in Python, using JAX to compile and run NumpPy-like samplers and variational methods on CPUs, GPUs, and TPUs. The library integrates well with probabilistic programming languages by working directly with the (un-normalized) target log density function. BlackJAX is intended as a collection of low-level, composable implementations of basic statistical 'atoms' that can be combined to perform well-defined Bayesian inference, but also provides high-level routines for ease of use. It is designed for users who need cutting-edge methods, researchers who want to create complex sampling methods, and people who want to learn how these work.

연구 동기 및 목표

• 블랙박스 샘플러와 함께 구조 인식형 베이지안 추론의 필요성을 제시합니다.
• 기본 구성요소로부터 추론 알고리즘을 구축하기 위한 구성 가능한 기능적 API를 소개합니다.
• 표준 파이프라인을 넘어 추론 방법을 커스터마이즈하고 확장하기 위한 저수준 API를 제공합니다.
• 비정규화된 로그-밀도 타깃과 최신 하드웨어와의 통합을 입증합니다.

제안 방법

• 각 반복이 사이드 이펙트 없이 새 상태를 업데이트하고 반환하는 기능적이고 상태 주도적인 API를 제안합니다.
• MH 수락/거부 단계, HMC/Langevin 다이나믹스, SMC, VI 도구, 템퍼링, 재샘플링 등 기본 추론 구성 요소 라이브러리를 제공합니다.
• 이 구성 요소의 조합을 통해 새로운 또는 기존 알고리즘을 형성할 수 있도록 하며, HMC의 윈도우/적응 스킴으로 예시를 제시합니다.
• 맞춤 커널 구축과 다수 입자 또는 커널의 병렬 실험을 위한 저수준 API를 제공합니다.
• 모든 구성 요소는 대상 로그-밀도 함수와 함께 작동하며 JAX를 통해 CPU/GPU/TPU 실행을 지원합니다.

실험 결과

연구 질문

• RQ1 베이지안 추론 방법을 재사용 가능하고 조합 가능한 구성 요소로 어떻게 분해할 수 있을까요?
• RQ2 기능적이며 사이드 이펙트가 없는 설계가 MCMC/SMC/VI 알고리즘의 병렬화와 구성 유연성을 개선할 수 있을까요?
• RQ3 BlackJAX가 비정규화된 로그-밀도 타깃과 확률 프로그래밍 워크플로우에 얼마나 효과적으로 통합될 수 있을까요?
• RQ4 고수준 API와 저수준 API를 모두 노출하는 것이 알고리즘 사용자 정의 및 실험에 어떤 영향을 미칠까요?

주요 결과

• 상태가 있는 구성 가능한 디자인으로 샘플링과 근사 추론을 위한 통합적이고 순수하게 기능적인 인터페이스를 제공합니다.
• 핵심 구성 요소의 여러 변형을 포함합니다(예: MH 수락/거부, 다양한 적분기를 갖춘 HMC, NUTS, 일반화된 HMC, 비가역 슬라이스 샘플링).
• 비정규화된 로그-밀도 타깃에서 작동함으로써 확률 프로그래밍 언어와의 상호 운용성을 입증합니다.
• 커널을 커스터마이즈하고 SMC에서 기본 커널을 혼합하거나 병렬 설정에서 사용하는 저수준 API를 제공합니다.
• JAX 기반 생태계와의 광범위한 호환성 및 관대한 Apache License v2.0를 주목하며, 오픈 소스 거버넌스와 포괄적인 테스트 스위트(99% 커버리지)를 제공합니다.
• 논문, 튜토리얼, 오픈 소스 강좌에의 채택을 통해 교육 및 연구에 미치는 영향을 기록합니다.