Support unbiased gradients through particle filtering via stop-gradient resampling.

This implements the approach proposed by Adam Scibior, Vaden Masrani, and Frank Wood in "Differentiable Particle
Filtering without Modifying the Forward Pass" (2021, https://arxiv.org/abs/2106.10314).

PiperOrigin-RevId: 384336880

Author: davmre

tensorflow_probability/python/experimental/mcmc/particle_filter.py

@@ -94,6 +94,11 @@ def _default_trace_fn(state, kernel_results):
     The default behavior resamples particles when the current effective
     sample size falls below half the total number of particles.
     Default value: `tfp.experimental.mcmc.ess_below_threshold`.
+  unbiased_gradients: If `True`, use the stop-gradient
+    resampling trick of Scibior, Masrani, and Wood [{scibor_ref_idx}] to
+    correct for gradient bias introduced by the discrete resampling step. This
+    will generally increase the variance of stochastic gradients.
+    Default value: `True`.
   rejuvenation_kernel_fn: optional Python `callable` with signature
     `transition_kernel = rejuvenation_kernel_fn(target_log_prob_fn)`
     where `target_log_prob_fn` is a provided callable evaluating
@@ -112,7 +117,7 @@ def _default_trace_fn(state, kernel_results):
 
 
 @docstring_util.expand_docstring(
-    particle_filter_arg_str=particle_filter_arg_str)
+    particle_filter_arg_str=particle_filter_arg_str.format(scibor_ref_idx=2))
 def infer_trajectories(observations,
                        initial_state_prior,
                        transition_fn,
@@ -122,6 +127,7 @@ def infer_trajectories(observations,
                        proposal_fn=None,
                        resample_fn=weighted_resampling.resample_systematic,
                        resample_criterion_fn=smc_kernel.ess_below_threshold,
+                       unbiased_gradients=True,
                        rejuvenation_kernel_fn=None,
                        num_transitions_per_observation=1,
                        seed=None,
@@ -224,6 +230,9 @@ def observation_fn(_, state):
       filtering and smoothing: Fifteen years later.
       _Handbook of nonlinear filtering_, 12(656-704), 2009.
       https://www.stats.ox.ac.uk/~doucet/doucet_johansen_tutorialPF2011.pdf
+  [2] Adam Scibior, Vaden Masrani, and Frank Wood. Differentiable Particle
+      Filtering without Modifying the Forward Pass. _arXiv preprint
+      arXiv:2106.10314_, 2021. https://arxiv.org/abs/2106.10314
 
   """
   with tf.name_scope(name or 'infer_trajectories') as name:
@@ -242,6 +251,7 @@ def observation_fn(_, state):
          proposal_fn=proposal_fn,
          resample_fn=resample_fn,
          resample_criterion_fn=resample_criterion_fn,
+         unbiased_gradients=unbiased_gradients,
          rejuvenation_kernel_fn=rejuvenation_kernel_fn,
          num_transitions_per_observation=num_transitions_per_observation,
          trace_fn=_default_trace_fn,
@@ -265,7 +275,7 @@ def observation_fn(_, state):
 
 
 @docstring_util.expand_docstring(
-    particle_filter_arg_str=particle_filter_arg_str)
+    particle_filter_arg_str=particle_filter_arg_str.format(scibor_ref_idx=1))
 def particle_filter(observations,
                     initial_state_prior,
                     transition_fn,
@@ -275,6 +285,7 @@ def particle_filter(observations,
                     proposal_fn=None,
                     resample_fn=weighted_resampling.resample_systematic,
                     resample_criterion_fn=smc_kernel.ess_below_threshold,
+                    unbiased_gradients=True,
                     rejuvenation_kernel_fn=None,  # TODO(davmre): not yet supported. pylint: disable=unused-argument
                     num_transitions_per_observation=1,
                     trace_fn=_default_trace_fn,
@@ -324,6 +335,12 @@ def particle_filter(observations,
       `trace_criterion_fn==None`, this is computed from the final step;
       otherwise, each Tensor will have initial dimension `num_steps_traced`
       and stacks the traced results across all steps.
+
+  #### References
+
+  [1] Adam Scibior, Vaden Masrani, and Frank Wood. Differentiable Particle
+      Filtering without Modifying the Forward Pass. _arXiv preprint
+      arXiv:2106.10314_, 2021. https://arxiv.org/abs/2106.10314
   """
 
   init_seed, loop_seed = samplers.split_seed(seed, salt='particle_filter')
@@ -356,7 +373,8 @@ def particle_filter(observations,
     kernel = smc_kernel.SequentialMonteCarlo(
         propose_and_update_log_weights_fn=propose_and_update_log_weights_fn,
         resample_fn=resample_fn,
-        resample_criterion_fn=resample_criterion_fn)
+        resample_criterion_fn=resample_criterion_fn,
+        unbiased_gradients=unbiased_gradients)
 
     # Use `trace_scan` rather than `sample_chain` directly because the latter
     # would force us to trace the state history (with or without thinning),

tensorflow_probability/python/experimental/mcmc/sample_sequential_monte_carlo.py

@@ -541,7 +541,7 @@ def smc_body_fn(stage, state, smc_kernel_result):
                 smc_kernel_result.particle_info.log_scalings, axis=0))
             )
       (resampled_state,
-       resampled_particle_info), _ = weighted_resampling.resample(
+       resampled_particle_info), _, _ = weighted_resampling.resample(
            particles=(state, smc_kernel_result.particle_info),
            log_weights=log_weights,
            resample_fn=resample_fn,

tensorflow_probability/python/experimental/mcmc/sequential_monte_carlo_kernel.py

@@ -140,6 +140,7 @@ def __init__(self,
                propose_and_update_log_weights_fn,
                resample_fn=weighted_resampling.resample_systematic,
                resample_criterion_fn=ess_below_threshold,
+               unbiased_gradients=True,
                name=None):
     """Initializes a sequential Monte Carlo transition kernel.
 
@@ -180,11 +181,23 @@ def __init__(self,
         default behavior is to resample particles when the effective
         sample size falls below half of the total number of particles.
         Default value: `tfp.experimental.mcmc.ess_below_threshold`.
+      unbiased_gradients: If `True`, use the stop-gradient
+        resampling trick of Scibior, Masrani, and Wood [{scibor_ref_idx}] to
+        correct for gradient bias introduced by the discrete resampling step.
+        This will generally increase the variance of stochastic gradients.
+        Default value: `True`.
       name: Python `str` name for ops created by this kernel.
+
+    #### References
+
+    [1] Adam Scibior, Vaden Masrani, and Frank Wood. Differentiable Particle
+        Filtering without Modifying the Forward Pass. _arXiv preprint
+        arXiv:2106.10314_, 2021. https://arxiv.org/abs/2106.10314
     """
     self._propose_and_update_log_weights_fn = propose_and_update_log_weights_fn
     self._resample_fn = resample_fn
     self._resample_criterion_fn = resample_criterion_fn
+    self._unbiased_gradients = unbiased_gradients
     self._name = name or 'SequentialMonteCarlo'
 
   @property
@@ -203,6 +216,10 @@ def propose_and_update_log_weights_fn(self):
   def resample_criterion_fn(self):
     return self._resample_criterion_fn
 
+  @property
+  def unbiased_gradients(self):
+    return self._unbiased_gradients
+
   @property
   def resample_fn(self):
     return self._resample_fn
@@ -234,7 +251,6 @@ def one_step(self, state, kernel_results, seed=None):
         proposal_seed, resample_seed = samplers.split_seed(seed)
 
         state = WeightedParticles(*state)  # Canonicalize.
-        num_particles = ps.size0(state.log_weights)
 
         # Propose new particles and update weights for this step, unless it's
         # the initial step, in which case, use the user-provided initial
@@ -266,19 +282,26 @@ def one_step(self, state, kernel_results, seed=None):
         # needed---but we're ultimately interested in adaptive resampling
         # for statistical (not computational) purposes, so this isn't a
         # dealbreaker.
-        resampled_particles, resample_indices = weighted_resampling.resample(
-            state.particles,
-            state.log_weights,
-            self.resample_fn,
+        [
+            resampled_particles,
+            resample_indices,
+            weights_after_resampling
+        ] = weighted_resampling.resample(
+            particles=state.particles,
+            # The `stop_gradient` here does not affect discrete resampling
+            # (which is nondifferentiable anyway), but avoids canceling out
+            # the gradient signal from the 'target' log weights, as described in
+            # Scibior, Masrani, and Wood (2021).
+            log_weights=tf.stop_gradient(state.log_weights),
+            resample_fn=self.resample_fn,
+            target_log_weights=(normalized_log_weights
+                                if self.unbiased_gradients else None),
             seed=resample_seed)
-        uniform_weights = tf.fill(
-            ps.shape(state.log_weights),
-            value=-tf.math.log(tf.cast(num_particles, state.log_weights.dtype)))
         (resampled_particles,
          resample_indices,
          log_weights) = tf.nest.map_structure(
              lambda r, p: tf.where(do_resample, r, p),
-             (resampled_particles, resample_indices, uniform_weights),
+             (resampled_particles, resample_indices, weights_after_resampling),
              (state.particles, _dummy_indices_like(resample_indices),
               normalized_log_weights))
 

tensorflow_probability/python/experimental/mcmc/weighted_resampling.py

@@ -34,7 +34,8 @@
 ]
 
 
-def resample(particles, log_weights, resample_fn, seed=None):
+def resample(particles, log_weights, resample_fn, target_log_weights=None,
+             seed=None):
   """Resamples the current particles according to provided weights.
 
   Args:
@@ -47,21 +48,46 @@ def resample(particles, log_weights, resample_fn, seed=None):
       Use 'resample_independent' for independent resamples.
       Use 'resample_stratified' for stratified resampling.
       Use 'resample_systematic' for systematic resampling.
+    target_log_weights: optional float `Tensor` of the same shape and dtype as
+      `log_weights`, specifying the target measure on `particles` if this is
+      different from that implied by normalizing `log_weights`. The
+      returned `log_weights_after_resampling` will represent this measure. If
+      `None`, the target measure is implicitly taken to be the normalized
+      log weights (`log_weights - tf.reduce_logsumexp(log_weights, axis=0)`).
+      Default value: `None`.
     seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
 
   Returns:
     resampled_particles: Nested structure of `Tensor`s, matching `particles`.
     resample_indices: int `Tensor` of shape `[num_particles, b1, ..., bN]`.
+    log_weights_after_resampling: float `Tensor` of same shape and dtype as
+      `log_weights`, such that weighted sums of the resampled particles are
+      equal (in expectation over the resampling step) to weighted sums of
+      the original particles:
+      `E [ exp(log_weights_after_resampling) * some_fn(resampled_particles) ]
+      == exp(target_log_weights) * some_fn(particles)`.
+      If no `target_log_weights` was specified, the log weights after
+      resampling are uniformly equal to `-log(num_particles)`.
   """
   with tf.name_scope('resample'):
     num_particles = ps.size0(log_weights)
+    log_num_particles = tf.math.log(tf.cast(num_particles, log_weights.dtype))
+
+    # Normalize the weights and sample the ancestral indices.
     log_probs = tf.math.log_softmax(log_weights, axis=0)
     resampled_indices = resample_fn(log_probs, num_particles, (), seed=seed)
-    resampled_particles = tf.nest.map_structure(
-        lambda x: mcmc_util.index_remapping_gather(  # pylint: disable=g-long-lambda
-            x, resampled_indices, axis=0),
-        particles)
-  return resampled_particles, resampled_indices
+
+    gather_ancestors = lambda x: (  # pylint: disable=g-long-lambda
+        mcmc_util.index_remapping_gather(x, resampled_indices, axis=0))
+    resampled_particles = tf.nest.map_structure(gather_ancestors, particles)
+    if target_log_weights is None:
+      log_weights_after_resampling = tf.fill(ps.shape(log_weights),
+                                             -log_num_particles)
+    else:
+      importance_weights = target_log_weights - log_probs - log_num_particles
+      log_weights_after_resampling = tf.nest.map_structure(
+          gather_ancestors, importance_weights)
+  return resampled_particles, resampled_indices, log_weights_after_resampling
 
 
 # TODO(b/153689734): rewrite so as not to use `move_dimension`.

tensorflow_probability/python/experimental/mcmc/weighted_resampling_test.py

@@ -23,6 +23,7 @@
 import tensorflow_probability as tfp
 from tensorflow_probability.python.experimental.mcmc.weighted_resampling import _resample_using_log_points
 from tensorflow_probability.python.experimental.mcmc.weighted_resampling import _scatter_nd_batch
+from tensorflow_probability.python.experimental.mcmc.weighted_resampling import resample
 from tensorflow_probability.python.experimental.mcmc.weighted_resampling import resample_deterministic_minimum_error
 from tensorflow_probability.python.experimental.mcmc.weighted_resampling import resample_independent
 from tensorflow_probability.python.experimental.mcmc.weighted_resampling import resample_stratified
@@ -275,6 +276,32 @@ def test_resample_using_extremal_log_points(self):
         log_probs_end, sample_shape, log_points_almost_one)
     self.assertAllEqual(indices, tf.fill([n], n - 1))
 
+  def resample_with_target_distribution(self):
+    particles = np.linspace(0., 500., num=2500, dtype=np.float32)
+    log_weights = tfd.Poisson(20.).log_prob(particles)
+
+    # Resample particles to target a Poisson(20.) distribution.
+    new_particles, _, new_log_weights = resample(
+        particles, log_weights,
+        resample_fn=resample_systematic,
+        seed=test_util.test_seed(sampler_type='stateless'))
+    self.assertAllClose(tf.reduce_mean(new_particles), 20., atol=1e-2)
+    self.assertAllClose(
+        tf.reduce_sum(tf.nn.softmax(new_log_weights) * new_particles),
+        20.,
+        atol=1e-2)
+
+    # Reweight the resampled particles to target a Poisson(30.) distribution.
+    new_particles, _, new_log_weights = resample(
+        particles, log_weights,
+        resample_fn=resample_systematic,
+        target_log_weights=tfd.Poisson(30).log_prob(particles),
+        seed=test_util.test_seed(sampler_type='stateless'))
+    self.assertAllClose(tf.reduce_mean(new_particles), 20., atol=1e-2)
+    self.assertAllClose(
+        tf.reduce_sum(tf.nn.softmax(new_log_weights) * new_particles),
+        30., atol=1.)
+
   def maybe_compiler(self, f):
     if self.use_xla:
       return tf.function(f, autograph=False, jit_compile=True)