Source-linked AI summary
Edward: A library for probabilistic modeling, inference, and criticism
Dustin Tran, Alp Kucukelbir, Adji B. Dieng, Maja Rudolph, Dawen Liang, David M. Blei
TL;DR
Probabilistic modeling needs flexible methods for representing diverse models, performing inference, and criticizing fit, especially with complex models and massive data. Edward addresses this through Box’s iterative workflow, broad modeling and inference infrastructure, and TensorFlow-backed scalable computation. The paper presents the library as enabling probabilistic models and algorithms at massive scale, while noting specific limitations in local inference and predictive-check interpretation.
Problem
Probabilistic modeling requires rich abstractions for broad model and inference classes, while researchers increasingly apply complex models to massive datasets.
Method
Edward implements Box’s iterative workflow with random-variable modeling abstractions, multiple inference algorithms, model-criticism methods, and TensorFlow-backed computation.
Results
Edward supports probabilistic modeling, inference, and criticism across broad model classes, with TensorFlow enabling distributed training and hardware acceleration.
Takeaways & Limitations
Edward provides interchangeable components for iteratively developing and evaluating probabilistic models across complex applications and large-scale computation.
Takeaways & Limitations
Edward does not support local inference within an individual random variable.
Abstract
from arXiv · showhide
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.
1 Introduction
Edward is a probabilistic modeling library designed around Box’s iterative cycle of modeling, inference, and model criticism. It combines broad model and inference abstractions with scalable TensorFlow-based computation.
- Edward’s interchangeable building blocks target rapid experimentation across probabilistic models, inference algorithms, and real-world applications.
- Edward organizes probabilistic modeling as an iterative loop: formulate a model, infer hidden structure from data, criticize fit, revise, and repeat.
- Edward provides random-variable abstractions for directed graphical models, stochastic neural networks, and programs with stochastic control flow.
- Edward supports stochastic and black box variational inference, Hamiltonian Monte Carlo, stochastic gradient Langevin dynamics, and infrastructure for developing new algorithms.
- Edward provides model-criticism methods based on scoring rules and predictive checks.
- Built on TensorFlow, Edward uses GPUs, distributed training, and automatic differentiation to support computation at scale.
2 Getting Started
Edward demonstrates probabilistic modeling with a Bayesian neural network: users simulate data, define a model, infer its latent variables, and inspect posterior draws. In the example, the inferred model captures the observed cosine relationship.
- The example uses a Bayesian neural network, defined as a neural network with a prior distribution on its weights.
- The tutorial simulates 50 observations with a cosine relationship before defining the model.
- Variational inference specifies normal approximations over the network’s weights and biases.
- Inference minimizes Kullback-Leibler divergence for 1000 iterations to infer latent variables from data.
- Posterior draws are used for graphical model checks that visualize how well sampled neural networks fit the data.
- The inferred model captures the cosine relationship between x and y in the observed domain.
3 Design
Edward’s design exposes interchangeable probabilistic-modeling components within Box’s loop and offers multiple data-reading modes. These choices support experimentation while accommodating data that may not fit in memory.
- Edward’s interchangeable building blocks enable rapid experimentation with probabilistic models.
- Box’s loop: Box’s loop cycles through building a model, reasoning from model and data, then criticizing, revising, and repeating.
- Box’s loop: A coin-flip example illustrates the loop by modeling independent flips, inferring hidden structure, and checking whether the model captures the phenomenon.
- Data: Edward supports preloaded data when it fits in memory, feeding for fine experimental control, and file-based pipelines when it does not.
- Data: Feeding represents data with TensorFlow placeholders supplied at runtime through feed_dict during inference updates.
3.2 Models
Edward represents probabilistic models as collections of random-variable objects parameterized by tensors, with compositional operations that support complex stochastic structure.
- A probabilistic model is a joint distribution p(x, z) over data x and latent variables z.
- Edward models random variables as objects parameterized by tensors, with object size determined by parameter dimensions.
- Edward supports scalar, vector, matrix, Dirichlet, and multivariate-normal random-variable constructions, with multivariate dimensions in the right-most parameter dimension.
- Random variables provide methods for log probabilities, means, and sampling, and associate graph tensors with individual samples.
- Compositional operations include arithmetic and TensorFlow operations, allowing random variables to participate in broader computational graphs.
Composing Random Variables
Edward builds models compositionally from random variables, covering graphical models, neural networks, mutable-state programs, and stochastic control flow.
- Composing Random Variables: Compositionality represents models as collections of random variables and provides fine control over modeling.
- Composing Random Variables: Edward outlines model classes including directed graphical models, neural networks, Bayesian nonparametrics, and probabilistic programs.
- Directed Graphical Models: Composing random variables implicitly defines directed edges in graphical models, as illustrated by a Beta-Bernoulli model.
- Directed Graphical Models: A random variable’s value can execute the graph to simulate a generative process, such as producing a binary vector of 50 elements.
- Composing Random Variables: TensorFlow state objects support fixed model parameters and discriminative programs whose feature inputs are supplied during training and testing.
- Neural Networks: Edward and high-level libraries construct stochastic neural networks and deep generative models with latent variables and neural-network-parameterized likelihoods.
- Neural Networks: Keras and TensorFlow Slim manage neural-network parameters as model parameters, so those parameters are not exposed for Bayesian prior distributions.
- Probabilistic Programs: Stochastic control flow creates dynamic conditional dependencies; computational graphs separate static structure for parallelism from dynamic structure handled by generic computations.
Developing Custom Random Variables
Edward supports custom random variables by combining its RandomVariable class with TensorFlow distributions and implementing sampling and probability methods as needed.
- Developing Custom Random Variables: Custom random variables inherit from Edward’s RandomVariable class and TensorFlow’s Distribution class.
- Developing Custom Random Variables: A custom implementation must provide methods such as _log_prob and _sample_n, whose template raises errors when they are unimplemented.
- Developing Custom Random Variables: The _sample_n method takes n and returns a tensor shaped as (n,) + batch_shape + event_shape.
- Developing Custom Random Variables: NumPy or SciPy sampling functions can be wrapped inside TensorFlow operations such as tf.py_func().
- Developing Custom Random Variables: Existing Edward random variables can be extended by overwriting a missing method, such as implementing custom Poisson sampling.
- Developing Custom Random Variables: The toy Poisson implementation does not correctly broadcast non-scalar parameters.
- Developing Custom Random Variables: For likelihood-only variables that are difficult to sample, supplying a value fixes the associated value and avoids the _sample_n error.
3.3 Inference
Edward frames inference as approximating a model’s posterior from observed data, while supporting configurable procedures and several specialized inference settings.
- Posterior inference: Inference takes latent variables with associated posterior variables and observed variables linked to training data.The posterior variables approximate the latent variables’ posterior distribution.
- Posterior inference: Edward adjusts qbeta and qz so their distribution approaches the posterior p(z, β | xtrain).
- Inference procedure: Inference supports fine control of the training procedure through configurable algorithm settings.
- Inference procedure: Inference algorithms expose a procedure that initializes update rules, repeatedly updates parameters, and finalizes computation.The run() method wraps this procedure.
- Other settings: Only a subset of inference algorithms supports estimation of model parameters.
- Other settings: Edward supports implicit prior samples, point-estimated model parameters, and conditional inference over selected posterior variables.Implicit prior samples use one prior sample; conditional inference fixes part of the posterior using other inferences.
Classes of Inference
Edward organizes inference around variational, Monte Carlo, and exact approaches, representing approximations and exposing algorithm-specific update procedures.
- Overview: Inference is broadly classified into variational inference, Monte Carlo, and exact inference.
- Variational inference: Variational inference searches a family of distributions for the member closest to the posterior, optimizing its parameters with respect to TensorFlow variables.
- Variational inference: MAP estimation uses PointMass approximations, with all probability mass concentrated at a point.MAP inherits from VariationalInference and uses a loss function, update rules, and TensorFlow optimizers.
- Monte Carlo: Monte Carlo approximates the posterior with samples represented as an empirical distribution.Markov chain Monte Carlo updates the current sample conditional on the previous sample; samplers may use gradients.
- Exact inference: Symbolic algebra on computational-graph nodes can uncover conjugacy and derive Gibbs, mean-field, and exact-inference updates.
Composing Inferences
Edward’s compositional design expresses inference as separate programs that can be combined into hybrid and message-passing algorithms, including distributed implementations.
- Compositionality: Compositionality enables fine control of inference by representing it as a collection of separate inference programs.
- Hybrid algorithms: Hybrid algorithms assign different inferences to latent variables, as in variational EM’s alternating E-step and M-step.The example uses approximate inference for local variables and an M-step for global variables.
- Hybrid algorithms: Conditional bindings let one inference update only part of the posterior while other variables remain fixed by other inferences.
- Message passing: Message passing operates on a posterior through a collection of local inferences, with local updates sharing a global posterior factor.The example alternates updates for local inferences while q(β) is shared.
- Scaling: With TensorFlow distributed training, compositionality supports distributed message passing across many workers and GPU acceleration through data and model parallelism.
- Message passing: Edward does not support local inferences within a single random variable.The limitation applies when all data points and cluster memberships are represented together rather than as separate variables.
Data Subsampling
Edward scales inference by subsampling data and operating on model subgraphs, reducing memory and per-iteration costs while supporting alternative parameterizations and online settings.
- Subsampling: Data subsampling updates inference using only part of the data, but only certain algorithms support it, including MAP, KLqp, and SGLD.
- Subgraphs: Subgraph inference is necessary when data and model do not fit in memory, with computation and memory per iteration independent of dataset size.
- Hierarchical models: A hierarchical model separates local variables zn for each data point from global variables β shared across data points.
- Memory complexity: Only M local variational parameters are stored in memory instead of N.
- Alternative approaches: Inference can use KLqp with variational factors, Empirical variables with SGLD, or an inference network to reduce memory complexity.
- Subgraphs: The subgraph procedure alternates global inference over β with local inference over a data subset’s z, scaling variables by N/M for unbiased stochastic gradients.
- Limitations: Subgraph inference does not apply when variational models must preserve dependencies across time steps, such as in time-series models.
3.4 Criticism
Edward supports model criticism through point-based evaluations and posterior predictive checks, helping assess fit and identify directions for model revision.
- Model criticism cannot establish that a model is true, but it can reveal where the model fails and guide revision.
- Point-based evaluations: Edward evaluates models with scalar metrics such as classification error and mean absolute error.Point-based evaluations compare predictions with observed labels or values and can use held-out data.
- Posterior predictive checks: Posterior predictive checks compare test statistics from replicated posterior-predictive data with the statistic computed on real data.The comparison can be numerical or graphical and uses discrepancy functions to assess model fit probabilistically.
- Posterior predictive checks: When the observed statistic falls in a low-probability region of the reference distribution, the check indicates poor model fit and suggests an area for improvement.
- Relationship between methods: Point-based evaluation is a special case of posterior predictive checking without a reference distribution, while PPCs provide probabilistic context.
- Uses and cautions: Edward supports PPCs for model comparison, selection, averaging, and hypothesis testing, but recommends many checks rather than a single binary decision.
4 End-to-end Examples
Edward’s examples pair probabilistic models with variational inference, posterior prediction, evaluation, and criticism. Bayesian neural networks capture nonlinear classification structure that logistic regression cannot represent in the example.
- Bayesian linear regression: The regression example simulates 500 training and test pairs with five-dimensional inputs, continuous outputs, linear dependence, and normally distributed noise.
- Bayesian linear regression: Bayesian linear regression models outputs as a linear function of inputs with latent weights and intercept, using known likelihood variance.
- Bayesian linear regression: Variational inference uses a fully factorized normal approximation and KL divergence; the evidence lower bound appears to converge in approximately 200 iterations.
- Bayesian linear regression: The inferred marginal coefficient posteriors are compared with simulated true coefficient values, and predictions have low mean squared error relative to output magnitude.
- Classification: The classification example uses 100 simulated two-dimensional points with a nonlinear decision boundary that challenges linear classifiers.
- Classification: Logistic regression fits a linear boundary, whereas a Bayesian neural network captures the nonlinear boundary and achieves better predictive accuracy metrics.