[논문 리뷰] Simple, Distributed, and Accelerated Probabilistic Programming

Dustin Tran, Matthew W. Hoffman외 4명|arXiv (Cornell University)|2018. 01. 01.

Machine Learning and Data Classification인용 수 30

한 줄 요약

이 논문은 텐서플로우에 통합된 경량이며 저수준의 확률적 프로그래밍 프레임워크를 소개한다. 이 프레임워크는 단일 추상화인 랜덤 변수에 기반하여 효율적인 모델 병렬 및 데이터 병렬 추론을 가능하게 한다. VAE 및 이미지 트랜스포머에서 256개의 TPUv2 코어를 대상으로 최적의 선형 속도 향상을 달성하였으며, 멀티 GPU NUTS 샘플링에서 스탠 대비 100배 빠른 성능을 보였다.

ABSTRACT

We describe a simple, low-level approach for embedding probabilistic programming in a deep learning ecosystem. In particular, we distill probabilistic programming down to a single abstraction—the random variable. Our lightweight implementation in TensorFlow enables numerous applications: a model-parallel variational auto-encoder (VAE) with 2nd-generation tensor processing units (TPUv2s); a data-parallel autoregressive model (Image Transformer) with TPUv2s; and multi-GPU No-U-Turn Sampler (NUTS). For both a state-of-the-art VAE on 64x64 ImageNet and Image Transformer on 256x256 CelebA-HQ, our approach achieves an optimal linear speedup from 1 to 256 TPUv2 chips. With NUTS, we see a 100x speedup on GPUs over Stan and 37x over PyMC3.

연구 동기 및 목표

• 확률적 프로그래밍을 단일 추상화인 랜덤 변수로 단순화함으로써 표현력과 성능를 손상시키지 않도록 한다.
• 최소한의 엔지니어링 오버헤드로 딥러닝 생태계 내에서 효율적이고 확장 가능한 추론을 가능하게 한다.
• 모델 병렬 및 데이터 병렬 확률적 모델에 대해 수천 개의 가속기에서 선형 속도 향상을 달성한다.
• 멀티 GPU No-U-Turn 샘플러(NUTS)를 통해 마르코프 체인 몬테카를로(MCMC) 샘플링을 이전 도구를 뛰어넘는 성능으로 가속화한다.

제안 방법

• 프레임워크는 복잡한 분포와 추론 메커니즘을 추상화하기 위해 단지 랜덤 변수를 핵심 추상화로 사용하여 확률적 프로그램을 모델링한다.
• 텐서플로우와 네이티브로 통합되어 자동 미분 및 분산 실행을 활용해 효율적인 계산을 수행한다.
• 모델 병렬 VAE 및 데이터 병렬 순차적 모델(예: 이미지 트랜스포머)은 TPUv2 하드웨어에서 이 추상화를 사용해 구현되었다.
• 장치 간 파라미터 분할과 기울기 동기화를 통해 멀티 GPU NUTS 추론이 가능해졌다.
• 최소한의 코드 변경으로 변분 추론과 MCMC 샘플링을 모두 지원한다.
• 1에서 256개의 TPUv2 코어에 걸쳐 효율적인 데이터 병렬 및 모델 병렬화를 통해 선형 확장이 달성되었다.

실험 결과

연구 질문

• RQ1 표현력이나 성능를 희생시키지 않고도 확률적 프로그래밍을 단일 추상화로 단순화할 수 있는가?
• RQ2 최소한의 인터페이스를 가진 저수준 인터페이스가 모델 병렬 및 데이터 병렬 환경에서 수천 대의 가속기에서 최적의 확장성을 달성할 수 있는가?
• RQ3 이러한 시스템은 기존의 확률적 프로그래밍 도구보다 MCMC 샘플링에서 뛰어난 성능을 달성할 수 있는가?
• RQ4 이 프레임워크는 분산 실행 환경에서 최신 기술인 VAE 및 트랜스포머 모델을 얼마나 잘 지원하는가?

주요 결과

• 64x64 ImageNet에서 최신 기술의 VAE에 대해 1에서 256개의 TPUv2 코어로 최적의 선형 속도 향상을 달성하였다.
• 256x256 CelebA-HQ에서 이미지 트랜스포머에 대해 최적의 선형 속도 향상을 달성하여 데이터 병렬 환경에서의 확장성을 입증하였다.
• 멀티 GPU NUTS 샘플링은 스탠 대비 100배, PyMC3 대비 37배의 속도 향상을 달성하였다.
• 동일한 저수준 추상화를 사용하여 변분 자동에코더와 순차적 모델 모두에 대해 효율적이고 확장 가능한 추론을 가능하게 하였다.
• 다양한 모델 아키텍처와 하드웨어 구성(예: TPU 및 멀티 GPU 설정)에서도 높은 성능를 유지하였다.
• 랜덤 변수 추상화의 단순성 덕분에 성능나 스케일링를 희생시키지 않고도 넓은 적용 범위를 가질 수 있었다.