Automatic Differentiation of Programs with Discrete Randomness
Anonymous Author(s)
AfﬁliationAddress
Abstract
Automatic differentiation (AD), a technique for constructing new programs whichcompute the derivative of an original program, has become ubiquitous throughoutscientiﬁc computing and deep learning due to the improved performance affordedby gradient-based optimization.
However, AD systems have been restricted to thesubset of programs that have a continuous dependence on parameters.
Programsthat have discrete stochastic behaviors governed by distribution parameters, such asﬂipping a coin with probability p of being heads, pose a challenge to these systemsbecause the connection between the result (heads vs tails) and the parameters ( p ) isfundamentally discrete.
In this manuscript we develop a new reparameterizationmethodology that allows for generating programs whose expectation is the derivative of the expectation of the original program.
We showcase how this methodgives an unbiased and low variance estimator which is as automated as traditional
AD mechanisms.
We demonstrate unbiased forward-mode AD of discrete-time
Markov chains, agent-based models such as Conway’s Game of Life, and unbiasedreverse-mode AD of a particle ﬁlter.
Introduction
Automatic differentiation (AD) is a technique for taking a mathematical program X ( p ) and generatinga new program (cid:101)
X
( p ) =
d X d p for computing the derivative [1, 2].
AD is widely used throughoutmachine learning and scientiﬁc computing due to the increased performance of gradient-basedoptimization techniques compared to derivative-free methods [3].
However, if X ( p ) returns the ﬂip ofa coin with probability p of receiving a 1 and probability 1 − p of receiving a 0 , it is clear that d X d p isnot deﬁned in the classical sense.
But when attempting to calibrate the parameter p to data, one maywish to ﬁt the model using the statistical quantities, e g. ﬁnd p such that the average of X ( p ) is closeto the average sum of N real-world coin ﬂips.
Given this use case, can one automatically construct aprogram (cid:101) X ( p ) that computes the derivative of the statistical quantities, i.e. E [ (cid:101) X ( p )]
=
d E [ X ( p )] d p ?
A naïve solution to this problem would be to use ﬁnite differences, i.e.:
[Equation]
Finite differences’ major issue in the context of stochastic programs is that this calculation does notcorrelate the calculation of E [ X ( p + ε )] with the calculation E [ X ( p )] .
This leads to a large variancein the ﬁnite-difference estimator [4] that goes to ∞ as ε → 0 .
This variance issue could be solvedif one could run the perturbed program with the same set of random numbers, i.e. sample from
X ( p + ε ) − X ( p ) as depicted in Figure 1. x
= X ( p )( z 1 )
X ( p + ε )( z 2 )
Pruning follows one alternative path
Smoothing ﬂattens onto primal
Y given X ( p ) =
x
We demonstrate how to automatically construct a stochastic program (cid:101) X ( p ) whose expectation satisﬁes
E [ (cid:101) X ( p )] =
d E [ X ( p )] dp .
We derive this through a technique which we term the stochastic derivative ,propagating the proportional probability of differing event outcomes due to inﬁnitesimal changes inp .
Our technique has several beneﬁts:
• Composability.
We derive an extension language-level AD to stochastic programs composed of any of a wide range of “elementary” stochastic programs, including discretedistributions such as Bernoulli, Geometric, and Poisson distributions, which may be freelymixed with continuous distributions such as Gaussian and exponential distributions.• Unbiasedness and low-variance.
The discrete structure of the original program is preservedwhile staying coupled, avoiding continuous relaxations which exhibit a variance-bias tradeoff and cannot be applied to highly discrete constructs.•
No tunable parameters.
A single run of the primal computation is transformed into anunbiased estimate of the derivative in a stateless fashion with no internal parameters to tune.• Speed.
Our technique’s computational cost is proportional to the number of steps in theoriginal program and is independent of how stochastic programs are composed
together.
We show the utility of our technique by demonstrating AD of stochastic simulations like inhomogenous random walks, AD of agent-based models such as the Game of Life, and the construction of anend-to-end differentiable particle sampler.
Additionally, we provide an open source implementationof the method, StochasticAD.jl , for readers to explore the technique on their own applications.
1.1 Related work
Gradient estimators for stochastic functions can be divided into three different classes [5]: scorefunction [6, 7, 8], measure-valued [9, 10], and pathwise [11, 12] gradient estimators.
Our methodoperates as a gradient estimator, but beyond that, the main feature is composability with arbitraryuser-written functions, enabling an automated mechanism.
When it comes to discrete randomness, thescore-function method is a popular general-purpose choice because it can be composed by stochasticcomputation graphs [13] and provides an unbiased estimate of the gradient.
However, the scorefunction method cannot be used to, e.g., differentiate a resampling step, and has other well-knownlimitations [5].
Importantly, the score-function method is not coupled to the primal evaluationand suffers from high variance, making gradient computation with discrete variables challenging.
A number of techniques (REBAR [14], RELAX [15]) have been introduced that rely on controlvariates for variance reduction [16].
Measure-valued derivations have no notion of intrinsic coupling,though coupling can be achieved by using common random numbers for certain distributions [17].
Adrawback of this coupling method is that its implementation may require nontrivial changes to theprimal evaluation code, which makes it less attractive for automation.
As an alternative direction,several recent works such as Gumbel-Softmax [18] have considered continuous relaxations of the program so that the reparameterization trick can then be applied.
However, such methods suffer avariance-bias tradeoff and furthermore are inapplicable to highly discrete programs which cannotbe continuously relaxed.
Our method extends the pathwise gradient estimator from continuous todiscrete programs while preserving all discrete structure in the original computatoin.
Smoothedperturbation analysis [19, 20] develops a mathematically equivalent object to the smoothed stochasticderivative based on conditional expectations, which is a special case of our formalism.
But suchideas have not previously been applied to construct a general-purpose AD algorithm via rigorouscomposition rules and computational structures.
While existing reverse-mode AD frameworks do not support unbiased differentiation of discreterandom programs, Storchastic [21] is a specialized framework for AD of stochastic computationgraphs, where the user can specify which estimator to use at each node, as well as any tunablehyperparameters.
Storchastic implements an exhaustive set of prior gradient estimation methods ateach sampling step.
However, previous reviews of these techniques revealed a performance scalingissue: the run time of the derivative estimate is exponential in the length of the largest chain ofstochastic nodes, a symptom of the way many prior gradient estimators compose [22].
In Section 3 wedemonstrate forward-mode and reverse-mode stochastic AD with identical computational complexityto their deterministic counterparts, alleviating these performance issues.
Composable derivatives of stochastic programs
In this section, we develop the notion of a stochastic derivative for programs containing discrete randomness.
We shall motivate this object as a natural generalization of the standard reparameterizationtrick, and present the key ideas underpinning the formalism.
In offering intuition, we make frequentuse of big-O notation for inﬁnitesimal asymptotics, where a quantity is O ( ε ) if bounded above inmagnitude by Cε for some real C and all sufﬁciently small ε .
Colloquially, we describe quantitiesthat are O ( ε ) as “inﬁnitesimal”.
2.1 Inﬁnitesimally perturbing a stochastic
program
Consider a stochastic program: a program which when provided with input p gives a random variable
X ( p ) as output.
To obtain a single sample of X ( p ) , which we call the primal evaluation , one shouldimagine a number z being randomly chosen from a sample space Z to produce a then deterministicoutput X ( p )( z ) .
For example, a Bernoulli random variable with probability p can be represented bypicking a uniform random number z ∈ [0 , 1] and deﬁning X ( p )( z ) =
[1 − p, 1] ( z ) where 1 Ω ( z ) is theindicator function for z ∈ Ω .
Note that parameterizations of many univariate distributions can beconstructed via the inversion method [23].
We deﬁne the stochastic differential d X ( ε ) as the change in the stochastic program X ( p ) w.r.t.
achange ε in its input p ,
[Equation]
Note that in the stochastic setting, d X ( ε ) is itself a random variable over the same sample space as
X ( p ) .
Crucially, X ( p + ε ) and X ( p ) share the same sample space (i.e. belong to a joint distribution),distinguishing Eq.
(2.1) from black-box ﬁnite differences.
For the full user-provided stochastic program, we are interested in computing ∂ p E [ X ( p )] , which isrelated to E [d X ( ε )] by
[Equation]
We thus expect d X ( ε ) to have “inﬁnitesimal” O ( ε ) expectation as ε approaches 0.
Let us considerthe form of d X ( ε ) for two illustrative examples: a program returning a sample from the continuousexponential distribution
Exp( p ) with scale p and a program returning a sample from the discrete
Bernoulli distribution Ber( p ) with probability p , both parameterized via the inversion method overthe sample space [0 , 1] (Fig. 2).
For the exponential distribution, X ( p )( z ) = − p log(1 − z ) isdifferentiable in p .
Thus, the range of d X ( ε ) is O ( ε ) for any ﬁxed z , as illustrated in Fig.
2a.
Incontrast, consider the Bernoulli distribution.
As shown in Fig.
2b, the output at a random z ∈ Zalmost surely does not change when the input probability p is perturbed by ε .
Speciﬁcally, d X ( ε )is 0 with probability 1 − ε ≈ 1 , and assumes the value 1 with probability
ε (Fig. 2c).
This is thefundamental challenge presented by discrete randomness: a O ( ε ) change in the input turns into
Output z
X ( p + ε )
X ( p ) dX ( ε )
O ( ε ) z
1 z(cid:29)
ε { { (a) X ( p ) ∼
Exp
( p ) .
Output z1 − p − ε 1 − p
X ( p + ε )
X ( p )  (b) X
( p ) ∼ Ber( p ) .
d X ( ε ) 1 − p − ε 1 − p
Output z
X ( p )( z ) = X ( p )( z ) =
1 { {  (c) X ( p ) ∼ Ber( p ) .
a ﬁnite perturbation in the output that is non-vanishing as ε → 0 .
The perturbation occurs with“inﬁnitesimal” O ( ε ) probability, contributing to the O ( ε ) expectation of d X ( ε ) .
2.2 Coupling to the primal
To produce low-variance estimates of the derivative of X ( p ) , it is important that the primal andderivative computations are coupled by a shared source of randomness.
To see this, let us understandwhat happens when they are entirely uncoupled.
A black-box ﬁnite difference approach with stepsize ε would independently sample z 1 and z 2 from the sample space (Fig. 2a), computing thederivative estimate 1 /ε · ( X ( p + ε )( z 2 ) − X ( p )( z 1 )) .
However, since the samples in the numeratorare independent, the variance of these samples is of order 1 /ε 2 → ∞ , so we are forced to pick aﬁnite step size and balance a variance-bias tradeoff.
In the case of continuous randomness, the reparameterization trick surmounts this issue by differentiating X ( p ) point-wise at each ﬁxed random sample z 1 (Fig. 2a) to form X (cid:48) ( p ) .
Mathematically, thereparameterization trick relies on the following interchange of limits and expectation,
[Equation]
But, considering a simple Bernoulli variable as in Fig.
2c, we see how this approach is ill-suited forthe discrete case!
As ε approaches 0, the stochastic differential d X ( ε ) is non-zero with inﬁnitesimal
O ( ε ) probability and thus d Xε is almost surely 0 as ε approaches 0.
This means X (cid:48) ( p ) is everywhere0, while the true derivative of E [ X ( p )] is ∂ p p = 1 .
The ﬁnite perturbation is neglected.
Thus, the reparameterization trick needs to be modiﬁed to handle discrete randomness.
Our insight isthat d X ( ε ) may in general be decomposed into an inﬁnitesimal component (i.e. of magnitude O ( ε ) )and a component consisting of ﬁnite perturbations that occur with inﬁnitesimal probability.
We cansample from these parts of the probability space separately , curing the degeneracy of ﬁnite changesoccurring with inﬁnitesimal probability.
As in stratiﬁed sampling, the expectation is then obtained asa weighted average of the separately sampled parts, cf.
Deﬁnition 2.1.
Suppose X ( p ) is a stochastic program with real input, whose output belongs to a
Euclidean space E .
We say that the triple of random variables ( δ, w, Y ) , with δ ∈ E , weight w ∈ Rand alternative outcome
Y ∈
E is a stochastic derivative of X at the input p if for all smooth boundedreal functions f ,
[Equation]
Theorem 2.2 (Simpliﬁed) .
Given a sufﬁciently regular stochastic program X ( p ) with derivative
X (cid:48) ( p ) , there exists a right stochastic derivative ( X (cid:48) ( p ) , w R , Y R ) of p with w R ≥ 0 .
Likewise thereexists a triple ( X (cid:48) ( p ) , w L , Y L ) with w L ≤ 0 giving a left stochastic derivative.
The components w and Y of the stochastic derivative form a correction to the output δ
= X (cid:48)
( p ) of thereparameterization trick, necessary to handle potential ﬁnite perturbations to the output.
Theorem 2.2shows existence of the stochastic derivative for a wide range of stochastic programs, subject totechnical assumptions given in our formal treatment (Appendix B).
Its proof is by construction, wherethe weight w can be understood as the derivative w.r.t.
ε of the probability that d X ( ε ) jumps by anon-inﬁnitesimal amount, as characterized by the following event,
[Equation]
0 and ε }
[Equation]
so that w is given as follows (separating the cases of right and left-handed limits):
[Equation]
The component Y R/L then has distribution given as the limit as ε → 0 + / − of the conditionaldistribution of X ( p + ε ) =
X ( p ) + d X ( ε ) given the event A ( ε ) and the outcome of X ( p ) .
In words,
Y represents the distribution of possible alternate values of X ( p ) , conditioned on the primal outputand on the fact that some ﬁnite perturbation occurs.
Example 2.3 (Stochastic derivative of Bernoulli variable) .
Suppose X ( p ) ∼ Ber( p ) , parameterizedvia the inversion method, and take ε > 0 .
As shown in Fig.
2c, d X is parameterized as,d X ( ε )( z ) =
(cid:26)if 1 − p − ε <
z < 1 − p , otherwise.
If X ( p ) = 1 , then z > 1 − p and d X ( ε ) is deterministically 0.
On the other hand, if X ( p ) is 0, thenz < 1 − p , so with probability ε/ (1 − p ) the stochastic differential d X ( ε ) assumes a value of 1,i.e.
the Bernoulli variable ﬂips from 0 to 1.
Thus, letting w R = 1 / (1 − p ) , Y R = 1 conditionallyon X ( p ) = 0 , and w R = 0 conditionally on X ( p ) = 1 , the triple (0 , w R , Y R ) is a right stochasticderivative of the program X ( p ) ∼ Ber( p ) .
We give a wide range of examples of left and right stochastic derivatives in Appendix A. It isof crucial importance that w and Y are random variables depending conditionally on the outputof X ( p ) .
Intuitively, this coupling allows us to consider only the smallest possible perturbationsto the program in the derivative computation (Fig. 1), and thereby achieve variance reductionwithout resorting to continuous relaxations.
For example, for a binomial variable X ( p ) ∼ Bin( n, p )parameterized via the inversion method (a natural parameterizaton for maximizing coupling), itholds that Y ∈ { X ( p ) − 1 , X ( p ) + 1 } .
We show in Example A.1 that the variance of our gradientestimator for such a binomial variable is O ( n ) , whereas the score function estimator has variance
O ( n 3 ) , justifying this intuition.
2.3 Composition of stochastic derivatives
Deﬁnition 2.1 deﬁnes the notion of stochastic derivative and Theorem 2.2 shows existence.
Themissing ingredient is composition , i.e. a stochastic derivative “chain rule”.
Because of the chain rule,we only need to hand-derive stochastic derivatives via Theorem 2.2 for each “elementary” stochasticprogram which makes up the full program: a stochastic derivative of the full program is obtainedautomatically by composition.
Proposition 2.4 (Chain rule) .
Consider independent stochastic programs X 1 and X 2 and theircomposition X 2 ◦ X 1 .
Suppose that X 1 has a stochastic derivative at p given by ( δ 1 , w 1 , Y 1 ) and X 2has a stochastic derivative ( δ 2 , w 2 , Y 2 ) in the direction δ 1 given conditionally on its input X 1 ( p ) , asconstructed by Theorem 2.2.
Then the stacked program [ X 1 ; X 2 ◦ X 1 ] has a stochastic derivativegiven by ( δ, w, Y )
where δ =
[ δ 1 ; | δ 1 | δ 2 ] and
[Equation]
with probability
[Equation]
with probability
[Equation]
and w = w 1 + | δ 1 | · w 2 .
Proposition 2.4 has two important special cases.
The ﬁrst is when X 2 is well-described by thereparameterization trick, i.e. w 2 = 0 .
In this case, Proposition 2.4 shows that the derivative δ 1of X 1 ( p ) simply propagates through X 2 via the chain rule.
In addition, the possible alternatevalues of X 1 ( p ) , as described by Y 1 , propagate through X 2 exactly as one would expect: the outputhas alternate values described by X 2 ( Y 1 ) , which are given the same weight w 1 (or “inﬁnitesimalprobability”) as the original perturbations.
The key is the case where X 2 is discrete random.
Here, an inﬁnitesimal perturbation to the input of
X 2 leads to potential ﬁnite perturbation(s) to the output, introducing another event with inﬁnitesimalprobability to be considered by the stochastic derivative.
This case is what links inﬁnitesimal changesto ﬁnite changes, necessitating the additional components w and Y of the stochastic derivative: without discrete randomness, w would always remain 0, recovering the traditional chain rule.
Crucially,as ε → 0 we may neglect the O ( ε 2 ) probability of two ﬁnite perturbations occurring simultaneously:in Eq.
(2.7), either X 1 is perturbed (and this perturbation propagates to X 2 ), or X 1 is not perturbedand a new perturbation occurs when evaluating X 2 .
2.4 Smoothed stochastic derivatives
For discrete random variables, the reparameterization trick is deterministically 0, while stochasticderivatives fully capture potential ﬁnite perturbations with inﬁnitesimal probability.
There exists amiddle ground between these methods: one may take the conditional expectation of the stochasticdifferential d X ( ε ) , such that the derivative has the correct expectation given the primal output, but aﬁnite perturbation has been replaced by an inﬁnitesimal one.
Deﬁnition 2.5 (Smoothed stochastic derivative) .
[Equation]
The key advantage of smoothed stochastic derivatives is that they immediately permit reverse-mode AD instead of forward-mode, as they have the same form as the traditional derivative.
However, theyenjoy more limited composition properties: they propagate exactly through smoothed functions fthat are linear over the conditional distribution of Y given X ( p ) via the standard chain rule, as weprove in Appendix B.5.
Due to the coupling of Y and X ( p ) , this is a much weaker requirement thanglobal linearity, which can lead to low-bias estimates: for example, for the program X ( p ) =
Geo( p ) 3smoothed stochastic derivatives give a derivative estimate with < 0 .
5% bias at p = 0 .
01 , even thoughthe cube function is highly non-linear on the inter-quartile range [28 , 137] of Geo( p ) .
Furthermore, aswe showcase in Section 3.4, smoothed stochastic derivatives lead to unbiased AD of the resamplingstep of a particle ﬁlter.
Automatic differentiation of stochastic programs
We develop StochasticAD.jl , a prototype package for stochastic AD based on our theory ofstochastic derivatives.
As discussed in Section 2.4, smoothed stochastic derivatives obey the usualchain rule, and thus can be used with existing AD infrastructure by supplying custom rules for discreterandom constructs.
We demonstrate unbiased reverse-mode AD of a particle ﬁlter via smoothing in
Section 3.4.
However, performing automatic differentiation with unsmoothed stochastic derivatives,which are unbiased in all cases, requires new innovation.
We develop a novel computational objectcalled a stochastic triple , introduced in Section 3.1 and showcased in Section 3.2 and Section 3.3,.
3.1 Educational toy example of stochastic triples
Traditionally, forward-mode AD has been made possible by dual numbers [2], which pair theprimal evaluation of a deterministic function f ( p ) with its derivative ∂ p f ( p ) .
Dual numbers can bepropagated through a program using the chain rule.
A useful alternative perspective of dual numbersis that they propagate an “inﬁnitesimal” perturbation ε to the input through the program, where ∂ p f ( p )is the coefﬁcient of the “dual” element ε .
Stochastic triples generalize dual numbers by including athird component ∆ s to describe discrete perturbations with inﬁnitesimal probability (Fig. 3, left).
In Fig.
3, we consider a toy program X ( p ) including discrete randomness.
We are interested inthe sensitivity of X ( p ) at p = 0 . 6 , and hence we provide the stochastic triple printed as "0.6 +
1 struct StochasticTriple value # primal evaluation δ # "infinitesimal" component Δ s
# component of discrete change # with "infinitesimal" # probability 7 end 1 using Distributions 2 function X(p) # p = {+0.
+} + ε a = p^2 # a = 0.36 + 1.2 ε b = rand(Binomial(10, p))
# b = 6 + (1 with probability 10.0 ε ) c
= 2 * b + 3 * rand(Bernoulli(p))
# c = 12 + (3 with probability 12.5 ε ) return c * rand(Normal(p, p^2))
9 end 1 julia> using StochasticAD 2 julia> st = stochastic_triple(X, 0.8) # sample a single stochastic triple at p = 0.
3 10.5 + 23.0 ε + (2.63 with probability 12.5 ε ) 4 julia> partial_expectation(st) # which produces a single derivative estimate...
6 julia> samples =
[derivative_estimate(X, 0.5) for i in 1:1000] # take many estimates!
7 julia> println("derivative of E[X(p)]: $(mean(samples)) ±
$(std(samples) / sqrt(1000))")
[Figure]
ε " as
First the is 36 +
2 ε this is the familiar that dual numbers propagate, via the chain rule.
But what happens when the triple is propagatedthrough the discrete and random Binomial variable?
The resultant stochastic triple "6 + (1 withprobability 10.0 ε )" is an integer with a component of discrete change, reﬂecting an inﬁnitesimalprobability of one more success in the Binomial.
Formally, we can interpret a printout "x + δε +( Δ with probability w ε )" as follows: x is a sample of the random variable X ( p ) describing theprimal evaluation, and δ , w , and x+ Δ are samples of the components δ , w , and Y , respectively, of thestochastic derivative of X at p .
Using our chain rule for stochastic derivatives (Proposition 2.4), we write rules for propagatingstochastic triples through functions via operator overloading [24], exploiting Julia’s multiple dispatchfeature [25].
When multiple discrete changes are possible, we pick one probabilistically: we callthis strategy pruning (see Fig. 1) and show its unbiasedness in Appendix B.4.
For example, in line6 of Fig. 3, right, we probabilistically choose between the perturbation to the Binomial and theperturbation to the Bernoulli, in this case picking the latter.
Thus, stochastic triples may efﬁcientlypropagate through the full toy function written in Fig.
The function derivative_estimatecreates a stochastic triple, propagates it, and collapses it into the derivative estimate δ + w Δ , formingan unbiased estimate of the derivative via Deﬁnition 2.1 (Fig. 3, bottom).
Thus, stochastic triplesenable the automatic construction of a program (cid:101) X
( p ) satisfying E
[ (cid:101) X ( p )]
=
d E [ X ( p )] d p .
3.2 Inhomogeneous random walk
We consider a Markovian random walk x 0 , . . .
, x n on Z ≥ 0 , with transition behavior dependent on aparameter p as follows,
[Equation]
probability
[Equation]
One can write a program that stochastically simulates this walk and applies an arbitrary non-linearfunction f to the output x n .
In practice, f may represent a loss or a likelihood estimate; in thistoy setting, we take f ( x ) = x 2 .
We are interested in studying the asymptotic behavior of thevariance of our automatically derived gradient estimator, and so we choose the initial value of p tobe n so that the transition function varies appreciably over the range of the walk for all n .
Fig.
4compares the asymptotic growth in variance of the score function estimator to that of stochastictriples.
Both estimators are unbiased with identical computational cost, but their variances differby an order of magnitude.
Empirically, we ﬁnd a polynomial ﬁt with leading term ∼ n 1 .
5 for thevariance of the score function, and a ﬁt with leading order ∼ n for the variance of stochastic triples.
X
[Figure]
for step in 1:n
[Figure]
i
[Figure]
Crucially, stochastic triples achieve this variance reduction while operating entirely in discrete space, automatically producing a gradient estimate that is provably unbiased.
3.3 Stochastic game of life
For our next example, we differentiate a program based on a stochastic version of John Conway’sgame of life, played on a two-dimensional board.
Such types of discrete stochastic examples arisein many applications.
For example, mixing machine learning with agent-based models found inepidemiology and sociological contexts [26] or rule-based models and discrete stochastic (Gillespie)simulations of discrete chemicals in systems biology [27, 28] require similar program constructs.
Inthe traditional game of life, a dead cell becomes alive when 3 of its neighbours are alive, while analive cell survives when 2 or 3 of its neighbours are alive.
In our stochastic version, each of theseevents instead has probability 95% , while their complementary events have probability 5% .
Consider a program that populates each cell of an N × N board with probability p , runs the stochasticgame of life for T time steps, and counts the number of alive cells n alive .
Our goal is to performa sensitivity analysis of the ﬁnal living population with respect to the initial living population, i.e.differentiate the expectation of n alive with respect to p .
Stochastic triples propagate fully throughthis program, leading to an unbiased estimate of the derivative, as we verify with black-box ﬁnitedifferences.
The board is depicted in Fig.
4 where the difference of the “alternative path” withinﬁnitesimal probability considered by stochastic triples in a single run of the program (recall Fig.
1)is highlighted.
This is a high-dimensional example (state space has dimension N 2 = 144 ) withfundamentally discrete structure, providing support for the algorithmic correctness and generality ofstochastic triples.
In particular, the program cannot be continuously relaxed by AD since it includesarray indexing; in the approach of stochastic triples, integer-valued quantities stay integers, andtherefore may be used in highly discrete constructs.
3.4 Differentiable particle ﬁlter
Based on the developed formalism, we now present an end-to-end differentiable particle samplercomputing a marginal log-likelihood L .
We consider a hidden Markov model with random states
X 1 , . . .
, X n as speciﬁed by a stochastic program X 1 ( θ ) giving the starting value depending onunknown parameters θ and consecutive states given by pointwise differentiable stochastic programs
X i ( x i − 1 , θ ) depending on the previous state x
i −
=
X i −
(Markov property) and the parameters.
In general, we allow probability densities p ( x 1 | θ ) of X 1 and transition probability densitiesp ( x i | x i − 1 , θ ) to depend arbitrarily on the parameters θ .
This latent process X 1 , . .
, X n is indirectlyobserved as the process Y 1 , . . .
, Y n with observations
Y i ( x i ) depending on x i
= X i which arespeciﬁed to have smooth conditional probability densities p ( y i | x i , θ ) depending only on x i and θ .
The likelihood of θ given the observations is in general not tractable, but a particle ﬁlter can be usedto compute an estimate of the likelihood.
As a concrete example, we consider the following systemwith a d -dimensional latent process,
[Equation]
i ∼ Normal(0 , Q ) , y i =
x i + ν i with ν i ∼ Normal(0 , R ) , where Q = 0 .
02 · 1 d × d , R = 0 .
01 · 1 d × d , x 1 ∼ Normal( µ, 0 .
001 · 1 d × d ) , µ ∼ Normal(0 , 1 d × d ) isa random initial position, and Φ is a d -dimensional rotation matrix.
We wish to differentiate withrespect to each entry of Φ , i.e. θ = vec(Φ) .
Given n observations y
=
Y 1 , . .
, y n = Y n of the hidden Markov model, a particle ﬁlter allows oneto estimate the likelihood L of these observations by simulating a number of “particles” followingthe dynamical model, and repeatedly performing resampling steps to pick the particles that match theobservations.
We give a more complete introduction to particle samplers in Appendix C, followingour terminology, but what is most important for our purposes is that the resampling step assigns animportance weight to each particle that is used in a purely linear way to produce the ultimate likelihoodestimate.
Therefore, we may differentiate each weight using a smoothed stochastic derivative, asformalized in Appendix B.6 and Appendix C.
Particle samplers are a powerful, generic approach to this problem, but our particular choice ofsystem allows the calculation of a ground-truth gradient of the log-likelihood with respect to Φ bydifferentiating the Kalman ﬁlter algorithm.
For our computations, we ﬁx the number of observationsn = 20 .
Fig.
5 visualizes the latent process, the observations, as well as the Kalman and the particleﬁlter trajectories in d = 2 dimensions.
Our estimator agrees with the ground-truth value from the
Kalman ﬁlter, unlike biased estimators [29, 30, 31, 32] that neglect the resampling step’ contributionto the derivative or perform it with entropy-regularized optimal transport (Fig. 5).
Furthermore,thanks to the similarity of smoothed stochastic derivatives and classical derivatives, we can alsouse reverse-mode instead of forward-mode AD.
This is convenient, as eventually reverse-mode ADbecomes superior than forward-mode AD for functions f from R n to R m with m (cid:29) n [2].
For ourstochastic program, we observe reverse-mode AD to perform better for more than ≈ 100 parameters.
Limitations and outlook
We have presented a method for unbiased AD of programs with discrete randomness, which, wehave argued, is a natural generalization of the often-used reparameterization trick.
However, morework will need to be done to turn our software demonstration StochasticAD.jl into an AD systemcapable of handling the full complexity of applications such as machine learning.
This includessupport for if statements with discrete random input; currently, such branches need to be rewrittenusing array indexing, which is supported.
Further, an interesting direction for future work is toautomatically handle constructs which deterministically turn a continuous random quantity intoa discrete one, such as an inequality comparison X(p) > 0 , and furthermore handle constructssuch as while loops based on such comparisons.
In a different direction, going beyond smoothingfor reverse-mode differentiation, and ideally achieving the unbiasedness of stochastic triples whileremaining coupled to the primal, is an important open problem for large-scale applications. [Table]
In the supplement, we: • provide a number of examples of stochastic derivatives and present hand-worked examples [Table] of how they compose (Appendix A),
[Table]
• prove our formal statements regarding stochastic derivatives, and state and prove formal
[Table]
statements regarding smoothed stochastic derivatives (Appendix B),
[Table]
• provide an introduction to the particle ﬁlter methodology, and show how smoothed stochasticderivatives leads to unbiased differentiation of the resampling step (Appendix C),
[Table]
• provide details about the hardware used for the experiments of the main text and the software
[Table] dependencies of StochasticAD.jl (Appendix D). A Examples of stochastic derivatives In this section, we present examples of stochastic derivatives for a number of stochastic programs,recalling Deﬁnition 2.1.
We consider both elementary programs (e.g. a Bernoulli variable) for which
Theorem 2.2 applies, as well as general programs (e.g. a two-step random walk).
Example A.1 (Integer-valued stochastic programs) .
The Bernoulli, Binomial, Geometric, and Poissondistributions are all discrete, and in fact (nonnegative) integer-valued.
Suppose X ( p ) follows one ofthese distributions, parameterized using the inversion method [23] over the interval [0 , 1] .
We presenta worked derivation of the stochastic derivative of X ( p ) .
Let X ( p ) have stochastic derivative ( δ, w, Y ) at p .
The program is integer-valued, and hence δ = 0 ,i.e. the point-wise derivative ∂ p X ( p )( z ) is 0 for all z .
Further, for small enough ε , d X ( ε ) ∈{− 1 , 0 , 1 } , and thus Y ∈ { X ( p ) − 1 , X ( p ) + 1 } .
This comes from the use of the inversion method:the only ﬁnite perturbations with non-zero probability switch to the next smallest or next largestdiscrete output of the program.
For a small change ε in the program input, let A + ( ε ) denote the event X ( p + ε ) =
x + 1 (switchingto the next largest output), and let A − ( ε ) denote the event X ( p − ε ) =
x − 1 (switching to the nextsmallest output).
Then, A ( ε ) is A − ( ε ) ∪ A + ( ε ) .
We derive,
[Equation]
We are interested in the derivative w + of this probability at ε = 0 ,
[Equation]
where [[ P ]] is the Iverson bracket of a condition P .
Informally, one can think of w + as the scaling factor of the “inﬁnitesimal” probability that X ( p + ε ) =
x + 1 , i.e. this probability is w + ε + O ( ε 2 ) .
Note that the sign of ε matters, where ε > 0 representsthe right stochastic derivative and ε < 0 the left stochastic derivative.
This is reﬂected throughthe inclusion of sgn( ε ) (an abuse of notation representing the direction of the limit) on the RHS of
Conditionally on X ( p ) =
x
[Equation]
w
[Equation]
w +
Ber( p ) x/p · [[ x = 1 and ε < 0]] 1 / (1 − p ) · [[ x = 0 and ε > 0]]
Bin( n, p ) n/p · [[ x > 0 and ε < 0]] ( n − x ) / (1 − p ) · [[ x < n and ε > 0]]
Geo( p ) x/ ( p (1 − p )) · [[ x > 0 and ε > 0]] ( x + 1) /p · [[ ε < 0]]
Pois( p ) x/p · [[ x > 0 and ε < 0]] [[ ε > 0]]
Eq.
(A.3): the magnitude of ε has cancelled out in the limit but a dependence on its sign remains.
Inan analogous fashion, we can derive
[Equation]
ε
[Equation]
p
[Equation]
x
[Equation]
1) > 0]] .
(A.4)
Putting these two cases together, we may write,
[Equation]
(A.5)
[Equation]
(A.6)
[Equation]
(A.7) and express the probability distribution of Y as
[Equation] [Equation]
[Equation]
2 lists the + and w − for the Bernoulli, Binomial, Poisson, and Geometric distributions,derived using the probability distributions given in Table 1.
Using w + and w − , we can easily expressw and Y as above.
Note that w + and w − for a Binomial variable may be derived from the compositionof Bernoulli variables, by writing the primal evaluation X ( p )
=
x as the sum of x Bernoulli successesand n − x Bernoulli failures.
Note that the stochastic derivative for a binomial variable has weight w of order n and thereforesamples from the estimator w · ( Y − X ( p )) of ∂ p E [ X ( p )] for X ( p ) ∼ Bin( n, p ) have variance oforder n .
In contrast, the score method gives an estimator
[Equation]
which is also unbiased, but its variance
[Equation]
is growing with cubic rate in n .
Example A.2 (Differentiably parameterized continuous stochastic programs) .
For continuous stochastic programs, our method reduces to the standard reparameterization trick.
Speciﬁcally, for a differentiably parameterized continuous random variable, the stochastic derivative w.r.t a parameter p ofthe distribution has the form ( X (cid:48) ( p ) , 0 , 0) .
For instance, the stochastic derivative of X ( p ) ∼ Exp( p )w.r.t.
p , the scale of the exponential distribution, is(cid:18) x p , 0 ,(cid:19) (A.10) conditionally on X ( p ) = x .
Such rules are already implemented in the Julia AD ecosystem in
Distributions.jl [1] and DistributionsAD.jl [2], and may be used without modiﬁcation asthe continuous case is a special case of our formalism.
Example
A.
(Categorical variable) .
Consider a categorical variable
Categorical a ( p 1 , p 2 , p 3 , . .
, p n ) assuming ﬁxed outputs a = a 1 , a 2 , . .
. , a n with probabilities p 1 , p 2 , p 3 , . . . , p n .
Assume that p 1 , p 2 , p 3 , . .
. , p n implicitly depend on a singleparameter p , to ﬁt our deﬁnition of stochastic derivative w.r.t.
a single parameter, and let
X ( p ) ∼ Categorical a ( p 1 , p 2 , p 3 , . .
. , p n ) .
The categorical variable is univariate and its discrete outputs do not depend on p , and hence itsstochastic derivative is of the form (0 , w, Y ) with Y ∈ { a x − 1 , a x +1 } .
Similar logic to Example A.1yields
[Equation]
12) with w and Y given from w + and w − as in Example A.1.
Note that a 1 , a 2 , . .
. , a n can be arbitrarydiscrete objects, such as arrays or strings, as the categorical variable may represent an intermediatevalue of the program which is ultimately converted into a number.
In handling such cases, onemust be careful to speak in terms of alternate outputs rather than perturbations to the output.
Ourformalism of stochastic derivatives can already accomodate this idea, as Y takes on the alternatevalues rather than the values of the perturbations.
We now consider some general stochastic programs.
Their stochastic derivatives derive automaticallyfrom composition and need not be hand-derived, but we do so for clarity.
Example A.4 (Bernoulli plus Exponential) .
Consider the program,
[Equation]
Then, X ( p ) has right (left) stochastic derivative,
[Equation]
conditionally on X 1 ( p ) =
x 1 and X 2 ( p ) =
x 2 , where w 1 and Y 1 are given conditionally on the
Bernoulli output X 1 ( p ) =
x 1 for ε > 0 ( ε < 0 ) by Table 2.
Example A.5 (Cubing a geometric) .
Consider the program,
X 1 ( p ) =
Geo( p
[Equation]
(A.17)
[Equation]
(A.18)
We know from Example A.1 that X 1 ( p ) has right (left) stochastic derivative (0 , w 1 , Y 1 ) with w 1and Y 1 given by Table 2 for ε > 0 ( ε < 0 ).
In particular, Y 1 ∈ { x 1 − 1 , x 1 + 1 } conditionally on
X 1 ( p ) = x 1 .
Thus, intuitively, when propagating through the cube function we only care aboutthe discretely spaced alternate values ( x 1 − 1) 3 and ( x 1 + 1) 3 ; the conventional derivative 3 x 21 isirrelevant as the input to the cube function is integer-valued.
Indeed, by Proposition 2.4 the stochasticderivative of X ( p ) reads(
[Equation]
31 ) , (A.19) where Y 31 ∈ { ( x 1 − 1) 3 , ( x 1 + 1) 3 } .
Example A.6 (Parameter-scaled Bernoulli) .
Consider the program,
X 1 ( p ) =
Ber( p ) , (A.20)
X ( p ) =
p · X 1 ( p ) .
(A.21)
By Proposition 2.4, the stochastic derivative of X ( p ) reads,( X ( p ) , w 1 , pY 1 ) , (A.22) where w 1 and Y 1 are given for X 1 ( p ) by Table 2.
Example A.7 (Two-step random walk) .
Consider the program,
[Equation]
[Equation]
p
[Equation]
Ber( p
[Equation]
(A
[Equation] [Equation]
This represents a two-step random walk with steps given by X 1 ( p ) and X 2 ( p ) , where there is atransition 0 → 1 with probability p , a transition 1 → 2 with probability 2 p , and self loops otherwise.
In this example, let us focus on a particular primal evaluation where X 1 ( p ) = 0 and X 2 ( p ) = 0 , sothat X ( p ) = 0 , and consider only the right stochastic derivative for simplicity.
Conditionally on
X 1 ( p ) = 0 , X 1 ( p ) has right stochastic derivative(cid:18)
[Equation]
(cid:19) , (A.26) as given by Table 2.
Now, conditionally on X 2 ( p ) = 0 , X 2 (0) has right stochastic derivative w.r.t.x 1 ,(cid:18)
[Equation]
− p
[Equation]
(cid:19) .
(A.27)
Let us now turn our attention to the stacked program [ X 1 ( p ); ( X 2 ◦ X 1 )( p )] .
By Proposition 2.4, thestacked program has right stochastic derivative (0 , w, Y 12 ) where conditionally on X 1 ( p ) = 0 and
X 2 (0) = 0 ,
Y= (cid:26) [1; X 2 (1)] with probability 1 / 2 , [0; 1] with probability 1 / 2 , (A.28) and w = 21 − p .
Note that conditionally on X 2 (0) = 0 , X 2 (1) has a chance 1 − 2 p 1 − p of also being 0, buta p 1 − p chance of ﬂipping to 1, so that the ﬁrst case expands into two events
Y
[Equation]
] with probability
[Equation]
[Equation]
with probability
[Equation] [Equation]
] with probability 1 /
[Equation]
(A.29)
Finally, the expression for X ( p ) may be thought of as a unary function of this stacked program, withstochastic derivative (0 , w, Y ) where conditionally on X 1 ( p ) = 0 and X 2 (0) =
Y = 
 with probability p 2(1 − p ) , with probability 1 − [Equation] p [Equation]
[Equation]
− p ) , with probability 1 / 2 .
(A.30)
B Proofs
B.1 Preliminaries
Let X ( p ) denote a stochastic program , formally deﬁned as a stochastic process with index set theﬁnite closed interval I =
[ a, b ] and values in the Euclidean space E .
Let p ∈
I be the parameter ofthe process.
As in the main text, we use the shorthand
[Equation]
Due to asymmetry of stochastic derivatives, a number of statements and propositions have equivalentforms for right and left stochastic derivatives.
In this case, we use the notation ε → 0 + / − to indicatethat statements and proofs for right and left stochastic derivatives can both be read off by choosingthe appropriate side of the limit for all expressions.
For clarity, we remark upon the (standard) notation,
[Equation]
(B.2) where U and V are both random variables deﬁned on a sample space Z .
In this setting, W is itselfalso a random variable deﬁned on Z .
In particular, for a particular sample z ∈ Z , W ( z ) is theconditional expectation E
[ A | B = B ( z )] .
In the language of σ -algebras, E [ U | V ] is equivalent to
E
[ U | σ ( V )] , where σ ( V ) is the sub- σ -algebra generated by V .
Finally, we use the notation S c to denote the complement of a set S .
B.2 Stochastic derivatives of elementary stochastic programs
In Theorem 2.2, w can be understood as the derivative of the probability that X ( p + ε ) and X ( p )vary by more than inﬁnitesimal amount.
The following event quantiﬁes this,
[Equation]
Outside the event A ( ε ) , the sensitivity is well-described by the reparameterization trick:
Proposition B.1.
If f : E → R is smooth with bounded derivative and X ( p ) is almost surelydifferentiable having derivative X (cid:48) ( p ) with bounded expectation, then
[Equation]
Proof.
Because f ◦ X is almost surely differentiable,
[Equation]
Therefore, it also holds that
[Equation]
Furthermore,
[Equation]
− f
[Equation]
X
[Equation]
p [Equation] |
[Equation]
By assumption E G < ∞ .
The proposition follows by bounded convergence.
We now give a technical condition under which Theorem 2.2 holds, in which we characterize w asthe derivative of the conditional probability of A ( ε ) given X ( p ) , and Y as realization of the weaklimit of the conditional distribution of X ( p + ε ) given X ( p ) and A ( ε ) .
Assumption B.2.
We assume that:
• For the quantity
[Equation]
the following limit converges in mean,
[Equation]
where E w < ∞ .
• The conditional distribution of X ( p + ε ) given X ( p ) and A ( ε ) converges in distribution to
a random variable Y .
Speciﬁcally, Y must satisfy
[Equation]
for all bounded smooth functions f : E → R with bounded derivative.
If that Z ∼ U(0 , 1] and X ( p ) = f ( p, Z ) is a continuous differentiable random variable or a discreterandom variable with f the left continuous inverse of its distribution function F p and if the boundariesof the interval ( a x ( p ) , b x ( p )]
= { z : f ( p, z ) = x } are differentiable in p for all values x with
P ( X ( p ) = x ) > 0 , Assumption B.2 holds.
Given Assumption B.2, we can now well-characterize thesensitivity of X ( p ) given that A ( ε ) occurs:
Proposition B.3.
If f : E → R is bounded and smooth with bounded derivative, then for the event Agiven by Eq.
(B.3) and the random variables w and Y constructed by Assumption B.2, it holds that
[Equation]
Proof.
Write
[Equation]
Tackling the ﬁrst term, by the tower property of conditional expectation,
[Equation]
which converges to E [ w · f ( Y )] − E [ w · f ( X ( p ))] as ε → 0 + / − by Assumption B.2 and the fact that
E [ 1 A ( ε ) | X ( p )]
= P ( A ( ε ) | X ( p )) =
w ( ε ) .
Here, the key manipulation has been to introduce a conditional expectation on X ( p ) , and hoist w , w ( ε ) , and f ( X ( p )) outside of it as they are deterministicfunctions of X ( p ) .
Now, we may bound the magnitude of the second term as(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:21)(cid:12)(cid:12)(cid:12)(cid:12) ≤
[Equation]
∞
[Equation]
E
[Equation]
( ε )
ε
[Equation]
w
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:21)(cid:21) ≤
[Equation]
(cid:12)(cid:12)(cid:12) (cid:21)
[Equation]
(cid:20)
[Equation]
( ε )
[Equation]
( ε ) (cid:12) (cid:12)(cid:12)(cid:12)
[Equation]
(cid:21)(cid:21)
[Equation] which approaches 0 as ε → 0 + / − by the convergence in mean of w ( ε )
/ε to w .
Putting the limits ofthese two terms together, we obtain the claim.
Now, recall
Deﬁnition 2.1.
Suppose X ( p ) is a stochastic program with real input, whose output belongs to a
Euclidean space E .
We say that the triple of random variables ( δ, w, Y ) , with δ ∈ E , weight w ∈ Rand alternative outcome
Y ∈
E is a stochastic derivative of X at the input p if for all smooth boundedreal functions f ,
[Equation]
We can now provide a proof of Theorem 2.2.
Theorem 2.2.
Given a stochastic program X ( p ) satisfying Assumption B.2 and almost surelydifferentiable having derivative X (cid:48) ( p ) with bounded expectation, there exists a stochastic derivative( X (cid:48) ( p ) , w R , Y R ) of p with w R ≥ 0 , which we call the right stochastic derivative.
Likewise thereexists a triple ( X (cid:48) ( p ) , w L , Y L ) with w L ≤ 0 giving a left stochastic derivative.
Proof.
Construct w R and Y R by Assumption B.2 applied with the right-sided limit ε → 0 + , andconstruct w L and Y L by Assumption B.2 applied with the left-sided limit ε → 0 − .
Now by
Proposition B.1 and Proposition B.3,
[Equation]
Thus, the triples ( X (cid:48) ( p ) , w R , Y R ) and ( X (cid:48) ( p ) , w L , Y L ) are stochastic derivatives of X ( p ) by Proposition B.3.
Note that w R ≥ 0 and w L ≤ 0 by Eq.
(B.6) and the non-negativity of w ( ε ) .
So( X (cid:48) ( p ) , w R , Y R ) is a right stochastic derivative of X ( p ) at p and ( X (cid:48)
( p ) , w L , Y L ) is a left stochasticderivative of X ( p ) at p , as desired.
B.3 Composition of stochastic derivatives
Recall,
Proposition 2.4 (Chain rule) .
Consider independent stochastic programs X 1 and X 2 and theircomposition X 2 ◦ X 1 .
Suppose that X 1 has a stochastic derivative at p given by ( δ 1 , w 1 , Y 1 ) and X 2has a stochastic derivative ( δ 2 , w 2 , Y 2 ) in the direction δ 1 given conditionally on its input X 1 ( p ) , asconstructed by Theorem 2.2.
Then the stacked program [ X 1 ; X 2 ◦ X 1 ] has a stochastic derivativegiven by ( δ, w, Y )
where δ =
[ δ 1 ; | δ 1 | δ 2 ] and
[Equation]
with probability
[Equation]
probability
[Equation]
and w = w 1 + | δ 1 | · w 2 .
Proof.
Let A 1 ( ε ) denote the event (B.3) for X 1 at p .
Since X 2 may have multivariate input, conditionally on X 1 ( p ) deﬁne the program X 2 ( X 1 ( p ) + δ 1 / | δ 1 | · p (cid:48) ) with real input p (cid:48) so that ( δ 2 , w 2 , Y 2 )is the stochastic derivative of this program for input p (cid:48)
= 0 , and let A 2 ( ε ) denotes the event (B.3) forthe program at p (cid:48)
= 0 .
Further, let
[Equation]
[Equation] B. [Equation]
The events
A 1 ( ε ) and (cid:101)
A 2 ( ε ) may be used to form a partition of the sample space:
[Equation]
B.
[Equation]
Informally, the ﬁrst event corresponds to the case of no large perturbation in either variable, the secondand third correspond to ﬁnite perturbations in X 1 and X 2 respectively, and the fourth corresponds toa ﬁnite perturbation in both.
As shorthands, deﬁne X
=
X 2 ◦ X 1 and ¯ X
=
[ X 1 ; X ] .
Our approachwill be to characterize the sensitivity of X ( p ) under the ﬁrst event using the reparameterizaton trick,under the second and third events using the argument of Proposition B.3, and ﬁnally show that thefourth event may be neglected.
Given the event A c 1 ( ε ) ∩ (cid:101)
A c 2 ( ε ) , we may bound | d X ( ε ) − X (cid:48)
( p ) | by ﬁrst bounding
| d X 1 ( ε ) | as O ( ε ) ,
[Equation]
(B.10) and then bounding | d X ( ε ) | for the composition X
=
X 2 ◦ X 1 ,
[Equation]
In addition to this bound, we have that the event A c 1 ( ε ) ∩ (cid:101)
A c 2 ( ε ) holds almost surely as ε → 0 + / − ,since X 1 and X 2 are almost surely differentiable.
Then, an analogous bounded convergence argumentas in Proposition B.1 applies, with the only difference being that the bound Eq.
(B.12) has accrued aprefactor in front of ε , that does not affect the asymptotic argument.
So,
[Equation]
for all smooth functions f : E → R with bounded derivative.
Now, for the events A 1 ( ε ) ∩ (cid:101)
A c 2 ( ε ) and
A c 1 ( ε ) ∩ (cid:101)
A 2 ( ε ) , we have
[Equation]
and
[Equation]
by the almost sure convergence of A c 1 ( ε ) and (cid:101)
A c 2 ( ε ) to 1 .
Furthermore, by the deﬁnitions of Y 1 and
Y 2 , and again using the almost sure convergence of A c 1 ( ε ) and (cid:101)
A c 2 ( ε ) to 1 , we have that the stackedprogram [ X 1 ; X ] converges in distribution to [ Y 1 ; X 2 ( Y 1 )] conditional on A 1 ( ε ) ∩ (cid:101)
A c 2 ( ε ) , X 1 ( p ) ,and X ( p ) , and converges in distribution to [ X 1 ; Y 2 ] conditional on A c 1 ( ε ) ∩ (cid:101)
A 2 ( ε ) , X 1 ( p ) , and X ( p ) .
Therefore, by similiar reasoning to Proposition B.3 we have
[Equation]
X
[Equation]
Y
[Equation]
¯ X
[Equation]
p ))
[Equation]
andlim
[Equation]
+
[Equation]
E
[Equation]
Finally, note that,
[Equation]
so by boundedness of f ,lim
[Equation]
All together, we have
[Equation]
so ( δ, w, Y ) is a stochastic derivative of ¯ X
=
[ X 1 ; X 2 ◦ X 1 ] by Deﬁnition 2.1.
B.4 Unbiasedness of pruning strategy
In the main text, we note that we can employ a pruning strategy so that we only ever track onealternative path (i.e. one sample from the stochastic derivative component Y ) and yet still obtain anunbiased estimate.
The following shows that the pruning method, whereby one chooses between twosamples of Y by picking with probability proportional to its weight, is indeed unbiased.
We prove by induction that the currently tracked alternative path is an unbiased choice amongstall possible alternative paths seen so far.
The base case, where the ﬁrst alternative path observedis followed, is trivial as it is the only choice.
Now, suppose that n alternative branches have beenobserved so far, where w = w 1 + w 2 + · · · + w n is the summed weight so far.
Suppose we observean ( n + 1) th branch, with weight w n +1 .
By our pruning strategy it is chosen with probability
[Equation]
while the j th path for some j ≤ n will have been chosen after this step with probability
[Equation]
as desired for unbiasedness.
Note that this enables us to prune online , i.e. without knowing the fullstructure of the computation a priori, which we exploit in StochasticAD.jl .
B.5 Smoothing
Deﬁnition B.4 (Smoothed stochastic derivative) .
For a stochastic process X ( p ) with right (left)stochastic derivative ( δ, w, Y ) at input p , the random variable (cid:101) δ is a right (left) smoothed stochasticderivative of X ( p ) at input p if
[Equation]
Using Deﬁnition 2.1, we can write the above in an alternative form that does not rely on the deﬁnitionof the stochastic derivative,
[Equation]
Smoothed stochastic derivatives propagate through functions that are locally linear over the range of
Y conditionally on X ( p ) , as we now formalize.
Note that this is a much weaker requirement thanglobal linearity over the full range of X ( p ) , as the range of Y conditional on X ( p ) is generally morerestricted (e.g. Y ∈ { X ( p ) − 1 , X ( p ) + 1 } for a binomial variable).
Proposition B.5.
Suppose that f is linear over the range of Y conditionally on X ( p ) =
x , for all x .
Then,
[Equation]
Proof.
For such an f , note that Eq.
(2.4) of Deﬁnition 2.1 simpliﬁes to
[Equation]
In most cases, local linearity will only hold approximately, leading to bias in the estimate producedby propagating smoothed stochastic derivatives.
However, in the case of a particle ﬁlter resamplingstep an exact estimate is produced, as we show in the following.
Let X ( p ) be a stochastic program approximated by the program (cid:101)
X ( p ) .
Assume we can sample from(cid:101)
X ( p ) .
Assuming absolute continuity, we may write expectations as E (cid:36) ( p ) f
( (cid:101) X ( p ))
=
E f
( X ( p ))
where (cid:36) ( p ) is the Radon-Nikodym derivative of the law of (cid:101) X
( p ) with respect to the law of X ( p ) ,evaluated in (cid:101) X
( p ) .
We thus consider the program which returns the pair of weight and value( (cid:36) ( p ) , (cid:101)
X ( p )) .
Proposition B.6.
If (cid:101) X ( p ) is a continuous program, reparameterized such that it is differentiablepointwise, and (cid:36) ( p ) has derivative ( δ, W, Y )
[Equation]
Proof.
As the function ( (cid:36), x ) (cid:55)→ (cid:36)g ( x ) is bilinear the statement follows from Proposition B.5.
Programs returning weighted samples can be monadically composed: If E (cid:36) i ( x i − 1 ) (cid:101)
X i ( x i − 1 ) =
E X i ( x i − 1 ) , i = 1 , 2 , then
[Equation]
) (cid:36) 1 ( p ) (cid:101) X 2 ( (cid:101)
X 1 ( p ))
(cid:105) by conditional independence.
Note that also the product (cid:36) 2 ( (cid:101) X 1 ( p )) (cid:36)
( p ) is linear in each of its arguments (cid:36) 2 ( p ) and (cid:36) 1 ( p ) .
C Formalism of the particle ﬁlter
Using a bootstrap particle sampler, we can approximate the posterior distributions of the statesgiven noisy and partial observations of the state of a hidden Markov model (Section 3.4) by a cloudof K weighted particles which we denote by x ( k ) i .
The particles evolve independently accordingto the program, x ( k ) i
=
X i ( x ( k )
i − 1 , θ ) .
Each particle has a corresponding likelihood weight (cid:36) ( k ) imeasuring how well the trajectory of the particle so far is matching the observations, updatedby (cid:36) ( k ) i = p
( y i | x ( k ) i , θ ) (cid:36) ( k ) i − 1 , such that the empirical measure (cid:80) k (cid:36) ( k ) i
δ x ( k ) i / (cid:80) k
(cid:36) ( k ) i of theweighted particles approximates the ﬁltering distribution p ( x i | y 1 , . . .
, y i , θ ) .
Most importantly forour application, the marginal likelihood of a parameter θ can be read off the weights at the last step,p ( y 1 , . . .
, y n | θ )
≈ (cid:80)
Kk =1 (cid:36) ( k ) n .
So far, the particles and weight trajectories are differentiable with respect to the parameter [3].
However, weight degeneracy, the collapse of all but a few weights, is a common problem.
Thestrategy to circumvent this issue is to include multinomial resampling steps, which are inherentlydiscrete, but differentiable with (smoothed) stochastic duals: The population of particles witha population of particles of equal weight ( x (cid:48)
( k ) i , (cid:36) (cid:48) i ) k ∈{ 1 ,...,K } , i.e. where the new particles are 673 chosen by sampling indices from the multinomial distribution with the normalized weight vectorw =
( w (1) , . . .
, w ( K ) ) =
( (cid:36) (1) i /(cid:36), . . .
, (cid:36) ( K ) i /(cid:36) ) .
After this step,
[Equation]
and (cid:36) :
= (cid:80) k
(cid:36) (cid:48)
( k ) i =
(cid:80) k (cid:36) ( k ) i .
Resampling is numerically attractive because unlikely particles arediscarded so numerical resources are not wasted on particles with vanishing weight.
One achieves asimilar effect by retaining each particle independently with probability proportional to its weight.
In ﬁrst approximation, a selected particle whose weight (cid:36) ( k ) decreases by h ( k ) has a chance ofh ( k ) (1 − w ( k ) ) (cid:36) ( k ) not being selected under the new weight.
We can disregard the effect 1 − w ( k ) on theother selected particles without incurring a bias, as shown below.
We thus propose to set the new weight of particle j spawned from a particle k to (cid:36)K + h ( k ) Kw ( k ) .
This isreadily accomplished using a formally differentiable weight function(cid:36) ( j ) =
(cid:36) newweight( w ( k ) ) /K
with deﬁnition:# primal new_weight(p::Real) =
# forward-mode AD rule function new_weight(p::ForwardDiff.
Dual{T}) where {T} ∂ _p = ForwardDiff.partials(p) value_p = ForwardDiff.value(p) value_p
=
max(1e-5, value_p)
ForwardDiff.
Dual{T}(1, ∂ _p / value_p) end # reverse-mode AD rule function ChainRulesCore.rrule(::typeof(new_weight), p) function new_weight_pullback( ∇ Ω ) return (ChainRulesCore.
NoTangent(), ∇ Ω / p) endreturn (one(p), new_weight_pullback) end
We now have:
Proposition C.1 (Resampling) .
Consider a vector of weights (cid:36) (1) , . .
, (cid:36) ( K ) with (cid:36) =
(cid:80) Kk =1 (cid:36) ( k )and a vector of particles x (1) , . . .
, x ( K ) as representation of µ = (cid:80)
Kk =1 w ( k ) f
( (cid:36) ( k ) )
where w ( j ) =(cid:36) ( j ) (cid:36) are normalized weights and f a functional of interest.
If J is a categorical random variablewith P ( J = j )
= w ( j ) , then(cid:36) E f
( x J )
= µ and for inﬁnitesimal differentials
h (1) , . .
, h ( K ) , ¯ h (1) , . . .
, ¯ h ( K ) for weights and particles,
E (cid:20)(cid:18) (cid:36) [Equation] ) (cid:36)
[Equation]
(cid:19)
[Equation]
[Equation]
This is as required, giving the unbiasedness of our rule for the resampling step of the particle ﬁlter.
D Implementation details
D.1 Experiment details
All computation times in Fig.
5 were measured on an Intel Xeon Platinum 8260 CPU and Julia version1.6.
Garbage collection times are included in the total run time.
We provide code and instructions torun the examples in the tutorials folder of StochasticAD.jl .
D.2 Software dependencies
StochasticAD.jl is implemented in the Julia Language [4] and uses internally
Distributions.jl and DistributionsAD.jl [2, 1],
ChainRulesCore.jl [5], and
ForwardDiff.jl [6].
The presented examples are using the additional packages: BenchmarkTools.jl [7], Zygote.jl [8],
GaussianDistributions.jl [9], and Plots.jl [10].
We refer the reader to the documentation of StochasticAD.jl for more detailed information on thepackage implementation.
Supplemental References
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section 3 and Section
for discussions on potentially large variance of the estimator and language-wide AD
support.
In addition, we provide a continuously updated list about the current limitations
of the StochasticAD.jl package in its README.
(c) Did you discuss any potential negative societal impacts of your work?
[No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Appendix.(b)
Did you include complete proofs of all theoretical results?
[Yes] See Appendix.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
[Table]
mental results (either in the supplemental material or as a URL)?
[Yes] See Appendix D
[Table]
and the README of StochasticAD.jl .
[Table]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes]
[Table]
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
[Table]
ments multiple times)?
[Yes] For the differentiable particle sampler, we have includedthe error on the mean value of the run times as error bars.
[Table]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix D.
[Table]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
[Table]
(a) If your work uses existing assets, did you cite the creators?
[Yes] We cite the Julia
[Table]
packages on which this work is based.
[Table]
(b) Did you mention the license of the assets?
[No] We only use open-source software
[Table]
packages.
[Table]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
[Table]
We have included our package in the supplemental material.
[Table]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
[Table]
using/curating?
[N/A] We do not use any previously collected data.
[Table]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
[Table]
information or offensive content?
[N/A]
[Table]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
[Table]
if
[Table]
[Table]
Did you wage
participants
the amount
[Table]
spent on participant compensation
[Table]