[논문 리뷰] tfp.mcmc: Modern Markov Chain Monte Carlo Tools Built for Modern Hardware

Junpeng Lao, Christopher Suter외 8명|arXiv (Cornell University)|2020. 02. 04.

Markov Chains and Monte Carlo Methods참고 문헌 15인용 수 17

한 줄 요약

이 논문은 TensorFlow Probability용으로 구축된 현대적인 마르코프 체인 몬테카를로(MCMC) 툴킷인 tfp.mcmc를 소개한다. 이 툴킷은 광범위한 데이터 병렬 처리, 벡터화된 계산, 그리고 조합 가능한 모듈식 커널을 통해 하드웨어 최적화된 고성능 MCMC 추론을 가능하게 한다. SIMD 및 다중 코어 실행을 활용하여 CPU, GPU, TPU 전반에서 막대한 병렬성을 실현하면서도, 간단한 대상 로그확률(TLP) 인터페이스를 통해 어떤 확률 모델과도 호환성을 유지한다.

ABSTRACT

Markov chain Monte Carlo (MCMC) is widely regarded as one of the most important algorithms of the 20th century. Its guarantees of asymptotic convergence, stability, and estimator-variance bounds using only unnormalized probability functions make it indispensable to probabilistic programming. In this paper, we introduce the TensorFlow Probability MCMC toolkit, and discuss some of the considerations that motivated its design.

연구 동기 및 목표

• 현대 하드웨어 기능인 SIMD, 다중 코어 CPU, GPU, TPU를 최대로 활용하여 확장 가능한 확률 추론을 위한 MCMC 프레임워크를 설계하는 것.
• 다중 체인에 걸쳐 벡터화된 연산을 사용하는 병렬 처리를 통해 효율적이고 막대한 병렬성을 가진 MCMC를 가능하게 하는 것.
• 도메인 특화 언어 없이도 특정 모델링 프레임워크에 밀접하게 결합되지 않은 모듈식이고 조합 가능한 API를 제공하여 복잡한 MCMC 전이 커널을 구축하는 것.
• 비정규화된 로그확률 함수와 자동 미분을 모두 지원하여 다양한 확률 모델과의 탄력적인 통합을 가능하게 하는 것.
• 조합 가능한 커널 구성 요소를 통해 적응형 단계 크기, 전처리, 초모수 조정과 같은 고급 MCMC 기법을 일반 목적 프레임워크 내에서 모듈식이고 조합 가능한 방식으로 구현하는 것.

제안 방법

• 프레임워크는 TensorFlow의 내장 배치 처리 및 자동 미분 기능을 활용하여 다중 MCMC 체인에서 동시에 병렬로 벡터화된 계산을 가능하게 한다.
• 대상 로그확률(TLP)을 파이썬 콜러블로 지정하는 기능적 인터페이스를 도입하여 MCMC 엔진과 모델 사양을 분리한다.
• 핵심 추상화는 전이 커널(TransitionKernel)으로, MCMC 전이 논리를 봉인하고 메트로폴리스-하스팅스와 같은 복잡한 알고리즘에서 비보정 제안을 위한 중첩을 지원한다.
• sample_chain 드라이버는 TensorFlow의 tf.while_loop와 XLA 컴파일을 활용하여 버닝 인, 샘플링, 추적을 조율하며, 하드웨어 가속과 저수준 최적화를 가능하게 한다.
• MetropolisHastings, SimpleStepSizeAdaptation, TransformedTransitionKernel 등의 조합 가능한 커널은 고급 MCMC 워크플로우의 모듈식 구축을 가능하게 한다.
• Bijectors를 통한 재매개변수화를 통해 상태 공간을 변환함으로써, 히브리지안 몬테카를로와 같은 알고리즘에서 샘플링 효율을 향상시킨다.

실험 결과

연구 질문

• RQ1 현대 하드웨어에서 막대한 병렬성을 활용하여 MCMC 추론을 어떻게 효율적으로 스케일링할 수 있는가?
• RQ2 다양한 하드웨어와 모델에서 작동하는 고성능이고 조합 가능한 MCMC 커널을 가능하게 하는 아키텍처 패턴은 무엇인가?
• RQ3 벡터화된 계산과 배치 처리를 MCMC에 원활하게 통합하면서도 모듈성이나 정확성을 훼손하지는 않는 방법은 무엇인가?
• RQ4 SIMD 및 다중 코어 실행을 사용해 수천 개의 병렬 MCMC 체인을 실행할 경우 성능 및 수렴에 어떤 영향을 미치는가?
• RQ5 일반 목적 프레임워크 내에서 적응형 MCMC 기법을 모듈식이고 조합 가능한 방식으로 어떻게 구현할 수 있는가?

주요 결과

• tfp.mcmc 라이브러리는 AVX512 벡터화를 활용해 단일 32코어 CPU에서 최대 1024개의 병렬 MCMC 체인을 실행할 수 있으며, 기존의 작업 병렬 처리 방식에 비해 수개의 주기적 속도 향상을 보였다.
• TensorFlow의 배치 문법을 통해 체인 간 벡터화된 계산이 네이티브로 지원되어 단일 tf.Tensor 입력으로 100개의 체인을 동시에 실행할 수 있다.
• 조합 가능한 커널 설계 덕분에 HMC에 적응형 단계 크기 또는 Bijectors를 통한 전처리 샘플링과 같은 복잡한 MCMC 알고리즘을 모듈식으로 구성할 수 있다.
• 프레임워크는 비정규화된 로그확률 함수와 자동 미분을 모두 지원하여, TensorFlow로 작성된 어떤 모델과도 원활한 통합이 가능하다.
• XLA 컴파일 및 TensorFlow의 저수준 API 사용 덕분에 하드웨어 가속 실행이 가능해져 GPU 및 TPU에서 성능 향상이著격히 향상되었다.
• 스트리밍 기대치 추정 및 다중 커널 드라이버와 같은 새로운 워크플로우를 지원하여 고급 진단 및 적응형 MCMC 전략을 가능하게 하였다.