Rejection sampling variational inference

edward/__init__.py

@@ -14,7 +14,7 @@
     HMC, MetropolisHastings, SGLD, SGHMC, \
     KLpq, KLqp, ReparameterizationKLqp, ReparameterizationKLKLqp, \
     ReparameterizationEntropyKLqp, ScoreKLqp, ScoreKLKLqp, ScoreEntropyKLqp, \
-    ScoreRBKLqp, WakeSleep, GANInference, BiGANInference, WGANInference, \
+    ScoreRBKLqp, RejectionSamplingKLqp, WakeSleep, GANInference, BiGANInference, WGANInference, \
     ImplicitKLqp, MAP, Laplace, complete_conditional, Gibbs
 from edward.models import RandomVariable
 from edward.util import check_data, check_latent_vars, copy, dot, \
@@ -52,6 +52,7 @@
     'ScoreKLKLqp',
     'ScoreEntropyKLqp',
     'ScoreRBKLqp',
+    'RejectionSamplingKLqp',
     'WakeSleep',
     'GANInference',
     'BiGANInference',

edward/inferences/__init__.py

@@ -42,6 +42,7 @@
     'ScoreKLKLqp',
     'ScoreEntropyKLqp',
     'ScoreRBKLqp',
+    'RejectionSamplingKLqp',
     'Laplace',
     'MAP',
     'MetropolisHastings',

edward/inferences/inference.py

@@ -123,7 +123,6 @@ def run(self, variables=None, use_coordinator=True, *args, **kwargs):
         Passed into `initialize`.
     """
     self.initialize(*args, **kwargs)
-
     if variables is None:
       init = tf.global_variables_initializer()
     else:
@@ -144,6 +143,7 @@ def run(self, variables=None, use_coordinator=True, *args, **kwargs):
 
     for _ in range(self.n_iter):
       info_dict = self.update()
+      print(info_dict)
       self.print_progress(info_dict)
 
     self.finalize()

edward/inferences/klpq.py

@@ -32,7 +32,7 @@ class KLpq(VariationalInference):
 
   with respect to $\\theta$.
 
-  In conditional inference, we infer $z` in $p(z, \\beta
+  In conditional inference, we infer $z$ in $p(z, \\beta
   \mid x)$ while fixing inference over $\\beta$ using another
   distribution $q(\\beta)$. During gradient calculation, instead
   of using the model's density

edward/inferences/klqp.py

@@ -6,7 +6,8 @@
 import tensorflow as tf
 
 from edward.inferences.variational_inference import VariationalInference
-from edward.models import RandomVariable
+from edward.models import RandomVariable, Gamma
+from edward.samplers import GammaRejectionSampler
 from edward.util import copy, get_descendants
 
 try:
@@ -616,6 +617,62 @@ def build_loss_and_gradients(self, var_list):
     return build_score_rb_loss_and_gradients(self, var_list)
 
 
+class RejectionSamplingKLqp(VariationalInference):
+
+    """
+    """
+
+    def __init__(self, latent_vars=None, data=None, rejection_sampler_vars=None):
+      """Create an inference algorithm.
+
+      # TODO: update me
+
+      Args:
+        latent_vars: list of RandomVariable or
+                     dict of RandomVariable to RandomVariable.
+          Collection of random variables to perform inference on. If
+          list, each random variable will be implictly optimized using a
+          `Normal` random variable that is defined internally with a
+          free parameter per location and scale and is initialized using
+          standard normal draws. The random variables to approximate
+          must be continuous.
+      """
+      if isinstance(latent_vars, list):
+        with tf.variable_scope(None, default_name="posterior"):
+          latent_vars_dict = {}
+          continuous = \
+              ('01', 'nonnegative', 'simplex', 'real', 'multivariate_real')
+          for z in latent_vars:
+            if not hasattr(z, 'support') or z.support not in continuous:
+              raise AttributeError(
+                  "Random variable {} is not continuous or a random "
+                  "variable with supported continuous support.".format(z))
+            batch_event_shape = z.batch_shape.concatenate(z.event_shape)
+            loc = tf.Variable(tf.random_normal(batch_event_shape))
+            scale = tf.nn.softplus(
+                tf.Variable(tf.random_normal(batch_event_shape)))
+            latent_vars_dict[z] = Normal(loc=loc, scale=scale)
+          latent_vars = latent_vars_dict
+          del latent_vars_dict
+      super(RejectionSamplingKLqp, self).__init__(latent_vars, data)
+      self.rejection_sampler_vars = rejection_sampler_vars
+
+    def initialize(self, n_samples=1, *args, **kwargs):
+      """Initialize inference algorithm. It initializes hyperparameters
+      and builds ops for the algorithm's computation graph.
+
+      Args:
+        n_samples: int, optional.
+          Number of samples from variational model for calculating
+          stochastic gradients.
+      """
+      self.n_samples = n_samples
+      return super(RejectionSamplingKLqp, self).initialize(*args, **kwargs)
+
+    def build_loss_and_gradients(self, var_list):
+      return build_rejection_sampling_loss_and_gradients(self, var_list)
+
+
 def build_reparam_loss_and_gradients(inference, var_list):
   """Build loss function. Its automatic differentiation
   is a stochastic gradient of
@@ -1127,3 +1184,90 @@ def build_score_rb_loss_and_gradients(inference, var_list):
   grads_vars.extend(model_vars)
   grads_and_vars = list(zip(grads, grads_vars))
   return loss, grads_and_vars
+
+
+def build_rejection_sampling_loss_and_gradients(inference, var_list, epsilon=None):
+    """
+    """
+    rej_samplers = {
+      Gamma: GammaRejectionSampler
+    }
+
+    rep = [0.0] * inference.n_samples
+    cor = [0.0] * inference.n_samples
+    base_scope = tf.get_default_graph().unique_name("inference") + '/'
+    for s in range(inference.n_samples):
+      # Form dictionary in order to replace conditioning on prior or
+      # observed variable with conditioning on a specific value.
+      scope = base_scope + tf.get_default_graph().unique_name("sample")
+      dict_swap = {}
+      for x, qx in six.iteritems(inference.data):
+        if isinstance(x, RandomVariable):
+          if isinstance(qx, RandomVariable):
+            qx_copy = copy(qx, scope=scope)
+            dict_swap[x] = qx_copy.value()
+          else:
+            dict_swap[x] = qx
+
+      p_log_prob = 0.
+      q_log_prob = 0.
+      r_log_prob = 0.
+
+      for z, qz in six.iteritems(inference.latent_vars):
+        # Copy q(z) to obtain new set of posterior samples.
+        qz_copy = copy(qz, scope=scope)
+        sampler = rej_samplers[qz_copy.__class__](density=qz)
+
+        if epsilon is not None:  # temporary
+          pass
+        else:
+          dict_swap[z] = qz_copy.value()
+          print('sample:', dict_swap[z])
+          epsilon = sampler.h_inverse(dict_swap[z])
+
+        dict_swap[z] = sampler.h(epsilon)
+        q_log_prob += tf.reduce_sum(
+          inference.scale.get(z, 1.0) * qz_copy.log_prob(dict_swap[z]))
+        r_log_prob += -tf.log(tf.gradients(dict_swap[z], epsilon))
+
+      for z in six.iterkeys(inference.latent_vars):
+        z_copy = copy(z, dict_swap, scope=scope)
+        p_log_prob += tf.reduce_sum(
+          inference.scale.get(z, 1.0) * z_copy.log_prob(dict_swap[z]))
+
+      for x in six.iterkeys(inference.data):
+        if isinstance(x, RandomVariable):
+          x_copy = copy(x, dict_swap, scope=scope)
+          p_log_prob += tf.reduce_sum(
+            inference.scale.get(x, 1.0) * x_copy.log_prob(dict_swap[x]))
+
+      rep[s] = p_log_prob
+      cor[s] = tf.stop_gradient(p_log_prob) * (q_log_prob - r_log_prob)
+
+    rep = tf.reduce_mean(rep)
+    cor = tf.reduce_mean(cor)
+    q_entropy = tf.reduce_sum([
+      tf.reduce_sum(qz.entropy())
+      for z, qz in six.iteritems(inference.latent_vars)])
+    reg_penalty = tf.reduce_sum(tf.losses.get_regularization_losses())
+
+    loss = -(rep + q_entropy - reg_penalty)
+
+    if inference.logging:
+      tf.summary.scalar("loss/reparam_objective", rep,
+                        collections=[inference._summary_key])
+      tf.summary.scalar("loss/correction_term", cor,
+                        collections=[inference._summary_key])
+      tf.summary.scalar("loss/q_entropy", q_entropy,
+                        collections=[inference._summary_key])
+      tf.summary.scalar("loss/reg_penalty", reg_penalty,
+                        collections=[inference._summary_key])
+
+    g_rep = tf.gradients(rep, var_list)
+    g_cor = tf.gradients(cor, var_list)
+    g_entropy = tf.gradients(q_entropy, var_list)
+
+    grad_summands = zip(*[g_rep, g_cor, g_entropy])
+    grads = [tf.reduce_sum(summand) for summand in grad_summands]
+    grads_and_vars = list(zip(grads, var_list))
+    return loss, grads_and_vars

edward/inferences/variational_inference.py

@@ -67,6 +67,8 @@ def initialize(self, optimizer=None, var_list=None, use_prettytensor=False,
 
     self.loss, grads_and_vars = self.build_loss_and_gradients(var_list)
 
+    self.grads_and_vars = grads_and_vars
+
     if self.logging:
       tf.summary.scalar("loss", self.loss, collections=[self._summary_key])
       for grad, var in grads_and_vars:
@@ -151,7 +153,9 @@ def update(self, feed_dict=None):
         feed_dict[key] = value
 
     sess = get_session()
-    _, t, loss = sess.run([self.train, self.increment_t, self.loss], feed_dict)
+    # _, t, loss = sess.run([self.train, self.increment_t, self.loss], feed_dict)
+    # TODO: delete me
+    _, t, loss, grads_and_vars_debug = sess.run([self.train, self.increment_t, self.loss, self.grads_and_vars], feed_dict)
 
     if self.debug:
       sess.run(self.op_check, feed_dict)
@@ -161,7 +165,7 @@ def update(self, feed_dict=None):
         summary = sess.run(self.summarize, feed_dict)
         self.train_writer.add_summary(summary, t)
 
-    return {'t': t, 'loss': loss}
+    return {'t': t, 'loss': loss, 'grads_and_vars_debug': grads_and_vars_debug}
 
   def print_progress(self, info_dict):
     """Print progress to output.

edward/samplers/__init__.py

@@ -0,0 +1,15 @@
+"""
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from edward.samplers.rejection import *
+
+from tensorflow.python.util.all_util import remove_undocumented
+
+_allowed_symbols = [
+    'GammaRejectionSampler',
+]
+
+remove_undocumented(__name__, allowed_exception_list=_allowed_symbols)

edward/samplers/rejection.py

@@ -0,0 +1,34 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+
+import tensorflow as tf
+
+
+class GammaRejectionSampler:
+
+  # As implemented in https://github.com/blei-lab/ars-reparameterization/blob/master/gamma/demo.ipynb
+
+  def __init__(self, density):
+    self.alpha = density.parameters['concentration']
+    self.beta = density.parameters['rate']
+
+  def h(self, epsilon):
+    a = self.alpha - (1. / 3)
+    b = tf.sqrt(9 * self.alpha - 3)
+    c = 1 + (epsilon / b)
+    d = a * c**3
+    return d / self.beta
+
+  def h_inverse(self, z):
+    a = self.alpha - (1. / 3)
+    b = tf.sqrt(9 * self.alpha - 3)
+    c = self.beta * z / a
+    d = c**(1 / 3)
+    return b * (d - 1)
+
+  @staticmethod
+  def log_prob_s(epsilon):
+    return -0.5 * (tf.log(2 * math.pi) + epsilon**2)