# Differentiable Particle Filtering
without Modifying the Forward Pass

Adam Ścibior1,2
<EMAIL_ADDRESS>
&Frank Wood1,2,3
<EMAIL_ADDRESS>

1University of British Columbia 2Inverted AI 3Mila

###### Abstract

Particle filters are not compatible with automatic differentiation due to the
presence of discrete resampling steps. While known estimators for the score
function, based on Fisher’s identity, can be computed using particle filters,
up to this point they required manual implementation. In this paper we show
that such estimators can be computed using automatic differentiation, after
introducing a simple correction to the particle weights. This correction
utilizes the stop-gradient operator and does not modify the particle filter
operation on the forward pass, while also being cheap and easy to compute.
Surprisingly, with the same correction automatic differentiation also produces
good estimators for gradients of expectations under the posterior. We can
therefore regard our method as a general recipe for making particle filters
differentiable. We additionally show that it produces desired estimators for
second-order derivatives and how to extend it to further reduce variance at
the expense of additional computation.

Differentiable Particle Filtering
without Modifying the Forward Pass

Adam Ścibior1,2<EMAIL_ADDRESS>Frank Wood1,2,3<EMAIL_ADDRESS>
1Inverted AI 2University of British Columbia 3Mila

##  1 Introduction

Sequential latent variable models can be learned by gradient ascent, applying
automatic differentiation (AD) to the log-marginal likelihood estimators
obtained with particle filters. Examples include robotic localization
(Jonschkowski et al., 2018), world-modeling in reinforcement learning (Igl et
al., 2018), and deep generative modeling of music and speech (Maddison et al.,
2017). Unfortunately, the discrete resampling steps allowing particle filters
to adaptively focus computational resources on promising paths cause problems
for automatic differentiation, with the usual algorithms producing gradient
estimators that suffer from excessive variance. In current practice, the
dependence of the resampling probabilities on model parameters is typically
ignored (Le et al., 2018a; Maddison et al., 2017; Naesseth et al., 2018),
introducing bias that was shown not to vanish asymptotically (Corenflos et
al., 2021) and impede learning.

Several recent articles (Corenflos et al., 2021; Zhu et al., 2020; Karkus et
al., 2018) have proposed continuous relaxations of the discrete resampling
steps to avoid this problem, with different trade-offs and varying degrees of
complexity, as discussed in detail in Section 6. At the same time, there exist
consistent estimators of the gradient of log-marginal likelihood that can be
obtained with usual discrete resampling (Poyiadjis et al., 2011), but until
now these could not be obtained by applying AD to the log-marginal likelihood
estimator produced by a particle filter.

In this work we show that a simple and computationally inexpensive
modification to an otherwise entirely standard “basic” particle filter
implementation allows AD to produce the estimators derived by Poyiadjis et al.
(2011). The modification, shown in Algorithm 1, can be thought of as stopping
the gradients of weights flowing into the resampling distribution, instead
directing them into the resulting weights. Crucially, it does not modify the
estimators produced by the particle filter on the forward pass, only changing
the gradients computed by AD. We also show that with our modification AD
produces sensible, easy to interpret, and consistent estimators for gradients
of expectations under the posterior, which are additionally unbiased if the
estimator obtained on the forward pass is unbiased. We can therefore regard
our modification as a general technique for making a particle filter
differentiable, without modifying its behavior on the forward pass. While
existing differentiable resampling mechanisms have their merits, the
simplicity of our approach is unmatched, so we advocate trying it first before
attempting more sophisticated alternatives.

Additionally, in Section 4 we show that our results extend to higher-order
derivatives computed by repeated AD and that the same method applies to
marginal particle filters (Klaas et al., 2012), recovering the more expensive
but lower variance estimator proposed by Poyiadjis et al. (2011). We also
briefly discuss the problem of jointly learning a proposal distribution with
the model, which is largely left for future work.

##  2 Preliminaries

Algorithm 1 Forward pass of the standard particle filter (left), compared with
our differentiable particle filter (right). The crucial modification required
by our method is highlighted in yellow, where the $\bot$ denotes the stop-
gradient operator. The application of stop-gradient inside the resampling
function is typically redundant, because automatic differentiation libraries
do not propagate gradients through samples from discrete distributions. We
choose the symbol $v$ to denote the weights of the differentiable particle
filter, to distinguish them from the traditionally computed weights in the
equations throughout the paper. Where $i$ is not bounded it denotes an
operation applied across all $i\in 1:N$, and we use colons to index sequences,
e.g. $x_{1:T}=(x_{1},\dots,x_{T})$.

function PF($p$, $q$, $N$, $T$)

$x_{0}^{i}\sim p(x_{0})$

$w_{0}^{i}=\overline{w}_{0}^{i}=\frac{1}{N}$

for $t\in 1:T$ do

if resampling condition then

$a^{1:N}_{t}\sim R(\overline{w}^{1:N}_{t-1})$

$\tilde{w}^{i}_{t}=\frac{1}{N}$

else

$a^{i}_{t}=i$

$\tilde{w}^{i}_{t}=\overline{w}^{i}_{t-1}$

$x^{i}_{t}\sim q_{\phi}(x_{t}|x^{a_{t}^{i}}_{t-1})$

$w^{i}_{t}=\tilde{w}^{i}_{t-1}\frac{p_{\theta}(x^{i}_{t},y_{t}|x^{a_{t}^{i}}_{t-1})}{q_{\phi}(x^{i}_{t}|x^{a_{t}^{i}}_{t-1})}$

$W_{t}=\sum_{i=1}^{N}w^{i}_{t}$

$\overline{w}^{i}_{t}=w^{i}_{t}/W_{t}$

$\hat{Z}_{\text{PF}}=\prod_{t=1}^{T}W_{t}$

function DPF($p$, $q$, $N$, $T$)

$x_{0}^{i}\sim p(x_{0})$

$v_{0}^{i}=\overline{v}_{0}^{i}=\frac{1}{N}$

for $t\in 1:T$ do

if resampling condition then

$a^{1:N}_{t}\sim R(\bot\left({\overline{v}^{1:N}_{t-1}}\right))$

$\tilde{v}^{i}_{t}=\frac{1}{N}$
$\displaystyle{\overline{v}^{a^{i}_{t}}_{t-1}/{\bot\left({\overline{v}^{a^{i}_{t}}_{t-1}}\right)}}$

else

$a^{i}_{t}=i$

$\tilde{v}^{i}_{t}=\overline{v}^{i}_{t-1}$

$x^{i}_{t}\sim q_{\phi}(x_{t}|x^{a_{t}^{i}}_{t-1})$

$v^{i}_{t}=\tilde{v}^{i}_{t-1}\frac{p_{\theta}(x^{i}_{t},y_{t}|x^{a_{t}^{i}}_{t-1})}{q_{\phi}(x^{i}_{t}|x^{a_{t}^{i}}_{t-1})}$

$V_{t}=\sum_{i=1}^{N}v^{i}_{t}$

$\overline{v}^{i}_{t}=v^{i}_{t}/V_{t}$

$\hat{Z}_{\text{DPF}}=\prod_{t=1}^{T}V_{t}$

We assume familiarity with automatic differentiation (Baydin et al., 2018),
including its application to stochastic computations (Schulman et al., 2016),
but provide some additional exposition to those and other topics in the
supplementary materials.

###  2.1 Particle filters

We assume familiarity with particle filters (Chopin and Papaspiliopoulos,
2020; Doucet and Johansen, 2009; Doucet et al., 2001), using notation defined
in Algorithm 1. Our results apply with any choice of the resampling scheme $R$
(Murray et al., 2016).

We assume that standard conditions for the existing convergence results
(Moral, 2004; Chopin and Papaspiliopoulos, 2020) are satisfied, specifically
for consistency of estimators of the log-marginal likelihood and expectations
under the posterior, and for unbiasedness of estimators of expectations under
the unnormalized posterior. For score function estimation with particle
filters, we use the estimators derived by Poyiadjis et al. (2011) using
Fisher’s identity, in particular

$\displaystyle\nabla_{\theta}\log
p_{\theta}(y_{1:T})\approx\sum_{i=1}^{N}\overline{w}_{T}^{i}\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T}),$ (1)

where $\tilde{x}_{1:t}^{i}$ is the ancestral lineage of a particle $i$ at time
$t$, satisfying the recursive equation
$\tilde{x}_{1:t}^{i}=(\tilde{x}_{1:t-1}^{a_{t}^{i}},x_{t}^{i}$). Our key
contribution is a construction that enables obtaining this estimator through
automatic differentiation of the $\log\hat{Z}$ estimate from a particle
filter, without affecting estimators produced on the forward pass.

###  2.2 The stop-gradient operator

The stop-gradient operator is a standard component of automatic
differentiation libraries, called stop_gradient in Tensorflow and detach in
PyTorch, producing an expression that evaluates to its normal value on the
forward pass but produces no gradient. In this section we introduce a minimal
formalism required to reason about gradient estimators produced by AD from
complex expressions involving stop-gradient.111Note that there doesn’t exist a
$\mathbb{R}\rightarrow\mathbb{R}$ function with the properties of stop-
gradient, so extra care needs to be taken when manipulating expressions
containing it. Following Foerster et al. (2018), we denote the stop-gradient
with a $\bot$, for example $\nabla\bot\left({x}\right)=0$.

We refer to expressions containing stop-gradient as _surrogate expressions_
and we need to distinguish between the results obtained by evaluating them on
the forward pass and the expressions obtained by applying AD. Following van
Krieken et al. (2021), we denote evaluation of an expression $E$ with an
overhead right arrow $\overrightarrow{E}$. Operationally, evaluation removes
all instances of stop-gradient, provided that no gradient operator is acting
on them. Gradients satisfy the usual chain rules, except that gradients of any
expression wrapped in stop-gradient are zero. The key equations are:

$\displaystyle\overrightarrow{f(E_{1},\dots,E_{n})}$
$\displaystyle=f(\overrightarrow{E_{1}},\dots,\overrightarrow{E_{2}})$ (2)
$\displaystyle\overrightarrow{\bot\left({E}\right)}$
$\displaystyle=\overrightarrow{E}$ (3)
$\displaystyle\nabla{\bot\left({E}\right)}$ $\displaystyle=0$ (4)
$\displaystyle\overrightarrow{\nabla{E}}$
$\displaystyle=\nabla{E}\quad\text{if}\ \bot\notin E,$ (5)

where $f$ is some fixed differentiable function and $\bot\notin E$ means that
stop-gradient is not applied anywhere in expression $E$, including its
subexpressions. Crucially, in general
$\nabla\overrightarrow{E}\neq\overrightarrow{\nabla E}$, and
$\overrightarrow{E_{1}}=\overrightarrow{E_{2}}$ _does not imply_
$\overrightarrow{\nabla E_{1}}=\overrightarrow{\nabla E_{2}}$.

We use this calculus to derive formulas for estimators obtained by algorithms
utilizing stop-gradient in automatic differentiation, which we generally
accomplish by using the rules above, along with the chain rule for the
gradient operator, to push the gradients inside, past any stop-gradient
operators, and then evaluating the resulting expression. Note that in this
calculus the gradient operator produces expressions rather than values, so
that we can model repeated automatic differentiation. Expressions without
stop-gradient can be equated with their usual interpretations. Throughout the
paper we are careful to avoid using integrals of expressions involving stop-
gradient, as we consider those undefined. We also present a short tutorial on
using this calculus in the appendix.

Surrogate expressions are commonly used when computing gradients of
computation involving non-reparameterizable random variables, which is
typically implemented using surrogates of the form
$f_{\theta}(x)+\bot\left({f_{\theta}(x)}\right)\log p_{\theta}(x)$. The
resulting gradient estimator is known in the reinforcement learning literature
as REINFORCE (Williams, 1992), and more broadly referred to as the score
function method or the likelihood ratio method (Fu, 2006). We refrain from
using those terms to avoid confusion with the score function
$\nabla_{\theta}\log p_{\theta}(y_{1:T})$ and the likelihood
$p_{\theta}(y_{1:T})$ associated with the sequential latent variable model.
Foerster et al. (2018) provide a method for constructing alternative
surrogates, called DiCE, which produces the same gradient estimators but has
the additional advantage of preserving the results of evaluating the forward
pass, as well as producing unbiased estimators for derivatives of higher
order. As an illustrative example, we use the calculus introduced above to
compute the gradient estimator produced by DiCE

$\displaystyle\overrightarrow{\nabla_{\theta}\left(\frac{p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}f_{\theta}(x)\right)}=$
(6) $\displaystyle
f_{\theta}(x)\overrightarrow{\frac{\nabla_{\theta}p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}}+\overrightarrow{\frac{p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}}\nabla_{\theta}f_{\theta}(x)=$
$\displaystyle f_{\theta}(x)\nabla_{\theta}\log
p_{\theta}(x)+\nabla_{\theta}f_{\theta}(x).$

Since it produces the same gradient estimators as REINFORCE, a.k.a. the score
function method, a.k.a. the likelihood ratio method, we use the name DiCE
throughout the paper to encapsulate all those methods.

##  3 Method

In this section we derive and analyze the gradient estimators obtained by
applying AD to two different versions of a particle filter, namely the
standard implementation (PF) and our modified version (DPF), contrasted in
Algorithm 1. To avoid further complicating notation, throughout the paper we
assume that resampling is performed at each time step. These two particle
filters are equivalent on the forward pass, in particular
$\overrightarrow{w_{t}^{i}}=\overrightarrow{v_{t}^{i}}$ and
$\overrightarrow{\hat{Z}_{\text{PF}}}=\overrightarrow{\hat{Z}_{\text{DPF}}}$,
and similarly $\bot\left({w_{t}^{i}}\right)=\bot\left({v_{t}^{i}}\right)$ and
$\bot\left({\hat{Z}_{\text{PF}}}\right)=\bot\left({\hat{Z}_{\text{DPF}}}\right)$.
The distinction is only material when computing gradients of such expressions,
when that is not the case we sometimes drop the subscript and simply write
$\hat{Z}$ to avoid clutter.

###  3.1 Estimating the score function

To provide a gentle exposition, we first consider simple importance sampling
in a non-sequential latent variable model, without any resampling or ancestral
variables. A simple way to obtain score function estimates is to apply
automatic differentiation to a program that estimates the log-marginal
likelihood with importance sampling, as is done in IWAE (Burda et al., 2016).
The forward pass of such a program defines the following objective

$\displaystyle\mathcal{L}_{\text{IWAE}}=\mathop{\mathbb{E}}_{x^{i}\sim
q_{\phi}(x)}\left[{\log\frac{1}{N}\sum_{i=1}^{N}\frac{p_{\theta}(x^{i},y)}{q_{\phi}(x^{i})}}\right].$
(7)

Automatic differentiation then computes gradients with respect to the model
parameters $\theta$ as

$\displaystyle\nabla_{\theta}\mathcal{L}_{\text{IWAE}}=\mathop{\mathbb{E}}_{x^{i}\sim
q_{\phi}(x)}\left[{\sum_{i=1}^{N}\overline{w}^{i}\nabla_{\theta}\log
p_{\theta}(x^{i},y)}\right],$ (8)

where $w^{i}=\frac{p_{\theta}(x^{i},y)}{q_{\phi}(x^{i})}$ and
$\overline{w}^{i}=w^{i}/\sum_{j=1}^{N}w^{j}$.

This is exactly the same estimator we would obtain by using Fisher’s identity
with a self-normalizing importance sampler to approximate the posterior.
However, the two methods are no longer equivalent if the proposal distribution
also depends on model parameters, such as when using the prior $p_{\theta}(x)$
as the proposal. Applying Fisher’s identity produces the same estimator as Eq.
8 whether $q$ depends on $\theta$ or not. With automatic differentiation,
since the sampling distribution depends on $\theta$ and assuming it can not be
reparameterized, we use DiCE and obtain the following gradient estimator

$\displaystyle\log\hat{Z}\ \sum_{i=1}^{N}\nabla_{\theta}\log
q_{\theta}(x^{i})+\sum_{i=1}^{N}\overline{w}^{i}\nabla_{\theta}\log\frac{p_{\theta}(x^{i},y)}{q_{\theta}(x^{i})},$
(9)

where $\hat{Z}=\sum_{i=1}^{N}w^{i}$. The first term inside the expectation is
the problematic one, and it tends to either be dropped, introducing bias, or
necessitates the use of baselines or other variance reduction methods. When
gradients of $q$ are stopped, the first term disappears, along with the
denominator in the second one, recovering the estimator from Eq. 8. We can
show this formally using the calculus introduced in Section 2.2. We introduce
a stop-gradient weight
$v^{i}={p_{\theta}(x^{i},y)}/{\bot\left({q_{\theta}(x^{i})}\right)}$ to
distinguish it from $w^{i}={p_{\theta}(x^{i},y)}/{q_{\theta}(x^{i})}$, noting
that they evaluate to the same value on the forward pass, that is
$\overrightarrow{v^{i}}=w^{i}$. Then

$\displaystyle\overrightarrow{\nabla_{\theta}\log\frac{1}{N}\sum_{i=1}^{N}\frac{p_{\theta}(x^{i},y)}{\bot\left({q_{\theta}(x^{i})}\right)}}=\overrightarrow{\nabla_{\theta}\log\frac{1}{N}\sum_{i=1}^{N}v^{i}}$
(10)
$\displaystyle\overrightarrow{\frac{\frac{1}{N}\sum_{i=1}^{N}\nabla_{\theta}v^{i}}{\frac{1}{N}\sum_{i=1}^{N}v^{i}}}=\overrightarrow{\sum_{i=1}^{N}\overline{v}^{i}\frac{\nabla_{\theta}v^{i}}{v^{i}}}=$
(11)
$\displaystyle\overrightarrow{\sum_{i=1}^{N}\overline{v}^{i}\nabla_{\theta}\log
v^{i}}=\sum_{i=1}^{N}\overline{w}^{i}\nabla_{\theta}\log p_{\theta}(x^{i},y).$
(12)

We emphasize that the use of this estimator is justified by Fisher’s identity,
independently of any stop-gradient considerations required to obtain it
through automatic differentiation. This lets us access known theoretical
results, in particular implying consistency under standard assumptions.

Those results suggests that _when the proposal distribution depends on the
model parameters_ , stop-gradient should be applied to it, in particular if it
can not be reparameterized. In IWAE, or sequential importance sampling, this
problem is usually avoided by having separate parameters for the proposal
distribution. In particle filters, however, even if the proposal distribution
uses separate parameters, the computation involves discrete ancestral indices,
which are sampled according to weights, which in turn depend on the model
parameters. Stopping those gradients changes the resampling distribution,
which requires multiplying the post-resampling weight with the ratio of
corresponding densities (Liu, 2004, Section 3.4.4), in this case
$\displaystyle{\overline{v}^{a^{i}_{t}}_{t-1}/{\bot\left({\overline{v}^{a^{i}_{t}}_{t-1}}\right)}}$.
With DPF as defined in Algorithm 1, we have

$\displaystyle\overrightarrow{\nabla_{\theta}\log\hat{Z}_{\text{DPF}}}=\sum_{i=1}^{N}\overline{w}_{T}^{i}\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T}),$ (13)

recovering the estimator from Eq. 1. Detailed calculations are presented in
the supplementary materials. This is known to be a consistent estimator of the
score function (Poyiadjis et al., 2011) and we obtain it by applying AD to a
log-likelihood estimator produced by a particle filter with minor correction
to its weights. Those corrections are easy to implement, their computational
overhead is small, and they do not affect the estimators produced on the
forward pass.

In contrast, applying AD to a standard PF implementation, using DiCE to
account for the dependence of resampling distribution on $\theta$, produces
the following gradient estimators (Le et al., 2018a, Eq. 31)

$\displaystyle\overrightarrow{\nabla_{\theta}\left(\prod_{t=1}^{T}\prod_{i=1}^{N}\frac{w_{t}^{a_{t}^{i}}}{\bot\left({w_{t}^{a_{t}^{i}}}\right)}\log\hat{Z}_{\text{PF}}\right)}=$
(14)
$\displaystyle\log\hat{Z}_{\text{PF}}\sum_{i=1}^{N}\sum_{t=1}^{T}\nabla_{\theta}\log\overline{w}_{t}^{i}+\nabla_{\theta}\log\hat{Z}_{\text{PF}}=$
(15) $\displaystyle\sum_{i=1}^{N}\sum_{t=1}^{T}\log\hat{Z}\
\nabla_{\theta}\log\overline{w}_{t}^{i}+\overline{w}_{t}^{i}\nabla_{\theta}\log
p_{\theta}(x_{t}^{i},y_{t}|x_{t-1}^{a_{t}^{i}}).$ (16)

The first term is typically dropped due to high-variance, introducing bias
which was shown not to disappear in the infinite particle limit (Corenflos et
al., 2021). Comparing with Eq. 13, this estimator does not track dependencies
across time beyond a single step.

###  3.2 Estimating expectations

While our primary goal is to obtain score function estimates by applying AD to
log-likelihood estimates, particle filters can also be used to estimate
expectations under filtering or smoothing distributions and we can apply AD to
estimate gradients of such expectations as well. In this section we show that
with the same DPF implementation, shown in Algorithm 1, AD produces sensible
and easy to interpret gradient estimators with desirable theoretical
properties. In the main text we only state the results, providing derivations
and proofs in the supplementary materials.

Consider a function $f_{\theta}:\mathcal{X}^{T}\rightarrow\mathcal{R}$,
differentiable in $\theta$ for all $x_{1:T}$. A particle filter yields a
consistent estimator for its expectation under the posterior

$\displaystyle\lim_{N\rightarrow\infty}\sum_{i=1}^{N}\overline{w}_{T}^{i}f_{\theta}(\tilde{x}_{1:T}^{i})=\mathop{\mathbb{E}}_{x\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{f_{\theta}(x_{1:T})}\right].$ (17)

Applying AD to DPF produces the following consistent estimator for the
gradient of this expectation

$\displaystyle\nabla_{\theta}\mathop{\mathbb{E}}_{x\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{f_{\theta}(x_{1:T})}\right]\approx\overrightarrow{\nabla_{\theta}\sum_{i=1}^{N}\overline{v}_{T}^{i}f_{\theta}(\tilde{x}_{1:T}^{i})}=$
$\displaystyle\sum_{i=1}^{N}\overline{w}_{T}^{i}\Big{(}\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})+$
$\displaystyle\quad\left(f_{\theta}(\tilde{x}_{1:T}^{i})-\overline{f}_{\theta}\right)\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})\Big{)},$ (18)

where
$\overline{f_{\theta}}=\sum_{i=1}^{N}\overline{w}_{T}^{i}f_{\theta}(\tilde{x}_{1:T}^{i})$.
The first term is the average value of the partial derivative of $f_{\theta}$
at the locations specified by surviving particles, while the second one can be
interpreted as a DiCE-like log-derivative correction under the posterior
distribution of the full ancestral lineage, with a variance-reducing baseline
provided by the average value of $f_{\theta}$ obtained on the forward pass.

In contrast, applying AD to the standard PF implementation with DiCE yields

$\displaystyle\overline{f}_{\theta}\sum_{t=1}^{T}\sum_{i=1}^{N}\nabla_{\theta}\log\overline{w}_{t-1}^{a_{t}^{i}}+\sum_{i=1}^{N}\overline{w}_{T}^{i}\Bigg{(}\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\
+$
$\displaystyle\left(f_{\theta}(\tilde{x}_{1:T}^{i})-\overline{f}_{\theta}\right)\nabla_{\theta}\log
p_{\theta}(x_{T},y_{T}|x_{T-1}^{a_{T}^{i}})\Bigg{)}.$ (19)

Once again, this estimators suffers from a credit assignment problem.
Gradients of resampling probabilities for particles from time steps other than
the last one are all given equal weight, regardless of how the particle
performed subsequently.

All the gradient estimators considered so far were biased, which is to be
expected given that they were obtained by differentiating biased estimators
from the forward pass. However, particle filters are known to produce unbiased
estimators for expectations under the _unnormalized_ posterior.
Differentiating through DPF, we obtain the following gradient estimator

$\displaystyle\nabla_{\theta}\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{Zf_{\theta}(x_{1:T})}\right]\approx\overrightarrow{\nabla_{\theta}\hat{Z}_{\text{DPF}}\sum_{i=1}^{N}\overline{v}_{T}^{i}f_{\theta}(\tilde{x}_{1:T}^{i})}=$
(20)
$\displaystyle\hat{Z}\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(f_{\theta}(\tilde{x}_{1:T}^{i})\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right).$
(21)

This very similar to Eq. 18, only without the baseline $\overline{f}_{\theta}$
and it is multiplied by the marginal likelihood estimate $\hat{Z}$. It is both
consistent and unbiased, as proven in the supplementary materials.

We compare it with the estimator obtained with PF and DiCE, which is also
consistent and unbiased, but suffers from credit assignment problems

$\displaystyle\nabla_{\theta}\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{Zf_{\theta}(x_{1:T})}\right]\approx$
$\displaystyle\hat{Z}\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})+\left(f_{\theta}(\tilde{x}_{1:T}^{i})-\overline{f}_{\theta}\right)\nabla_{\theta}\log
w_{T}^{i}\right)+$
$\displaystyle\quad\quad\hat{Z}\overline{f}_{\theta}\sum_{t=1}^{T}\sum_{i=1}^{N}\nabla_{\theta}\log\overline{w}_{t-1}^{a_{t}^{i}}.$
(22)

##  4 Extensions

###  4.1 Handling proposal distributions

So far we have only considered gradients with respect to model parameters
$\theta$. However, it is often desirable to jointly learn separate parameters
$\phi$ of the proposal distributions. Doing this requires reparameterizing
$q_{\phi}$, which we assume to be given by
$x_{t}^{i}=h_{\phi}(\epsilon_{t}^{i},x_{t-1}^{a_{t}^{i}})$ with
$\epsilon_{t}^{i}\sim\mathcal{N}(0,1)$. This induces a reparameterization
across the entire ancestral lineage
$\tilde{x}_{1:T}^{i}=\tilde{h}_{\phi}(\tilde{\epsilon}_{1:T}^{i})$, which
provides convenient notation. We then compute the result of applying automatic
differentiation to $\log\hat{Z}_{\text{DPF}}$

$\displaystyle\overrightarrow{\nabla_{\phi}\log\hat{Z}_{\text{DPF}}}=\sum_{i=1}^{N}\overline{w}_{T}^{i}\nabla_{\phi}\log\frac{p_{\theta}(\tilde{h}_{\phi}(\tilde{\epsilon}_{1:T}^{i}),y_{1:T})}{q_{\phi}(\tilde{h}_{\phi}(\tilde{\epsilon}_{1:T}^{i}))}.$
(23)

This is very similar to the gradients $\nabla_{\phi}\mathcal{L}_{\text{IWAE}}$
(Burda et al., 2016, Eq. 14), the only difference being that the trajectories
are sampled from the particle filter rather then directly from the proposal.
We can therefore reasonably expect those gradients to provide learning signal
for $\phi$, although it is unclear whether it would be better than
$\overrightarrow{\nabla_{\phi}\log\hat{Z}_{\text{PF}}}$ (Le et al., 2018a;
Maddison et al., 2017; Naesseth et al., 2018) due to particle degeneracy
problems. Our initial experiments in this regard, presented in Section 5.2,
were inconclusive, so we leave this question for future work.

In some settings, particularly when using the bootstrap particle filter, the
proposal is derived from model parameters. For the reasons discussed in
Section 3.1, in such settings we recommend applying stop-gradient to the
proposal throughout the program. This was done by Le et al. (2018a, Section
5.1) without theoretical justification and we do it in our experiments as
well.

###  4.2 Reducing variance with additional computation

Algorithm 2 The marginal particle filter in the usual formulation (left) and
in our differentiable formulation (right). When estimating the score function,
DPF2 reduces variance from quadratic to linear in $T$. The computational cost
increases from linear to quadratic in $N$.

function PF2($p$, $q$, $N$, $T$)

$x_{0}^{i}\sim p(x_{0})$

$w_{0}^{i}=\overline{w}_{0}^{i}=\frac{1}{N}$

for $t\in 1:T$ do

$x^{i}_{t}\sim\sum_{j=1}^{N}\overline{v}_{t-1}^{j}q_{\phi}(x_{t}|x^{a_{t}^{i}}_{t-1},y_{t})$

$w_{t}^{n}=\frac{1}{N}\frac{\sum_{i=1}^{N}\overline{w}_{t-1}^{i}\
p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})}{\sum_{i=1}^{N}\overline{w}_{t-1}^{i}\
q_{\phi}(x_{t}^{n}|x_{t-1}^{i},y_{t})}$

$W_{t}=\sum_{i=1}^{N}w^{i}_{t}$

$\overline{w}^{i}_{t}=w^{i}_{t}/W_{t}$

$\hat{Z}_{\text{PF2}}=\prod_{t=1}^{T}W_{t}$

function DPF2($p$, $q$, $N$, $T$)

$x_{0}^{i}\sim p(x_{0})$

$v_{0}^{i}=\overline{v}_{0}^{i}=\frac{1}{N}$

for $t\in 1:T$ do

$x^{i}_{t}\sim\sum_{j=1}^{N}\bot\left({\overline{v}_{t-1}^{j}}\right)q_{\phi}(x_{t}|x^{a_{t}^{i}}_{t-1},y_{t})$

$v_{t}^{n}=\frac{1}{N}$ $\frac{\sum_{i=1}^{N}\overline{v}_{t-1}^{i}\
p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})}{\sum_{i=1}^{N}\bot\left({\overline{v}_{t-1}^{i}}\right)\
q_{\phi}(x_{t}^{n}|x_{t-1}^{i},y_{t})}$

$V_{t}=\sum_{i=1}^{N}v^{i}_{t}$

$\overline{v}^{i}_{t}=v^{i}_{t}/V_{t}$

$\hat{Z}_{\text{DPF2}}=\prod_{t=1}^{T}V_{t}$

Poyiadjis et al. (2011) have shown that the variance of score function
estimators from Eq. 1 grows quadratically with sequence length, but can be
reduced to linear using what is now known as the marginal particle filter
(Klaas et al., 2012). We obtain this reduced-variance estimator with AD by
stopping the gradients of the mixture component weights in the proposal
distribution. This modification is shown in Algorithm 2.

Poyiadjis et al. (2011, Eq. 20) define their reduced-variance score function
estimator as

$\displaystyle\nabla_{\theta}\log
p_{\theta}(y_{1:T})\approx\sum_{i=1}^{N}\overline{w}_{T}^{i}\bar{\alpha}_{T}^{i},$
(24)

using the marginal particle filter and the following recursive definition of
$\bar{\alpha}$

$\displaystyle\bar{\alpha}_{t+1}^{i}=\frac{1}{\sum_{j=1}^{N}\overline{w}_{t+1}^{j}p_{\theta}(x_{t+1}^{i}|x_{t}^{j})}$
(25)
$\displaystyle\quad\sum_{j=1}^{N}\overline{w}_{t+1}^{j}p_{\theta}(x_{t+1}^{i}|x_{t}^{j})\left(\bar{\alpha}_{t}^{j}+\nabla_{\theta}\log
p_{\theta}(x_{t+1}^{i},y_{t+1}|x_{t}^{j})\right).$

In the supplementary materials we show that
$\overrightarrow{\nabla_{\theta}\log\hat{Z}_{\text{DPF2}}}$ produces the same
estimator.

The marginal particle filter avoids issues with particle degeneracy, at the
expense of not producing joint samples $\tilde{x}_{1:T}^{i}$. This means that
it is not suitable for estimating expectations of arbitrary functions under
the smoothing density, but on the flipside it achieves linear scaling of
variance with $T$ for estimating the score function. Its computational cost is
quadratic in the number of particles, since the mixture densities need to be
evaluated for each particle separately. As shown by Poyiadjis et al. (2011),
this trade-off is always favorable for sufficiently large $T$.

Lai et al. (2021) propose to use the marginal particle filter to jointly learn
proposal parameters $\phi$. Regardless of whether stop-gradients are applied,
the key consideration is how to reparameterize sampling from the mixture
proposal. Lai et al. (2021) consider two methods. The first draws the mixture
component first from a discrete distribution, which introduces bias in the
gradients. This bias can be avoided with the stop-gradient corrections of
DPF2, but that potentially introduces problems with particle degeneracy, much
like in the normal particle filter. The second method is to reparameterize the
mixture distribution using cumulative distribution functions (Graves, 2016).
In this method there is still a choice of whether to stop gradients and
reparameterize a mixture with fixed weights and it is not clear whether that
would perform better or worse.

###  4.3 Higher-order derivatives

While in most applications we are interested in first-order derivatives, we
can apply automatic differentiation repeatedly to obtain derivatives of higher
orders. Poyiadjis et al. (2011) show that the Hessian of the log-marginal
likelihood can be approximated with a particle filter using Louis’ identity

$\displaystyle\nabla_{\theta}^{2}\log p_{\theta}(y_{1:T})\approx$ (26)
$\displaystyle-\sum_{i=1}^{N}\sum_{j=1}^{N}\overline{w}_{T}^{i}\overline{w}_{T}^{j}\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})^{T}$
$\displaystyle+\sum_{i=1}^{N}\overline{w}_{T}^{i}\Bigg{(}\nabla_{\theta}^{2}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})+$
$\displaystyle\quad\quad\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},,y_{1:T})\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},,y_{1:T})^{T}\Bigg{)},$

where superscript $T$ denotes transpose. This is exactly the estimator
obtained by applying AD twice to DPF, specifically from
$\overrightarrow{\nabla_{\theta}^{2}\log\hat{Z}_{\text{DPF}}}$, as we show in
the supplementary materials. Furthermore, Hessian-vector products can be
computed with reduced memory requirements in the usual fashion.

We can similarly derive higher-order derivatives for expectations under the
posterior. We conjecture that DPF with automatic differentiation provides
unbiased estimators for derivatives of any order of expectations under the
unnormalized posterior, similarly to how DiCE produces unbiased estimators for
derivatives of any order, leaving the proof for future work.

##  5 Experiments

In our experiments we focus on the task of model learning using the score
function estimators given by
$\overrightarrow{\nabla_{\theta}\log\hat{Z}_{\text{DPF}}}$. We show that DPF
outperforms PF on model learning, while the results for jointly learning the
model and the proposal are open for interpretation.

We use the following baselines: sequential importance sampling without
resampling (SIS), biased particle filter ignoring the gradients resulting from
resampling (PF), and a particle filter with DiCE for ancestral variables (PF-
SF). We refer to our method as DPF-SGR, for “stop-gradient resampling”, to
distinguish it from other versions of differentiable particle filterning
proposed in the literature. We borrow the experimental setup from Le et al.
(2018a), using the code provided by those authors with our modifications
needed to implement DPF-SGR. An additional experiment on a stochastic
volatility model, including comparisons with the method proposed by Corenflos
et al. (2021), is provided in the supplementary materials.

###  5.1 Linear Gaussian State Space Model

In the first experiment we use a Linear Gaussian State Space Model (LGSSM)
with a bootstrap proposal exactly as specified by Le et al. (2018a), with two
global parameters, a one dimensional latent variable and 200 time steps.
Figure 1 in our paper corresponds exactly to Figure 2 from Le et al. (2018a),
with DPF-SGR and PF-SF added. We find that DPF-SGR offers a noticeable
improvement over the standard PF in the low particle regime, while in the high
particle regime DPF-SGR achieves significantly faster convergence at the
expense of a minimal decrease in the final performance.

Figure 1: Estimates for the test set log-marginal likelihood throughout the
training of the LGSSM model. The bottom plot is the zoomed in version of the
top one, showing improvement of DPF-SGR over PF.

###  5.2 Variational Recurrent Neural Network

To jointly learn the model and the proposal, we train a variational recurrent
neural network (VRNN) modeling the movement of agents in a 2D world where the
obvservations are provided in the form of partially occluded images. We use
the setting introduced by Le et al. (2018a), with 32x32 images and 40 time
steps. Our Figure 2 corresponds exactly to Figure 4 of Le et al. (2018a),
except that we show performance both on the training and test sets. While PF
achieves better performance on the training set, this is largely a result of
overfitting, with DPF-SGR learning slower but ultimately achieving better test
set performance. It is unclear to us what general conclusions can be drawn
from this experiment regarding the suitability of DPF-SGR for learning the
proposal, other than that it warrants further investigation.

Figure 2: Log-marginal likelihood estimates for training and test set
throughout the training of the VRNN model. Unlike DPF-SGR, SIS and PF show
clear signs of overfitting.

##  6 Related Work

Due to high variance of estimators obtained by applying DiCE to the ancestral
variables, the gradients associated with sampling those variables are
typically ignored (Le et al., 2018a; Maddison et al., 2017; Naesseth et al.,
2018). In particular, Jonschkowski et al. (2018) construct an algorithm they
call Differentiable Particle Filtering, which explicitly ignores the
propagation of gradients through the resampling steps. They suggest adding
auxiliary gradients to weights in an attempt to mitigate the resulting bias as
a direction of future work, which is exactly what we provide here.

There are several existing approaches to making resampling differentiable, all
of which modify its behavior on the forward pass. Karkus et al. (2018)
resample with probabilities proportional to $\alpha w^{i}+\frac{1-\alpha}{N}$,
which effectively interpolates between no resampling and no gradient bias
forcing the users to make an undesirable trade-off. We get the best of both
worlds, retaining full advantages of resampling while allowing the gradients
to propage. Corenflos et al. (2021) replace traditional resampling with a
continuous transformation of particles based on optimal transport equations.
This is computationally intensive, biased, introduces additional
hyperparameters, and can lead to undesirable behavior, in the extreme cases
even putting the resampled particles outside of the support of the prior. On
the other hand, it avoids problems with particle degeneracy while achieving
full reparameterization. Zhu et al. (2020) modify particles in a similar
fashion using learned resampling networks, with similar disadvantages, noting
that in practice propagating gradients through them leads to instabilities so
they avoid it.

##  7 Discussion

We have demonstrated that applying AD to DPF recovers known estimators for the
score function, which enables maximum likelihood model learning by gradient
ascent. We note that those estimators target the gradients of log-marginal
likelihood directly, rather than the gradients of some lower bound to it. For
this reason it’s unclear whether such gradients are suitable for learning the
proposal, since the marginal likelihood does not depend on the proposal. The
topic of constructing objectives for proposal learning is somewhat
controversial, and it has in particular been shown that tightening lower
bounds to log-marginal likelihood eventually impedes proposal learning
(Rainforth et al., 2019). Generally, it is neither necessary nor sufficient to
optimize a lower bound to log-marginal likelihood to provide good learning
signal to the proposal. For example, Le et al. (2018b) suggest using a
separate objective for proposal learning with alternating optimization, which
is competitive but not clearly a superior approach. Since our score function
estimators contain bias dependent on the proposal, they might be able to
provide a learning signal for it, as indicated in Sections 4.1 and 5.2. We
believe the question warrants further investigation.

The stop-gradient operator is relatively rarely used in theoretical formulas
and many authors prefer to avoid it. We believe it presents an opportunity to
write equations more closely resembling the implementation, reducing the scope
for translation errors. Ultimately, stop-gradient is just a different notation
for familiar concepts, in particular
$\nabla_{\theta}\frac{p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}$
expresses the same thing as
$\nabla_{\theta^{\prime}}\frac{p_{\theta^{\prime}}(x)}{p_{\theta}(x)}\big{\rvert}_{\theta^{\prime}=\theta}$,
only making it easier to build compound expressions. We believe that using
this operator more broadly can drastically simplify certain derivations and
proofs.

## Acknowledgements

We thank Tuan Anh Le, Maximilian Igl, Adrien Corenflos, and James Thornton for
sharing the code for their experimental setups which we used to conduct the
experiments described in this paper, and Vaden Masrani for helping to adapt
this code to our needs. We also thank Marcin Tomczak, Wilder Lavington and the
anonymous reviewers for providing feedback on a draft of this paper.

We acknowledge the support of the Natural Sciences and Engineering Research
Council of Canada (NSERC), the Canada CIFAR AI Chairs Program, and the Intel
Parallel Computing Centers program. This material is based upon work supported
by the United States Air Force Research Laboratory (AFRL) under the Defense
Advanced Research Projects Agency (DARPA) Data Driven Discovery Models (D3M)
program (Contract No. FA8750-19-2-0222) and Learning with Less Labels (LwLL)
program (Contract No.FA8750-19-C-0515). Additional support was provided by
UBC’s Composites Research Network (CRN), Data Science Institute (DSI) and
Support for Teams to Advance Interdisciplinary Research (STAIR) Grants. This
research was enabled in part by technical support and computational resources
provided by WestGrid (https://www.westgrid.ca/) and Compute Canada
(www.computecanada.ca).

## References

  * Abadi et al. [2015] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng.  TensorFlow: Large-scale machine learning on heterogeneous systems, 2015.
  * Baydin et al. [2018] Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind.  Automatic differentiation in machine learning: A survey.  _arXiv:1502.05767 [cs, stat]_ , February 2018.
  * Burda et al. [2016] Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov.  Importance Weighted Autoencoders.  _arXiv:1509.00519 [cs, stat]_ , November 2016.
  * Chopin and Papaspiliopoulos [2020] Nicolas Chopin and Omiros Papaspiliopoulos.  _An Introduction to Sequential Monte Carlo_.  Springer Series in Statistics. Springer International Publishing, 2020.  ISBN 978-3-030-47844-5.  doi: 10.1007/978-3-030-47845-2.
  * Corenflos et al. [2021] Adrien Corenflos, James Thornton, Arnaud Doucet, and George Deligiannidis.  Differentiable Particle Filtering via Entropy-Regularized Optimal Transport.  _arXiv:2102.07850 [cs, stat]_ , February 2021.
  * Doucet and Johansen [2009] Arnaud Doucet and Adam M Johansen.  A tutorial on particle filtering and smoothing: Fifteen years later.  _Handbook of nonlinear filtering_ , 12(656-704):3, 2009.
  * Doucet et al. [2001] Arnaud Doucet, Nando De Freitas, and Neil Gordon.  An introduction to sequential Monte Carlo methods.  In _Sequential Monte Carlo methods in practice_ , pages 3–14. Springer, 2001.
  * Farquhar et al. [2019] Gregory Farquhar, Shimon Whiteson, and Jakob Foerster.  Loaded DiCE: Trading off Bias and Variance in Any-Order Score Function Estimators for Reinforcement Learning.  _arXiv:1909.10549 [cs, stat]_ , September 2019.
  * Fearnhead and Clifford [2003] Paul Fearnhead and Peter Clifford.  On-Line Inference for Hidden Markov Models via Particle Filters.  _Journal of the Royal Statistical Society. Series B (Statistical Methodology)_ , 65(4):887–899, 2003.  ISSN 1369-7412.
  * Foerster et al. [2018] Jakob Foerster, Gregory Farquhar, Maruan Al-Shedivat, Tim Rocktäschel, Eric Xing, and Shimon Whiteson.  DiCE: The Infinitely Differentiable Monte Carlo Estimator.  In _International Conference on Machine Learning_ , pages 1529–1538. PMLR, July 2018.
  * Fu [2006] Michael C. Fu.  Chapter 19 Gradient Estimation.  In Shane G. Henderson and Barry L. Nelson, editors, _Handbooks in Operations Research and Management Science_ , volume 13 of _Simulation_ , pages 575–616. Elsevier, January 2006.  doi: 10.1016/S0927-0507(06)13019-4.
  * Glasserman [2003] Paul Glasserman.  _Monte Carlo Methods in Financial Engineering_.  Stochastic Modelling and Applied Probability. Springer-Verlag, New York, 2003.  ISBN 978-0-387-00451-8.  doi: 10.1007/978-0-387-21617-1.
  * Graves [2016] Alex Graves.  Stochastic Backpropagation through Mixture Density Distributions.  _arXiv:1607.05690 [cs]_ , July 2016.
  * Griewank [1989] Andreas Griewank.  On Automatic Differentiation.  In _Mathematical Programming: Recent Developments and Applications_ , pages 83–108. Kluwer Academic Publishers, Amsterdam, 1989\.
  * Igl et al. [2018] Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson.  Deep variational reinforcement learning for POMDPs.  In _International Conference on Machine Learning_ , pages 2117–2126. PMLR, 2018.
  * Jankowiak and Obermeyer [2018] Martin Jankowiak and Fritz Obermeyer.  Pathwise Derivatives Beyond the Reparameterization Trick.  _arXiv:1806.01851 [cs, stat]_ , July 2018.
  * Jonschkowski et al. [2018] Rico Jonschkowski, Divyam Rastogi, and Oliver Brock.  Differentiable Particle Filters: End-to-End Learning with Algorithmic Priors.  _arXiv:1805.11122 [cs, stat]_ , May 2018.
  * Karkus et al. [2018] Peter Karkus, David Hsu, and Wee Sun Lee.  Particle Filter Networks with Application to Visual Localization.  In _Conference on Robot Learning_ , pages 169–178. PMLR, October 2018.
  * Kingma and Welling [2014] Diederik P. Kingma and Max Welling.  Auto-Encoding Variational Bayes.  _arXiv:1312.6114 [cs, stat]_ , May 2014.
  * Klaas et al. [2012] Mike Klaas, Nando de Freitas, and Arnaud Doucet.  Toward Practical N2 Monte Carlo: The Marginal Particle Filter.  _arXiv:1207.1396 [cs, stat]_ , July 2012.
  * Lai et al. [2021] Jinlin Lai, Daniel Sheldon, and Justin Domke.  Variational Marginal Particle Filters.  _arXiv:2109.15134 [cs, stat]_ , September 2021.
  * Le et al. [2018a] Tuan Anh Le, Maximilian Igl, Tom Rainforth, Tom Jin, and Frank Wood.  Auto-Encoding Sequential Monte Carlo.  _arXiv:1705.10306 [stat]_ , April 2018a.
  * Le et al. [2018b] Tuan Anh Le, Adam R Kosiorek, N Siddharth, Yee Whye Teh, and Frank Wood.  Revisiting reweighted wake-sleep.  2018b.
  * Liu [2004] Jun S. Liu.  Theory of Sequential Monte Carlo.  In Jun S. Liu, editor, _Monte Carlo Strategies in Scientific Computing_ , Springer Series in Statistics, pages 53–77. Springer, New York, NY, 2004.  ISBN 978-0-387-76371-2.  doi: 10.1007/978-0-387-76371-2_3.
  * Maddison et al. [2017] Chris J. Maddison, John Lawson, George Tucker, Nicolas Heess, Mohammad Norouzi, Andriy Mnih, Arnaud Doucet, and Yee Teh.  Filtering Variational Objectives.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Mao et al. [2019] Jingkai Mao, Jakob Foerster, Tim Rocktäschel, Maruan Al-Shedivat, Gregory Farquhar, and Shimon Whiteson.  A Baseline for Any Order Gradient Estimation in Stochastic Computation Graphs.  In _Proceedings of the 36th International Conference on Machine Learning_ , pages 4343–4351. PMLR, May 2019.
  * Moral [2004] Pierre Del Moral.  _Feynman-Kac Formulae: Genealogical and Interacting Particle Systems with Applications_.  Probability and Its Applications. Springer-Verlag, New York, 2004\.  ISBN 978-0-387-20268-6.  doi: 10.1007/978-1-4684-9393-1.
  * Murray et al. [2016] Lawrence M Murray, Anthony Lee, and Pierre E Jacob.  Parallel resampling in the particle filter.  _Journal of Computational and Graphical Statistics_ , 25(3):789–805, 2016.
  * Naesseth et al. [2018] Christian Naesseth, Scott Linderman, Rajesh Ranganath, and David Blei.  Variational Sequential Monte Carlo.  In _International Conference on Artificial Intelligence and Statistics_ , pages 968–977. PMLR, March 2018.
  * Neal and Hinton [1998] Radford M. Neal and Geoffrey E. Hinton.  A View of the Em Algorithm that Justifies Incremental, Sparse, and other Variants.  In Michael I. Jordan, editor, _Learning in Graphical Models_ , NATO ASI Series, pages 355–368. Springer Netherlands, Dordrecht, 1998.  ISBN 978-94-011-5014-9.  doi: 10.1007/978-94-011-5014-9_12.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  PyTorch: An Imperative Style, High-Performance Deep Learning Library.  _arXiv:1912.01703 [cs, stat]_ , December 2019.
  * Poyiadjis et al. [2011] George Poyiadjis, Arnaud Doucet, and Sumeetpal Singh.  Particle approximations of the score and observed information matrix in state space models with application to parameter estimation.  _Biometrika_ , 98(1):65–80, 2011.  ISSN 0006-3444.
  * Rainforth et al. [2019] Tom Rainforth, Adam R. Kosiorek, Tuan Anh Le, Chris J. Maddison, Maximilian Igl, Frank Wood, and Yee Whye Teh.  Tighter Variational Bounds are Not Necessarily Better.  _arXiv:1802.04537 [cs, stat]_ , March 2019.
  * Ranganath et al. [2013] Rajesh Ranganath, Sean Gerrish, and David M. Blei.  Black Box Variational Inference.  _arXiv:1401.0118 [cs, stat]_ , December 2013.
  * Rezende et al. [2014] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra.  Stochastic Backpropagation and Approximate Inference in Deep Generative Models.  _arXiv:1401.4082 [cs, stat]_ , May 2014.
  * Rumelhart et al. [1986] David E. Rumelhart, Geoffrey E. Hinton, and Ronald J. Williams.  Learning representations by back-propagating errors.  _Nature_ , 323(6088):533–536, October 1986.  ISSN 1476-4687.  doi: 10.1038/323533a0.
  * Schulman et al. [2016] John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel.  Gradient Estimation Using Stochastic Computation Graphs.  _arXiv:1506.05254 [cs]_ , January 2016.
  * van Krieken et al. [2021] Emile van Krieken, Jakub M. Tomczak, and Annette ten Teije.  Storchastic: A Framework for General Stochastic Automatic Differentiation.  _arXiv:2104.00428 [cs, stat]_ , May 2021.
  * Wengert [1964] R. E. Wengert.  A simple automatic derivative evaluation program.  _Communications of the ACM_ , 7(8):463–464, August 1964.  ISSN 0001-0782.  doi: 10.1145/355586.364791.
  * Williams [1992] Ronald J. Williams.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  _Machine Learning_ , 8(3):229–256, May 1992.  ISSN 1573-0565.  doi: 10.1007/BF00992696.
  * Zhu et al. [2020] Michael Zhu, Kevin Murphy, and Rico Jonschkowski.  Towards Differentiable Resampling.  _arXiv:2004.11938 [cs, stat]_ , April 2020.

Supplementary Materials

##  8 Stop-gradient tutorial

The key equations for manipulating expressions involving stop-gradient are

$\displaystyle\overrightarrow{f(E_{1},\dots,E_{n})}$
$\displaystyle=f(\overrightarrow{E_{1}},\dots,\overrightarrow{E_{2}})$ (27)
$\displaystyle\overrightarrow{\bot\left({E}\right)}$
$\displaystyle=\overrightarrow{E}$ (28)
$\displaystyle\nabla{\bot\left({E}\right)}$ $\displaystyle=0$ (29)
$\displaystyle\overrightarrow{\nabla{E}}$
$\displaystyle=\nabla{E}\quad\text{if}\ \bot\notin E.$ (30)

Crucially, it is generally the case that $\overrightarrow{\nabla
E}\neq\nabla\overrightarrow{E}$, that is different expressions involving stop-
gradient can always evaluate to the same value on the forward pass, but
produce different gradients. For example,

$\displaystyle\nabla_{x}\overrightarrow{x+\bot\left({x}\right)^{2}+x\bot\left({x}\right)^{2}}=\nabla_{x}\left(\overrightarrow{x}+\overrightarrow{\bot\left({x}\right)^{2}}+\overrightarrow{x\bot\left({x}\right)^{2}}\right)=\nabla_{x}\left(x+\overrightarrow{\bot\left({x}\right)}^{2}+\overrightarrow{x}\overrightarrow{\bot\left({x}\right)^{2}}\right)=$
(31)
$\displaystyle\nabla_{x}\left(x+x^{2}+xx^{2}\right)=\nabla_{x}\left(x+x^{2}+x^{3}\right)=1+2x+3x^{2}.$
(32)

On the other hand,

$\displaystyle\overrightarrow{\nabla_{x}\left(x+\bot\left({x}\right)^{2}+x\bot\left({x}\right)^{2}\right)}=\overrightarrow{1+2\bot\left({x}\right)\nabla_{x}\bot\left({x}\right)+\left(\nabla
x\right)\bot\left({x}\right)^{2}+x\left(\nabla_{x}\bot\left({x}\right)^{2}\right)}=$
(33)
$\displaystyle\overrightarrow{1+2\bot\left({x}\right)0+\bot\left({x}\right)^{2}+x0}=\overrightarrow{1+\bot\left({x}\right)^{2}}=\overrightarrow{1}+\overrightarrow{\bot\left({x}\right)^{2}}=1+\overrightarrow{\bot\left({x}\right)}^{2}=1+x^{2}.$
(34)

The chain rule for standard functions works normally, as does repeated
differentiation. Moreover, it is possible to interleave differentiation with
evaluation. For example,

$\displaystyle\nabla_{x}\overrightarrow{\nabla_{x}\exp\left(\bot\left({x}\right)+\log
x\right)}=\nabla_{x}\overrightarrow{\exp\left(\bot\left({x}\right)+\log
x\right)\nabla_{x}\left(\bot\left({x}\right)+\log x\right)}=$ (35)
$\displaystyle\nabla_{x}\overrightarrow{\exp\left(\bot\left({x}\right)+\log
x\right)\left(\nabla_{x}\bot\left({x}\right)+\nabla_{x}\log
x\right)}=\nabla_{x}\overrightarrow{\exp\left(\bot\left({x}\right)+\log
x\right)\left(0+\frac{1}{x}\right)}=$ (36)
$\displaystyle\nabla_{x}\overrightarrow{\exp\left(\bot\left({x}\right)+\log
x\right)\frac{1}{x}}=\nabla_{x}\left(\overrightarrow{\exp\left(\bot\left({x}\right)+\log
x\right)}\overrightarrow{\frac{1}{x}}\right)=$ (37)
$\displaystyle\nabla_{x}\left(\exp\overrightarrow{\left(\bot\left({x}\right)+\log
x\right)}\frac{1}{x}\right)=\nabla_{x}\left(\exp\left(x+\log
x\right)\frac{1}{x}\right)=$ (38)
$\displaystyle\frac{1}{x}\nabla_{x}\exp\left(x+\log x\right)+\exp\left(x+\log
x\right)\nabla_{x}\frac{1}{x}=$ (39) $\displaystyle\frac{1}{x}\exp\left(x+\log
x\right)\left(1+\frac{1}{x}\right)-\frac{1}{x^{2}}\exp\left(x+\log
x\right)=\frac{1}{x}\exp\left(x+\log x\right).$ (40)

All the usual properties of standard functions hold, for example both addition
and multiplication are associative and commutative, and the logarithm of a
product is a sum of logarithms. We do not consider limits or integrals of
expressions involving stop-gradient, leaving them undefined.

The key identities showing up repeatedly in this paper are

$\displaystyle\overrightarrow{\nabla\frac{E}{\bot\left({E}\right)}}=\overrightarrow{\frac{\nabla
E}{\bot\left({E}\right)}}=\frac{\nabla E}{E}=\log E,$ (41)

and

$\displaystyle\overrightarrow{\nabla\frac{\bot\left({E}\right)}{E}}=\overrightarrow{-\frac{\bot\left({E}\right)}{E^{2}}\nabla
E}=-\overrightarrow{\frac{\bot\left({E}\right)}{E^{2}}}\overrightarrow{\nabla
E}=-\frac{E}{E^{2}}\nabla E=-\frac{\nabla E}{E}=-\log E,$ (42)

both of which hold assuming that $E$ does not contain any stop-gradient terms.

Equipped with those, we can see what DiCE evaluates to

$\displaystyle\overrightarrow{\nabla_{\theta}\left(\frac{p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}f_{\theta}(x)\right)}=\overrightarrow{\nabla_{\theta}\left(\frac{p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}\right)f_{\theta}(x)+\frac{p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}\nabla_{\theta}f_{\theta}(x)}=$
(43) $\displaystyle
f_{\theta}(x)\overrightarrow{\frac{\nabla_{\theta}p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}}+\overrightarrow{\frac{p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}}\nabla_{\theta}f_{\theta}(x)=f_{\theta}(x)\nabla_{\theta}\log
p_{\theta}(x)+\nabla_{\theta}f_{\theta}(x).$ (44)

##  9 Detailed derivations

###  9.1 Fisher’s identity with stop-gradient

The IWAE [Burda et al., 2016] objective is

$\displaystyle\mathcal{L}_{\text{IWAE}}=\mathop{\mathbb{E}}_{x^{i}\sim
q_{\phi}}\left[{\log\hat{Z}}\right]=\mathop{\mathbb{E}}_{x^{i}\sim
q_{\phi}(x)}\left[{\log\frac{1}{N}\sum_{i=1}^{N}w^{i}}\right]=\mathop{\mathbb{E}}_{x^{i}\sim
q_{\phi}(x)}\left[{\log\frac{1}{N}\sum_{i=1}^{N}\frac{p_{\theta}(x^{i},y)}{q_{\phi}(x^{i})}}\right].$
(45)

Automatic differentiation then computes gradients with respect to the model
parameters $\theta$ as follows

$\displaystyle\nabla_{\theta}\mathcal{L}_{\text{IWAE}}$
$\displaystyle=\mathop{\mathbb{E}}_{x^{i}\sim
q_{\phi}(x)}\left[{\nabla_{\theta}\log\frac{1}{N}\sum_{i=1}^{N}w^{i}}\right]=\mathop{\mathbb{E}}_{x^{i}\sim
q_{\phi}(x)}\left[{\frac{\nabla_{\theta}\sum_{i=1}^{N}w^{i}}{\sum_{i=1}^{N}w^{i}}}\right]$
(46) $\displaystyle=\mathop{\mathbb{E}}_{x^{i}\sim
q_{\phi}(x)}\left[{\frac{\sum_{i=1}^{N}w^{i}\nabla_{\theta}\log
w^{i}}{\sum_{i=1}^{N}w^{i}}}\right]=\mathop{\mathbb{E}}_{x^{i}\sim
q_{\phi}(x)}\left[{\sum_{i=1}^{N}\overline{w}^{i}\nabla_{\theta}\log
w^{i}}\right]$ (47) $\displaystyle=\mathop{\mathbb{E}}_{x^{i}\sim
q_{\phi}(x)}\left[{\sum_{i=1}^{N}\overline{w}^{i}\nabla_{\theta}\log
p_{\theta}(x^{i},y)}\right].$ (48)

This is exactly the same estimator we would obtain by using Fisher’s identity
with a self-normalizing importance sampler to approximate the posterior

$\displaystyle\nabla_{\theta}\log p_{\theta}(y)$
$\displaystyle=\mathop{\mathbb{E}}_{x\sim
p_{\theta}(x|y)}\left[{\nabla_{\theta}\log
p_{\theta}(x,y)}\right]=\mathop{\mathbb{E}}_{x\sim
q_{\phi}(x)}\left[{\frac{p_{\theta}(x|y)}{q_{\phi}(x)}\nabla_{\theta}\log
p_{\theta}(x,y)}\right]$ (49) $\displaystyle=\mathop{\mathbb{E}}_{x\sim
q_{\phi}(x)}\left[{\frac{w}{p_{\theta}(y)}\nabla_{\theta}\log
p_{\theta}(x,y)}\right]=\mathop{\mathbb{E}}_{x^{i}\sim
q_{\phi}(x)}\left[{\frac{1}{N}\sum_{i=1}^{N}\frac{w^{i}}{p_{\theta}(y)}\nabla_{\theta}\log
p_{\theta}(x^{i},y)}\right]$ (50)
$\displaystyle\approx\mathop{\mathbb{E}}_{x^{i}\sim
q_{\phi}(x)}\left[{\sum_{i=1}^{N}\overline{w}^{i}\nabla_{\theta}\log
p_{\theta}(x^{i},y)}\right].$ (51)

However, the two methods are no longer equivalent if the proposal distribution
$q_{\theta}$ also depends on model parameters, such as when using the prior
$p_{\theta}(x)$. Furthermore, suppose it is not reparameterizable, so we apply
DiCE to obtain the gradient estimators

$\displaystyle\nabla_{\theta}\mathop{\mathbb{E}}_{x^{i}\sim
q_{\theta}(x)}\left[{\log\frac{1}{N}\sum_{i=1}^{N}\frac{p_{\theta}(x^{i},y)}{q_{\theta}(x^{i})}}\right]=\mathop{\mathbb{E}}_{x^{i}\sim
q_{\theta}(x)}\left[{\overrightarrow{\nabla_{\theta}\left(\prod_{i=1}^{N}\frac{p_{\theta}(x^{i})}{\bot\left({q_{\theta}(x^{i})}\right)}\log\frac{1}{N}\sum_{i=1}^{N}\frac{p_{\theta}(x^{i},y)}{q_{\theta}(x^{i})}\right)}}\right]=$
(52) $\displaystyle\mathop{\mathbb{E}}_{x^{i}\sim
p_{\theta}(x))}\left[{\overrightarrow{\left(\nabla_{\theta}\prod_{i=1}^{N}\frac{p_{\theta}(x^{i})}{\bot\left({q_{\theta}(x^{i})}\right)}\right)\log\frac{1}{N}\sum_{i=1}^{N}\frac{p_{\theta}(x^{i},y)}{q_{\theta}(x^{i})}+\left(\prod_{i=1}^{N}\frac{p_{\theta}(x^{i})}{\bot\left({q_{\theta}(x^{i})}\right)}\right)\sum_{i=1}^{N}\overline{w}^{i}\nabla_{\theta}\log\frac{p_{\theta}(x^{i},y)}{q_{\theta}(x^{i})}}}\right]=$
(53) $\displaystyle\mathop{\mathbb{E}}_{x^{i}\sim
q_{\theta}(x)}\left[{\log\hat{Z}\ \sum_{i=1}^{N}\nabla_{\theta}\log
p_{\theta}(x^{i})+\sum_{i=1}^{N}\overline{w}^{i}\nabla_{\theta}\log\frac{p_{\theta}(x^{i},y)}{q_{\theta}(x^{i})}}\right].$
(54)

###  9.2 Score function estimation with DPF

Let
$L_{t}^{i}=\frac{p_{\theta}(\tilde{x}_{1:t}^{i},y_{1:t})}{q_{\phi}(\tilde{x}_{1:t}^{i})}$,
which can be thought of as the importance weight accumulated along the full
ancestral lineage of a particle $i$ at time $t$ and satisfies the recursion
$L_{t}^{i}=Nw_{t}^{i}L_{t-1}^{a_{t}^{i}}$

First, we derive a formula for the modified weights at given time steps

$\displaystyle v_{t}^{i}$
$\displaystyle=\frac{p_{\theta}(x_{t}^{i},y_{t}|x_{t-1}^{a_{t}^{i}})}{q_{\phi}(x_{t}^{i}|x_{t-1}^{a_{t}^{i}})}\
\tilde{v}_{t}^{i}=w_{t}^{i}\frac{\overline{v}_{t-1}^{a_{t}^{i}}}{\bot\left({\overline{v}_{t-1}^{a_{t}^{i}}}\right)}=w_{t}^{i}\frac{v_{t-1}^{a_{t}^{i}}}{\bot\left({v_{t-1}^{a_{t}^{i}}}\right)}\frac{\bot\left({V_{t-1}}\right)}{V_{t-1}}.$
(55)

This allows us to derive and solve the following recursive equation

$\displaystyle\frac{v_{t}^{i}}{\bot\left({v_{t}^{i}}\right)}=\frac{w_{t}^{i}}{\bot\left({w_{t}^{i}}\right)}\frac{v_{t-1}^{a_{t}^{i}}}{\bot\left({v_{t-1}^{a_{t}^{i}}}\right)}\frac{\bot\left({V_{t-1}}\right)}{V_{t-1}}=\frac{L_{t}^{i}}{\bot\left({L_{t}^{i}}\right)}\prod_{j=1}^{t-1}\frac{\bot\left({V_{j}}\right)}{V_{j}}.$
(56)

Plugging this back into Eq. 55, and noting that
$\bot\left({v_{t}^{i}}\right)=\bot\left({w_{t}^{i}}\right)$, we obtain the
following formula

$\displaystyle
v_{t}^{i}=w_{t}^{i}\frac{L_{t-1}^{a_{t}^{i}}}{\bot\left({L_{t-1}^{a_{t}^{i}}}\right)}\prod_{j=1}^{t-1}\frac{\bot\left({V_{j}}\right)}{V_{j}}=\bot\left({v_{t}^{i}}\right)\frac{L_{t}^{i}}{\bot\left({L_{t}^{i}}\right)}\prod_{j=1}^{t-1}\frac{\bot\left({V_{j}}\right)}{V_{j}}$
(57)

With that, we can derive an alternative form of $\hat{Z}_{\text{DPF}}$

$\displaystyle\hat{Z}_{\text{DPF}}$
$\displaystyle=\prod_{t=1}^{T}V_{t}=\left(\prod_{t=1}^{T-1}V_{t}\right)\sum_{i=1}^{N}v_{T}^{i}=\left(\prod_{t=1}^{T-1}\bot\left({V_{t}}\right)\right)\sum_{i=1}^{N}\bot\left({v_{T}^{i}}\right)\frac{L_{T}^{i}}{\bot\left({L_{T}^{i}}\right)}$
(58)
$\displaystyle=\left(\prod_{t=1}^{T-1}\bot\left({V_{t}}\right)\right)\sum_{i=1}^{N}\bot\left({v_{T}^{i}}\right)\frac{L_{T}^{i}}{\bot\left({L_{T}^{i}}\right)}\frac{\bot\left({V_{T}}\right)}{\sum_{i=1}^{N}\bot\left({v_{T}^{i}}\right)}$
(59)
$\displaystyle=\left(\prod_{t=1}^{T}\bot\left({V_{t}}\right)\right)\frac{\sum_{i=1}^{N}\bot\left({v_{T}^{i}}\right)\frac{L_{T}^{i}}{\bot\left({L_{T}^{i}}\right)}}{\sum_{i=1}^{N}\bot\left({v_{T}^{i}}\right)}=\bot\left({\hat{Z}_{\text{DPF}}}\right)\sum_{i=1}^{N}\bot\left({\overline{v}_{T}^{i}}\right)\frac{L_{T}^{i}}{\bot\left({L_{T}^{i}}\right)}.$
(60)

This form allows us to evaluate the gradient computed by automatic
differentiation

$\displaystyle\overrightarrow{\nabla_{\theta}\hat{Z}_{\text{DPF}}}=\overrightarrow{\hat{Z}_{\text{DPF}}}\sum_{i=1}^{N}\overrightarrow{\overline{v}_{T}^{i}}\overrightarrow{\frac{\nabla_{\theta}L_{T}^{i}}{\bot\left({L_{T}^{i}}\right)}}=\overrightarrow{\hat{Z}_{\text{DPF}}}\sum_{i=1}^{N}\overrightarrow{\overline{v}_{T}^{i}}\nabla_{\theta}\log
L_{T}^{i}.$ (61)

The estimator for the score is obtained using the chain rule for the logarithm

$\displaystyle\overrightarrow{\nabla_{\theta}\log\hat{Z}_{\text{DPF}}}=\frac{\overrightarrow{\nabla_{\theta}\hat{Z}_{\text{DPF}}}}{\overrightarrow{\hat{Z}_{\text{DPF}}}}=\sum_{i=1}^{N}\overrightarrow{\overline{v}_{T}^{i}}\nabla_{\theta}\log
L_{T}^{i}=\sum_{i=1}^{N}\overline{w}_{T}^{i}\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T}).$ (62)

###  9.3 Estimators of expectations

Consider a function $f_{\theta}:\mathcal{X}^{T}\rightarrow\mathcal{R}$,
differentiable in $\theta$ for all $x_{1:T}$. A particle filter yields a
consistent estimator for its expectation under the posterior

$\displaystyle\lim_{N\rightarrow\infty}\sum_{i=1}^{N}\overline{w}_{T}^{i}f_{\theta}(\tilde{x}_{1:T}^{i})=\mathop{\mathbb{E}}_{x\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{f_{\theta}(x_{1:T})}\right].$ (63)

Differentiating through the usual PF implementation, we obtain the following
gradient estimator

$\displaystyle\nabla_{\theta}\mathop{\mathbb{E}}_{x\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{f_{\theta}(x_{1:T})}\right]\approx\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(f_{\theta}(\tilde{x}_{1:T}^{i})\nabla_{\theta}\log\overline{w}_{T}^{i}+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right)=$
(64) $\displaystyle-\nabla_{\theta}\log
W_{T}\sum_{i=1}^{N}\overline{w}_{T}^{i}f_{\theta}(\tilde{x}_{1:T}^{i})+\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(f_{\theta}(\tilde{x}_{1:T}^{i})\nabla_{\theta}\log
w_{T}^{i}+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right)=$ (65)
$\displaystyle-\nabla_{\theta}\log
W_{T}\overline{f}_{\theta}+\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(f_{\theta}(\tilde{x}_{1:T}^{i})\nabla_{\theta}\log
w_{T}^{i}+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right)=$ (66)
$\displaystyle-\overline{f}_{\theta}\sum_{i=1}^{N}\overline{w}_{T}^{i}\nabla_{\theta}\log
w_{T}^{i}+\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(f_{\theta}(\tilde{x}_{1:T}^{i})\nabla_{\theta}\log
w_{T}^{i}+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right)=$ (67)
$\displaystyle\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(\left(f_{\theta}(\tilde{x}_{1:T}^{i})-\overline{f}_{\theta}\right)\nabla_{\theta}\log
p_{\theta}(x_{T},y_{T}|x_{T-1}^{a_{T}^{i}})+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right).$
(68)

Furthermore, employing DiCE on ancestral variables results in the following
correction

$\displaystyle\overline{f}_{\theta}\sum_{t=1}^{T}\sum_{i=1}^{N}\nabla_{\theta}\log\overline{w}_{t-1}^{a_{t}^{i}}=\overline{f}_{\theta}\sum_{t=1}^{T}\sum_{i=1}^{N}\nabla_{\theta}\left(\log
w_{t-1}^{a_{t}^{i}}-\log W_{t-1}\right)=$ (69)
$\displaystyle\overline{f}_{\theta}\sum_{t=1}^{T}\sum_{i=1}^{N}\left(\nabla_{\theta}\log
w_{t-1}^{a_{t}^{i}}-\sum_{j=1}^{N}\overline{w}_{t-1}^{i}\nabla_{\theta}\log
w_{t-1}^{i}\right).$ (70)

When differentiating through the implementation of the same estimator obtained
by DPF, we obtain the following gradient estimator, using Eq. 57 and 62

$\displaystyle\nabla_{\theta}\mathop{\mathbb{E}}_{x\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{f_{\theta}(x_{1:T})}\right]\approx\overrightarrow{\nabla_{\theta}\sum_{i=1}^{N}\overline{v}_{T}^{i}f_{\theta}(\tilde{x}_{1:T}^{i})}=$
(71)
$\displaystyle\overrightarrow{\nabla_{\theta}\sum_{i=1}^{N}\bot\left({\overline{v}_{T}^{i}}\right)\frac{L_{T}^{i}}{\bot\left({L_{T}^{i}}\right)}\frac{\bot\left({\hat{Z}_{\text{DPF}}}\right)}{\hat{Z}_{\text{DPF}}}f_{\theta}(\tilde{x}_{1:T}^{i})}=$
(72)
$\displaystyle\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(f_{\theta}(\tilde{x}_{1:T}^{i})\nabla_{\theta}\log
L_{T}^{i}+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right)-\overline{f}_{\theta}\overrightarrow{\nabla_{\theta}\log\hat{Z}_{\text{DPF}}}=$
(73)
$\displaystyle\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(\left(f_{\theta}(\tilde{x}_{1:T}^{i})-\overline{f}_{\theta}\right)\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right).$
(74)

For the expectations under the unnormalized posterior, differentiating through
DPF produces the following gradient estimator

$\displaystyle\nabla_{\theta}\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{Zf_{\theta}(x_{1:T})}\right]\approx\overrightarrow{\nabla_{\theta}\hat{Z}_{\text{DPF}}\sum_{i=1}^{N}\overline{v}_{T}^{i}f_{\theta}(\tilde{x}_{1:T}^{i})}=$
(75)
$\displaystyle\overrightarrow{\nabla_{\theta}\hat{Z}_{\text{DPF}}\sum_{i=1}^{N}\bot\left({\overline{v}_{T}^{i}}\right)\frac{L_{T}^{i}}{\bot\left({L_{T}^{i}}\right)}\frac{\bot\left({\hat{Z}_{\text{DPF}}}\right)}{\hat{Z}_{\text{DPF}}}f_{\theta}(\tilde{x}_{1:T}^{i})}=$
(76)
$\displaystyle\hat{Z}\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(f_{\theta}(\tilde{x}_{1:T}^{i})\nabla_{\theta}\log
L_{T}^{i}+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right)=$ (77)
$\displaystyle\hat{Z}\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(f_{\theta}(\tilde{x}_{1:T}^{i})\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right).$
(78)

We compare it with the estimator obtained with PF and DiCE, which is also
unbiased

$\displaystyle\nabla_{\theta}\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{Zf_{\theta}(x_{1:T})}\right]\approx$ (79)
$\displaystyle\overline{f}_{\theta}\nabla_{\theta}\hat{Z}_{\text{PF}}+\hat{Z}_{\text{PF}}\nabla_{\theta}\sum_{i=1}^{N}\overline{w}_{T}^{i}f_{\theta}(\tilde{x}_{1:T}^{i})+\hat{Z}_{\text{PF}}\overline{f}_{\theta}\sum_{t=1}^{T}\sum_{i=1}^{N}\nabla_{\theta}\log\overline{w}_{t-1}^{a_{t}^{i}}=$
(80)
$\displaystyle\hat{Z}\overline{f}_{\theta}\sum_{t=1}^{T}\nabla_{\theta}\log
W_{t}+\hat{Z}\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})+\left(f_{\theta}(\tilde{x}_{1:T}^{i})-\overline{f}_{\theta}\right)\nabla_{\theta}\log
w_{T}^{i}\right)+$ (81)
$\displaystyle\quad\hat{Z}\overline{f}_{\theta}\sum_{t=1}^{T}\sum_{i=1}^{N}\left(\nabla_{\theta}\log
w_{t-1}^{a_{t}^{i}}-\sum_{j=1}^{N}\overline{w}_{t-1}^{i}\nabla_{\theta}\log
w_{t-1}^{i}\right)=$ (82)
$\displaystyle\hat{Z}\overline{f}_{\theta}\sum_{t=1}^{T}\nabla_{\theta}\log
W_{t}+\hat{Z}\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})+\left(f_{\theta}(\tilde{x}_{1:T}^{i})-\overline{f}_{\theta}\right)\nabla_{\theta}\log
w_{T}^{i}\right)+$ (83)
$\displaystyle\quad\hat{Z}\overline{f}_{\theta}\sum_{t=1}^{T}\sum_{i=1}^{N}\left(\nabla_{\theta}\log
w_{t-1}^{a_{t}^{i}}-\sum_{j=1}^{N}\overline{w}_{t-1}^{i}\nabla_{\theta}\log
w_{t-1}^{i}\right)=$ (84)
$\displaystyle\hat{Z}\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})+\left(f_{\theta}(\tilde{x}_{1:T}^{i})-\overline{f}_{\theta}\right)\nabla_{\theta}\log
w_{T}^{i}\right)+$ (85)
$\displaystyle\quad\hat{Z}\overline{f}_{\theta}\sum_{t=1}^{T}\sum_{i=1}^{N}\left(\nabla_{\theta}\log
w_{t-1}^{a_{t}^{i}}+\overline{w}_{t}^{i}\nabla_{\theta}\log
w_{t}^{i}-\sum_{j=1}^{N}\overline{w}_{t-1}^{i}\nabla_{\theta}\log
w_{t-1}^{i}\right).$ (86)

###  9.4 Gradients for proposal learning

The derivation leading up to Eq. 61 is valid for $\nabla_{\phi}$. Then we
compute the gradient of $\log L_{T}^{i}$ assuming reparameterization

$\displaystyle\overrightarrow{\nabla_{\phi}\log\hat{Z}_{\text{DPF}}}=\frac{\overrightarrow{\nabla_{\phi}\hat{Z}_{\text{DPF}}}}{\overrightarrow{\hat{Z}_{\text{DPF}}}}=\sum_{i=1}^{N}\overrightarrow{\overline{v}_{T}^{i}}\nabla_{\phi}\log
L_{T}^{i}=\sum_{i=1}^{N}\overline{w}_{T}^{i}\nabla_{\phi}\log\frac{p_{\theta}(\tilde{h}_{\phi}(\tilde{\epsilon}_{1:T}^{i}),y_{1:T})}{q_{\phi}(\tilde{h}_{\phi}(\tilde{\epsilon}_{1:T}^{i}))}.$
(87)

###  9.5 Score estimation with marginal particle filters

Poyiadjis et al. [2011, Eq. 20] define their score function estimator as

$\displaystyle\nabla_{\theta}\log
p_{\theta}(y_{1:T})\approx\sum_{i=1}^{N}\overline{w}_{T}^{i}\bar{\alpha}_{T}^{i},$
(88)

using the marginal particle filter and the following recursive definition of
$\alpha$

$\displaystyle\bar{\alpha}_{t+1}^{i}$
$\displaystyle=\frac{\sum_{j=1}^{N}\overline{w}_{t+1}^{j}p_{\theta}(x_{t+1}^{i}|x_{t}^{j})\left(\bar{\alpha}_{t}^{j}+\nabla_{\theta}\log
p_{\theta}(x_{t+1}^{i},y_{t+1})\right)}{\sum_{j=1}^{N}\overline{w}_{t+1}^{j}p_{\theta}(x_{t+1}^{i}|x_{t}^{j})}$
(89)
$\displaystyle=\frac{\sum_{j=1}^{N}\overline{w}_{t+1}^{j}p_{\theta}(x_{t+1}^{i},y_{t+1}|x_{t}^{j})\left(\bar{\alpha}_{t}^{j}+\nabla_{\theta}\log
p_{\theta}(x_{t+1}^{i},y_{t+1})\right)}{\sum_{j=1}^{N}\overline{w}_{t+1}^{j}p_{\theta}(x_{t+1}^{i},y_{t+1}|x_{t}^{j})}.$
(90)

We show that the same result is obtained by differentiating
$\log\hat{Z}_{\text{DPF2}}$. To find the corresponding recursion, we expand
the definition of $v_{t}^{i}$ in the following formula

$\displaystyle\overrightarrow{\nabla_{\theta}\frac{v_{t}^{i}}{\bot\left({v_{t}^{i}}\right)}\prod_{k=1}^{t-1}\frac{V_{k}}{\bot\left({V_{k}}\right)}}=$
(91)
$\displaystyle\overrightarrow{\nabla_{\theta}\frac{\sum_{i=1}^{N}\overline{v}_{t-1}^{i}\
p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})}{\bot\left({\sum_{i=1}^{N}\overline{v}_{t-1}^{i}\
p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})}\right)}\frac{\bot\left({\sum_{i=1}^{N}\bot\left({\overline{v}_{t-1}^{i}}\right)\
q_{\phi}(x_{t}^{n}|x_{t-1}^{i},y_{t})}\right)}{\sum_{i=1}^{N}\bot\left({\overline{v}_{t-1}^{i}}\right)\
q_{\phi}(x_{t}^{n}|x_{t-1}^{i},y_{t})}\prod_{k=1}^{t-1}\frac{V_{k}}{\bot\left({V_{k}}\right)}}=$
(92)
$\displaystyle\overrightarrow{\nabla_{\theta}\frac{\sum_{i=1}^{N}\bot\left({\overline{v}_{t-1}^{i}}\right)\
p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})\frac{\overline{v}_{t-1}^{i}}{\bot\left({\overline{v}_{t-1}^{i}}\right)}\prod_{k=1}^{t-1}\frac{V_{k}}{\bot\left({V_{k}}\right)}}{\bot\left({\sum_{i=1}^{N}\overline{v}_{t-1}^{i}\
p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})}\right)}}=$ (93)
$\displaystyle\frac{\sum_{i=1}^{N}\overline{w}_{t-1}^{i}\overrightarrow{\nabla_{\theta}\left(p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})\frac{v_{t-1}^{i}}{\bot\left({v_{t-1}^{i}}\right)}\prod_{k=1}^{t-2}\frac{V_{k}}{\bot\left({V_{k}}\right)}\right)}}{\sum_{i=1}^{N}\overline{w}_{t-1}^{i}p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})}=$
(94)
$\displaystyle\frac{\sum_{i=1}^{N}\overline{w}_{t-1}^{i}\left(\nabla_{\theta}p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})+p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})\overrightarrow{\nabla_{\theta}\frac{v_{t-1}^{i}}{\bot\left({v_{t-1}^{i}}\right)}\prod_{k=1}^{t-2}\frac{V_{k}}{\bot\left({V_{k}}\right)}}\right)}{\sum_{i=1}^{N}\overline{w}_{t-1}^{i}p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})}=$
(95)
$\displaystyle\frac{\sum_{i=1}^{N}\overline{w}_{t-1}^{i}p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})\left(\nabla_{\theta}\log
p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})+\overrightarrow{\nabla_{\theta}\frac{v_{t-1}^{i}}{\bot\left({v_{t-1}^{i}}\right)}\prod_{k=1}^{t-2}\frac{V_{k}}{\bot\left({V_{k}}\right)}}\right)}{\sum_{i=1}^{N}\overline{w}_{t-1}^{i}p_{\theta}(x_{t}^{n},y_{t}|x_{t-1}^{i})}.$
(96)

This is exactly the same recursive relationship as defined by Eq. 90, so we
have
$\overrightarrow{\nabla_{\theta}\frac{v_{t}^{i}}{\bot\left({v_{t}^{i}}\right)}\prod_{k=1}^{t-1}\frac{V_{k}}{\bot\left({V_{k}}\right)}}=\bar{\alpha}_{t}^{i}$.
Differentiating through DPF2 produces the following expression

$\displaystyle\overrightarrow{\nabla_{\theta}\log\hat{Z}_{\text{DPF2}}}=\frac{\overrightarrow{\nabla_{\theta}\hat{Z}_{\text{DPF2}}}}{\hat{Z}}=\frac{1}{\hat{Z}}\overrightarrow{\nabla_{\theta}\prod_{t=1}^{T}V_{t}}=\frac{1}{\hat{Z}}\overrightarrow{\nabla_{\theta}\sum_{i=1}^{N}v_{T}^{i}\prod_{t=1}^{T-1}V_{t}}=$
(97)
$\displaystyle\frac{1}{\hat{Z}}\overrightarrow{\nabla_{\theta}\sum_{i=1}^{N}\bot\left({\overline{v}_{T}^{i}V_{T}}\right)\frac{v_{T}^{i}}{\bot\left({v_{T}^{i}}\right)}\prod_{t=1}^{T-1}V_{t}}=\frac{1}{\hat{Z}}\overrightarrow{\nabla_{\theta}\bot\left({\hat{Z}}\right)\sum_{i=1}^{N}\bot\left({\overline{v}_{T}^{i}}\right)\frac{v_{T}^{i}}{\bot\left({v_{T}^{i}}\right)}\prod_{t=1}^{T-1}\frac{V_{t}}{\bot\left({V_{t}}\right)}}=$
(98)
$\displaystyle\sum_{i=1}^{N}\overline{w}_{T}^{i}\overrightarrow{\nabla_{\theta}\frac{v_{T}^{i}}{\bot\left({v_{T}^{i}}\right)}\prod_{t=1}^{T-1}\frac{V_{t}}{\bot\left({V_{t}}\right)}}=\sum_{i=1}^{N}\overline{w}_{T}^{i}\bar{\alpha}_{T}^{i},$
(99)

recovering Eq. 88.

###  9.6 Hessian of log-marginal likelihood with DPF

While in most applications we are interested in first-order derivatives, we
can apply automatic differentiation repeatedly to obtain derivatives of higher
orders. Poyiadjis et al. [2011] show that the Hessian of the log-marginal
likelihood can be approximated with a particle filter using Louis’ identity

$\displaystyle\nabla_{\theta}^{2}\log p_{\theta}(y_{1:T})=-\nabla_{\theta}\log
p_{\theta}(y_{1:T})\nabla_{\theta}\log p_{\theta}(y_{1:T})^{T}+$ (100)
$\displaystyle\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{\nabla_{\theta}^{2}\log
p_{\theta}(x_{1:T},y_{1:T})+\nabla_{\theta}\log
p_{\theta}(x_{1:T},y_{1:T})\nabla_{\theta}\log
p_{\theta}(x_{1:T},y_{1:T})^{T}}\right]\approx$ (101)
$\displaystyle-\sum_{i=1}^{N}\sum_{j=1}^{N}\overline{w}_{T}^{i}\overline{w}_{T}^{j}\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})^{T}+$ (102)
$\displaystyle\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(\nabla_{\theta}^{2}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})+\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},,y_{1:T})\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},,y_{1:T})^{T}\right),$ (103)

where superscript $T$ denotes transpose.

We can derive the estimator produced by applying automatic differentiation
twice to DPF using Eq. 60

$\displaystyle\overrightarrow{\nabla_{\theta}^{2}\log\hat{Z}_{\text{DPF}}}=\overrightarrow{\nabla_{\theta}\frac{\nabla_{\theta}\hat{Z}_{\text{DPF}}}{\hat{Z}_{\text{DPF}}}}=\overrightarrow{\nabla_{\theta}\left(\frac{\bot\left({\hat{Z}_{\text{DPF}}}\right)}{\hat{Z}_{\text{DPF}}}\sum_{i=1}^{N}\bot\left({\overline{v}_{T}^{i}}\right)\frac{\nabla_{\theta}L_{T}^{i}}{\bot\left({L_{T}^{i}}\right)}\right)}=$
(104)
$\displaystyle-\overrightarrow{\nabla_{\theta}\log\hat{Z}_{\text{DPF}}}\sum_{i=1}^{N}\overline{w}_{T}^{i}\nabla_{\theta}\log
L_{T}^{i}+\sum_{i=1}^{N}\overline{w}_{T}^{i}\overrightarrow{\frac{\nabla_{\theta}^{2}L_{T}^{i}}{\bot\left({L_{T}^{i}}\right)}}=$
(105)
$\displaystyle-\sum_{i=1}^{N}\sum_{j=1}^{N}\overline{w}_{T}^{i}\overline{w}_{T}^{j}\nabla_{\theta}\log
L_{T}^{i}(\nabla_{\theta}\log
L_{T}^{j})^{T}+\sum_{i=1}^{N}\overline{w}_{T}^{i}\overrightarrow{\nabla_{\theta}\left(\frac{L_{T}^{i}}{\bot\left({L_{T}^{i}}\right)}(\nabla_{\theta}\log
L_{T}^{i})^{T}\right)}=$ (106)
$\displaystyle-\sum_{i=1}^{N}\sum_{j=1}^{N}\overline{w}_{T}^{i}\overline{w}_{T}^{j}\nabla_{\theta}\log
L_{T}^{i}(\nabla_{\theta}\log
L_{T}^{j})^{T}+\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(\nabla_{\theta}\log
L_{T}^{i}(\nabla_{\theta}\log L_{T}^{i})^{T}+\nabla_{\theta}^{2}\log
L_{T}^{i}\right),$ (107)

which is exactly the same as Eq. 103. Furthermore, Hessian-vector products can
be computed with reduced memory requirements in the usual fashion.

##  10 Proofs

Consistency of the DPF estimator for the expectation under the posterior (Eq.
74) can be shown as follows

$\displaystyle\lim_{N\rightarrow\infty}\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(\left(f_{\theta}(\tilde{x}_{1:T}^{i})-\overline{f}_{\theta}\right)\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right)=$
(108)
$\displaystyle\lim_{N\rightarrow\infty}\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(\left(f_{\theta}(\tilde{x}_{1:T}^{i})-\overline{f}_{\theta}\right)\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i}|y_{1:T})+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right)=$
(109) $\displaystyle\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{\left(f_{\theta}(x_{1:T})-\mathop{\mathbb{E}}_{x^{\prime}_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{f_{\theta}(x^{\prime}_{1:T})}\right]\right)\nabla_{\theta}\log
p_{\theta}(x_{1:T}|y_{1:T})+\nabla_{\theta}f_{\theta}(x_{1:T})}\right]=$ (110)
$\displaystyle\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{f_{\theta}(x_{1:T})\nabla_{\theta}\log
p_{\theta}(x_{1:T}|y_{1:T})+\nabla_{\theta}f_{\theta}(x_{1:T})}\right]=\nabla_{\theta}\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{f_{\theta}(x_{1:T})}\right].$ (111)

Similarly for the consistency of the DPF estimator for the expectation under
the unnormalized posterior (Eq. 21)

$\displaystyle\lim_{N\rightarrow\infty}\mathop{\mathbb{E}}_{PF}\left[{\hat{Z}\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(f_{\theta}(\tilde{x}_{1:T}^{i})\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right)}\right]=$
(112) $\displaystyle\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{p_{\theta}(y_{1:T})\left(f_{\theta}(x_{1:T})\nabla_{\theta}\log
p_{\theta}(x_{1:T},y_{1:T})+\nabla_{\theta}f_{\theta}(x_{1:T})\right)}\right]=$
(113) $\displaystyle\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{p_{\theta}(y_{1:T})\left(f_{\theta}(x_{1:T})\nabla_{\theta}\log
p_{\theta}(x_{1:T}|y_{1:T})+\nabla_{\theta}f_{\theta}(x_{1:T})\right)+f_{\theta}(x_{1:T})\nabla_{\theta}p_{\theta}(y_{1:T})}\right]=$
(114) $\displaystyle\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{p_{\theta}(y_{1:T})f_{\theta}(x_{1:T})\nabla_{\theta}\log
p_{\theta}(x_{1:T}|y_{1:T})+\nabla_{\theta}\left(f_{\theta}(x_{1:T})p_{\theta}(y_{1:T})\right)}\right]=$
(115) $\displaystyle\nabla_{\theta}\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{p_{\theta}(y_{1:T})f_{\theta}(x_{1:T})}\right].$
(116)

######  Theorem 1.

Eq. 74 is an unbiased estimator for the expectation under the unnormalized
posterior, that is

$\displaystyle\mathop{\mathbb{E}}_{x_{t}^{i}\sim
q_{\phi}(x_{t}|x_{t-1}^{a_{t}^{i}},y_{t}),\ a_{t}^{1:N}\sim
R(\overline{w}_{t-1}^{1:N})}\left[{\hat{Z}\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(f_{\theta}(\tilde{x}_{1:T}^{i})\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T})+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:T}^{i})\right)}\right]=$
(117) $\displaystyle\nabla_{\theta}\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{p_{\theta}(y_{1:T})f_{\theta}(x_{1:T})}\right].$
(118)

###### Proof.

The proof is by induction on $T$ and the inductive hypothesis is precisely the
theorem for any function $f$.

The inductive step is as follows. Let
$g(x_{1:T+1})=f_{\theta}(x_{1:{T+1}})\nabla_{\theta}\log
p_{\theta}(x_{1:{T+1}},y_{1:{T+1}})+\nabla_{\theta}f_{\theta}(x_{1:{T+1}})$
and $h(x_{1:T})=\mathop{\mathbb{E}}_{x_{T+1}\sim
p_{\theta}(x_{T+1}|y_{T+1},x_{T}^{i})}\left[{p_{\theta}(y_{T+1}|x_{T}^{i})f_{\theta}(x_{1:{T+1}})}\right]$.
We first show that

$\displaystyle\mathop{\mathbb{E}}_{x_{T+1}\sim
p_{\theta}(x_{T+1}|y_{T+1},x_{T})}\left[{p_{\theta}(y_{T+1}|x_{T})g(x_{1:T+1})}\right]=$
(119) $\displaystyle\mathop{\mathbb{E}}_{x_{T+1}\sim
p_{\theta}(x_{T+1}|y_{T+1},x_{T})}\left[{p_{\theta}(y_{T+1}|x_{T})\left(f_{\theta}(x_{1:{T+1}})\nabla_{\theta}\log
p_{\theta}(x_{1:{T+1}},y_{1:{T+1}})+\nabla_{\theta}f_{\theta}(x_{1:{T+1}})\right)}\right]=$
(120) $\displaystyle\mathop{\mathbb{E}}_{x_{T+1}\sim
p_{\theta}(x_{T+1}|y_{T+1},x_{T})}\left[{p_{\theta}(y_{T+1}|x_{T})f_{\theta}(x_{1:{T+1}})}\right]\nabla_{\theta}\log
p_{\theta}(x_{1:{T}},y_{1:{T}})+$ (121)
$\displaystyle\mathop{\mathbb{E}}_{x_{T+1}\sim
p_{\theta}(x_{T+1}|y_{T+1},x_{T})}\left[{p_{\theta}(y_{T+1}|x_{T})\left(f_{\theta}(x_{1:{T+1}})\nabla_{\theta}\log
p_{\theta}(x_{{T+1}},y_{{T+1}}|x_{T})+\nabla_{\theta}f_{\theta}(x_{1:{T+1}})\right)}\right]=$
(122) $\displaystyle h(x_{1:T})\nabla_{\theta}\log
p_{\theta}(x_{1:{T}},y_{1:{T}})+\mathop{\mathbb{E}}_{x_{T+1}}\Bigg{[}p_{\theta}(y_{T+1}|x_{T})f_{\theta}(x_{1:{T+1}})\nabla_{\theta}\log
p_{\theta}(x_{{T+1}}|y_{{T+1}},x_{T})+$ $\displaystyle\quad\quad
p_{\theta}(y_{T+1}|x_{T})f_{\theta}(x_{1:{T+1}})\frac{\nabla_{\theta}p_{\theta}(y_{T+1}|x_{T})}{p_{\theta}(y_{T+1}|x_{T})}+p_{\theta}(y_{T+1}|x_{T})\nabla_{\theta}f_{\theta}(x_{1:{T+1}})\Bigg{]}=$
(123) $\displaystyle h(x_{1:T})\nabla_{\theta}\log
p_{\theta}(x_{1:{T}},y_{1:{T}})+\mathop{\mathbb{E}}_{x_{T+1}\sim
p_{\theta}(x_{T+1}|y_{T+1},x_{T})}\Big{[}p_{\theta}(y_{T+1}|x_{T})f_{\theta}(x_{1:{T+1}})\nabla_{\theta}\log
p_{\theta}(x_{{T+1}}|y_{{T+1}},x_{T})+$
$\displaystyle\quad\quad\nabla_{\theta}\left(p_{\theta}(y_{T+1}|x_{T})f_{\theta}(x_{1:{T+1}})\right)\Big{]}=$
(124) $\displaystyle h(x_{1:T})\nabla_{\theta}\log
p_{\theta}(x_{1:{T}},y_{1:{T}})+\nabla_{\theta}\mathop{\mathbb{E}}_{x_{T+1}\sim
p_{\theta}(x_{T+1}|y_{T+1},x_{T})}\left[{p_{\theta}(y_{T+1}|x_{T})f_{\theta}(x_{1:{T+1}})}\right]=$
(125) $\displaystyle h(x_{1:T})\nabla_{\theta}\log
p_{\theta}(x_{1:{T}},y_{1:{T}})+\nabla_{\theta}h(x_{1:T}).$ (126)

Then

$\displaystyle\mathop{\mathbb{E}}_{x_{t}^{1:N},a_{t}^{1:N}}\left[{\left(\prod_{t=1}^{T+1}W_{t}\right)\sum_{i=1}^{N}\overline{w}_{T+1}^{i}\left(f_{\theta}(\tilde{x}_{1:{T+1}}^{i})\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:{T+1}}^{i},y_{1:{T+1}})+\nabla_{\theta}f_{\theta}(\tilde{x}_{1:{T+1}}^{i})\right)}\right]=$
(127)
$\displaystyle\mathop{\mathbb{E}}_{x_{t}^{1:N},a_{t}^{1:N}}\left[{\left(\prod_{t=1}^{T+1}W_{t}\right)\sum_{i=1}^{N}\overline{w}_{T+1}^{i}g(\tilde{x}_{1:T+1}^{i})}\right]=$
(128)
$\displaystyle\mathop{\mathbb{E}}_{x_{1:T}^{1:N},a_{1:T}^{1:N}}\left[{\left(\prod_{t=1}^{T}W_{t}\right)\mathop{\mathbb{E}}_{a_{T+1}^{1:N},x_{T+1}^{1:N}}\left[{W_{T+1}\sum_{i=1}^{N}\overline{w}_{T+1}^{i}g(\tilde{x}_{1:T+1}^{i})}\right]}\right]=$
(129)
$\displaystyle\mathop{\mathbb{E}}_{x_{1:T}^{1:N},a_{1:T}^{1:N}}\left[{\left(\prod_{t=1}^{T}W_{t}\right)\mathop{\mathbb{E}}_{a_{T+1}^{1:N}}\left[{\mathop{\mathbb{E}}_{x_{T+1}^{1:N}}\left[{\sum_{i=1}^{N}w_{T+1}^{i}g(\tilde{x}_{1:T+1}^{i})}\right]}\right]}\right]=$
(130)
$\displaystyle\mathop{\mathbb{E}}_{x_{1:T}^{1:N},a_{1:T}^{1:N}}\left[{\left(\prod_{t=1}^{T}W_{t}\right)\mathop{\mathbb{E}}_{a_{T+1}^{1:N}}\left[{\frac{1}{N}\sum_{i=1}^{N}\mathop{\mathbb{E}}_{x_{T+1}\sim
q_{\phi}(x_{T+1}|x_{T}^{a_{T+1}^{i}})}\left[{\frac{p_{\theta}(x_{T+1},y_{T+1}|x_{T}^{a_{T+1}^{i}})}{q_{\phi}(x_{T+1}|x_{T}^{a_{T+1}^{i}})}g(\tilde{x}_{1:T}^{a_{T+1}^{i}},x_{T+1})}\right]}\right]}\right]=$
(131)
$\displaystyle\mathop{\mathbb{E}}_{x_{1:T}^{1:N},a_{1:T}^{1:N}}\left[{\left(\prod_{t=1}^{T}W_{t}\right)\mathop{\mathbb{E}}_{a_{T+1}^{1:N}}\left[{\frac{1}{N}\sum_{i=1}^{N}\mathop{\mathbb{E}}_{x_{T+1}\sim
p_{\theta}(x_{T+1}|y_{T+1},x_{T}^{a_{T+1}^{i}})}\left[{p_{\theta}(y_{T+1}|x_{T}^{a_{T+1}^{i}})g(\tilde{x}_{1:T}^{a_{T+1}^{i}},x_{T+1})}\right]}\right]}\right]=$
(132)
$\displaystyle\mathop{\mathbb{E}}_{x_{1:T}^{1:N},a_{1:T}^{1:N}}\left[{\left(\prod_{t=1}^{T}W_{t}\right)\sum_{i=1}^{N}\overline{w}_{T}^{i}\mathop{\mathbb{E}}_{x_{T+1}\sim
p_{\theta}(x_{T+1}|y_{T+1},x_{T}^{i})}\left[{p_{\theta}(y_{T+1}|x_{T}^{i})g(\tilde{x}_{1:T}^{i},x_{T+1})}\right]}\right]=$
(133)
$\displaystyle\mathop{\mathbb{E}}_{x_{1:T}^{1:N},a_{1:T}^{1:N}}\left[{\left(\prod_{t=1}^{T}W_{t}\right)\sum_{i=1}^{N}\overline{w}_{T}^{i}\left(h(x_{1:T})\nabla_{\theta}\log
p_{\theta}(x_{1:{T}},y_{1:{T}})+\nabla_{\theta}h(x_{1:T})\right)}\right]=$
(134) $\displaystyle\nabla_{\theta}\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{p_{\theta}(y_{1:T})h(x_{1:T})}\right]=\nabla_{\theta}\mathop{\mathbb{E}}_{x_{1:T+1}\sim
p_{\theta}(x_{1:T+1}|y_{1:T+1})}\left[{p_{\theta}(y_{1:T+1})f_{\theta}(x_{1:T+1})}\right],$
(135)

where the penultimate equality is the inductive hypothesis for $T$ and $h$.
This completes the proof.

∎

##  11 Additional experiments

In Table 1 we provide numerical results for the experiments presented in the
main text. We also include an additional experiment, with a comparison to the
differentiable particle filter based on optimal transport (DPF-OT) proposed by
Corenflos et al. [2021].

N | Method | LGSSM
---|---|---
| PF-SF | -475.73 $\pm$ 55.69
10 | SIS | -2742.99 $\pm$ 279.97
| PF | -300.46 $\pm$ 2.62
| DPF-SGR (us) | -292.70 $\pm$ 1.28
| PF-SF | -1038.85 $\pm$ 176.02
10k | SIS | -343.65 $\pm$ 15.61
| PF | -290.04 $\pm$ 0.00
| DPF-SGR (us) | -290.08 $\pm$ 0.02

N | Method | VRNN
---|---|---
| – | –
10 | SIS | -356.22 $\pm$ 0.08
| PF | -355.98 $\pm$ 0.13
| DPF-SGR (us) | -355.82 $\pm$ 0.14
| – | –
20 | SIS | -356.41 $\pm$ 0.50
| PF | -355.97 $\pm$ 0.16
| DPF-SGR (us) | -355.74 $\pm$ 0.02

Table 1: Best values for the test set log-marginal likelihood $\log Z$ for $N$
particles in the LGSSM/VRNN experiments.

###  11.1 Stochastic Volatility Model

Figure 3: Left is the L1 error in parameter estimation throughout the training
of the stochastic volatility model. Right is the average wall clock execution
time for a single forward and backward pass for each particle filter variant.

In this experiment we learn the parameters of a stochastic volatility model, a
popular financial model for which particle filters are known to perform well.
We use the implementation provided in the code repository associated with the
paper of Corenflos et al. [2021]222https://github.com/JTT94/filterflow, which
allows us to reuse the DPF-OT settings provided by the original authors. The
model is described by the following equations:

$\displaystyle x_{0}$
$\displaystyle\sim\mathcal{N}\left(0,\frac{\sigma^{2}_{x}}{(1-\phi^{2})}\right)$
$\displaystyle x_{t+1}$ $\displaystyle=\mu\cdot(1-\phi)+\phi\cdot
x_{t}+\sigma_{x}\cdot\eta_{t},$ $\displaystyle y_{t}$
$\displaystyle=\epsilon_{t}\cdot\exp\left(\frac{x_{t}}{2}\right),$
$\displaystyle\eta_{t}$ $\displaystyle\sim\mathcal{N}(0,1),t=1,...,T,$
$\displaystyle\epsilon_{t}$ $\displaystyle\sim\mathcal{N}(0,1).$

We generate synthetic data with T = 100, $\mu$ = 2, $\phi$ = 0.9, $\sigma_{x}$
= 1, $\theta_{1:3}=\\{\mu,\phi,\sigma_{x}\\}$ and train the model for 500
epochs using Adam with a learning rate of $0.01$. Resampling is triggered when
the effective sample size falls beneath $N/2$, and for all methods we use the
bootstrap proposal like in the LGSSM experiment. Figure 3 shows progression of
the parameter estimation error $|\hat{\mu}-\mu|$ throughout training and the
execution time of different algorithms. DPF-SGR produces results at least on
par with DPF-OT, while executing almost as fast as PF.

##  12 Additional background

Particle filters [Doucet and Johansen, 2009, Chopin and Papaspiliopoulos,
2020, Doucet et al., 2001] are an extremely popular family of algorithms for
state estimation in non-linear state space models. Their versatility, ease of
implementation, and well-understood theoretical behavior make them a popular
choice in a wide variety of application domains, including robotics, ecology,
and finance. The crucial advantage of particle filters is that they can
adaptively focus computation on more promising trajectories, which is usually
accomplished through discrete resampling steps. Since particle filters also
provide estimates of the log-marginal likelihood, they can be used to
facilitate model learning.

It has been known for a long time [Wengert, 1964, Griewank, 1989] that
programs computing derivatives can be automatically obtained from programs
computing differentiable functions, the corresponding algorithms being known
as “automatic differentiation” (AD). A specialized version of this technique
is used in deep learning under the name “backpropagation” [Rumelhart et al.,
1986], more recently becoming popular across all of machine learning [Baydin
et al., 2018], accelerated by the availability of open-source libraries such
as PyTorch [Paszke et al., 2019] and TensorFlow [Abadi et al., 2015]. In
recent years there has been an explosion of interest in the machine learning
community in computing gradients of expressions involving sampling random
variables, motivated primarily by problems in deep generative modeling [Neal
and Hinton, 1998] and reinforcement learning [Williams, 1992]. While generic
methods for extending AD to handle stochastic computations have been developed
[Schulman et al., 2016, Foerster et al., 2018], the variance of resulting
gradient estimators is still problematic[Rezende et al., 2014] where the
reparameterization trick [Kingma and Welling, 2014, Ranganath et al., 2013]
does not apply, in particular for discrete variables.

Building on those developments, it is becoming increasingly popular to learn
sequential latent variable models by gradient ascent using estimators derived
from the application of AD to the log-marginal likelihood estimators obtained
with particle filters. Examples include robotic localization [Jonschkowski et
al., 2018], world-modeling in reinforcement learning [Igl et al., 2018], and
deep generative modeling of music and speech [Maddison et al., 2017]. The main
obstacle is the discrete resampling steps, for which existing gradient
estimators suffer from excessive variance, leading to the dependence of the
resampling probabilities on model parameters typically being ignored [Le et
al., 2018a, Maddison et al., 2017, Naesseth et al., 2018], introducing bias
that was shown not to vanish asymptotically [Corenflos et al., 2021].

Several recent articles[Corenflos et al., 2021, Zhu et al., 2020, Karkus et
al., 2018] have proposed continuous relaxations of particle filter discrete
resampling steps to avoid this problem, with different trade-offs and varying
degrees of complexity.

Alternatively, there exist consistent estimators of the gradient of log-
marginal likelihood that can be obtained with an unmodified pass of a particle
filter with usual discrete resampling [Poyiadjis et al., 2011], but until now
these could not be obtained by applying AD to the log-marginal likelihood
estimator produced by a particle filter.

###  12.1 Particle filters

Here we briefly review the basic particle filter algorithm, summarized in
Algorithm 1, referring the reader to [Chopin and Papaspiliopoulos, 2020,
Doucet and Johansen, 2009, Doucet et al., 2001] for more details. For a given
sequence of observations $y_{1:T}$, particle filters approximate the filtering
distributions $p_{\theta}(x_{t}|y_{1:t})$ and the marginal likelihood
$p_{\theta}(y_{1:T})$ of a state space model with the following factorization
of the joint

$\displaystyle
p_{\theta}(x_{0:T},y_{1:T})=p_{\theta}(x_{0})\prod_{t=1}^{T}p_{\theta}(x_{t}|x_{t-1})p_{\theta}(y_{t}|x_{t}).$
(136)

The algorithm starts by sampling $N$ particles independently from an initial
distribution $x_{0}^{i}\sim p_{\theta}(x_{0})$ and setting their weights to be
uniform $w_{0}^{i}=\frac{1}{N}$. Then, at each iteration it advances one step
forward in time, by first advancing each particle according to a proposal
distribution $q_{\phi}(x_{t+1}|x_{t})$, then accumulating the importance
sampling ratio $\frac{p_{\theta}(x_{t+1}|x_{t})}{q_{\phi}(x_{t+1}|x_{t})}$ and
the likelihood $p_{\theta}(y_{t+1}|x_{t+1})$ in the corresponding weight,
saving the sum of weights and normalizing them before proceeding to the next
time step. We assume that $q$ and $p$ have the same support and that the
likelihoods are always non-zero, which guarantees that all the weights are
always positive. We also assume that both $p$ and $q$ are differentiable with
respect to their parameters.

As the weights accumulate more terms, they become increasingly unbalanced and
it is more efficient to focus more computational resources to particles with
higher weights. This is accomplished via a resampling step, which in our
formulation is optionally performed at the beginning of each time step. The
ancestral indices for each particle are chosen from a resampling distribution
$R(\overline{w}_{t-1}^{1:N})$, which in the simplest multinomial resampling
picks ancestral indices independently from a categorical distribution, but
other variants like stratified and systematic resampling are usually favored
in practice [Murray et al., 2016]. We assume post-resampling weights are
always uniform for simplicity, although non-uniform resampling schemes also
exist [Fearnhead and Clifford, 2003]. Note that we write $\tilde{w}$ for the
weights obtained after resampling, which in a standard particle filter are
always $\frac{1}{N}$, $w$ for weights after the likelihood was accumulated at
a given time step, and $\overline{w}$ for the subsequently normalized weights.

Regardless of the specific details of the particle filter algorithm used, the
marginal likelihood $Z=p_{\theta}(y_{1:T})$ is approximated by
$\hat{Z}=\prod_{t=1}^{T}W_{t}$. This estimator is known to be unbiased and
under weak assumptions consistent. The estimator for the log-marginal
likelihood $\log p_{\theta}(y_{1:T})$ given by $\log\hat{Z}$ is consistent but
biased. Similarly, a particle filter produces consistent but biased estimators
for expectations under the posterior, while the estimators for the
expectations under the unnormalized posterior are additionally unbiased. We
refer the reader to Moral [2004] and Chopin and Papaspiliopoulos [2020] for
detailed convergence results.

###  12.2 Score function estimation with particle filters

Maximum likelihood learning of generative models by gradient ascent can be
performed using the score function, which is the gradient of the log-marginal
likelihood $\nabla_{\theta}\log p_{\theta}(y_{1:T})$. While the exact value of
the score function is in general intractable, it can be approximated using any
method for sampling from the posterior distribution with the help of Fisher’s
identity

$\displaystyle\nabla_{\theta}\log
p_{\theta}(y_{1:T})=\nabla_{\theta}\log\mathop{\mathbb{E}}_{x_{1:T}\sim
p(x_{1:T})}\left[{p_{\theta}(x_{1:T},y_{1:T})}\right]=$ (137)
$\displaystyle\mathop{\mathbb{E}}_{x_{1:T}\sim
p_{\theta}(x_{1:T}|y_{1:T})}\left[{\nabla_{\theta}\log
p_{\theta}(x_{1:T},y_{1:T})}\right].$ (138)

In the particle filtering context, the approximate samples from the posterior
can be obtained by tracing the ancestral lineages of the surviving particles
form the final time step. Let $\tilde{x}_{1:t}^{i}$ be the ancestral lineage
of a particle $i$ at time $t$, that is
$\tilde{x}_{1:t+1}^{i}=(\tilde{x}_{1:t}^{a_{t+1}^{i}},x_{t+1}^{i})$ and
$\tilde{x}_{1}^{i}=x_{1}^{i}$. Poyiadjis et al. [2011] propose a consistent
score function estimator of the form

$\displaystyle\nabla_{\theta}\log
p_{\theta}(y_{1:T})\approx\sum_{i=1}^{N}\overline{w}_{T}^{i}\nabla_{\theta}\log
p_{\theta}(\tilde{x}_{1:T}^{i},y_{1:T}).$ (139)

They also show how to compute this estimator without having to store full
ancestral lineages and that the variance of this estimator can be further
reduced using the marginal particle filter.

In this work we show how to obtain this estimator through automatic
differentiation of $\log\hat{Z}$ produced by a particle filter, without
modifying the estimators obtained on the forward pass.

###  12.3 The stop-gradient operator

The stop-gradient operator is a standard component of automatic
differentiation libraries, called stop_gradient in Tensorflow and detach in
PyTorch, producing an expression that evaluates to its normal value on the
forward pass but has zero gradient. While stop-gradient does not correspond to
any formal mathematical function, we need to include it in our formulas to
formally reason about programs involving stop-gradient and how automatic
differentiation acts upon them. Following Foerster et al. [2018], we denote
the stop-gradient with a $\bot$, for example $\nabla\bot\left({x}\right)=0$.

Expressions involving the stop-gradient operator should be understood as
expressions in a formal calculus, which correspond to programs written using
automatic differentiation libraries. Following established terminology, we
sometimes refer to them as _surrogate expressions_. For such expressions we
need to distinguish between the results obtained by evaluating them on the
forward pass and the expressions obtained by applying the gradient operator.
Following van Krieken et al. [2021], we denote evaluation of an expression $E$
with an overhead right arrow $\overrightarrow{E}$. Operationally, evaluation
removes all instances of stop-gradient, provided that no gradient operator is
acting on them. Gradients satisfy the usual chain rules, except that gradients
of any expression wrapped in stop-gradient are zero. Formally, the key
equations are:

$\displaystyle\overrightarrow{f(E_{1},\dots,E_{n})}$
$\displaystyle=f(\overrightarrow{E_{1}},\dots,\overrightarrow{E_{2}})$ (140)
$\displaystyle\overrightarrow{\bot\left({E}\right)}$
$\displaystyle=\overrightarrow{E}$ (141)
$\displaystyle\nabla{\bot\left({E}\right)}$ $\displaystyle=0$ (142)
$\displaystyle\overrightarrow{\nabla{E}}$
$\displaystyle=\nabla{E}\quad\text{if}\ \bot\notin E,$ (143)

where $f$ is some differentiable function and $\bot\notin E$ means that stop-
gradient is not applied anywhere in expression $E$, including its
subexpressions. We use this calculus to derive formulas for the estimators
obtained by algorithms utilizing stop-gradient in automatic differentiation,
which is generally accomplished by using the rules above, along with the chain
rule for the gradient operator, to push the gradients inside, past any stop-
gradient operators, and then evaluating the resulting expression. Note that in
this calculus the gradient operator produces expressions rather than values,
so that we can model repeated automatic differentiation. Expressions without
stop-gradient can be equated with their usual interpretations. Throughout the
paper we are careful to avoid using integrals of expressions involving stop-
gradient, as we consider those undefined. We also present a short tutorial on
using this calculus in the appendix.

In previous work [Foerster et al., 2018, van Krieken et al., 2021], the stop-
gradient operator has been used to show that certain expressions produce
unbiased estimators for derivatives of arbitrary order under repeated
automatic differentiation. In this work we consider a different family of
expressions involving stop-gradient, showing that under automatic
differentiation they produce score function estimators corresponding to the
application of Fisher’s identity and self-normalizing importance sampling.

###  12.4 Automatic differentiation of stochastic computation

When using AD for computation involving stochastic choices, it is typically
desirable to obtain gradient estimators that in expectation equal the gradient
of the expected value of the original computation. The generic methods for
accomplishing this generally fall into two categories. Consider a computation
with a single random variable, $\mathcal{L}=\mathop{\mathbb{E}}_{x\sim
p_{\theta}(x)}\left[{f_{\theta}(x)}\right]$. If $p_{\theta}$ is a continuous
distribution, we can express $x$ as a deterministic function $h$ of some
random variable $\epsilon$ with a fixed distribution independent of $\theta$,
such as $\mathcal{N}(0,1)$. With that, the gradient commutes with the
expectation

$\displaystyle\nabla_{\theta}\mathcal{L}=\nabla_{\theta}\mathop{\mathbb{E}}_{\epsilon\sim\mathcal{N}(0,1)}\left[{f_{\theta}(h(\epsilon,\theta))}\right]=\mathop{\mathbb{E}}_{\epsilon\sim\mathcal{N}(0,1)}\left[{\nabla_{\theta}f_{\theta}(h(\epsilon,\theta))}\right],$
(144)

at which point AD can be applied for a fixed $\epsilon$. This technique is
broadly known as the pathwise derivative [Glasserman, 2003], while in the
machine learning it’s usually being referred to as the reparameterization
trick [Kingma and Welling, 2014, Ranganath et al., 2013]. It is easy to apply
for some distributions such as Gaussian and uniform, but in general finding a
suitable $\epsilon$ and $h$ that can be efficiently computed is difficult and
an active research area [Jankowiak and Obermeyer, 2018, Graves, 2016]. When
applicable, this method tends to produce gradient estimators with relatively
low variance.

An alternative, typically used when reparemeterization is not possible, is
based on the following identity

$\displaystyle\nabla_{\theta}\mathop{\mathbb{E}}_{x\sim
p_{\theta}(x)}\left[{f_{\theta}(x)}\right]=\mathop{\mathbb{E}}_{x\sim
p_{\theta}(x)}\left[{\nabla_{\theta}f_{\theta}(x)+f_{\theta}(x)\nabla_{\theta}\log
p_{\theta}(x)}\right].$ (145)

The two gradients inside the expectation are typically computed with AD
jointly for a fixed $x$ by forming a surrogate expression
$f_{\theta}(x)+\bot\left({f_{\theta}(x)}\right)\log p_{\theta}(x)$. In the
reinforcement learning literature this method is known as REINFORCE [Williams,
1992], and more broadly referred to as the score function method or the
likelihood ratio method [Fu, 2006]. We refrain from using those terms to avoid
confusion with the score function $\nabla_{\theta}\log p_{\theta}(y_{1:T})$
and the likelihood $p_{\theta}(y_{1:T})$ associated with the sequential latent
variable model. Instead, we refer to this method as DiCE [Foerster et al.,
2018], which produces the same gradient estimator using the following
surrogate expression

$\displaystyle\overrightarrow{\nabla_{\theta}\left(\frac{p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}f_{\theta}(x)\right)}=$
(146) $\displaystyle
f_{\theta}(x)\overrightarrow{\frac{\nabla_{\theta}p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}}+\overrightarrow{\frac{p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}}\nabla_{\theta}f_{\theta}(x)=$
(147) $\displaystyle f_{\theta}(x)\nabla_{\theta}\log
p_{\theta}(x)+\nabla_{\theta}f_{\theta}(x),$ (148)

while also remaining unbiased under repeated AD. Incidentally, the term
$\frac{p_{\theta}(x)}{\bot\left({p_{\theta}(x)}\right)}$, which Foerster et
al. [2018] call “the magic box”, is precisely the AD-centric expression of the
likelihood ratio from which the likelihood ratio method derives its name. The
rules for constructing the surrogate expression in the presence of multiple
random variables and complex dependencies were given by Schulman et al.
[2016], with Foerster et al. [2018] deriving equivalent rules for DiCE. While
DiCE is broadly applicable, the resulting gradient estimators tend to have
high variance, which to some extent can be mitigated through a careful use of
baselines [Mao et al., 2019, Farquhar et al., 2019].