Automatic Differentiation of Programs with Discrete Randomness
Abstract
Automatic differentiation (AD), a technique for constructing new programs which compute the derivative of an original program, has become ubiquitous throughout scientific computing and deep learning due to the improved performance afforded by gradient-based optimization.
However, AD systems have been restricted to the subset of programs that have a continuous dependence on parameters.
Programs that have discrete stochastic behaviors governed by distribution parameters, such as flipping a coin with probability p of being heads, pose a challenge to these systems because the connection between the result (heads vs tails) and the parameters ( p ) is fundamentally discrete.
In this paper we develop a new reparameterizationbased methodology that allows for generating programs whose expectation is the derivative of the expectation of the original program.
We showcase how this method gives 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 unbiased reverse-mode AD of a particle filter.
Our code is available at https://github.com/gaurav-arya/StochasticAD.jl .
Introduction
Automatic differentiation (AD) is a technique for taking a mathematical program X ( p ) and generating a new program (cid:101)
X
( p ) =
d X d p for computing the derivative [1, 2].
AD is widely used throughout machine learning and scientific computing due to the increased performance of gradient-based optimization techniques compared to derivative-free methods [3].
However, if X ( p ) returns the flip of a coin with probability p of receiving a 1 and probability 1 − p of receiving a 0 , it is clear that d X d p is not defined in the classical sense.
But when attempting to calibrate the parameter p to data, one may wish to fit the model using statistical quantities, e.g. find p such that the average of X ( p ) is close to the average sum of N real-world coin flips.
Given this use case, can one automatically construct a program (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 finite differences, i.e.:
[Equation]
Pruning follows one alternative path
[Figure]
X ( p + ε )( ω 2 )
Smoothing ﬂattens onto
[Figure]
x = X ( p )( ω 1 )
y ∼ Y | X ( p ) =
x
Finite differences’ major issue in the context of stochastic programs is that this calculation does not correlate the calculation of E [ X ( p + ε )] with the calculation of E [ X ( p )] .
This leads to a large variance in the finite-difference estimator [4] that goes to ∞ as ε → 0 .
This issue of unbounded variance could be solved if one could run the perturbed program with the same set of random numbers, i e. directly estimate E [ X ( p + ε ) − X ( p )] , by a method that is well-posed in the limit of ε → 0 (Fig. 1).
We demonstrate how to automatically construct a new stochastic program (cid:101) X ( p ) whose expectation satisfies E [ (cid:101) X ( p )]
=
d E [ X ( p )] d p .
We derive this through a technique which we term the stochastic derivative , propagating the proportional probability of differing event outcomes due to infinitesimal changes in p .
Our technique has the following design goals:
• Composability.
We differentiate stochastic programs
X ( p ) which are themselves composed of many “elementary” stochastic programs, including samples from discrete distributions such as the Bernoulli, Poisson, and Geometric, and samples from continuous distributions.
These may be composed (chained, added, etc.)
arbitrarily, with computational cost independent of how they are composed.
• Unbiasedness.
The program’s discrete structure is preserved, leading to a provably unbiased estimator that avoids continuous relaxations and tunable or learned internal parameters.
• Low variance.
We consider correlated paths through the computation that are linked by the smallest possible perturbation (Fig. 1), generalizing the widely-used pathwise gradient estimator for continuous randomness [5] to the discrete case.
We show the utility of our technique by demonstrating forward-mode AD of stochastic simulations like inhomogeneous random walks and agent-based models such as the Game of Life.
To achieve O (1) computational overhead for these applications, we combine stochastic derivatives with an online strategy we call pruning .
We also demonstrate a straightforward way to perform reverse-mode
AD from our approach via smoothing , allowing us to derive from first principles the biased but empirically successful straight-through gradient estimator [6] as a special case and to construct an unbiased end-to-end reverse-differentiable particle filter, recovering a technique discovered in [7].
We provide an open-source implementation of 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 [8, 9, 10], measure-valued [11, 12], and pathwise [13, 14] gradient estimators.
When it comes to discrete randomness, the score-function method is a popular general-purpose choice because it is unbiased and can be composed through stochastic computation graphs [15].
However, the score-function method does not search for correlated paths through the computation and thus suffers from high variance, making gradient computation with discrete variables challenging.
A number of techniques (REBAR [16], RELAX [17]) have been introduced that rely on control variates for variance reduction [18].
Measure-valued derivatives also have no notion of intrinsic coupling, though coupling can be achieved using common random numbers for certain distributions [19].
As an alternative direction, several recent works such as Gumbel-Softmax [20] consider continuous relaxations of discrete programs so that a pathwise gradient estimator, based on the “reparameterization trick” [14], can then be applied.
However, such methods face a variance-bias tradeoff and are inapplicable to discrete programs that cannot be continuously relaxed.
Our stochastic derivatives also extend the pathwise gradient estimator to discrete programs but do so unbiasedly.
The field of smoothed perturbation analysis [21, 22] develops a mathematically equivalent object to our smoothed stochastic derivative based on conditional expectations, which is a special case of our formalism, and has also considered a randomized approach similar to our pruning technique in the context of generalized semi-Markov processes [23].
However, these ideas have not previously been applied to construct a general-purpose AD method via rigorous composition rules and the algorithms and data structures to realize the approach automatically.
Although our method can be used to hand-derive a gradient estimator, the main feature is composability through user-written functions, enabling an automated mechanism.
While mainstream AD frameworks do not support unbiased differentiation of discrete random programs, Storchastic [24] is a specialized framework for AD of stochastic computation graphs [15] where the user can specify which estimator to use at each node, as well as any tunable hyperparameters.
Storchastic implements an exhaustive set of prior gradient estimation methods at each sampling step.
However, the runtime of the derivative estimate is in general exponential in the length of the largest chain of stochastic nodes, an artifact of the way many prior gradient estimators compose [24, 25].
In Section 3 we demonstrate stochastic AD that matches the computational complexity of deterministic AD even when discrete random functions are chained together, 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 pathwise gradient estimator for continuous random programs, and present the key ideas underpinning the formalism.
When describing infinitesimal asymptotics, we say a function g ( ε ) is O ( ε ) if | g ( ε ) | ≤ C | ε | for some real C and all sufficiently small ε .
Colloquially, we describe quantities that are O ( ε ) as “infinitesimal”.
2.1 Infinitesimally perturbing a stochastic program
We are interested in differentiating stochastic programs, formally defined below.
A stochastic program
X ( p ) can be thought of as a map from an input p to a random variable X ( p ) .
Here, X ( p ) can represent either an “elementary” program such as a draw from a Bernoulli distribution, or the full user-provided stochastic program represented using many elementary programs, e.g. a simulation of a random walk.
Definition 2.1.
A stochastic program X ( p ) is a stochastic process with values in a Euclidean space
E , whose index set I is either an open subset of a Euclidean space or a closed real interval.
Let Ω be the sample space, equipped with a probability distribution P .
To sample X ( p ) at input p ∈
I , which we call the primal evaluation , one should imagine a sample ω being randomly chosen from Ω according to P to produce an output X ( p )( ω ) ∈ E .
Note that P is independent of p and X ( p ) is a map Ω → E ; such a formulation has been called the “reparameterization trick” [14].
For example, a Bernoulli distribution Ber( p ) can be represented by choosing a uniform random ω ∈ [0 , 1] and defining X ( p )( ω ) =
1 [1 − p, 1] ( ω ) , where 1 S is the indicator function for a set S (see Fig. 2b).
At fixed p ∈
I we define the differential d X ( ε ) , which is itself a stochastic program:
[Equation]
X ( p ) .
Let us now restrict our attention to the case where I is a closed real interval, so that p, ε ∈ R .
The sensitivity of stochastic programs Z with more general index sets can be understood at an input u by studying at p =
the directional perturbation X ( p ) =
Z ( u + p v ) in a direction v .
Output Ω
X ( p + ε )
X ( p ) d X ( ε )
O ( ε ) ω 1 ω { { Not O ( ε ) (a) X ( p ) ∼ Exp( p ) .
Output Ω1 − p − ε 1 − p
X ( p + ε )
X ( p )  (b) X
( p ) ∼ Ber( p ) .
d X ( ε ) 1 − p − ε 1 − p
Output Ω
X ( p )( ω ) = X ( p )( ω ) =
1 { {  (c) X ( p ) ∼ Ber( p ) .
For the full user-provided stochastic program, we are interested in computing d E [ X ( p )]
d p , which is related to E [d X ( ε )] by d E [ X ( p )]
d p
[Equation]
E
[Equation]
] ε .
We thus expect d X ( ε ) to have “infinitesimal” O ( ε ) expectation as ε approaches 0.
Let us consider the form of d X ( ε ) for two elementary programs: a program returning a sample from the continuous
Exponential distribution Exp( p ) with scale p and a program returning a sample from the discrete
Bernoulli distribution Ber( p ) with success probability p .
In both cases, the program is parameterized via the inversion method [26] over the sample space Ω
= [0 , 1] , which chooses X ( p )( ω ) to be nondecreasing in ω .
For X ( p ) ∼ Exp( p ) , X ( p )( ω ) = − p log(1 − ω ) is differentiable in p .
Thus, the range of d X ( ε ) is O ( ε ) for any fixed ω , as illustrated in Fig.
2a.
In contrast, consider X ( p ) ∼ Ber( p ) , with X ( p )( ω ) = 1 [1 − p, 1] ( ω ) .
As shown in Fig.
2b, as ε →
the output of X ( p ) at a random ω ∈ Ω almost surely does not change in X ( p + ε ) .
Specifically, d X ( ε ) is 0 with probability 1 − ε ≈ 1 , and assumes the value 1 with probability
ε (Fig. 2c).
This is the fundamental challenge presented by discrete randomness: a O ( ε ) change in the input turns into a finite perturbation to the output, where “finite” means that it is non-vanishing as ε → 0 .
The perturbation occurs with “infinitesimal” 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 and derivative computations are coupled by a shared source of randomness.
To see this, let us understand what happens when they are entirely uncoupled.
A black-box finite difference approach with step size ε would independently sample ω 1 and ω 2 from Ω (Fig.
2a, red), computing the derivative estimate ( X ( p + ε )( ω 2 ) − X ( p )( ω 1 )) /ε .
However, since the samples are independent, the variance of the estimator is of order 1 /ε 2 , so we are forced to pick a finite ε to balance a variance-bias tradeoff.
This motivates using the same random sample for the primal and derivative computations (Fig.
2a, blue).
For continuous randomness, taking the limit of this approach as ε → 0 leads to the widely-used pathwise gradient estimator δ , given as the almost sure limit of d X ( ε ) /ε as ε → 0 (i.e. the pointwise derivative of X ( p ) at each fixed ω ), so that
[Equation]
But, considering a simple Bernoulli variable X ( p ) ∼ Ber( p ) as in Fig.
2c, we see how this approach is ill-suited for the discrete case!
As ε approaches 0, the differential d X ( ε ) is non-zero with infinitesimal O ( ε ) probability.
This means that δ is almost surely 0, while the true derivative of
E [ X ( p )]
= p is 1 .
The finite perturbation is neglected.
Thus, the pathwise gradient estimator needs to be modified to handle discrete randomness.
The issue with the interchange of limit and expectation in (2.3) is that d X ( ε ) /ε is unbounded in ε in the presence of a finite perturbation.
This motivates explicitly considering the event of a large jump in d X ( ε ) , which may be characterized by the event A B ( ε ) = {| d X ( ε ) | > B | ε |} for some chosen bound B .
The event A B ( ε ) has O ( ε ) probability, which is vanishingly small, but its contribution to the derivative estimate cannot be neglected because it contains finite perturbations.
This motivates sampling from this part of the probability space separately .
Formally, we introduce a random weight w ∈ R and alternate value Y ∈ E that characterize the sensitivity of X ( p ) when the probability space is restricted to A B ( ε ) [given by the r.h.s.
of (2.4) below], forming the “stochastic derivative”:
Definition 2.2 (Stochastic derivative) .
Suppose X ( p ) ∈
E is a stochastic program with index set I a closed real interval.
We say that the triple of random variables ( δ, w, Y ) , with w ∈ R and Y ∈ E , is a right (left) stochastic derivative of X at the input p ∈
I if d X ( ε ) /ε
→ δ almost surely as ε → 0 , and there is an integrable (i.e. of bounded expectation) random variable B > | δ | such that for all bounded functions f : E → R with bounded derivative it holds almost surely that
[Equation]
p
[Equation]
(cid:21)
[Equation]
2.4) with limit taken from above (below), where P ( A B ( ε ) | X ( p )) /ε is dominated by an integrable random variable for all ε > 0 ( ε < 0 ).
A stochastic derivative may be collapsed into an unbiased estimator of the derivative of E [ X ( p )] .
Proposition 2.3 (Unbiasedness) .
If ( δ, w, Y ) is a stochastic derivative of X ( p ) at p , it holds that
[Equation]
Proof sketch.
With f as identity, by Definition 2.2 the sensitivity of X ( p ) over A B ( ε ) is given by w ( Y − X ( p )) .
On the other hand, given A cB ( ε ) it holds that | d X ( ε ) /ε
| ≤ B : a dominated convergence argument shows that the sensitivity of X ( p ) is then given by its almost-sure derivative δ .
Theorem 2.4 shows the existence of the stochastic derivative for a stochastic program, subject to technical assumptions given in our formal treatment (Appendix B).
Theorem 2.4 (Existence, simplified) .
Given a sufficiently regular stochastic program X ( p ) with index set I a closed interval [ a, b ] ⊂ R , there exists a right stochastic derivative ( δ, w R , Y R ) with w R ≥ 0 at any p ∈
[ a, b ) and a left stochastic derivative ( δ, w L , Y L ) with w L ≤ 0 at any p ∈ ( a, b ] .
Proof sketch.
The proof is by construction: at a high level, w is the derivative of the probability of a large jump, while Y follows the distribution of the possible jumps, conditional on a jump happening.
Specifically, Y has distribution given as the limit as ε → 0 of the conditional distribution of X ( p + ε ) =
X ( p )+d X ( ε ) given the event A B ( ε ) and the outcome of X ( p ) .
The weight w is given as the derivative w.r.t.
ε of the probability P
( A B ( ε ) |
X ( p ))
that d X ( ε ) jumps by a non-infinitesimal amount, conditional on X ( p ) .
Essentially, since P ( A B ( ε ) | X ( p ))
≈
wε , multiplying by w bridges the gap between conditioning on A B ( ε ) [recall Y is constructed conditional on A B ( ε ) ] and simply restricting the probability space to A B ( ε ) [i.e. multiplying by 1 A B ( ε ) as in (2.4)].
Example 2.5 (Right stochastic derivative of Bernoulli variable) .
Suppose X ( p ) ∼ Ber( p ) , parameterized via the inversion method, and take ε > 0 .
As shown in Fig.
2c, d X ( ε ) is parameterized as,
[Equation]
Given the event X ( p ) = 1 , we have that ω ≥ 1 − p and d X ( ε ) is deterministically 0.
On the other hand, given X ( p ) = 0 , we have ω < 1 − p , so with probability ε/ (1 − p ) the differential d X ( ε ) assumes a value of 1, i.e. the Bernoulli variable flips from 0 to 1.
Thus, we may construct a right stochastic derivative (0 , w R , Y R ) of X ( p ) by letting w R = 1 / (1 − p ) , Y R = 1 conditionally on
X ( p ) = 0 and w R
=
Y R = 0 conditionally on X ( p ) = 1 .
For a concrete example, with p = 0 . 6 , we have w R = 2 .
5 , Y R = 1 conditionally on X ( p ) = 0 ; StochasticAD.jl prints this as "
+ ( with probability 2.5 ε )" (as explained further in Section 3.1).
We give a number of examples of stochastic derivatives in Appendix A.
It is of crucial importance that w and Y depend conditionally on the output of X ( p ) .
In particular, given a primal evaluation X ( p ) = x , the forms of w and Y depend only on the distribution of the perturbed program on the set { ω : X ( p )( ω ) = x } ⊂ Ω , which elegantly generalizes the coupling achieved by the pathwise gradient estimator δ at each fixed ω ∈ Ω .
Intuitively, this coupling allows us to consider only the smallest possible perturbations to the program in the derivative computation (Fig. 1), and thereby achieve variance reduction without resorting to continuous relaxations.
For example, for a binomial variable X ( p ) ∼ Bin( n, p ) parameterized via the inversion method (a natural parameterization for maximizing coupling), it holds that Y ∈ { X ( p ) − 1 , X ( p ) + 1 } .
We show in Example A.1 that the variance of our gradient estimator 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
Proposition 2.3 shows that the derivative estimate produced by stochastic derivatives is correct in expectation.
But this is insufficient to ensure composition , i.e. a stochastic derivative “chain rule”.
While Definition 2.2 requires composition through deterministic test functions f to enforce a sufficiently strict definition, Theorem 2.6 provides a general-purpose composition result through any program which has a stochastic derivative, as well as multidimensional programs with directional stochastic derivatives (i.e. stochastic derivatives of directional perturbations of the program).
Theorem 2.6 (Chain rule, simplified) .
Consider independent stochastic programs X 1 and X 2 and their composition X 2 ◦ X 1 .
Suppose that X 1 has a right (left) stochastic derivative at p ∈ R given by ( δ 1 , w 1 , Y 1 ) , and X 2 has a right stochastic derivative ( δ 2 , w 2 , Y 2 ) in the direction ˆ δ 1 =
δ 1 / | δ 1 | given conditionally on its input X 1 ( p ) .
Then, under assumptions to ensure integrability, the stacked program [ X 1 ; X 2 ◦ X 1 ] has a right (left) stochastic derivative ( δ, w, Y ) at p where δ =
[ δ 1 ; | δ 1 | δ 2 ] ,
[Equation]
and w = w 1 + | δ 1 | · w 2 .
Proof sketch.
Let A 1 ( ε ) be the event of a jump in X 1 and A 2 ( ε ) be the event of a jump in X 2 when its input is perturbed by εδ 1 .
Given A c 1 ( ε ) ∩ A c 2 ( ε ) (a jump in neither), a dominated convergence argument implies that X 2 ◦ X 1 does not jump either.
Given A 1 ( ε ) , Definition 2.2 yields that the sensitivity of X 2 ◦ X 1 is described by alternate values Y 1 with weight w 1 , while given A 2 ( ε ) the sensitivity of X 2 ◦ X 1 is described by alternate values Y 2 with weight | δ 1 | w 2 .
We may then form w as the sum of these weights and Y as a weighted distribution over the two cases.
Crucially, we may neglect the sensitivity of X 2 ◦ X 1 given A 1 ( ε ) ∩ A 2 ( ε ) (a jump in both), as this event has probability O ( ε 2 ) .
2.4 Smoothed stochastic derivatives
The reparameterization trick neglects finite perturbations, while stochastic derivatives precisely capture all possible finite perturbations.
There exists a middle ground between these methods: one may take a conditional expectation on X ( p ) so that finite perturbations have been “smoothed” into infinitesimal ones.
(E.g. "0 + (1 with probability 2.5 ε )" becomes "
+ 2.5 ε " .)
Definition 2.7 (Smoothed stochastic derivative) .
For a stochastic program X ( p ) with a right (left) stochastic derivative ( δ, w, Y ) at input p , a right (left) smoothed stochastic derivative (cid:101) δ of X at input p is given as
[Equation]
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 0.36 + 1.2 ε b = rand(Binomial(10, p))
6 + (1 with probability 10.0 ε ) c
= 2 * b + 3 * rand(Bernoulli(p))
12 + (3 with probability 12.5 ε ) return a * c * rand(Normal(b, a)) 9 end 1 julia> using StochasticAD 2 julia> st = stochastic_triple(X, 0.6)
# sample a single stochastic triple at p = 0.
3 27.11 + 94.32 ε + (6.78 with probability 12.5 ε ) 4 julia> derivative_contribution(st) # which produces a single derivative estimate...
6 julia> samples =
[derivative_estimate(X, 0.6) for i in 1:1000] # take many estimates!
7 julia> println("d/dp of E[X(p)]: $(mean(samples)) ±
$(std(samples) / sqrt(1000))") 8 d/dp of E[X(p)]: 204.63 ± 1.
Smoothed stochastic derivatives easily permit reverse-mode AD instead of forward-mode, as they have the same form as the usual derivative.
However, they enjoy more limited composition properties: they propagate exactly through differentiable functions f that are linear over the conditional distribution of
Y given X ( p ) via the standard chain rule, as we prove in Appendix B.6.
Due to the coupling of Y and
X ( p ) , this is a much weaker requirement than global linearity, which can lead to low-bias estimates: for example, for the program X ( p ) =
Geo( p ) 3 smoothed stochastic derivatives give a derivative estimate with < 0 .
5% bias at p = 0 .
01 , even though the cube function is highly non-linear on the inter-quartile range [28 , 137] of Geo( p ) .
Smoothed stochastic derivatives recover the widely-used straight-through gradient estimator [6] as a special case, as we work out in Example A.8.
Automatic differentiation of stochastic programs
We develop StochasticAD.jl , a prototype package for stochastic AD based on our theory of stochastic derivatives.
As discussed in Section 2.4, smoothed stochastic derivatives obey the usual chain rule, and thus can be used with existing AD infrastructure by supplying custom rules for discrete random constructs, and we do so for a particle filter 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 object called 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
Forward-mode AD is often implemented with dual numbers [2], which pair the primal evaluation of a deterministic function f ( p ) with its derivative dd p f ( p ) .
Dual numbers can be propagated through a program using the chain rule.
A useful alternative perspective of dual numbers is that they propagate an “infinitesimal” perturbation ε to the input through the program, where dd p f ( p ) is the coefficient of the “dual” element ε .
Stochastic triples generalize dual numbers by including a third component ∆ s to describe finite perturbations with infinitesimal probability (Fig. 3, left).
In Fig.
3, we consider a toy program X ( p ) including discrete randomness.
We are interested in the derivative of E [ X ( p )] at p = 0 . 6 , and hence we provide the stochastic triple printed as "0.
+ ε " as input.
First, the triple is squared and becomes "0.36 + 1.2 ε " : this is the familiar way that dual numbers propagate, via the chain rule.
But what happens when the triple is propagated
[Table]
through the discrete and random Binomial variable?
The resultant stochastic triple "6 + (1 with probability 10.0 ε )" is an integer with a component of discrete change, reflecting an infinitesimal probability of one more success in the Binomial Bin(10 , 0 . 6) .
We can in fact understand why the probability is "10 ε " by representing the Binomial as the sum of 10 Bernoulli variables, each with probability 0 .
Since 4 of the Bernoulli’s have an output of 0, they each have a probability "2.5 ε " of switching to 1 (recall Example 2.5), and thus there is in total a "10 ε " probability that the output of the Binomial increases by 1 (rigorously, we have applied Theorem 2.6).
Formally, we can interpret a printout "x + δε + ( Δ with probability w ε )" as follows: x is a sample of the random variable X ( p ) describing the primal evaluation, and δ , w , and x+ Δ are samples of the components δ , w , and Y , respectively, of the stochastic derivative of X at p .
Using our chain rule for stochastic derivatives (Theorem 2.6), we write rules for propagating stochastic triples through functions via operator overloading [27], exploiting Julia’s multiple dispatch feature [28].
When multiple discrete changes are possible, we pick one probabilistically: we call this strategy pruning (recall Fig. 1) and show its unbiasedness in Appendix B.5.
For example, in line 6 of Fig. 3, right, we probabilistically choose between the perturbation to the Binomial and the perturbation to the Bernoulli, in this case picking the latter.
To handle causal relationships between perturbations as in the first case of Eq.
(2.7), we associate each perturbation with a tag to avoid erroneously pruning between two perturbations that occur simultaneously.
Thus, stochastic triples can efficiently propagate through the full toy function written in Fig.
The function derivative_estimate creates a stochastic triple, propagates it, and collapses it into the derivative estimate δ + w Δ , forming an unbiased estimate of the derivative via Proposition 2.3 (Fig. 3, bottom).
3.2 Inhomogeneous random walk
We consider a Markovian random walk x 0 , . . .
, x n on Z ≥ 0 , with transition behavior dependent on a parameter p as follows,
[Equation]
We consider a program that stochastically simulates this walk and applies an arbitrary non-linear function f to the output x n .
In practice, f may represent a loss or a likelihood estimate; in this toy setting, we take f ( x ) = x 2 .
We are interested in studying the asymptotic behavior of the variance of our automatically derived gradient estimator, and so set p = n so that the transition function varies appreciably over the range of the walk for all n .
We find that the stochastic triple estimator has asymptotically lower variance than the score function estimator (Fig. 4b).
Crucially, stochastic triples achieve this variance reduction while remaining 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 stochastic version of John Conway’s game of life, played on a two-dimensional board.
In the traditional game of life, a dead cell becomes alive when 3 of its neighbors are alive, while a living cell survives when 2 or 3 of its neighbors are alive.
In our stochastic version, each of these events instead has probability 95% , while their complementary events have probability 5% .
Such types of discrete stochastic examples arise in many applications.
For example, mixing machine learning with agent-based models found in epidemiology and sociological contexts [29] or rule-based models and discrete stochastic (Gillespie) simulations of discrete chemicals in systems biology [30, 31] requires similar program constructs.
Consider a program that populates each cell of an N × N board with probability p , runs the stochastic game of life for T time steps, and counts the number of living cells n living .
We perform a sensitivity analysis of the final living population with respect to the initial living population, i.e. differentiate the expectation of n living with respect to p .
Stochastic triples propagate fully through the program, leading to an unbiased estimate of the derivative, as we verify with black-box finite differences.
The final board is depicted in Fig.
4c, along with the difference to the alternative final board chosen by pruning.
This is a high-dimensional example (the state space has dimension N 2 = 625 ) with fundamentally discrete structure, providing support for the algorithmic correctness and generality of stochastic triples.
In particular, the program cannot directly be continuously relaxed since it includes array indexing; in the approach of stochastic triples, integer-valued quantities stay integers.
3.4 Particle filter
As a final example, we consider a hidden Markov model with random latent states X 1 , . . .
, X n , observations y 1 ∼ Y 1 , . . .
, y n ∼ Y n , and parameters θ .
The likelihood L = p ( y 1 , . .
, y n | θ ) is in general not tractable, but a particle filter can be used to compute an estimate of L .
Here, we assume familiarity with particle filters [32, 33] but provide a short exposition with focus on the resampling step and experimental details in Appendix C. We assume the latent states to be continuous random, but discete randomness enters through the resampling step.
To differentiate the particle filter resampling step, we provide a stochastic derivative formulation of the returned particles and importance weights (Appendix C).
Importantly, L can be expressed by the sum of the weights at the last step, and the weight assigned to each particle in the resampling steps is used in a purely linear way.
Smoothed stochastic derivatives permit unbiased reverse-mode AD in this case.
Our formulated approach, though derived very differently, is equivalent to the particle filter AD scheme developed in [7] implementing the first estimator derived in [34], as we show in
Appendix C. Our particular choice of system allows the calculation of a ground-truth gradient of L by differentiating the Kalman filter algorithm [35].
Fig.
5a visualizes the latent process and observations, and the Kalman and particle filter trajectories.
Our estimator agrees with the Kalman filter derivative, unlike biased estimators [36, 37, 38] that neglect the contribution of the resampling step or perform it with entropy-regularized optimal transport (Fig. 5b).
For our program, we observe reverse-mode
AD to perform better for more than ≈ 100 parameters (Fig. 5c).
However, we find that the variance increases more rapidly with the number of steps and dimension as compared to biased estimators, suggesting that there is room for improvement in the coupling approach [7, 34].
Limitations and outlook
We have presented a method for unbiased AD of programs with discrete randomness, which, we have argued, is a natural generalization of pathwise gradient estimators based on the reparameterization trick to the discrete case.
However, more work will need to be done to turn our software demonstration
StochasticAD.jl into an AD system capable of handling the full complexity of applications such as machine learning.
A useful improvement would be better support for discrete constructs such as if statements with discrete random input; currently, such branches need to be rewritten using array indexing, which is supported.
Further, an interesting direction for future work is to automatically handle functions which deterministically turn a continuous random quantity into a discrete one, such as an inequality comparison X(p) > 0 or a Bernoulli variable implicitly implemented as rand() < p , and also to handle constructs such as while loops based on such functions.
We also expect future work to focus on further variance reduction.
Our method’s variance depends on the degree to which the primal and derivative computations can be coupled : while we present a natural method of coupling for a number of elementary stochastic programs and their compositions, the design space is rich when it comes to challenging examples such as the game of life or the resampling step of a particle filter.
(This design space is reflected in the fact that the form of the stochastic derivative depends not only on a program’s probability distribution, but also on the way it is parameterized.)
Furthermore, the pruning operation can introduce additional variance: for example, the derivative estimate automatically produced for the program B 1 + 2 B 2 with i.i.d.
B i ∼ Bin( n, p ) has variance O ( n 2 ) due to the pruning between the +1 and +2 perturbations, even though O ( n ) variance can be achieved for Bin(3 n, p ) .
Smoothing does not face this issue but accrues a bias through non-linear functions.
Pruning and smoothing may be thought of as the simplest ways to construct an AD algorithm from stochastic derivatives, lying on opposite ends of the design space: we anticipate future work to address their suboptimalities and ideally form unbiased estimators for discrete random programs that fully close the variance gap to their continuous counterparts [14, 39].
Finally, going beyond smoothing for reverse-mode AD, and ideally achieving unbiasedness while remaining coupled to the primal, is an important open problem for large-scale applications.
Acknowledgments and disclosure of funding
We thank Alan Edelman, Guillaume Dalle, and the anonymous reviewers for their feedback, Emile van
Krieken for helpful discussions regarding composing gradient estimators, and Simeon Schaub for help with the package.
We acknowledge the MIT SuperCloud and Lincoln Laboratory Supercomputing Center for providing HPC resources that have contributed to the research results reported within this paper.
This material is based upon work supported by the National Science Foundation OAC1835443, SII-2029670, ECCS-2029670, OAC-2103804, and PHY-2021825; the Advanced Research
Projects Agency-Energy DE-AR0001211 and DE-AR0001222; the Defense Advanced Research Projects Agency (DARPA) HR00112290091; the United States Artificial Intelligence Accelerator FA8750-19-2-
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section 3 and Sectionfor 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] 2.
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 experimental results (either in the supplemental material or as a URL)?
[Yes] See Appendix D and the README of StochasticAD.jl .
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] For the differentiable particle sampler, we have included the error on the mean value of the run times as error bars.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix D. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] We cite the Julia packages on which this work is based.
(b) Did you mention the license of the assets?
[No] We only use open-source software packages.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] We have included our package in the supplemental material.
(d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] We do not use any previously collected data.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
[Table]
In the supplement
[Table]
we:
• provide a number of examples of stochastic derivatives and present hand-worked examples of how they compose (Appendix A), • prove our formal statements regarding stochastic derivatives, and state and prove formal statements regarding smoothed stochastic derivatives (Appendix B), • provide an introduction to the particle filter methodology, and show how smoothed stochastic derivatives leads to unbiased differentiation of the resampling step (Appendix C), • provide details about the hardware used for the experiments of the main text and the software 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.
Example A.1 (Integer-valued stochastic programs) .
The Bernoulli, Binomial, Geometric, and Poisson distributions are all discrete, and in fact (nonnegative) integer-valued.
Suppose X ( p ) follows one of these distributions, parameterized using the inversion method [26] over the interval [0 , 1] .
We present a worked derivation of the stochastic derivative of X ( p ) , following the construction used in
Theorem 2.4.
We construct a stochastic derivative ( δ, w, Y ) at input p satisfying Definition 2.2.
Since X ( p ) is parameterized via the inversion method, d X ( ε ) ∈ {− 1 , 0 , 1 } for small enough ε .
Thus, for small enough ε , the conditional distribution of d X ( ε ) given X ( p ) =
x is:
P
[Equation]
X
[Equation]
p
[Equation]
P
[Equation]
X
[Equation]
p + ε ) =
x + 1 and X ( p ) = x )
P
[Equation]
X
[Equation]
p ) = x ) , (A.1)
P (d X ( ε ) = − 1 | X ( p ) = x ) =
P ( X ( p + ε ) =
x − 1 and X ( p ) = x )
P ( X ( p ) = x ) , (A.2) and d X ( ε ) = 0 otherwise.
According to the inversion method, X ( p )( ω ) = x for ω between
P ( X ( p ) ≤ x − 1) and P ( X ( p ) ≤ x ) .
Taking the derivative of the above quantities as ε → 0 therefore yields:
[Equation]
[Equation] [Equation]
and
[Equation] [Equation]
p
[Equation]
(A.6) where the ε included on the right hand sides is an abuse of notation to express their dependence on the direction of the limit, and 1 { C } = 1 { C } for a condition C .
By the differentiability of
Conditionally on X ( p ) =
x
X ( p ) w − w +
Ber( p ) − 1 /p · 1 { x = 1 and ε < 0 } 1 / (1 − p ) · 1 { x = 0 and ε > 0 }
Bin( n, p ) − x/p · 1 { x > 0 and ε < 0 } ( n − x ) / (1 − p ) · 1 { x < n and ε > 0 }
Geo( p ) x/ ( p (1 − p )) · 1 { x > 0 and ε > 0 } − ( x + 1) /p · 1 { ε < 0 }
Pois( p ) − x/p · 1 { x > 0 and ε < 0 } 1 { ε > 0 }
P (d X ( ε ) ̸ = 0) , d X ( ε ) /ε almost surely approaches 0 , so δ = 0 .
A simple choice for the bound B is 1 : then | d X ( ε ) | > Bε = ε ⇐⇒ d X ( ε ) ∈ {− 1 , 1 } .
We may now form w as the almost sure limit of
P ( A B ( ε ) | X ( p )) /ε
=
P (d X ( ε ) ̸ = 0 | X ( p )) and Y as the limit of the conditional distribution of
X ( p ) + d X ( ε ) given X ( p ) and d X ( ε ) ̸
= 0 .
By the above, w is given conditionally on X ( p ) =
x as w = w + + w − , (A.7) while the probability distribution of Y is given conditionally on X ( p ) =
x as
[Equation]
Table 2 lists the weights w + and w − for the Bernoulli, Binomial, Poisson, and Geometric distributions, derived using the parameterizations given in Table 1.
Using w + and w − , we can easily express w and
Y as above.
Note that w + and w − for a Binomial variable X ( p ) ∼ Bin( n, p ) may be derived via representation as the sum of Bernoulli variables, with the primal evaluation X ( p ) =
x corresponding to the sum of x Bernoulli successes and n − x Bernoulli failures.
The stochastic derivative for a Binomial variable Bin( n, p ) has weight w of order n and Y deterministically either 1 or − 1 depending on the chosen direction of the derivative.
Therefore, the derivative estimator w ( Y − X ( p )) has variance of order 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.
Specifically, for a differentiably parameterized continuous stochastic program X ( p ) , the stochastic derivative w.r.t a parameter p of the distribution has the form ( δ, 0 , 0) .
To satisfy Definition 2.2, one only has to select a bound
B such that | d X ( ε ) | <
B | ε | is satisfied almost surely: a natural choice is B = | δ | + 1 , so that
A cB ( ε ) = {| d X ( ε ) | ≤ ε | δ | + ε } simply requires that the higher-order correction to the derivative linearization be smaller than ε when applied at a point at distance ε from p .
For instance, the stochastic derivative of X ( p ) ∼ Exp( p ) w.r.t.
p , the scale of the Exponential distribution, is
[Equation]
conditionally on X ( p ) =
x . Such rules are already implemented in the Julia AD ecosystem in
Distributions.jl [1] and DistributionsAD
Example
A.
(Categorical variable) .
Consider a categorical variable
Categorical a ( p 1 , p 2 , p 3 , . .
. , p n ) assuming fixed outputs a = a 1 , a 2 , . .
. , a n with probabilities p 1 , p 2 , p 3 , . . . , p n > 0 .
Assume that p 1 , p 2 , p 3 , . .
. , p n implicitly depend on a single parameter p , to fit our definition of stochastic derivative w.r.t.
a single parameter, and let
X ( p ) ∼ Categorical a ( p 1 , p 2 , p 3 , . .
. , p n ) .
Suppose X ( p ) is parameterized via the inversion method with outputs ordered as a 1 , a 2 , . . .
a n .
Similar logic to Example A.1 then yields that its stochastic derivative is of the form (0 , w, Y ) with
Y ∈ { a x − 1 , a x +1 } , and
[Equation]
In theory, a 1 , a 2 , . . .
, a n can be arbitrary discrete objects such as arrays or strings, as the categorical variable could represent an intermediate value of the program that is ultimately converted into a number.
(Interpreted strictly, our formalism requires an embedding into Euclidean space; one can imagine a trivial embedding a i (cid:55)→ i which is then provided to a learned embedding.)
In handling such cases, one must be careful to speak in terms of alternate outputs rather than perturbations to the output.
Our formalism of stochastic derivatives can already accommodate this idea, as Y takes on the alternate values rather than the values of the perturbations.
We now consider some general stochastic programs.
Their stochastic derivatives arise automatically from composition and need not be hand-derived, but we do so to illustrate Theorem 2.6.
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
[Equation]
We know from Example A.1 that X 1 ( p ) has right (left) stochastic derivative (0 , w 1 , Y 1 ) with wand 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 about the discretely spaced alternate values ( x 1 − 1) 3 and ( x 1 + 1) 3 ; the conventional derivative 3 x 21 is irrelevant as the input to the cube function is integer-valued.
Indeed, by Theorem 2.6 the stochastic derivative of X ( p ) reads (
[Equation]
[Equation]
21) where Y 31 ∈ { ( x 1 − 1) 3 , ( x 1 + 1) 3 } .
Example A.6 (Parameter-scaled Bernoulli) .
Consider the program,
X
[Equation]
( p
[Equation]
Ber( p ) , (A.22)
[Equation]
(A.23)
By Theorem 2.6, the stochastic derivative of X ( p ) reads, (
[Equation]
( p
[Equation]
w
[Equation]
, pY 1 ) , (A.24) 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]
This represents a two-step random walk with steps given by X 1 ( p ) and X 2 ( p ) , where there is a transition 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 , so that 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)
[Equation]
(A.28) as given by Table 2.
Now, conditionally on X 2 ( p ) = 0 , X 2 (0) has right stochastic derivative w.r.t.
x 1 ,
[Equation]
(A
[Equation] [Equation]
Let us now turn our attention to the stacked program [ X 1 ( p ); ( X 2 ◦ X 1 )( p )] .
By Theorem 2.6, the stacked 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.30) 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, but a p 1 − p chance of flipping to 1, so that the first case expands into two events
[Equation]
Finally, the expression for X ( p ) may be thought of as a unary sum function operating on this stacked program, with stochastic derivative (0 , w, Y ) where conditionally on X 1 ( p ) = 0 and X 2 (0) =
[Equation]
with probability
[Equation]
Finally, we present an example using smoothed stochastic derivatives to rederive a popular gradient estimator.
Example A.8 (Recovering the straight-through gradient estimator of [6]) .
The straight-through gradient estimator formally assigns a derivative of 1 to the hard-thresholding function HT = 1 [0 , ∞ ) , such that a Bernoulli variable Ber( p ) ∼ HT( p − U [0 , 1]) is assigned a derivative of 1.
Now, let (cid:101) δ L and (cid:101) δ R be the left and right smoothed stochastic derivatives of Ber( p ) .
By Table 2 and Definition 2.7, for the right-sided ( ε > 0 ) case,
[Equation]
, (A
[Equation]
33) and for the left-sided ( ε < 0 ) case, (cid:102)
δ
[Equation]
= 1 /p · 1 { x = 1 } .
(A.34)
By linearity of expectation applied to Definition 2.7, any affine combination of (cid:102) δ L and (cid:102) δ R is also a valid smoothed stochastic derivative.
In particular,
[Equation]
is a valid smoothed stochastic derivative of a Bernoulli variable, which explains why the straightthrough estimator provides a low-bias estimate.
Note that a constant-valued affine combination of the left and right smoothed stochastic derivatives is not possible in general, e.g. it is not possible for
Geometric and Poisson random variables, in which cases smoothed stochastic derivatives generalize the straight-through estimator (mixing (cid:101) δ R and (cid:101) δ L may still be useful in these cases to reduce bias.)
B Proofs B.1 Preliminaries
Recall from the main text,
Definition 2.1.
A stochastic program X ( p ) is a stochastic process with values in a Euclidean space
E , whose index set I is either an open subset of a Euclidean space or a closed real interval.
Throughout the formalism, we let X ( p ) denote such a stochastic program, where I =
[ a, b ] ⊂ R is a closed interval.
As noted in the main text, it is sufficient to consider this case because sensitivities of stochastic programs Z with more general index sets can be understood for an input u by studying at p =
the directional perturbation X ( p ) =
Z ( u + p v ) in a direction v , where X ( p ) is then a stochastic program with index set a closed interval containing 0.
As in the main text, we use the shorthand
[Equation]
A number of statements and propositions can be stated for both right and left stochastic derivatives (which are generally not equal).
In this case, we often use the notation ε → 0 + / − to indicate that statements and proofs for right and left stochastic derivatives can both be read off by choosing the appropriate side of the limit for all expressions, where objects such as w have different definitions depending on the chosen reading ε → 0 + or ε → 0 − .
For clarity, we remark upon the (standard) notation,
[Equation]
where U and V are both random variables defined on a sample space Ω .
In this setting, W is itself also a random variable defined on Ω .
In particular, for a particular sample ω ∈ Ω , W ( ω ) is the conditional expectation E
[ U | V = V ( ω )] .
In the language of σ -algebras, E [ U | V ] is equivalent to
E
[ U | σ ( V )] , where σ ( V ) is the sub- σ -algebra generated by V .
B.2 Unbiasedness of stochastic derivatives
We first prove a general result regarding when a class of events A ( ε ) allows for the pathwise gradient estimator δ to be applied when the probability space is restricted to A c ( ε ) .
Proposition B.1.
Suppose d X ( ε ) /ε
→ δ almost surely, and | d X ( ε ) | ≤ Bε holds given the event
A c ( ε ) , where B is integrable.
Further, suppose A ( ε ) almost surely does not hold as ε → 0 + / − .
Then, for any function f : E → R with bounded derivative,
[Equation]
Proof.
Since 1 A c ( ε ) → 1 and d X ( ε ) /ε
→ δ almost surely, it holds by the chain rule that
[Equation]
almost surely as ε → 0 + / − .
Furthermore,
[Equation]
By assumption E G < ∞ .
The proposition follows by the dominated convergence theorem applied to conditional expectations.
As in the main text, define
[Equation]
and recall
Definition 2.2 (Stochastic derivative) .
Suppose X ( p ) ∈
E is a stochastic program with index set I a closed real interval.
We say that the triple of random variables ( δ, w, Y ) , with w ∈ R and Y ∈ E , is a right (left) stochastic derivative of X at the input p ∈
I if d X ( ε ) /ε
→ δ almost surely as ε → 0 , and there is an integrable random variable B > | δ | such that for all bounded functions f : E → R with bounded derivative it holds almost surely that
[Equation]
with limit taken from above (below), where P ( A B ( ε ) | X ( p )) /ε is dominated by an integrable random variable for all ε > 0 ( ε < 0 ).
Using Proposition B.1, we may show that outside the event A B ( ε ) , the sensitivity is well-described by the pathwise gradient estimator, allowing us to prove Proposition 2.3.
Proposition 2.3 (Unbiasedness) .
If ( δ, w, Y ) is a stochastic derivative of X ( p ) at p , it holds that
[Equation]
Proof.
Let B be the associated bound.
Since A cB ( ε ) holds almost surely as ε → 0 + / − and d X ( ε ) /ε
≤ Bε given A cB ( ε ) , Proposition B.1 applied to A cB ( ε ) with f as identity implies
[Equation]
while Definition 2.2 also applied with f as identity gives
[Equation]
Summing Eq.
(B.8) and Eq.
(B.9), we have
[Equation]
and applying the tower property of conditional expectations
[Equation]
we obtain
[Equation]
as desired
[Equation]
B.3 Construction of stochastic derivatives for elementary programs
We now give a technical condition under which a stochastic derivative can be constructed for an elementary stochastic program, in which we characterize w as the derivative of the conditional probability of A B ( ε ) given X ( p ) , and Y as realization of the weak limit of the conditional distribution of X ( p + ε ) given X ( p ) and A B ( ε ) .
Assumption B.2.
We assume
• there exists an integrable δ such that d X ( ε ) /ε
→ δ almost surely as
ε → 0 , and that we may find an integrable random bound B > | δ | such that
• for the quantity
[Equation]
w ( ε ) /ε is dominated in ε by an integrable random variable and converges almost surely to a random variable w as ε → 0 + / − .
• the conditional distribution of X ( p + ε ) given X ( p ) and A B ( ε ) converges in distribution to the law of a random variable Y as ε → 0 + / − .
Specifically, Y must satisfy
[Equation]
for
all bounded continuous
functions
f
E
→ R .
Given Assumption B.2, we may construct a stochastic derivative for X ( p ) .
Theorem 2.4.
Given a stochastic program X ( p ) satisfying Assumption B.2 for all p
∈
I
=
[ a, b ] ⊂ R , there exists a right stochastic derivative ( δ, w R , Y R ) with w R ≥ 0 at any p ∈
[ a, b ) and a left stochastic derivative ( δ, w L , Y L ) with w L ≤ 0 at any p ∈ ( a, b ] .
Proof.
Construct ( δ, w, Y ) via Assumption B.2 (taking the limit from above for the case of right stochastic derivatives, and the limit from below for the case of left stochastic derivatives).
For bounded f with bounded derivative, write
[Equation]
Note that E [ 1 A B ( ε ) | X ( p )]
= P ( A B ( ε ) | X ( p )) =
w ( ε ) .
Thus, as ε → 0 + / − the first term approaches
[Equation]
by Assumption B.2.
Now, using the X ( p ) -measurability of w ( ε ) and w , we may bound the magnitude of the second term as
[Equation]
which approaches 0 almost surely as ε → 0 + / − by the almost sure convergence of w ( ε )
/ε to w .
Thus,
[Equation]
Additionally, P ( A B ( ε ) | X ( p )) /ε is dominated in ε by an integrable random variable and d X ( ε ) /ε
→ δ almost surely by assumption.
So ( δ, w, Y ) is a stochastic derivative of X at p .
Finally, note that w ( ε ) /ε ≥ 0 for ε ≥ 0 while w ( ε ) /ε ≤ 0 for ε ≤ 0 , by the non-negativity of w ( ε ) .
So applying Assumption B.2 with limits from above indeed produces a right stochastic derivative, while applying Assumption B.2 with limits from below produces a left stochastic derivative.
B.4 Composition of stochastic derivatives
We now provide a proof of the chain rule for stochastic derivatives.
Theorem 2.6 (Chain rule) .
Consider independent stochastic programs X 1 and X 2 and their composition X 2 ◦ X 1 .
Suppose that X 1 has a right (left) stochastic derivative at p ∈ R given by ( δ 1 , w 1 , Y 1 ) with bound B 1 , and X 2 has a right stochastic derivative ( δ 2 , w 2 , Y 2 ) in direction ˆ δ 1 =
δ 1 / |
δ 1 | with bound B 2 given conditionally on its input X 1 ( p ) , where with d X 2 ( v ) =
X 2 ( X 1 ( p )+ v ) − X
( X 1 ( p ))
the event (cid:101)
A 2 ( ε ) =
{| d X 2 ( v )
| > B 2 | v | for some v satisfying | v | ≤ ε } (B.23) of a jump in some input direction v has P (cid:16) (cid:101)
A 2 ( ε ) (cid:12)(cid:12)(cid:12) X 1 ( p ) , X 2 ( p ) (cid:17) (cid:14) ε dominated in ε by an integrable random variable W .
Then, if B 1 B 2 and | δ 1 | W are integrable, the stacked program [ X 1 ; X 2 ◦ X 1 ] has a right (left) stochastic derivative at p given by ( δ, w, Y ) where δ =
[ δ 1 ; | δ 1 | δ 2 ] ,
Y = 
 [ Y 1 ; X 2 ( Y 1 )] with probability w
[Equation]
w 2 , [ X 1 ( p ); Y 2 ] with probability | [Equation] w
[Equation]
(2.7) and w = w 1 + | δ 1 | · w 2 , with associated bound B = B 1 + B 1 B 2 .
Proof.
We prove the case of right stochastic derivatives, as the case of left stochastic derivatives is analogous.
Let X =
X 2 ◦ X 1 as a shorthand.
Now, for ε ≥ 0 consider the events
[Equation]
Intuitively, the events
A 1 ( ε ) , A 2 ( ε ) , and A ( ε ) represent large jumps in X 1 , X 2 , and X respectively.
Throughout the proof, we let f denote a bounded function with bounded derivative.
Step 1: No large jump in X 1 or X 2 implies no large jump in X .
Given A c 1 ( ε ) and (cid:101)
A c 2 ( ε | δ 1 | ) , we have
[Equation]
[Equation]
which implies A c ( ε ) (no jump in output).
We now remark that (cid:101)
A 2 ( ε | δ 1 | ) and A 2 ( ε ) may be interchanged when considering sensitivities.
Specifically,
[Equation]
by applying the dominated convergence argument of Proposition B.1 for the directional perturbation ε (cid:55)→ X 2 ( X 1 ( p ) + εδ 1 ) to both of the event classes
A 2 ( ε ) and (cid:101)
A 2 ( ε | δ 1 | ) , noting that both A 2 ( ε ) and (cid:101)
A 2 ( ε | δ 1 | ) almost surely do not hold as ε → 0 and that (cid:101)
A c 2 ( ε | δ 1 | ) implies A c 2 ( ε ) .
Step 2: A large jump in both X 1 and X 2 has negligible probability.
Note that A 1 ( ε ) and (cid:101)
A 2 ( ε | δ 1 | ) are independent conditional on X 1 ( p ) .
Therefore, in the limit, we have
[Equation]
almost surely as ε → 0 + by the domination of P ( A 1 ( ε )
| X 1 ( p ))
(cid:14) ε by an integrable random variable and that P (cid:16) (cid:101)
A 2 ( ε | δ 1 | ) (cid:12)(cid:12)(cid:12) X 1 ( p ) , X ( p ) (cid:17) → 0 .
By boundedness of f this implies
[Equation]
so the sensitivity of X is negligible under A 1 ( ε ) ∩ (cid:101)
A 2 ( ε | δ 1 | ) .
Step 3: Sensitivity of X ( p ) given a large jump in X 1 is described by w 1 , Y 1 .
The random variable
[Equation]
is X 1 ( p ) , X ( p ) measurable and almost surely continuous, and the function f in the characterization
Eq.
(2.4) can be taken as bounded and almost surely continuous in X 1 ( p ) by an extension of the
Portmanteau theorem [40].
It follows, as X ( p ) and X 1 ( p + ε ) are independent given X 1 ( p ) , that
[Equation]
so that
[Equation]
Step 4: Sensitivity of X ( p ) given a large jump in X 2 is described by w 2 , Y 2 .
Using Eq.
(B.27), given (cid:101)
A c 2 ( ε | δ 1 | ) ∩ A c 1 ( ε ) it holds that
[Equation]
is dominated by an integrable random variable.
Therefore, since d X 1 ( ε ) /ε
→ δ 1 almost surely,
[Equation]
43) by dominated convergence.
Putting this together with Eq.
(B.31), by definition of w 2 and Y 2 , for any bounded f with bounded derivative we have
[Equation]
where we freely neglect (or add back in) 1 A 1 ( ε ) ∩ (cid:101)
A 2 ( ε ) and 1 A 1 ( ε ) ∩
A 2 ( ε ) ≤ 1 A 1 ( ε ) ∩ (cid:101)
A 2 ( ε ) inside the conditional expectation by Eq.
(B.35).
Step 5: ( δ, w, Y ) is a stochastic derivative of the stacked program [ X 1 ( p ); X ( p )] .
By the almost sure convergences d X 1 ( ε ) /ε
→ δ 1 and d (cid:98) X 2 ( ε ) /ε
→ δ 2 , the chain rule implies d X ( ε ) /ε
→ | δ 1 | δ 2 almost surely, where | δ 1 | δ 2 ≤ B 1 B 2 is integrable by assumption.
So δ =
[ δ 1 ; | δ 1 | δ 2 ] is an almost sure derivative of [ X 1 ( p ); X ( p )] .
Given A ( ε ) , | d X ( ε ) | > B 1 B 2 ε , and given A 1 ( ε ) , | d X 1 ( ε ) | > B 1 ε .
So we can choose the bound
B = B 1 B 2 + B 1 for the stacked program, so that A B ( ε ) ⊆
A 1 ( ε ) ∪ A ( ε ) .
Now, since A ( ε ) ⊆
A 1 ( ε ) ∪ (cid:101)
A 2 ( ε | δ 1 | ) ,
P ( A ( ε ) | X 1 ( p ) , X ( p ))
[Equation]
where P ( A 1 ( ε ) | X 1 ( p ) , X ( p )) (cid:14) ε and | δ 1 | W are integrable by assumption.
Therefore,
P ( A B ( ε ) | X 1 ( p ) , X ( p )) is also dominated by an integrable random variable, as desired.
Now, let f be operating on the stacked program, and for convenience write f ( x 1 ; x 2 ) =
f
([ x 1 ; x 2 ]) .
As a shorthand, let X ( p ) =
[ X 1 ( p ); X ( p )] .
With w and Y as given in the statement,
[Equation]
We conclude that ( δ, w, Y ) is a valid stochastic derivative of X ( p ) =
[ X 1 ( p ); X ( p )] .
B.5 Unbiasedness of pruning strategy
In the main text, we note that we can employ a pruning strategy so that we only ever track one alternative path (i.e. one sample from the stochastic derivative component Y ) and yet still obtain an unbiased estimate.
The following shows that the pruning method, whereby one chooses between two samples of Y by picking one with probability proportional to its weight, is indeed unbiased.
We prove by induction that the currently tracked alternative path is an unbiased choice amongst all possible alternative paths seen so far.
The base case, where the first alternative path observed is followed, is trivial as it is the only choice.
Now, suppose that n alternative branches have been observed so far, where w = w 1 + w 2 + · · · + w n is the summed weight so far.
Suppose we observe an ( n + 1) th branch, with weight w n +1 .
By our pruning strategy it is chosen with probability
[Equation]
[Equation]
while the j th path for some j ≤ n will have been chosen after this step with probability
[Equation]
[Equation]
as desired for unbiasedness.
Note that this enables us to prune online , i.e. without knowing the full structure of the computation a priori, which we exploit in StochasticAD.jl .
B.6 Smoothing
Recall from the main text,
Definition 2.7 (Smoothed stochastic derivative) .
For a stochastic program X ( p ) with a right (left) stochastic derivative ( δ, w, Y ) at input p , a right (left) smoothed stochastic derivative (cid:101) δ of X at input p is given as
[Equation]
Using Definition 2.2, we can write the above in an alternative form that does not rely on the definition of the stochastic derivative,
[Equation]
as given in [21, 22].
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 than global linearity over the full range of X ( p ) , as the range of Y conditional on X ( p ) is generally more restricted (e.g. Y ∈ { X ( p ) − 1 , X ( p ) + 1 } for a binomial variable).
Proposition B.3.
Suppose that f is linear over the range of Y conditionally on X ( p ) =
x , for all x .
Then,
[Equation]
Proof.
Note that ( f ′ ( X ( p )) · δ, w, f ( Y )) is a stochastic derivative of f ◦ X by Theorem 2.6.
Now, by Proposition 2.3 and the local linearity of f , we have the simplification,
[Equation]
In most cases, local linearity will only hold approximately, leading to bias in the estimate produced by propagating smoothed stochastic derivatives.
However, in the case of a particle filter resampling step an exact estimate is produced, as we show in the following.
C Formalism of the particle filter C.1 Hidden Markov model
Let us consider a hidden Markov model with random states X 1 , . . .
, X n as specified by a stochastic program X 1 ( θ ) giving the starting value depending on parameters θ and consecutive states given by pointwise differentiable stochastic programs
X i ( x i − 1 , θ ) depending on the previous state x i − 1 ∼
X i −
(Markov property) and θ .
In general, we allow continuous probability densities p ( x 1 | θ ) of Xand continuous transition probability densities p ( x i | x i − 1 , θ ) to depend arbitrarily on θ .
This latent process X 1 , . .
, X n is indirectly observed as the process y 1 ∼ Y 1 , . . .
, y n ∼ Y n with n observations
Y i ( x i ) depending on x i ∼ X i which are specified to have smooth conditional probability densities p ( y i | x i , θ ) depending only on x i and θ .
As a concrete, special case, we consider the following linear Gaussian state-space model with a d -dimensional latent process,
[Equation]
C
[Equation]
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 ) is a random initial position, and Φ is a d -dimensional rotation matrix.
Here, the parameters θ are defined by the entries of Φ , i.e. θ = vec(Φ) .
We use a particle filter to compute an estimate of the likelihood L = p ( y 1 , . . .
, y n | θ ) .
C.2 Differentiating a particle filter with a resampling step
Given n observations y 1 ∼ Y 1 , . . .
, y n ∼ Y n of the hidden Markov model defined via Eqs.
(C.1) and (C.2), a bootstrap particle filter allows us to approximate the posterior distributions of the states and the likelihood L of these observations by propagating a cloud of K weighted particles.
We denote the k th particle by x ( k ) i .
Each particle evolves independently according to the stochastic program [Eqs.
(C.1) and (C.2)] and carries a weight ϖ ( k ) i measuring how well the trajectory of the particle so far is matching the observations.
This importance weight is updated by ϖ ( k ) i =
[Equation]
such that the empirical measure (cid:80) k ϖ ( k ) i δ x ( k ) i / (cid:80)
k ϖ ( k ) i of the weighted particles approximates the filtering distribution p ( x i | y 1 , . . .
, y i , θ ) .
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.
Our goal is to discard unlikely particles, so that numerical resources are not wasted on particles with vanishing weight.
The strategy to accomplish this goal is to include resampling steps, where we pick the particles that best match the observations.
However, such resampling steps present discrete randomness, where the particle population is resampled according to the particle weights ϖ ( k ) i to form a new population ( x ′ )
( k ) i with equal weight ϖ ′
i =
1 /K · (cid:80)
Kk =1 ϖ ( k ) i .
To differentiate the resampling step, we need to describe how perturbing the weight distribution ϖ ( k ) i provided to the resampling procedure affects the resampled particles and weights.
Importantly, the marginal likelihood of a parameter θ can be read off using the weights at the last step
[Equation]
≈
K (cid:88) k
= ϖ ( k
[Equation]
n .
(C.3)
The strategies for resampling vary, but what they have in common is that the marginal distribution of each resampled particle x ′ ( k ) i is a multinomial distribution over the original particles with weights given by the normalized weight vector ( ϖ (1)
i /ϖ i , . . .
, ϖ ( K ) i /ϖ i ) with ϖ i = (cid:80)
Kk =1 ϖ ( k ) i .
Thus, a weighted sample of the marginal distribution of each resampled particle can be obtained by repeating the following procedure until obtaining a weight of ϖ i :
Sample an integer k uniformly from 1 to K . 2.
Return the particle x ( k ) i with assigned weight ϖ i Ber( ϖ ( k )
i /ϖ i ) .
Conditional on the assigned weight being ϖ i , the returned particle obeys the multinomial distribution.
The key insight is that the alternate possibility of a resampled particle x ( k ) i not being chosen can be written instead as the alternate possibility of its weight ϖ ( k ) i changing to 0.
Recall from Example 2.5 that there is an asymmetry between the left and right stochastic derivative of the Bernoulli distribution Ber( ϖ ( k ) i /ϖ ) conditioned on the output 0 or 1 being chosen.
In this case, it is convenient to take the left stochastic derivative (0 , w ( k ) i , Y ( k ) i ) with Y ( k ) i
= 0 and w ( k ) i =
− ϖ ki
/ϖ
i of the particle’s weight ϖ i because it is 0 when the assigned weight is 0, meaning that this case has no influence on the derivative estimate.
Therefore, particles that are not resampled [and thus do not contribute to the primal value, cf.
Eq.
(C.3)] do (also) not contribute to the derivative computation.
The fact that only the weights have a stochastic derivative, but not the particles imposes the following setting:
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 ϖ ( p ) f
( (cid:101) X ( p ))
=
E f
( X ( p )) where ϖ ( p ) is the Radon-Nikodym derivative [4] 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 ( ϖ ( p ) , (cid:101)
X ( p )) .
(C.4)
Proposition C.1.
If (cid:101) X ( p ) is a continuous program, reparameterized such that it is differentiable pointwise, and ϖ ( p ) has a smoothed stochastic derivative (cid:101) δ ,
[Equation]
Proof.
By assumption,
[Equation]
As the function ( ϖ, x ) (cid:55)→ ϖf ( x ) is linear in ϖ , the statement follows from Proposition B.3.
Therefore, we can replace the stochastic derivative of ϖ Ber( ϖ ( k ) i /ϖ ) with its smoothed stochastic derivative w ( k ) i (1 − Y ( k ) i ) =
ϖ ki /ϖ i obtained using Table 2.
This is convenient, as smoothed stochastic derivatives permit forward- and reverse-mode whereas reverse-mode AD becomes superior than forward-mode AD for functions f from R n to R m with m ≫ n [2].
Since the weights are used in a purely linear fashion, Proposition B.3 and Proposition C.1 guarantee that the derivative estimator is unbiased.
Numerically, we accomplish this in our code with a formally differentiable weight function new_weight(p) whose primal value is always 1 but whose derivative is the left smoothed stochastic derivative of Ber( p ) for primal output 1,
[Equation]
so that a particle x ( k ) i has weight given as
[Equation]
In [7], an equivalent expression is derived by different means, using the stop-gradient operator ⊥ :
[Equation]
where ⊥ is formally assigned a derivative of 0.
D Implementation details D.1 Experiment details
All computation times in Fig.
5 were measured on an Intel Xeon Platinum 8260 CPU and Julia version 1.6.
Garbage collection times are included in the total run time.
We provide code and instructions to run 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 [1, 2],
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 the package implementation.
E Funding details
This material is based upon work supported by the National Science Foundation under grant no.
OAC-1835443, grant no. SII-2029670, grant no. ECCS-2029670, grant no.
OAC-2103804, and grant no. PHY-2021825.
We also gratefully acknowledge the U.S. Agency for International Development through Penn State for grant no.
S002283-USAID.
The information, data, or work presented herein was funded in part by the Advanced Research Projects Agency-Energy (ARPA-E), U.S. Department of Energy, under Award Number DE-AR0001211 and DE-AR0001222.
This material is based upon work supported by the Defense Advanced Research Projects Agency (DARPA) under Agreement No
HR00112290091.
We also gratefully acknowledge the U.S. Agency for International Development through Penn State for grant no.
S002283-USAID.
The views and opinions of authors expressed herein do not necessarily state or reflect those of the United States Government or any agency thereof.
This material was supported by The Research Council of Norway and Equinor ASA through
Research Council project "308817 - Digital wells for optimal production and drainage".
Research was sponsored by the United States Air Force Research Laboratory and the United States Air Force
Artificial Intelligence Accelerator and was accomplished under Cooperative Agreement Number
FA8750-19-2-1000.
The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the
United States Air Force or the U.S. Government.
The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein.
We also acknowledge financial support from the NCCR QSIT funded by the Swiss National Science
Foundation (Grant No. 51NF40-185902), as well as from the Swiss National Science Foundation individual grant (Grant No. 200020_200481).