Edward2 Documentation

Bayesian Neural Network (BNN) layers extend deterministic Keras layers by placing prior distributions on their weights and biases. They are designed as drop-in replacements that maintain the same constructor arguments and input/output Tensor shapes as their deterministic counterparts.

Key design principles:

• Estimators as Layers: Because computing the integrals over weights/biases is often intractable, each estimation method is implemented as a distinct Layer. For example, ed.layers.Conv2DFlipout and ed.layers.Conv2DReparameterization are different estimators for a convolutional layer. Gradients work automatically with tf.GradientTape.
• Parameter Distributions: Use Keras' kernel_initializer and bias_initializer to specify distributions over parameters. Use ed.initializers for Bayesian-specific additions.
• Regularization: Use Keras' kernel_regularizer and bias_regularizer to specify regularizers like the KL penalty. Use ed.regularizers for Bayesian-specific additions.

# Example of swapping a deterministic layer for a Bayesian one
model = tf.keras.Sequential([
  ed.layers.Conv2DFlipout(32, (3, 3), activation='relu'),
  tf.keras.layers.MaxPooling2D((2, 2)),
  # ...
])

Edward2 provides reversible layers that allow for invertible neural networks, enabling the propagation of uncertainty from input to output. These layers are useful for transformations like log-normal distributions or high-dimensional flow-based models.

Key Design Points

• Inversion: Reversible layers implement a reverse method to perform the inverse computation of the standard call. They may also provide a log_det_jacobian method. You can use ed.layers.Reverse(layer) to create a new layer that swaps the forward and reverse computations of the provided input layer.
• Propagating Uncertainty: While deterministic layers are Tensor-in/Tensor-out, reversible layers can also take a RandomVariable as input and return a transformed RandomVariable based on the call, reverse, and log_det_jacobian logic.

Example: Discrete Autoregressive Flow

To implement a discrete flow over sequences, you can compose ed.layers.DiscreteAutoregressiveFlow with other layers like ed.layers.MADE within a tf.keras.Sequential model.

sequence_length, vocab_size = ...

# Define the model.
flow = tf.keras.Sequential([
  ed.layers.DiscreteAutoregressiveFlow(ed.layers.MADE(vocab_size, hidden_dims=[256, 256])),
  ed.layers.DiscreteAutoregressiveFlow(ed.layers.MADE(vocab_size, hidden_dims=[256, 256], order='right-to-left')),
  ed.layers.DiscreteAutoregressiveFlow(ed.layers.MADE(vocab_size, hidden_dims=[256, 256])),
])
base = ed.Categorical(logits=tf.Variable(tf.random.normal([sequence_length, vocab_size]))

# Specify custom loss function and run training loop. Or use model.compile and
# model.fit.
def loss_fn(features):
  whitened_features = flow.reverse(features)
  # In this example, we don't include log-det-jacobian as in continuous flows.
  # Discrete flows don't require them.
  return -tf.reduce_mean(base.distribution.log_prob(whitened_features))

In Edward, models were typically written inline by composing random variables. In Edward2, the recommended pattern is to write models as functions.

In this functional pattern:

• Function Inputs: Represent what the probabilistic program conditions on (the $x$ in $p(y|x)$).
• Function Outputs: Represent what the probabilistic program is over (the $y$ in $p(y|x)$).

Best Practice: Always provide a name argument to all random variables (e.g., ed.Gamma(..., name="w2")). This ensures cleaner names in the computational graph and facilitates model manipulation.

def deep_exponential_family(data_size, feature_size, units, shape):
  """A multi-layered topic model over a documents-by-terms matrix."""
  w2 = ed.Gamma(0.1, 0.3, sample_shape=[units[2], units[1]], name="w2")
  w1 = ed.Gamma(0.1, 0.3, sample_shape=[units[1], units[0]], name="w1")
  w0 = ed.Gamma(0.1, 0.3, sample_shape=[units[0], feature_size], name="w0")

  z2 = ed.Gamma(0.1, 0.1, sample_shape=[data_size, units[2]], name="z2")
  z1 = ed.Gamma(shape, shape / tf.matmul(z2, w2), name="z1")
  z0 = ed.Gamma(shape, shape / tf.matmul(z1, w1), name="z0")
  x = ed.Poisson(tf.matmul(z0, w0), name="x")
  return x

In Edward2, RandomVariables are used to define the structure of a probabilistic model. Each random variable rv contains a rv.distribution (a TensorFlow Distribution instance) which provides methods like log_prob and sample.

Key Behaviors:

• Sampling: By default, instantiating a random variable creates a sampling operation. The number of samples is controlled by the sample_shape argument. If a value argument is provided during instantiation, no sampling operation is created.
• Interoperability: Random variables interoperate with TensorFlow ops; operations performed on a random variable act on its sample.
• Accessing Samples: You can access the underlying sample via the random variable object itself (e.g., x + y where x is a RandomVariable).

import edward2 as ed
import tensorflow as tf

# Basic instantiation
normal_rv = ed.Normal(loc=0., scale=1.)

# Accessing log_prob
log_prob = normal_rv.distribution.log_prob(1.231)

# Interoperating with TF ops
x = ed.Normal(loc=tf.zeros(2), scale=tf.ones(2))
y = 5.
result = x + y  # Operates on the sample

Stochastic output layers add stochasticity to the output of a model. Given a Tensor input, they perform deterministic computations and return an ed.RandomVariable. This is useful for models where the output has a known, tractable distribution, such as Variational Autoencoders (VAEs).

• Output Dimensionality: An optional units argument determines output dimensionality via a trainable linear projection. If omitted, the layer maintains the input shape without projection.
• Inference: Unlike BNN or GP layers, testing with stochastic output layers does not strictly require Monte Carlo averaging unless the model contains other stochastic layers (like BNN/GP layers).

# Example: Variational Autoencoder (VAE) structure
encoder = tf.keras.Sequential([
  tf.keras.layers.Conv2D(128, 5, 1, padding='same', activation='relu'),
  # ...
  ed.layers.Normal(name='latent_code'),
])

decoder = tf.keras.Sequential([
  # ...
  ed.layers.Categorical(name='image'),
])

# Loss function using log_prob and KL divergence
def loss_fn(features):
  encoding = encoder(features)
  nll = -decoder(encoding).log_prob(features)
  kl = encoding.distribution.kl_divergence(ed.Normal(0., 1.).distribution)
  return tf.reduce_mean(nll + kl)

Gaussian Process (GP) layers represent distributions over functions by specifying function values at different inputs, rather than distributions over weights.

Key features:

• Estimators: Like BNNs, GP integration is handled via specific layers. ed.layers.GaussianProcess provides exact integration, while ed.layers.SparseGaussianProcess uses inducing variable approximations for scalability.
• Type Signature: GP layers maintain typical Keras arguments. For example, ed.layers.GaussianProcess(units) acts as a Bayesian nonparametric extension of tf.keras.layers.Dense(units).
• Arguments: Instead of an activation function, GP layers use mean and covariance function arguments (defaulting to the zero function and squared exponential kernel, respectively).
• Regularization: Use kernel_regularizer and bias_regularizer to specify regularizers like the KL penalty.

# Define the model.
model = tf.keras.Sequential([
  tf.keras.layers.Flatten(),
  ed.layers.SparseGaussianProcess(256, num_inducing=512),
  ed.layers.SparseGaussianProcess(256, num_inducing=512),
  ed.layers.SparseGaussianProcess(3, num_inducing=512),
])
predictions = model(features)

In Edward, variational inference is handled by high-level inference classes (like ed.KLqp) that manage the relationship between random variables and the model. In Edward2, inference is modularized and manual. You must define your own variational approximation (often as another Edward2 program) and use ed.tape() and ed.condition() to compute the Evidence Lower Bound (ELBO).

Key changes:

• Edward: Use an inference class and call .run() or manual .update() loops.
• Edward2: Use ed.tape() to record the model's forward pass, ed.condition() to inject variational samples into the model, and manually compute the KL-divergence between the variational and prior distributions to derive the ELBO.

# Edward2 pattern for ELBO computation
@tf.function
def train_step(bag_of_words, step):
  with tf.GradientTape() as tape:
    # 1. Sample from variational distribution
    qw2, qw1, qw0, qz2, qz1, qz0 = deep_exponential_family_variational()

    # 2. Record forward pass with ed.tape()
    with ed.tape() as model_tape:
      with ed.condition(w2=qw2, w1=qw1, w0=qw0, z2=qz2, z1=qz1, z0=qz0):
        posterior_predictive = deep_exponential_family(data_size, feature_size, units, shape)

    # 3. Compute log-likelihood
    log_likelihood = posterior_predictive.distribution.log_prob(bag_of_words)

    # 4. Compute KL-divergence using the recorded tape
    kl = 0.
    for rv_name, variational_rv in [("z0", qz0), ("z1", qz1), ("z2", qz2), 
                                    ("w0", qw0), ("w1", qw1), ("w2", qw2)]:
      kl += tf.reduce_sum(variational_rv.distribution.kl_divergence(
          model_tape[rv_name].distribution))

    elbo = tf.reduce_mean(log_likelihood - kl)
    loss = -elbo

  # 5. Standard TensorFlow optimization
  gradients = tape.gradient(loss, trainable_variables)
  optimizer.apply_gradients(zip(gradients, trainable_variables))
  return loss

In Edward, MCMC is implemented via inference classes like ed.HMC that manage state and sampling. In Edward2, MCMC is treated as a transition kernel (a function that maps one state to the next).

To implement MCMC in Edward2:

1. Define a target log-probability function using ed.make_log_joint_fn.
2. Create a function representing the log-joint probability of the model given fixed hyperparameters and observations.
3. Apply a transition kernel (e.g., a NUTS sampler) iteratively over the states.

# Edward2 pattern for MCMC
# 1. Create log-joint function
log_joint = ed.make_log_joint_fn(deep_exponential_family)

# 2. Define target log-prob function
def target_log_prob_fn(w2, w1, w0, z2, z1, z0):
  return log_joint(data_size, feature_size, units, shape, 
                   w2=w2, w1=w1, w0=w0, z2=z2, z1=z1, z0=z0, x=bag_of_words)

# 3. Apply transition kernel
for _ in range(num_samples):
  [states, target_log_prob, grads] = no_u_turn_sampler.kernel(
      target_log_prob_fn=target_log_prob_fn,
      current_state=current_state,
      step_size=step_size,
      current_target_log_prob=target_log_prob,
      current_grads_target_log_prob=grads)