[Paper Review] TensorFlow Distributions

Joshua V. Dillon, Ian Langmore+8 more|arXiv (Cornell University)|Nov 28, 2017

Gaussian Processes and Bayesian Inference41 references240 citations

TL;DR

TensorFlow Distributions provides two core abstractions—Distributions and Bijectors—for fast, differentiable probabilistic programming in TensorFlow, enabling modular construction of complex high-dimensional distributions and transformations and powering backends like Edward.

ABSTRACT

The TensorFlow Distributions library implements a vision of probability theory adapted to the modern deep-learning paradigm of end-to-end differentiable computation. Building on two basic abstractions, it offers flexible building blocks for probabilistic computation. Distributions provide fast, numerically stable methods for generating samples and computing statistics, e.g., log density. Bijectors provide composable volume-tracking transformations with automatic caching. Together these enable modular construction of high dimensional distributions and transformations not possible with previous libraries (e.g., pixelCNNs, autoregressive flows, and reversible residual networks). They are the workhorse behind deep probabilistic programming systems like Edward and empower fast black-box inference in probabilistic models built on deep-network components. TensorFlow Distributions has proven an important part of the TensorFlow toolkit within Google and in the broader deep learning community.

Motivation & Objective

• Enable end-to-end differentiable probabilistic programming within the TensorFlow ecosystem.
• Provide fast, numerically stable sampling, log-densities, and statistics for a large set of distributions.
• Support batching, automatic differentiation, and accelerator (GPU/TPU) compatibility.
• Offer composable transformations (Bijectors) to build complex distributions efficiently.
• Integrate with higher-level tools like Edward and TensorFlow Estimator to enable scalable research and production use.

Proposed method

• Introduce two abstractions: Distribution (60+ distributions with fast sampling and log_prob) and Bijector (22 composable, differentiable transformations).
• Define shape semantics (sample, batch, event) to enable vectorized operations and broadcasting.
• Implement sampling via device-specific C++ kernels with end-to-end differentiation via reparameterization where possible.
• Provide higher-order distributions (functions of distributions) and distribution functionals (e.g., entropy, KL).
• Utilize TransformedDistribution and Chain/Inver t bijectors for modular composition and efficient density computation.
• Incorporate caching for transformed samples and log-determinants to speed up importance sampling and variational inference.

Experimental results

Research questions

• RQ1 How can we design a probability library that is fast, differentiable, and scalable for deep probabilistic programming?
• RQ2 How can distributions and transformations be composed to express rich, high-dimensional probabilistic models?
• RQ3 What abstractions (Distributions and Bijectors) enable modular, reusable construction of complex models while preserving numerical stability?
• RQ4 How can such a library integrate with the broader TensorFlow ecosystem and accelerator hardware?
• RQ5 What are the practical benefits and limitations of higher-order distributions and distribution functionals in deep learning contexts?

Key findings

• The library provides about 60 distributions with fast sampling and log-density computations plus 22 composable bijectors.
• Distributions and Bijectors enable modular, end-to-end differentiable construction of complex models such as VAEs, autoregressive flows, and PixelCNN-based architectures.
• Shape semantics (sample, batch, event) enable idiomatic vectorization and broadcasting across large tensors.
• Sampling is implemented with device-specific kernels and supports reparameterization for valid backpropagation through stochastic nodes.
• Bijectors automatically cache transformations and log-determinants, improving efficiency for sampling-based inference and variational methods.
• TensorFlow Distributions integrates with TensorFlow components (layers, data pipelines, serving, visualization) and serves as the backend for Edward.