Source-linked AI summary
Deep Probabilistic Programming
Dustin Tran, Matthew D. Hoffman, Rif A. Saurous, Eugene Brevdo, Kevin Murphy, David M. Blei
TL;DR
Probabilistic programming has rich compositional models but often treats inference as a black box. Edward makes inference compositional and first-class, supporting flexible inference methods and efficient TensorFlow execution, including 35x and 6x speedups over Stan and PyMC3.
Problem
Existing probabilistic programming languages often treat inference as a black box, limiting compositional inference that reuses the model’s representation.
Method
Edward represents models as collections of random variables and inference as parameter updates, supporting composable methods and model-reusing constructions such as GANs.
Results
Edward expands probabilistic programming toward the flexibility and computational efficiency of deep learning, including 35x and 6x runtime speedups over Stan and PyMC3.
Takeaways & Limitations
Edward serves as a research platform for combining inference methods, developing variational models and GANs, and designing new probabilistic models and algorithms.
Takeaways & Limitations
Edward’s flexible inference for programs with complex control flow and recursion remains an open challenge, and extending it to dynamic computational graphs may sacrifice performance.
Abstract
from arXiv · showhide
We propose Edward, a Turing-complete probabilistic programming language. Edward defines two compositional representations---random variables and inference. By treating inference as a first class citizen, on a par with modeling, we show that probabilistic programming can be as flexible and computationally efficient as traditional deep learning. For flexibility, Edward makes it easy to fit the same model using a variety of composable inference methods, ranging from point estimation to variational inference to MCMC. In addition, Edward can reuse the modeling representation as part of inference, facilitating the design of rich variational models and generative adversarial networks. For efficiency, Edward is integrated into TensorFlow, providing significant speedups over existing probabilistic systems. For example, we show on a benchmark logistic regression task that Edward is at least 35x faster than Stan and 6x faster than PyMC3. Further, Edward incurs no runtime overhead: it is as fast as handwritten TensorFlow.
1 INTRODUCTION
Edward addresses limited compositionality in probabilistic-programming inference by treating inference alongside modeling as a first-class, composable representation.
- The design is motivated by deep learning’s compositional structure, where users connect layers without manually implementing testing or gradient-based inference.
- Probabilistic programming lets users specify generative probabilistic models as programs and compile them into inference procedures.
- Existing probabilistic-programming languages often treat inference engines as model-independent black boxes.
- Edward introduces compositional representations for both random variables and inference in a Turing-complete probabilistic programming language.
2 RELATED WORK
Related probabilistic-programming systems explore efficient or composable inference, but Edward combines model reuse with compositional inference and supports richer inference constructions.
- Probabilistic-programming languages typically trade expressive language features against computational efficiency of inference.
- Venture and Anglican build fast program-specific inference over local program fragments but do not support programmable posterior approximations, inference models, or data subsampling.
- WebPPL supports amortized inference but does not reuse the model representation, instead annotating programs and using helper functions.
- Edward composes inference within modeling and modeling within inference, including variational models.
3 COMPOSITIONAL REPRESENTATIONS FOR PROBABILISTIC MODELS
Edward represents probabilistic models with random variables embedded in TensorFlow computational graphs, enabling compositional neural, stochastic, control-flow, and data-dependent programs.
- Edward’s design integrates probabilistic representations with computational graphs and keeps the representation reusable during inference.
- Random variables provide methods for density computation and sampling while associating each variable with a tensor representing a sample on the computational graph.
- All computation on the graph can combine random variables with deep neural networks, mathematical operations, third-party libraries, and complex stochastic structure.
- Example: Beta-Bernoulli Model: A Beta-Bernoulli example uses a latent probability shared across 50 binary data points, with fetching x generating a 50-element binary vector.
- Example: Variational Auto-Encoder: A variational auto-encoder uses random variables for both its probabilistic and variational models, with neural networks defining image generation and posterior approximation.
- Examples: Bayesian RNN and Stochastic Control Flow: A Bayesian RNN supports variable-length sequences, while stochastic control flow exposes static dependence structure for model parallelism with GPUs and batch training.
4 COMPOSITIONAL REPRESENTATIONS FOR INFERENCE
Edward represents inference compositionally, allowing posterior approximations and parameter updates to be specified independently of the model. This design supports variational, Monte Carlo, adversarial, message-passing, and stochastic optimization methods while preserving computational-graph efficiency.
- Inference representation: Inference is represented as a composable class that associates latent variables with posterior variables and observed variables with their data realizations.Inference solves an optimization that adjusts posterior and model parameters toward the posterior.
- Efficiency and extensibility: Edward’s inference computations use the same computational graph as handwritten model-specific code, so the runtime is the same.The framework also supports compositional inference programs, including variational EM, exact EM, contrastive divergence, pseudo-marginal methods, Gibbs sampling, and message passing.
- Inference representation: Edward separates inference from a distinct model block, enabling partial-model inference, shared inference components, and posterior reuse in new models.The paper gives layer-wise training, multi-task learning, and Bayesian updating as examples.
- Classes of inference: Variational inference builds an approximating family in the computational graph, while MAP uses PointMass variables and a negative log joint-density loss.The example variational family uses Normal variables for β and Categorical variables for z.
- Classes of inference: Monte Carlo represents the approximation as an empirical distribution over samples and updates one sample at a time using sampler-specific rules.Samplers may use gradients, as in Hamiltonian Monte Carlo, or graph structure, as in sequential Monte Carlo.
- Classes of inference: GANInference optimizes parameters inside generator and discriminator neural networks using random noise and real-data probability outputs.The generator maps noise to real-valued data, while the discriminator predicts whether inputs are real.
- Data subsampling: Stochastic optimization subsamples M data points, estimates the log joint density unbiasedly, and scales local-variable densities to support massive datasets.Edward represents only a subgraph of the full model to avoid materializing the full model and excessive memory consumption.
5 EXPERIMENTS
Edward’s experiments demonstrate flexibility by applying diverse inference methods and novel combinations to shared models, and efficiency through GPU acceleration without runtime overhead.
- Flexibility: Edward makes it easy to compare multiple inference algorithms on the same model, supporting experimentation with composable methods.The experiments include several gradient estimators and hierarchical variational models for a VAE on binarized MNIST.
- Flexibility: Hierarchical variational models with a normalizing flow prior produced similar results to latent-space normalizing flows and better results than IWAEs.Gradient estimators typically reached the same optima at different convergence rates, with score-function gradients slowest.
- Flexibility: Novel combinations included HVMs with the IWAE objective, GAN-based decoder optimization, and Rényi divergence on the decoder.GAN-based optimization prevented log-likelihood calculation, while Rényi divergence did not directly optimize log-likelihood and performed poorly.
- Efficiency: 35x speedup over Stan and 6x speedup over PyMC3 were achieved by Edward on GPU for large-scale logistic regression.The benchmark used 100 HMC iterations on the Covertype dataset; the comparison involved Stan on one CPU and PyMC3 on twelve CPUs.
- Efficiency: The speedup stems from fast, GPU-parallelized matrix multiplication when calculating the model’s log-likelihood.The authors expect similar speedups for models whose bottleneck is also matrix multiplication, such as deep neural networks.
- Efficiency: Edward has no runtime overhead and runs as fast as handwritten TensorFlow.The inference computational graphs are the same for Edward and the handwritten code.
- Probability Zoo: The Probability Zoo extends Edward with a community repository of pretrained probability models and their posteriors.Its design is inspired by Caffe’s model zoo and Forest’s probabilistic-program examples.
6 DISCUSSION: CHALLENGES & EXTENSIONS
Edward combines compositional probabilistic models and inference to target flexibility and efficiency, while identifying unresolved challenges for complex control flow and dynamic graphs.
- Contributions: Edward is a Turing-complete probabilistic programming language with compositional representations for probabilistic models and inference.Its stated goal is to approach the flexibility and computational efficiency of traditional deep learning.
- Flexibility: Edward supports composable inference methods, variational inference, generative adversarial networks, and finely controlled inference algorithms.
- Efficiency: Edward leverages computational graphs for fast, parallelizable computation, scaling to massive data without runtime overhead over handwritten code.
- Challenges: Complex control flow and recursion remain open challenges because flexible inference strategies for such programs are not yet enabled.
- Challenges: Expanding Edward to dynamic computational graph frameworks remains unresolved because greater programming flexibility may sacrifice performance.
A.2 LATENT DIRICHLET ALLOCATION
Figure 11 presents a latent Dirichlet allocation model.
- A.2 LATENT DIRICHLET ALLOCATION: Figure 11 is labeled as latent Dirichlet allocation.
- A.2 LATENT DIRICHLET ALLOCATION: The referenced example is attributed to Blei et al. (2003).
- A.2 LATENT DIRICHLET ALLOCATION: The supplied passage identifies the example by figure number and model name.
A.3 GAUSSIAN MATRIX FACTORIZATIONN
Figure 12 presents a Gaussian matrix factorization model.
- A.3 GAUSSIAN MATRIX FACTORIZATIONN: Figure 12 is labeled as Gaussian matrix factorization.
- A.3 GAUSSIAN MATRIX FACTORIZATIONN: The supplied passage identifies the example by its figure number and model name.
- A.3 GAUSSIAN MATRIX FACTORIZATIONN: The figure is presented as the Gaussian matrix factorization example.
A Dirichlet process mixture model is written as follows:
Edward implements a Dirichlet process mixture with minibatch subgraphs and separate global and local inference. Scaling and reinitialization support unbiased stochastic updates across batches.
- Model: The DirichletProcess random variable draws N samples from a Normal base distribution and uses a stochastic while loop.The resulting samples have shape determined by the base distribution, with mu shaped (N, D).
- Scalable inference: The subgraph setting enables data subsampling when the full data and model do not fit in memory.Per-iteration computational and memory complexity are independent of dataset size.
- Scalable inference: The example uses a 10,000,000-point dataset with dimension D=2, five clusters, and minibatches of size M=128.
- Variational model: The variational model stores only M local parameters for γ_m rather than N.
- Inference: KLqp performs separate global inference over β and local inference over z, using a TensorFlow placeholder to change each batch.
- Inference: Scaling updates by N/M enables unbiased stochastic-gradient estimates, while local factors are reinitialized for each new batch.
C.1 VARIATIONAL AUTO-ENCODER
The VAE example defines a two-dimensional latent-variable model that generates MNIST digits through a neural network, with batch training producing samples periodically.
- The model uses a two-dimensional Normal latent variable z with batch size M = 100.
- A 256-unit ReLU layer transforms z before a Bernoulli output layer models 28 × 28 MNIST pixels.
- The complete VAE script generates MNIST digits after every 1000 updates during batch training.
C.2 PROBABILISTIC MODEL FOR WORD EMBEDDINGS
The word-embedding example performs minibatch-based probabilistic modeling, while its unregularized objective recovers negative sampling and its figure illustrates an exponential-family embedding.
- The example subsamples data by defining priors, conditional likelihoods, and variational embeddings only for each minibatch.TensorFlow variables store embeddings for the full vocabulary, and placeholders select the vectors needed for each minibatch.
- TensorFlow variables contain embedding vectors for the entire vocabulary, while placeholders select the correct variational parameters for each minibatch.
- Without priors, the objective for fixed positive and negative labels is identical to negative sampling.The positive labels are 1s and the negative labels are 0s.
- The exponential-family embedding example uses MAP to maximize the total sum of conditional log-likelihoods and log-priors.