[Paper Review] Simple, Distributed, and Accelerated Probabilistic Programming

Dustin Tran, Matthew W. Hoffman+4 more|arXiv (Cornell University)|Jan 1, 2018

Machine Learning and Data ClassificationComputer Science30 citations

TL;DR

This paper introduces a lightweight, low-level probabilistic programming framework embedded in TensorFlow, centered on a single abstraction—random variables—enabling efficient model- and data-parallel inference. It achieves optimal linear speedup across 256 TPUv2 chips for VAEs and Image Transformers, and delivers 100x speedup over Stan with multi-GPU NUTS sampling.

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.

Motivation & Objective

• To simplify probabilistic programming by reducing it to a single abstraction: the random variable.
• To enable efficient, scalable inference in deep learning ecosystems using minimal engineering overhead.
• To achieve linear speedup across thousands of accelerators for both model-parallel and data-parallel probabilistic models.
• To accelerate Markov Chain Monte Carlo sampling with multi-GPU No-U-Turn Sampler (NUTS) beyond existing tools.

Proposed method

• The framework models probabilistic programs using only random variables as the core abstraction, abstracting away complex distributions and inference mechanisms.
• It integrates natively with TensorFlow, leveraging its automatic differentiation and distributed execution for efficient computation.
• Model-parallel VAEs and data-parallel autoregressive models (e.g., Image Transformer) are implemented using this abstraction on TPUv2 hardware.
• Multi-GPU NUTS inference is enabled through parameter sharding and gradient synchronization across devices.
• The system supports both variational inference and MCMC sampling with minimal code changes.
• Linear scaling is achieved via efficient data and model parallelism across 1 to 256 TPUv2 chips.

Experimental results

Research questions

• RQ1 Can probabilistic programming be simplified to a single abstraction without sacrificing expressiveness or performance?
• RQ2 Does a minimalistic, low-level interface enable optimal scaling across thousands of accelerators in both model- and data-parallel settings?
• RQ3 Can such a system achieve superior performance in MCMC sampling compared to existing probabilistic programming tools?
• RQ4 To what extent does the framework support state-of-the-art models like VAEs and Transformers under distributed execution?

Key findings

• The framework achieves optimal linear speedup from 1 to 256 TPUv2 chips for a state-of-the-art VAE on 64x64 ImageNet.
• It achieves optimal linear speedup for an Image Transformer on 256x256 CelebA-HQ, demonstrating scalability in data-parallel settings.
• Multi-GPU NUTS sampling achieves a 100x speedup over Stan and a 37x speedup over PyMC3 on GPUs.
• The system enables efficient, scalable inference for both variational auto-encoders and autoregressive models using the same low-level abstraction.
• The approach maintains high performance across diverse model architectures and hardware configurations, including TPUs and multi-GPU setups.
• The simplicity of the random variable abstraction enables broad applicability without sacrificing performance or scalability.