Source-linked AI summary
TensorFlow Distributions
Joshua V. Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, Rif A. Saurous
TL;DR
Deep learning creates new opportunities for probabilistic programming in perceptual and scientific applications. TensorFlow Distributions addresses this opportunity with modular, differentiable probability abstractions and has seen broad adoption across probabilistic programming and statistical systems.
Problem
Deep learning and deep generative models create new opportunities for probabilistic programming across perceptual and scientific applications.
Method
TensorFlow Distributions provides Distributions for sampling, log densities, and statistics, and Bijectors for composable volume-tracking transformations with caching.
Results
TensorFlow Distributions has been adopted in production and research systems, serves as Edward’s backend, and underlies related systems including Greta, Pyro, and ZhuSuan.
Takeaways & Limitations
The two abstractions support modular construction of rich deep probabilistic models and differentiable functions of samples.
Takeaways & Limitations
Distributions do not implement approximate inference or approximations of properties and statistics.
Abstract
from arXiv · showhide
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.
1 Introduction
TensorFlow Distributions addresses missing modern backend support for probabilistic programming by providing efficient, composable probability distributions and transformations integrated with TensorFlow. Its examples show modular construction of increasingly rich variational autoencoders, while its stated scope excludes universal and approximate-inference functionality.
- Motivation: Deep probabilistic programming lacked backend libraries combining probability operations with batching, automatic differentiation, accelerators, compiler optimization, and composability.The paper identifies these capabilities as missing despite their necessity for sampling, log densities, statistics, and probabilistic-program manipulation.
- Illustration: A few code changes transform a standard MNIST VAE into an architecture using a PixelCNN++ decoder and autoregressive flows for the encoder and prior.The standard version uses a Bernoulli decoder, fully factorized Gaussian encoder, and Gaussian prior, whereas the richer version uses a convolutional encoder and prior with autoregressive flows.
- Applications: TensorFlow Distributions supports modular compositions such as pixelCNN applications and higher-level abstractions including Edward and TF Estimator.The figures demonstrate composing fast, idiomatic modules to express rich, deep structure.
- Contributions: TensorFlow Distributions (r1.4) defines 56 Distributions for sampling, log densities, and statistics, plus 22 Bijectors for composable volume-tracking transformations.The library is designed as an efficient, composable TensorFlow system for manipulating probability distributions.
- Contributions: TensorFlow integration provides compatibility with neural-network layers, data pipelines, distributed serving, visualization, graph operations, differentiation, batching, and accelerators.The ecosystem includes CPUs, GPUs, and TPUs, along with XLA and device-specific kernel optimizations.
- Non-Goals: The library does not aim for universality or approximate inference: distributions must have expected-polynomial-time sampling and log-probability implementations, and approximations are excluded.The stated non-goals exclude distributions such as Multinomial-LogisticNormal and Monte Carlo entropy approximation, with explicit exceptions for some analytically tractable cases.
2 Related Work
TensorFlow Distributions is positioned relative to R, SciPy, and Stan Math as an object-oriented, differentiable library tailored to deep probabilistic programming. It adds tensor-oriented execution, accelerator support, and design innovations beyond API compatibility.
- R provides a comprehensive, user-contributed distribution collection, whereas TensorFlow Distributions uses object-oriented Distribution objects and fast, differentiable operations.
- TensorFlow Distributions models its API after SciPy while adding arbitrary tensor-dimensional vectorization, computational-graph operations, automatic differentiation, and accelerator execution.
- Unlike Stan Math, TensorFlow Distributions focuses on deep probabilistic programming through bijectors, shape semantics, higher-order distributions, and distribution functionals.
3 Distributions
Distributions provide object-oriented, analytically specified probability-distribution operations with broadcasting, shape semantics, and reparameterization support for differentiable computation.
- Distribution abstraction: TensorFlow Distributions offers approximately 60 distributions with methods for sampling, log density, and statistics.The library defines a common Distribution interface and supports additional functionals such as cdf, quantile, mean, variance, and entropy.
- Distribution abstraction: Distribution instances expose standardized parameters, validation controls, dtypes, and TensorFlow-style broadcasting.Alternative parameterizations are handled through mutually exclusive arguments or separate distribution classes.
- Distribution contracts: Distribution methods are designed to be efficiently computable, vectorized, statistically consistent, analytical except for sampling, and fixed in their properties.The library favors parameterizations such as MultivariateNormalTriL that avoid additional factorizations.
- Shape semantics: Sample, batch, and event shapes distinguish repeated draws, independently parameterized distributions, and the dimensions of one draw.Combining these shapes in one Tensor enables vectorized computation and broadcasting, while dynamic sample and batch ranks are disallowed.
- Reparameterization: Reparameterized samples are smooth functions of parameters and parameter-free randomness, enabling gradients through expectations of sample-dependent functions.For smooth enough reparameterized samples, Monte Carlo estimates of both expected loss and its parameter gradient are justified.
4 Bijectors
Bijectors separate deterministic transformations from stochastic distributions, providing composable, differentiable mappings that automate transformed densities, shape changes, and volume corrections.
- Motivation: Bijectors isolate deterministic transformations from Distributions, enabling modular construction, inherited methods, and algebraic relationships among random variables.This separation makes distributions easier to design, implement, and validate while avoiding a combinatorial collection of specialized transformations.
- Definition: The Bijector API supports differentiable bijections and selected non-injective maps, with TransformedDistribution expressing new distributions from base distributions and transforms.For diffeomorphisms, transformed densities use the inverse transformation and the absolute inverse-Jacobian determinant.
- Definition: Forward, inverse, and inverse_log_det_jacobian operations respectively transform samples, recover inputs for log_prob, and correct for volume changes.TransformedDistribution uses these operations to automate sampling and density evaluation.
- Composition: Chain and Invert compose or reverse Bijectors, supporting constructions such as autoregressive flows, matrix-shaped logit-Normals, and unconstrained parameterizations.The framework also supports transformations such as softplus inversion for Gamma distributions used by algorithms operating on real-valued spaces.
- Caching: Bijector caching reuses forward and inverse computation, benefiting settings where inversion is slow, unstable, or unavailable.For an InverseAutoregressiveFlow, caching reduces overall complexity from quadratic to linear in event size.
- Non-injective transformations: Smooth-covering Bijectors represent non-injective transformations by returning all inverse preimages and enable half-distributions such as half-Cauchy.These distributions allocate probability mass over only the positive half-plane of the original distribution.
5 Applications
TensorFlow Distributions supports diverse probabilistic applications by composing distributions and bijectors into density estimators, autoregressive image models, stochastic recurrent networks, and TPU-based regression.
- Kernel Density Estimators: KDEs represent unknown distributions nonparametrically and can be constructed flexibly as MixtureSameFamily models.Changing the callable permits alternative distribution-based kernels and bootstrap sampling schemes.
- Bijectors: Bijector caching is not currently supported for smooth coverings.
- PixelCNN: PixelCNN models images as fully visible autoregressive distributions whose callable returns per-time-step distributions for likelihood training and sampling.The model handles batches of 32 × 32 × 3 RGB images from Small ImageNet.
- Stochastic recurrent neural networks: Stochastic recurrent neural networks model sequences with random hidden states, nonlinear dynamics, and neural-network parameterizations of latent variables and likelihoods.Edward supplies approximate inference for the generative program, while inference remains a higher-level concern than TensorFlow Distributions.
- TensorFlow Estimator API: Multivariate regression can model heteroscedastic noise through MultivariateNormalTriL parameterized by neural layers.Using TPUEstimator provides training, evaluation, and prediction across diverse hardware and network topologies.
6 Discussion
TensorFlow Distributions provides a broad set of distribution and transformation primitives for differentiable computation, and the library has gained adoption across probabilistic programming ecosystems.
- Library scope: The library provides 56 distributions with numerically stable sampling, log-density, and statistics methods, alongside 22 composable bijectors with volume tracking and caching.
- Adoption: TensorFlow Distributions has seen widespread adoption inside and outside Google, including use in Edward and Greta and as a design basis for PyTorch, Pyro, and ZhuSuan.
- Future directions: Planned extensions include distribution support properties, exponential-family structure, broader TPU compatibility, and native SparseTensor support.Examples include unnormalized_log_prob and log_normalizer methods and compatibility with rejection- and whileloop-based sampling.