[논문 리뷰] Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro

Du Phan, Neeraj Pradhan외 1명|arXiv (Cornell University)|2019. 12. 24.

Parallel Computing and Optimization Techniques참고 문헌 15인용 수 226

한 줄 요약

이 논문은 NumPyro를 소개하는데, NumPy 기반 확률 프로그래밍 라이브러리로 JAX 변환을 사용하는 조합 가능한 효과 핸들러를 통해 엔드-투-엔드 JIT 컴파일과 상당한 속도 향상을 달성하며, 기존 구현을 능가하는 반복적 NUTS 샘플러를 포함한다.

ABSTRACT

NumPyro is a lightweight library that provides an alternate NumPy backend to the Pyro probabilistic programming language with the same modeling interface, language primitives and effect handling abstractions. Effect handlers allow Pyro's modeling API to be extended to NumPyro despite its being built atop a fundamentally different JAX-based functional backend. In this work, we demonstrate the power of composing Pyro's effect handlers with the program transformations that enable hardware acceleration, automatic differentiation, and vectorization in JAX. In particular, NumPyro provides an iterative formulation of the No-U-Turn Sampler (NUTS) that can be end-to-end JIT compiled, yielding an implementation that is much faster than existing alternatives in both the small and large dataset regimes.

연구 동기 및 목표

• Pyro와 유사한 효과 핸들러를 JAX 기반 백엔드와 통합하여 NumPyro의 모델링 및 추론 인터페이스를 확장할 수 있음을 증명한다.
• 조합 가능한 변환(jit, grad, vmap)이 Pyro-호환 모델링 언어의 품질을 유지하면서 추론 하위 루틴을 가속하는지 보여준다.
• XLA를 활용한 CPU 및 GPU 속도 향상을 위한 엔드-투-엔드 JIT 컴파일된 NUTS 구현을 반복적 방식으로 제시한다.
• 벡터화된 하위 루틴으로 배치 추론 및 예측을 가능하게 하는 vmap의 이점을 설명한다.
• 작고 큰 데이터셋에서의 성능 개선을 정량화하기 위해 NumPyro를 Stan 및 Pyro와 벤치마킹한다.

제안 방법

• NumPyro 내에서 seed, trace, condition과 같은 Pyro 유사 효과 핸들링 추상화를 도입하여 JAX 백엔드에서 조합 가능한 프로그램 변환을 가능하게 한다.
• 추론 하위 루틴을 구성하고 최적화하기 위해 functional하고 추적 가능한 코드로 JAX 변환(jit, grad, vmap)을 활용한다.
• NUTS BuildTree 하위 루틴을 엔드-투-엔드 JIT 컴파일과 XLA를 통한 연산자 융합 개선을 위해 반복형 형태로 변환한다.
• prior 샘플링, posterior predictive 샘플링, 로그 가능도 계산과 같은 일반적인 추론 작업을 배치 처리하기 위해 vmap을 사용한다.
• 효과 핸들러를 JAX 변환과 통합하여 Pyro 호환 모델링 인터페이스를 유지하면서 가속을 얻는다.

실험 결과

연구 질문

• RQ1 Pyro 스타일의 효과 핸들러를 JAX 변환과 효과적으로 조합하여 NumPyro의 확률 프로그래밍 워크로드를 가속할 수 있는가?
• RQ2 NumPyro에서 NUTS와 같은 추론 하위 루틴의 엔드-투-엔드 JIT 컴파일이 Pyro 및 Stan에 비해 어떤 성능 이점을 제공하는가?
• RQ3 vmap을 통한 벡터화가 효과 핸들러와 어떻게 상호작용하여 확장 가능한 배치 추론 및 예측을 가능하게 하는가?
• RQ4 반복적이고 JIT-컴파일된 NUTS 구현이 정확성을 유지하면서 다양한 크기의 데이터셋에서 상당한 속도 향상을 제공하는가?
• RQ5 대규모 모델에서 CPU, GPU, TPU 백엔드 사용 시 NumPyro의 조합 가능한 변환을 활용한 실용적 이득은 무엇인가?

주요 결과

• NumPyro의 반복적 NUTS 구현은 테스트된 모델에서 Pyro 및 Stan보다 훨씬 빠르며, 특정 작업에서 Pyro 대비 약 340x, Stan 대비 약 6x의 속도향상을 보고한다.
• XLA를 통한 엔드-투-엔드 JIT 컴파일은 추론 하위 루틴의 완전 융합 및 최적화된 실행을 가능하게 하여 상당한 속도향상을 가져온다.
• vmap을 사용한 추론 벡터화는 priors 샘플링, posterior predictive 샘플링, 로그 가능도 계산의 배치를 효율적으로 가능하게 한다.
• NumPyro의 프레임워크는 Pyro 호환 모델링 언어를 유지하면서 CPU, GPU, TPU 백엔드 전반에 걸쳐 JAX 변환을 활용한 가속을 제공한다.
• 엔드-투-엔드 가속 이점은 소규모 CPU 환경과 대규모 GPU 환경 모두에서 관찰되며, JAX와 함께 컴포지션 가능한 효과의 보다 폭넓은 적용 가능성을 시사한다.