API and Documentation

Composing Inferences

Core to Edward’s design is compositionality. Compositionality enables fine control of inference, where we can write inference as a collection of separate inference programs.

We outline how to write popular classes of compositional inferences using Edward: hybrid algorithms and message passing algorithms. We use the running example of a mixture model with latent mixture assignments z, latent cluster means beta, and observations x.

Hybrid algorithms

Hybrid algorithms leverage different inferences for each latent variable in the posterior. As an example, we demonstrate variational EM, with an approximate E-step over local variables and an M-step over global variables. We alternate with one update of each (Neal & Hinton, 1993).

from edward.models import Categorical, PointMass

qbeta = PointMass(params=tf.Variable(tf.zeros([K, D])))
qz = Categorical(logits=tf.Variable(tf.zeros[N, K]))

inference_e = ed.VariationalInference({z: qz}, data={x: x_data, beta: qbeta})
inference_m = ed.MAP({beta: qbeta}, data={x: x_data, z: qz})
...
for _ in range(10000):
  inference_e.update()
  inference_m.update()

In data, we include bindings of prior latent variables (z or beta) to posterior latent variables (qz or qbeta). This performs conditional inference, where only a subset of the posterior is inferred while the rest are fixed using other inferences.

This extends to many algorithms: for example, exact EM for exponential families; contrastive divergence (Hinton, 2002); pseudo-marginal and ABC methods (Andrieu & Roberts, 2009); Gibbs sampling within variational inference (Wang & Blei, 2012); Laplace variational inference (Wang & Blei, 2013); and structured variational auto-encoders (Johnson, Duvenaud, Wiltschko, Datta, & Adams, 2016).

Message passing algorithms

Message passing algorithms operate on the posterior distribution using a collection of local inferences (Koller & Friedman, 2009). As an example, we demonstrate expectation propagation. We split a mixture model to be over two random variables x1 and x2 along with their latent mixture assignments z1 and z2.

from edward.models import Categorical, Normal

N1 = 1000  # number of data points in first data set
N2 = 2000  # number of data points in second data set
D = 2  # data dimension
K = 5  # number of clusters

# MODEL
beta = Normal(loc=tf.zeros([K, D]), scale=tf.ones([K, D]))
z1 = Categorical(logits=tf.zeros([N1, K]))
z2 = Categorical(logits=tf.zeros([N2, K]))
x1 = Normal(loc=tf.gather(beta, z1), scale=tf.ones([N1, D]))
x2 = Normal(loc=tf.gather(beta, z2), scale=tf.ones([N2, D]))

# INFERENCE
qbeta = Normal(loc=tf.Variable(tf.zeros([K, D])),
               scale=tf.nn.softplus(tf.Variable(tf.zeros([K, D]))))
qz1 = Categorical(logits=tf.Variable(tf.zeros[N1, K]))
qz2 = Categorical(logits=tf.Variable(tf.zeros[N2, K]))

inference_z1 = ed.KLpq({beta: qbeta, z1: qz1}, {x1: x1_train})
inference_z2 = ed.KLpq({beta: qbeta, z2: qz2}, {x2: x2_train})
...
for _ in range(10000):
  inference_z1.update()
  inference_z2.update()

We alternate updates for each local inference, where the global posterior factor \(q(\beta)\) is shared across both inferences (Gelman et al., 2017).

With TensorFlow’s distributed training, compositionality enables distributed message passing over a cluster with many workers. The computation can be further sped up with the use of GPUs via data and model parallelism.

This extends to many algorithms: for example, classical message passing, which performs exact local inferences; Gibbs sampling, which draws samples from conditionally conjugate inferences (Geman & Geman, 1984); expectation propagation, which locally minimizes \(\text{KL}(p || q)\) over exponential families (Minka, 2001); integrated nested Laplace approximation, which performs local Laplace approximations (Rue, Martino, & Chopin, 2009); and all the instantiations of EP-like algorithms in Gelman et al. (2017).

In the above, we perform local inferences split over individual random variables. At the moment, Edward does not support local inferences within a random variable itself. We cannot do local inferences when representing the random variable for all data points and their cluster membership as x and z rather than x1, x2, z1, and z2.

References