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

@@ -616,6 +616,63 @@ def build_loss_and_gradients(self, var_list):
     return build_score_rb_loss_and_gradients(self, var_list)
 
 
+# TODO: you can probably make another base class that implements a `sample` method?
+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,106 @@ 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):
+    """
+    """
+    p_log_prob = [0.0] * inference.n_samples
+    q_log_prob = [0.0] * inference.n_samples
+    r_log_prob = [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
+
+      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)
+
+        # Of course, this will evaluate to `True`. We just do this as a simple first pass.
+        if 'rsvi':
+          # --- RSVI
+
+          # Get variable shortnames
+          qz_class = qz.__class__
+          epsilon_likelihood = inference.rejection_sampler_vars[qz_class]['epsilon_likelihood']
+          reparam_func = inference.rejection_sampler_vars[qz_class]['reparam_func']
+          m = inference.rejection_sampler_vars[qz_class]['m']
+          alpha = qz.parameters['concentration']
+          beta = qz.parameters['rate']
+
+          # Sample
+
+          # TODO: pass in the real `qalpha` and `qbeta`
+          # TODO: pass in the real `qz`
+          epsilon = epsilon_likelihood.value()
+          sample = reparam_func(epsilon, alpha, beta)
+          eps_prob = epsilon_likelihood.prob(epsilon)
+          qz_prob = qz.prob(sample)
+          random_uniform = tf.random_uniform([])
+
+          # We need this line. However, let's just accept for now.
+          # if random_uniform * m * eps_prob <= qz_prob:
+
+          # RSVI ---
+        else:
+          z = qz_copy.value()
+
+        dict_swap[z] = sample
+
+        q_log_prob[s] += tf.reduce_sum(
+            inference.scale.get(z, 1.0) * qz_copy.log_prob(dict_swap[z]))
+        r_log_prob[s] += tf.reduce_sum(
+            inference.scale.get(z, 1.0) * epsilon_likelihood.log_prob(dict_swap[z]))
+
+      for z in six.iterkeys(inference.latent_vars):
+        z_copy = copy(z, dict_swap, scope=scope)
+        p_log_prob[s] += 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[s] += tf.reduce_sum(
+              inference.scale.get(x, 1.0) * x_copy.log_prob(dict_swap[x]))
+
+    p_log_prob = tf.reduce_mean(p_log_prob)
+    q_log_prob = tf.reduce_mean(q_log_prob)
+    r_log_prob = tf.reduce_mean(r_log_prob)
+
+    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())
+
+    if inference.logging:
+      tf.summary.scalar("loss/p_log_prob", p_log_prob,
+                        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])
+
+    loss = -(p_log_prob + q_entropy - reg_penalty)
+
+    # RSVI gradient components
+    model_grad = tf.gradients(p_log_prob, sample)[0]
+    q_entropy_grad = tf.gradients(q_entropy, var_list)
+    g_rep = [model_grad * grad for grad in tf.gradients(sample, var_list)]
+    g_cor = [p_log_prob * grad for grad in tf.gradients(q_log_prob - r_log_prob, var_list)]
+    grad_summands = zip(*[g_rep, g_cor, q_entropy_grad])
+
+    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

@@ -9,6 +9,7 @@
 
 from edward.inferences.inference import Inference
 from edward.models import RandomVariable
+from edward.optimizers import KucukelbirOptimizer
 from edward.util import get_session, get_variables
 
 
@@ -67,6 +68,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:
@@ -110,6 +113,14 @@ def initialize(self, optimizer=None, var_list=None, use_prettytensor=False,
         optimizer = tf.train.FtrlOptimizer(learning_rate)
       elif optimizer == 'rmsprop':
         optimizer = tf.train.RMSPropOptimizer(learning_rate)
+      elif optimizer == 'kucukelbir':
+        optimizer = KucukelbirOptimizer(
+          t=0.1,
+          delta=10e-3,
+          eta=1e-1,
+          s_n=tf.Variable([0., 0.], trainable=False),
+          n=tf.Variable(0., trainable=False)
+        )
       else:
         raise ValueError('Optimizer class not found:', optimizer)
     elif not isinstance(optimizer, tf.train.Optimizer):
@@ -151,7 +162,10 @@ 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 +175,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/optimizers/__init__.py

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

edward/optimizers/sgd.py

@@ -0,0 +1,35 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+class KucukelbirOptimizer:
+
+  """
+  Used for RSVI (Rejection-Sampling Variational Inference).
+
+  # TODO: add me
+  """
+
+  def __init__(self, t, delta, eta, s_n, n):
+    self.t = t
+    self.delta = delta
+    self.eta = eta
+    self.s_n = s_n
+    self.n = n
+
+  def apply_gradients(self, grads_and_vars, global_step=None):
+    self.n = tf.assign_add(self.n, 1.)
+    ops = []
+    for i, (grad, var) in enumerate(grads_and_vars):
+      updated_s_n = self.s_n[i].assign( (self.t * grad**2) + (1 - self.t) * self.s_n[i] )
+
+      p_n_first = self.eta * self.n**(-.5 + self.delta)
+      p_n_second = (1 + tf.sqrt(updated_s_n[i]))**(-1)
+      p_n = p_n_first * p_n_second
+
+      updated_var = var.assign_add(-p_n * grad)
+      ops.append(updated_var)
+    return ops