[Paper Review] Edward: A library for probabilistic modeling, inference, and criticism

Dustin Tran, Alp Kucukelbir+4 more|arXiv (Cornell University)|Oct 31, 2016

Machine Learning and Data Classification1 references221 citations

TL;DR

Edward is a TensorFlow-based library that enables iterative probabilistic modeling, scalable inference, and model criticism for a wide range of probabilistic models and algorithms.

ABSTRACT

Probabilistic modeling is a powerful approach for analyzing empirical information. We describe Edward, a library for probabilistic modeling. Edward's design reflects an iterative process pioneered by George Box: build a model of a phenomenon, make inferences about the model given data, and criticize the model's fit to the data. Edward supports a broad class of probabilistic models, efficient algorithms for inference, and many techniques for model criticism. The library builds on top of TensorFlow to support distributed training and hardware such as GPUs. Edward enables the development of complex probabilistic models and their algorithms at a massive scale.

Motivation & Objective

• Provide a flexible language of random variables to define broad probabilistic models (directed graphs, neural networks, nonparametrics, probabilistic programs).
• Offer scalable inference algorithms (variational inference, Monte Carlo, HMC, Langevin dynamics) and a framework to develop new algorithms.
• Incorporate model criticism tools (scoring rules, predictive checks) to assess and iterate on models.
• Integrate with TensorFlow for GPU/ distributed computation and automatic differentiation.
• Demonstrate end-to-end workflows on standard statistics and ML tasks (e.g., Bayesian linear regression, classification).

Proposed method

• Edward builds models by composing random variables (directed graphical models, stochastic neural networks, probabilistic programs).
• Inference is implemented as modular classes (e.g., Variational Inference, Monte Carlo, exact inference) with support for black-box and model-specific methods.
• The system enables data input modes (preloaded, feeding, and file reading) and supports model parameters and conditional inference.
• Criticism uses predictive checks and scoring rules to assess model fit and guide revision.
• Edward runs on TensorFlow to leverage GPUs, distributed training, and automatic differentiation.
• It supports custom random variables and compositional inference, including hybrid and message-passing algorithms.

Experimental results

Research questions

• RQ1 How can probabilistic models be expressed, inferred, and critiqued within a unified software framework?
• RQ2 What inference algorithms (variational, Monte Carlo, exact) are effective across diverse model families in Edward?
• RQ3 How can flexible composition of models and inferences facilitate rapid experimentation at scale?
• RQ4 What design choices enable probabilistic programs with stochastic control flow to be trained efficiently on large data?
• RQ5 How can model criticism be integrated into the iterative Box loop to drive model revision?

Key findings

• Edward provides a broad language of random variables enabling directed graphs, neural networks, Bayesian nonparametrics, and probabilistic programs.
• A modular inference framework supports variational, Monte Carlo, and exact methods, including hybrid and message-passing algorithms.
• The library emphasizes compositionality, enabling conditional inference and hybrid inference workflows for complex posteriors.
• Edward is implemented on TensorFlow, enabling GPU acceleration, distributed training, and automatic differentiation.
• The paper demonstrates end-to-end examples such as Bayesian linear regression and neural network classification to illustrate workflows.
• Edward supports custom random variables and stochastic control flow to model dynamic computational graphs.