Automatic Differentiation of Programs with Discrete Randomness
Authors: Gaurav Arya
Moritz Schauer
Frank Schäfer
Christopher Vincent Rackauckas
Venue: NeurIPS 2022 Accept (2022)
Keywords: stochastic methods
automatic differentiation
reparameterization trick
discrete randomness
gradient based inference
differentiable stochastic programming
chain rule
compositionality

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 reparameterization-based 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 package is available at https://github.com/gaurav-arya/StochasticAD.jl.

Automatic Differentiation of Programs with Discrete Randomness
Gaurav Arya
Moritz Schauer
Massachusetts Institute of Technology, USA aryagdmit. edu
Chalmers University of Technology, Sweden
University of Gothenburg, Sweden smori Seerrttt@caallmess.co
Frank Schäfer
Chris Rackauckas
Massachusetts Institute Technology, USA University of Basel, Switzerland franksch@min edu
Massachusetts Institute Technology, USA Julia Computing Inc., USA
Pumas-Al Inc., USA crackauc@mit.co edu
Abstract
Automatic differentiation (AD), technique for constructing 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 continuous dependence that have discrete stochastic behaviors governed by distribution parameters, such as flipping coin with probability p of being heads, pose challenge these systems because the connection between the result (heads vs tails) and the parameters (p) develop
parameters. Programs
is fundamentally discrete. In this paper based methodology that allows for generating programs whose expectation 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 traditional AD mechanisms. We demonstrate unbiased forward-mode AD of discrete- time Markov chains, and unbiased reverse-mode AD of particle filter. Our code package available https: //tpt/gggttttttt/// com/ htt///garav-vee artaaa/occccsssii........
reparameterization-
models such as Conway's 's Game of Life,
1 Introduction
Automatic differentiation (AD) is technique for taking mathematical program ( and generating
new program (()) dx for computing the derivative [1 21.. AD is widely used throughout dp
machine learning and scientific computing due to the increased performance of gradient-based optimization techniques compared to derivative-fii free methods (3]. However, if X (p) returns the flip of coin with probability of receiving 1 and probability
of receiving a O, it is clear that 4X dp is not defined in the classical sense. But when attempting to calibrate the parameter to data, one may wish 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 program (p) that computes the derivative of the statistical quantities, i.e. E[X(p)) dz()(p)) dp
naïve solution to this problem would be to use finite differences, i.e.: dE[X(p) E[X :)] E(X(p))
(1.1)
dp
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

(|X(p) = = ((p)(w))
Pruning follows one alternative path
Primal
Smoothing flattens onto primal
State space
Uncoupled path
Steps
X(p + (W2)
Figure 1: Qualitative sketch of our method and comparison to finite differences. The primal computa- tion (solid black line) samples X (p) with random number sequence w1. Black-box finite differences samples the perturbed program X + =) with independent random number sequence W2 (bottom) for some finite choice of e. In contrast, the component Y of the stochastic derivative of the program (Section?) 2 considers the effect of the minimal possible perturbations (gray lines diverging from primal path) to the original program which could stem from a truly infinitesimal change in the input p.
Finite differences' major issue the context of stochastic programs is that this calculation does not correlate the calculation of [[[X [X(+++2 [)] with the calculation of E[X (p)]. This leads to large variance in the finite-difference estimator 4 that goes to oo as - 0. This issue of unbounded variance could be solved if one could run the perturbed program with the same set of random numbers and directly estimate {[ + E) - X (p)] by method that is well-posoe in the limit of & - (Fig. I
We demonstrate how to automatically construct new stochastic program (()) whose expectation satisfies [p)] dE[X(p]] We derive this through technique which we term the stochastic
dp
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 (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 distribu- tions. These may be composed (chained, added, etc.) arbitrarily, with computational cost independent of how they are composed.
Unbiasedness. The program's discrete structure preserved. leading 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. IT. 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 ((1) computational overhead for these applications, combine stochastic derivatives with online strategy we call pruning. We also demonstrate straightforward way to perform reverse-mode AD from our approach via smoothing, allowing empirically successful straight-through gradient estimator (6] as a special case and to construct an unbiased end-to-end difers--iifeerentiable particle filter, recovering technique discovered in (7).
derive from first principles the biased but
We provide an open-source implementation of the method, StochasticAD. jj. 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 [ score-function | measure-valued [2], and pathwise [4] gradient estimators. When it comes discrete randomness, the score-function method is a popular general- 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 (e.g. REBAR [6), RELAX (77) have been introduced that rely on control variates for variance reduction [8]. Measure-valuce 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, Gumbel-Softmax 20 considers a continuous relaxation of discrete programs so that pathwise gradient estimator, based on the "reparameterization trick" [14], can be applied. However, such methods face a bias-variance 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 conceptual starting point of our approach is finite differences with common random numbers [21 22 23 whose ideas have also been extended by direct optimization [24] 25), but crucially we show how to take the exact limit of step size - 0 even in the discrete case. The field of smoothed perturbation analysis 26 27 develops mathematically equívalent 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 our pruning technique in the context of generalized semi-Markov processes [28]. However, these ideas have not previously been applied construct general purppss 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 compos- ability through user-written functions, enabling an automated mechanism. While mainstream AD frameworks do not support unbiased differentiation of discrete random programs, Storchastic [29] is 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 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 (29 33). 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.
2 Composable derivatives of stochastic programs
In this section, we develop the notion of stochastic derivative for programs containing discrete randomness. We shall motivate this object 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 function g(s) is OEs) if (g(e) S Clal for some real C and all sufficiently small e. Colloquially, we describe quantities that are O(E) as "infinitesimal..
2.1 Infinitesimally perturbing a stochastic program
We are interested in differentiating stochastic programs, formally defined below. A stochastic program can be thought of as a map from an input p to a random variable (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 a random walk.
Definition 2.1. A stochastic program X (p) stochastic process with values in a Euclidean space E, whose index set I is either an open subset of Euclidean space or a closed real interval.
Let Ω be the sample space, equipped with a probability distribution P. To sample (p) at input E 1. which we call the primal evaluation, one should imagine sample w being randomly chosen from I according to PP to produce an output X (w) E E. Note that IP is independent of p and ((p) is map 92 > E; such a formulation has been called the "reparameterization trick" [14]. For example,
a Bernoulli distribution Ber(p) can be represented by choosing uniform random w E [0. 0] and
defining X (w) (1-p.1)(4)) (4:), where 1s is the indicator function for a set S (see Fig. 2b
At fixed p E I we define the differential dxx which is itself a stochastic program:
dX(x) X + (p).
(2.1)
Let us now restrict our attention to the case where / is a closed real interval, so that P. E R. The sensitivity of stochastic programs Z with more general index sets can be understood at an input by studying p 0 the directional perturbation X X(p) Z(u + pv) in direction
3

---

Output
Output
Output
X(p + E)
which X()l(w) =0 X(plw)=
1
dX(e)
(())
Not (e))
X(p +5)
X(p)
dX((=)
Ω
Ω
0
W2
W1
P 1-p (b) (p) ~ Ber(p).
0 1- 1-p
(a) (p) Exp(p).
(c) (p) Ber(p).
Figure 2: Illustrations of d (E) for a small finite assuming > 0 for simplicity. (a) The (es) range of (x) for continuous X (p). The blue bracket indicates sample from dx) while the red bracket indicates a sample obtained by black-box finite differences. (b) The original program X (p) and the perturbed program (p + s) for discrete (Bernoulli) (p). (c) Only those samples w below the hatched area cause non-zero change in output. The hatched area equals E[dX (E)), the expected size of change for randomly drawn w.
For the full user-provided stochastic program, we are interested in computing related to [s)] by
dF X
dp
(p)] lim EldX(=))
(2.2)
dp
8-0
We thus expect (s) to have "infinitesimal" (E)) expectation as E approaches O. Let us consider
the form of d (e) for two elementary programs: a program returning sample from the continuous Exponential distribution Exp(p) with scale p and a program returning sample from the discrete Bernoulli distribution Ber(p) with success probability p. both cases, the program parameterized via the inversion method [31] over the sample space f [O, 1]. which chooses [P to be uniform over
[0, 1] and X (w) to be non-decreasing in w. For X (p) ~ Exp(p), X (w) = -p ---lloog log w) is
differentiable in p. Thus, the range of (e) is Oss) for any fixed w, illustrated in Fig.
2a
In contrast, consider X (p) Ber(p), with (D) [p-).j((w). As shown in Fig. 2b
as E + 0 Specifically, x 1, and assumes the value with probability (Fig. 2ch This is the
the output of (p) at random w E Ω almost surely does not change in X(p + e).
(e is with probability fundamental challenge presented by discrete randomness: a (()) change in the input turns into a finite perturbation the output, where "finite" means that non-vanishing as - 0. The perturbation occurs with "infinitesimal" (s)) probability, contributing to the (ss) expectation of (e).
2.2 Coupling to the primal
To produce low-variance estimates of the derivative of (p),, it is important that the primal and derivative computations are coupled by a shared source of randomness. see this, let us understand what happens when they are entirely uncoupled. A black-box finite difference approach with step size e would independently sample wi and we from (Fig. [2 red), computing the derivative estimate (X(p + (W2) X x(p) (p)(w) /e. However, since the samples are independent, the variance of the estimator is order 1/E", so we are forced to pick a finite E to balance a bias- variance 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 E + 0 leads to the widely-used pathwise gradient estimator ò, given as the almost sure limit of dx (s)/s derivative of X(p) at each fixed w)), so that
0 (i.c. the pointwise
dL[ lim [(d LE lim dX(e)) dp
E [5]
(2.3)
E
But, considering a simple Bernoulli variable ((p) Ber(p) as in Fig. 2c we see how this approach
is ill-suited for the discrete case! As approaches O, the differential dxX(s) is non-zero with

---

infinitesimal (E)) probability. This means that ò is almost surely O, while the true derivative of [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 dx (e)/e is unbounded in e in the presence of finite perturbation. This motivates explicitly considering the event of a large jump in (e), as characterized by the event Ag(e) {IdX(e)) > B for a chosen random bound B > 10). The event (e) has (E) 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 random weight w E R and alternate value Y E that characterize the sensitivity of (p) when the probability space restricted to Ag (E) [as given by the r.h.s. of (2.4) below], forming the "stochastic derivative":
Definition 2.2 (Stochastic derivative). Suppose (()) E is stochastic program with index set I closed real interval. We say that the triple of random variables (6. w. Y), with w E IR and Y E E, is a right (left) stochastic derivative of X at the imput p if d (x)/== - ô almost surely as - 0. and there is an integrable (i.e. of bounded expectation) random variable >> |o] such that for all bounded functions f: F IR with bounded derivative it holds almost surely that
f(X + 8)) f(X())) LAs(e)
E [ (S(Y) - ((x()))) X()) lim E --+0+/-
X(O)
(2.4)
with limit taken from above (below), where IP (AB(s) X())) /e is dominated by an integrable random variable for all (5 < 0).
stochastic derivative may be collapsed into an unbiased estimator of the derivative of (p)). Proposition 2.3 (Unbiasedness). If (ō, w, Y) is stochastic derivative of X (p) p. it holds that
dE (p) = E|6 + w (Y-(())))) dp
(2.5)
Proof sketch. With f as identity, by Definition 2.2 the sensitivity of X (p) over AB(E) is given by w(Y (p)). Given the complement event (E), it holds that [dX(e)/Ell s B: a dominated convergence argument shows that the sensitivity of (p) is then given by its almost-sure derivative 8.
Theorem 2.4 shows the existence of the stochastic derivative for 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) C R, there exists right stochastic derivative (5, wn, Yn) with wn 2 0 at any p E (a, b) and a left stochastic derivative (6, Wr. YL) with wt S 0 at any p E (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 (m++)= = (p) +dX(() given the event Ag(s) and the outcome of (p).. The weight w is given as the derivative w.r.t. E of the probability IP (AB(E) (x)) that dxx (=) jumps by a non non-infinitesimal amount, conditional on X (p). Essentially, since P (AB(E) X X()) x ws, multiplying by w bridges the gap between conditioning on Ag (a) [recall Y is constructed conditional on As()) and simply restricting the probability space to Ag(e) [i.e. multiplying by An(=) in (.44)
Example 2.5 (Right stochastic derivative of Bernoulli variable). Suppose X(p) Ber(p), parame-
terized via the inversion method, and take & > O. As shown in Fig. 2c dx (E) is parameterized as,
dx (e) = {1 if p- - < 1-p.
(2.6)
otherwise.
Given the event (() = 1, we have that 2 p and (s) is deterministically O. On the other hand, given (p) 0, we have w < 1 - p. so with probability 8/(1 - p) the differential d (e) assumes a value of 1, i.c. the Bernoulli variable flips from 0 to 1. Thus, we may construct right
5

---

stochastic derivative WR, Yx) of (p) by letting wn = 1/(1-)), YR = conditionally on
(p) = 0 and wR = Yn 0 conditionally on (p) 1. For concrete example, with p 0.6,
we have wn 2.5, Yn = 1 conditionally on (p) = O; StochasticAD. prints this as "O (1
with probability 2. 55) (as explained further in Section 3.1
We give number of examples of stochastic derivatives in Appendix A It is of crucial importance that w and Y depend conditionally on the output of (p).. In particular, given a primal evalua- tion X (p) Ii, the forms of wo and Y depend only the distribution of the perturbed program on the set {w (p) = x C S2 which elegantly generalizes the coupling achieved by the pathwise gradient estimator ó at each fixed w S2. Intuitively, this coupling allows us consider only the smallest possible perturbations to the program in the derivative computation (Fig. [ and thereby achieve variance reduction without resorting continuous relaxations. For example, for binomial variable X (p) ~ Bin(n, p) parameterized via the inversion method (a natural parameterization for maximizing coupling), holds that Y E (X(p) 1, (() + 1]. We show in Example A..1 that our gradient estimator for such a binomial variable has variance of order n, whereas the score function estimator has variance of order n², justifying this intuition.
2.3 Composition of stochastic derivatives
Proposition shows that the derivative estimate produced by stochastic derivatives is correct in expectation. But this is insufficient to ensure composition, i.c. stochastic derivative "chain rule". While Definition reqrurss composition through deterministic test functions f to enforce a sufficiently strict definition, Theorem 2.eerremmm22ppovaeee provides a ceneral-purpons 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 X1 and X2 and their composition X₂ X1- Suppose that X₂ has a right (left) stochastic derivative p E R given by (δ1, w1, Y1), and X2 has right stochastic derivative (82, W2, Y2) in the direction ő, 81/161 given conditionally on its input X1 (I)). Then, under regularity and integrability assumptions. the stacked program [X1; X2 X1] has a right (left) stochastic derivative (ő, w. Y) at p where d [ō1: (61 62).
[Yr X2(Y)]
with probability
W1
W1 +
8 102
Y
86 W2
(2.7)
[p)
with probability
W1
+ |
and W1 + 181 /w2.
Proof sketch. Let A₁ (E) be the event of a jump in X, and A2(e) be the event of jump in X₂ when its input is perturbed by eδ, Given Af(E) n A5()) (a jump in neither), a dominated convergence argument implies that X₂ o X1 does not jump either. Given (s)s) Definition yields that the sensitivity of X₂ o X, is described by alternate values Y1 with weight w1. while given A₂(e) the sensitivity of X2 o Xi is described by alternate values Y2 with weight (6 (22. 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 X2 X1 given (e) A2(e) (a jump in both), as this event has probability OEE2). This prevents a combinatorial explosion in the complexity of Y.
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.5E) becomes "O 2.5E".)
Definition 2.7 (Smoothed stochastic derivative). For stochastic program (()) with right (left) stochastic derivative (ō, w. Y) at input p, a right (left) smoothed stochastic derivative 3 of X at input p is given as
σ = E (6 w( x(p)) [x)].
(2.8)
6

---

using Distributions function X(p) 0.6
struct StochasticIriple
value prámal evaluation
=p"2 3. + 1.2e
5 "enfénstesimal" component
b rand (Binomial p))
áa # component of discrete change
6 (1 with probability 10.0t)
$ with "infénétesima!"
3 (Bernoullll (p))
probabality
12 + (3 with probability 12.5e)
end
return a * rand(Normal(b, a))
end
ijulia> using StochasticAD
jjulia> pecchassti__rrppeee.. 0.6) sample a single stochastic triple p 0.6
27.11 94.32e (6.78 with probability 12.5e)
+julia> prrriaative_contributiio
which produces single derivative estimate...
179.04
jjulia> samples - Lieivaaivv_essiimttkk.k 0.6) for i in 1:1000] # take many estimates!
jjulia> printIn("d/dd of E[X(p)]: (mean (samples))) ((std(ammles) sqr((1)))))))
d/dp E[X())]: 204.63 1.25
Figure 3: Left: Stochastic triple structure (simplified). Right: A toy program X (p), using discrete distributions Bin(n, p), Ber(p), and the continuous normal distribution N(H, o); the used stochastic derivatives are given in Appendix Highlights show intermediate values during a single derivative estimate. Bottom: Differentiating X (p)]; printout float precision reduced for clarity.
Smoothed stochastic derivatives easily permit reverse-mode AD instead of forward- mode, as they have the same form 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 Appendix B.6 Due to the coupling of Y and (p), this is much weaker requirement than global linearity, which can lead to low-bias estimates: for example, for the program (p) Gco(p) 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 special case, we work out in Example A.8
3 Automatic differentiation of stochastic programs
We develop StochasticAD. jjl, 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 However, performing automatic differentiation with unsmoothed stochastic derivatives, which are unbiased in all cases, requires new innovation. We develop novel computational object called a stochastic triple, introduced in Section and
and showcased in Section 3.2 and Section 3K
3.1 Educational toy example of stochastic triples
Forward-mode AD often implemented with dual numbers 21 which pair the primal evaluation of deterministic function (p) with its derivative f Dual numbers be propagated through 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 dp (p) is the coefficient of the "dual" element Stochastic triples generalize dual numbers by including third component As to describe finite perturbations with infinitesimal probability (Fig. left).
In Fig. 3 we consider a toy program X(p) including discrete randomness. We are interested in the derivative of (p)) at p 0.6, and hence we provide the stochastic triple printed as "0. 6 E" as input. First, the triple squared and becomes "0.36 1.2e": this is the familiar way that dual numbers propagate, via the chain rule. But what happens when the triple is propagated
7

---

(a)
(b)
standard deviation
score-function Score- Conre-tnncion CV
80-
X ==0
triples
for step in 1:n
60
i - rand (aatoooaaaa (probs (X)))
40-
X [i]
end
20-
return f(X)
20
40 60
11 100
80
Figure 4: Automatic differentiation of discrete-time Markov processes. (a) Code snipped for a ID random walk, which can be automatically differentiated by StochasticAD. j1; probs (x) gives the transition probabilities at value X and steps [i] gives the step size for the ith transition. (b) The variance of unbiased gradient estimates of the random walk program using stochastic triples and the score function, which is applied both without control variate (CV) and with pre-computed batch-average CV. (c) The final board for one run of the stochastic Game of Life with N = 25 and = 10, where the "+" signs represent additional living cells (white) in the stochastic alternative path, and the "X" signs represent additional dead cells (grey).
through the discrete and random Binomial variable? The resultant stochastic triple "6 + (1 with probability (0.0g) " is integer with component of discrete change, reflecting infinitesimal probability of one more success the Binomial Bin (10, 0.6). We can in fact understand why the probability is " "10s" by representing the Binomial the sum of 10 Bernoulli variables, each with probability 0.6. Since 4 of the Bernoulli's have an output of O, they each have a probability "2 5e" of switching to I (recall Example 2.5), and thus there is in total " "10s" probability that the output of the Binomial increases by (rigorously, we have applied Theorem 2.6). Formally, we can interpret printout "x 8 + (A with probability wE) " follows: is sample of the random variable X (p) describing the primal evaluation, and 8, w, and x+A are samples of the components 8, 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 stochas- tic triples through functions via operator overloading [32] exploiting Julia's 's multiple dispatch feature (33] When multiple discrete changes are possible, pick one probabilistically:
call this strategy pruning (recall Fig. [T) and show its unbiasedness in Appendix B.5 For example, in line 6 of Fig. right, 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), associate each perturbation with tag to avoid erroneously pruning between two perturbations that occur simultaneously. Thus, stochas- tie triples can efficiently propagate through the full toy function written in Fig. 3 The function srrivvatveeeestmmlt creates a stochastic triple, propagates it, and collapses it into the derivative estimate 8 + wd, forming an unbiased estimate of the derivative via Proposition 2.3 (Fig. bottom).
3.2 Inhomogeneous random walk
In n on Zzo. with transition behavior dependent on a We consider a Markovian random walk ...,,,,, parameter p as follows,
In 1 + 1 with probability exp
#n-1
In
D0 = 0.
(3.1)
Fn-1 with probability - exp
Øn-1
We consider program that stochastically simulates this walk and applies an arbitrary non-linear function f to the output Zy. In practice, ff may represent a loss likelihood estimate; in this toy setting, we take f(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 n. We find that the stochastic triple estimator has asympiotically 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.
8

---

(8)
1.0
(b)
0.020
(c)
relative frequency
smoothing blased
primal reverrs Al
0.5
0.015
run times
-0.5
-1.0
0.005
-1.
0.000
1.0
-0.5
0.0
0.5
1.0
2.5
100
100
200
10'
10'
10
derivative estimate
parameters
Figure 5: Unbiased differentiable particle sampler using smoothed stochastic derivatives. (a) Latent path (blue) and observations (orange) for a particular realization the process given by the hidden Markov model (see Appendix
[ for details) in d = 2 dimensions, including the recursively applied Kalman filter estimator (red) and the particles of the particle filter (going from blue at n = 1 to pink
20). (b) The value of the derivative of log with respect to the first parameter for = 2 calculated by differentiating the Kalman filter is marked by a black line. The blue distribution corresponds to the derivative computed with smoothed stochastic derivatives. The orange distribution represents the biased derivative approach, where the resampling step is not differentiated. The means of the two distributions are highlighted by line in the respective color. In both cases, we use 1000 samples. (c) Averaged run times in ms comparing forward- and reverse-mode AD for increasing numbers of parameters.
3.3 Stochastic Game of Life
For our next example, differentiate 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 alive, while living cell survives when 2 or 3 of its neighbors stochastic version, each of these events instead has probability 95%, while their complementary events have probability 5%. Such types of discrete stochastic programs arise in many applications. For example, mixing machine learning with agent- models found in epidemiology and sociological rule-based models and discrete stochastic (Gillespie) simulations in systems biology 36 requires similar program constructs.
alive. In
contexts 34 35
Consider a program that populates each cell of an N x N board with probability p. runs the stochastic Game of Life for T time steps, and counts the number of living cells nivinng We perform a sensitivity analysis of the final living population with respect to the initial living population, i.e. differentiate the expectation of Vivvvng with respect to p. Stochastic triples propagate fully through the program, leading to an unbiased estimate the derivative, as we verify with black-box finite differences. An example 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 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 X1, X... observations y/1 ~ Yy...., . Vn
Y,. and parameters 0. The likelihood p(v1..... Un ()) is in general not tractable, but a particle filter can used to compute an estimate of L. Here, we assume familiarity with particle filters [38 39) but provide short exposition with focus the resampling step and experimental details in Appendix C We assume the latent states to be continuous random, but discrete randomness enters through the resampling step.
To differentiate the particle filter resampling step, we provide 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 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 [4]] as we show in Appendix C Our particular choice of system allows the calculation of a ground-truth gradient of & by differentiating the Kalman filter algorithm [41] 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 42 43 44 that neglect the contribution of the resampling step or perform with entropy-regularizee optimal transport (Fig. 5b). For our program, we observe reverse- mode AD to perform better for more than = 100 parameters (Fig. 5c). However, 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 40
find that the variance
4 Limitations and outlook
We have presented method for unbiased AD of programs with discrete randomness, which, we have argued, is a natural generalization of pathwise gradient estimators based on the reparametcrizatioo trick to the discrete case. However, more work will need to be done to turn our software demonstration StochasticAD . j1 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, interesting direction for future work is to automatically handle functions which deterministically turn a continuous random quantity into a discrete one, such as inequality comparison (p) 0 or a Bernoulli variable implemented implicitly as rand () < and also to handle constructs such as while loops based on such functions. It would also be interesting explore synergies of our approach with ADEV (55). Haskell-based framework for provably correct stochastic AD developed concurrently with our work, which formulates gradient estimation strategies using Haskell's continuation passing style.
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 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 partiele filter. (This design space is reflected in the fact that the form of the stochastic derivative depends not only on program's '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 B1 + 2B₂ with i.i.d. B, ~ Bin(n, p) has variance On²) due to the pruning between the +1 and +2 perturbations, even though stochastic derivatives give gradient estimators with variance (nn) for each of B3 and B2. Smoothing does not face this issue but accrues 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 [1]].6] Finally, going beyond smoothing for reverse- 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 Kricken 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 the research results reported within this paper. This material is based upon work supported by the National Science Foundation OAC- 1835443, SII-2029670, BCCC-202960
OAC-2103804, and PHY-2021825 the Advanced Research Projects Agency-Energy DE-AR0001211 and DE-ARO001 2200 the Defense Advanced Research Projects Agency (DARPA) HR001 12290091; the United States Artificial Intelligence Accelerator 1A8750-19-2-10000 the Chalmers AI Research Centre; and the Swiss National Science Foundation 51 115940-185902 The views and opinions of authors expressed herein do not necessarily state or reflect those of the United States Govermment any agency thereof.
10

---

References
[1] Andreas Griewank. On automatic differentiation. Mathematical Programming: recent develop- ments and applications, 6(6):83-107, 1989.
[2] Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: survey. Journal of Machine Leurning Research, 18:1-43, 2018.
[3] George Corliss, Christele Faure, Andreas Griewank, Laurent Hascoet, and Uwe Naumann. Automatic differentiation of algorithms: from simulation to optimization. Springer Science & Business Media, 2002.
[4] Søren Asmussen and Peter W Glynn. Stochastic simulation: algorithms and analysis, volume 57. Springer. 2007.
[5] Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte Carlo gradient estimation in machine learning- Mach. Learn. Res., 21(132):1 2020.
[6] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. a38.3: 2013.
[7] Adam Scibior and Frank Wood. Differentiable particle filtering without modifying the forward
pass. arXiv:2106. 10314, 2021.
[8] Peter W Glynn. Likelihood ratio gradient estimation for stochastic systems. Communications of the ACM, 33(10))75-84, 1990.
[9] Ronald Williams. Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Machine learning. 8(3):229-256, 1992.
[10] Jack PC Kleijnen and Reuven Rubinstein. Optimization and sensitivity analysis of computer simulation models by the score function method. European Journal of Operational Research, 88(3):41-4277, 1996.
[11] Gh Pflug. Sampling derivatives of probabilities. Computing. 42(4): -288, 1989.
[12] Bernd Heidergott and FJ Vázquez-Abad. Measure-valued differentiation for Markov chains. Journal of Optimization Theory and Applications, 136(2): 181-)):8872999 2008.
[13] Paul Glasserman and Yu-Chi Ho. Gradient estimation via perturbation analysis, volume 1 116. Springer Science & Business Media, 1991.
[14] Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. arXiv: 1312.61312.6114 2013.
[15] John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel. Gradient estimation using stochastic computation graphs. Advances in Neural Information Processing Systems, 28, 2015.
[16] George Tucker, Andriy Mnih, Chris J Maddison, John Lawson, and Jascha Sohl-Dickstein. REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models. Advances in Neural Information Processing Systems, 30, 2017.
[17] Will Grathwohl, Dami Choi, Yuhuai Wu, Geoffrey Roeder, and David Duvenaud. Back- propagation through the void: Optimizing control variates for black-box gradient estimation. arXiv: 1.101111..0023... 2017.
[18] Peter W Glynn and Roberto Szechtman. Some new perspectives on the method of control variates. In Monte Carlo and Quasi-Monte Carlo Methods 2000, pages 27-49. Springer, 2002. [19] Bernd Heidergott, Felisa J Wázquez-Abad, and Warren Volk-Makarewicz. Sensitivity estimation for Gaussian systems. European Journal of Operational Research, 187(1):193-207, 2008
[20] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with Gumbel-Softmax.
In International Conference on Learning Representations, 2017.
[21] David Anderson. An efficient finite difference method for parameter sensitivities of continuous time markov chains. SIAM Journal Numerical Analysis, 50(5):2237-2258, 2012.
[22] Vo Hong Thanh, Roberto Zunino, and Corrado Priami. Efficient finite-difference method for computing sensitivities of biochemical reactions. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences. 474(2118):0180330 2018.
11

---

[23] Paul Glasserman and David D Yao. Some guidelines and guarantees for common random numbers. Management Science, 38(6):884 908. 1992.
[24] Tamir Hazan, Joseph Keshet, and David McAllester. Direct loss minimization for structured prediction. In J. Lafferty, C. Williams, J. Saylor R. Zemel, and A. Culotta, editors, Advances in Neural Information Processing Systems, volume 23. Curran Associates, Inc., 2010.
[25] Guy Lorberbom, Chris J Maddison, Nicolas Heess, Tamir Hazan, and Daniel Tarlow. Direct policy gradients: Direct optimization of policies in discrete action spaces. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 18076-18086. Curran Associates, Inc., 2020.
[26] Wei-Bo Gong and Yu-Chi Ho. Smoothed (conditional) perturbation analysis of discrete event
dynamical systems. IEEE Transactions on Automatic Control, 32(10):858- 866, 1987.
[27] Paul Glasserman and W-B Gong. Smoothed perturbation analysis for class of discrete-event systems. IEEE Transactions on Automatic Control, 35(1 11:11188881 1990.
[28] Michael Fu and Jian- Dian-Qiang Hu. Conditional Monte Carlo: Gradient estimation and optimiza-
lion applications, volume 392. Springer Science & Business Media, 2012.
[29] Emile van Krieken, Jakub Mikolaj Tomczak, and Annette Ten Teije. Storchastic: A framework for general stochastic automatic differentiation. In A. Beygelzimer, Y. Dauphin, P. Liang, and
J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[30] Emile van Krieken, Jakub Tomezak, and Annette Ten Teije. Openreview: Stor-
chastic: A framework for general stochastic automatic differentiation. URL wttp://wwweeeeeeeooommmdid.-aeee
[31] Luc Devroye. Non-Uniform Random Variate Generation. Springer-Verlag, 1986.
[32] Jarrett Revels, Miles Lubin, and Theodore Papamarkou. Forward-mode automatic differentiation
in Julia, arXiv: [cs.MS], 2016.
[33] Jeff Bezanson, Alan Edelman, Stefan Karpinski, and Viral B Shah. Julia: A fresh approach to numerical computing. SIAM Review, (9(11)65-996 98, 2017.
[34] Eric Silverman, Umberto Gostoli, Stefano Picascia. Jonatan Almagor, Mark McCann, Richard Shaw, and Claudio Angione. Situating agent- modelling in population health research. Emerging Themes in Epidemiology. 18(1):10, Jul 2021.
[35] Ayush Chopra, Alexander Rodríguez. Jayakumar Subramanian, Balaji Krishnamurthy, B Aditya Prakash, and Ramesh Raskar. Differentiable agent-based epidemiology. arXiv preprint
arXiv:2207.09714, 2022.
[36] Lily A Chylek, Leonard A Harris, Chang-Shunp Tung, James R Faeder, Carlos F Lopez, and William S Hlavacek. Rule-based modeling: computational approach for studying biomolecular site dynamics in cell signaling systems. Wiley Interdiscip. Rev. Syst. Biol. Med., 6(1):13-36, January 2014.
[37] James R Faeder, Michael L Blinov, and William S Hlavacek. Rule-based modeling of biochemi- cal systems with BioNetGen. Methods Mol. Biol., 500:113-167, 2009.
[38] Nicolas Chopin and Omiros Papaspiliopoolos. An introduction to sequential Monte Carlo. Springer, 2020.
[39] Arnaud Doucet and Adam M Johansen. A tutorial partiele filtering and smoothing: Fifteen years later. Oxford Handbook of nonlinear filtering, 12(656-704), 2011.
[40] George Poyiadjís, Arnaud Doucet, and Sumeetpal S. Singh. Particle approximations of the score and observed information matrix in state space models with application to parameter estimation. Biometrika, 98(1 (111)))6-88)))) 80,000000000 02 201
[41] R. E. Kalman. A New Approach to Linear Filtering and Prediction Problems. Journal of Basic Engineering, 82(())35-455 03 1960.
[42] Rico Jonschkowski, Divyam Rastogi, and Oliver Brock. Differentiable particle filters: End-to- end learning with algorithmic priors. / / /222, 2018.
[43] Adrien Corenflos, James Thornton, George Deligiannidis, and Arnaud Doucet. Differentiable particle filtering via mntropy-regulare
optimal transport. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 2100-2111. PMLR, 18-24 Jul 2021.
12

---

[44] Conor Rosato, Lee Devlin, Vincent Beraud, Paul Horridge, Thomas B. Schön, and Simon Maskell. Efficient learning of the parameters of non-lines models using differentiable resampling in particle filters. IEEE Transactions on Signal Processing, 70:3676-3692, 2022.
[45] Alexander Lew, Mathieu Huot, Sam Staton, and Vikash K Mansinghka. Adev: Sound automatic differentiation of expected values of probabilistic programs. arXiv preprini arXiv:2212.06336 2022.
[46] Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Representatioss, 2017.
[47] Fernando Mazzone. A characterization of almost everywhere continuous functions. Real Analysis Exchange, 21(1):317-319, 1995.
Checklist
1. 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 4
(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 experi- mental results (either in the supplemental material or as a URL)? [Yes] See Appendix
and the provided code.
(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 the random seed after running experi- ments multiple times)? [No] Our examples were small-scale, we could take enough samples to achieve minimal error.
(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.
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "kjWuDdEsZtp", "reviewer": "Reviewer_kBKn", "summary": "The paper introduces a technique named “stochastic derivative” that helps construct a stochastic program whose expectation is the derivative of the expectation of a given stochastic program, which might include discrete random variables. The technique is then shown to give unbiased estimators and low variance (compared to the baseline score-function method) through three examples: inhomogeneous random walk, stochastic game of life, and particle filter.", "questions": "+ It would be helpful if the authors give a list of parameterizations of distributions used in the main paper, like Exp, Bin.\n+ Line 87: I wonder what is the definition of Z?\n+ Line 105: Why the value of dX(z) is O(eps), I thought this quantity is -eps * log(1-z) which is not bounded when z near 1?\n+ Figure 2: The red dashed line indicates a sample obtained by black-box finite differences. What does O(1) imply at this dashed line?\n+ Lines 119-120: Why the variance of those independent samples z1, z2 is of order infinity? In addition, it might be more consistent with the introduction section to use X(p + eps) - X(p), rather than X(p + eps/2) − X(p − eps/2).\n+ Line 145: Does f need to be bounded? In later examples, I’m seeing f(x) can be x^2 or x^3, which are not bounded.\n+ Line 152: typo: infinitely. What does “infinitely close to 1” mean? Why is the importance weight close to 1 here?\n+ Line 162: typo: occurring\n+ Lines 160-164: I don’t understand this paragraph. Could the author add more details? What does “event of two finite perturbations occurring simultaneously” mean?\n+ Line 172: typo: parameterized\n+ Line 173: please provide which section in the supplemental that we should look at\n+ Figure 4: how can we interpret the output in this figure?\n", "limitations": "Nothing to report.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, I think this is an important paper which introduces a novel technique but needs a major revision for clarity.\n\nThe paper is technically sound. I really like that the authors provide detailed motivations and illustrations for each idea and want to thank the authors for that. However, the structure is a bit hard to follow. For example, Section 2 talks about “composable stochastic derivatives” but I only get what “stochastic derivatives” means until Theorem 2.3. Or the section 2.4 talks about Smoothing but lacks background for why it is needed (is this related to Smoothed perturbation analysis mentioned in the introduction?)\n\nI’m not familiar with Julia but would like to reproduce some results in the Demonstrations section in Python to better understand the technique. However, the paper lacks an outline on how to design that stochastic AD system.\n\nThe proposed approach is only compared to the score-function method in the inhomogeneous random walk example, showing that using stochastic triples the variance of unbiased gradient estimates. Could the author also provide empirical comparisons to the score-function method in other examples?\n\n**Updated:** I'm happy to increase my score given that the writing quality has been much improved in the last revision.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "avpAlEw2_7T", "reviewer": "Reviewer_x6td", "summary": "This paper proposes a novel approach for both forward- and reverse-mode automatic differentiation of probabilistic programs with discrete RVs. Since the derivative $\\frac{dX}{dp}$ of a program $X$ with respect to the parameters $p$ of a discrete distribution is not well-defined, the authors circumvent the problem by finding the derivative of the expectation of $X$, i.e. $\\frac{d\\mathbb{E}(X)}{dp}$.\nThe technique results in an unbiased estimator, exhibiting lower variance with respect to prior approaches. \n", "questions": "\n1- In 4.1, what value(s) of $p$ were used? Why not showing in Fig.5  (center) the standard deviation over multiple values of $p$?\n\n2- In 4.2 the program is almost deterministic with $p=0.05$. What happens as $p$ approaches $0.5$?\n\n3- Why not applying the score-function method to the tasks 4.2 and 4.3? It is also not clear to me why some approaches mentioned in the Related Work section were not considered (REBAR, RELAX, or the reparametrization trick with Gumbel-softmax).\n\n4- In Fig.6 (center), why not reporting both forward and reverse mode AD, possibly using the same colors in Fig.6 (right)?\n\n5- The paper mentions some additional complications when dealing with multivariate functions. Why not reporting experiments in that sense?\n\n---\n\n**Minor comments**\n\n- Fig.1: I don't get it. $X'(p + \\epsilon)$ with random number sequence\n  z(2) (bottom) in the caption is $X(p + \\Delta p)$ in the figure. What\n  is the relation between $\\Delta p$ and $\\epsilon$?\n---\n\"This variance issue can be alleviated if one could run the perturbed program with the same set of random numbers, i.e. directly compute $\\mathbb{E}[X(p + \\epsilon) − X(p)]$\" \n\n- Shoudn't it be 'solved' rather than 'alleviated'?\n---\n- Def. 2.2 \"right (left) derivative\" 'stochastic' is missing.\n---\n- When is a stochastic program \"sufficiently regular\"?\n---\n\n\"It is convenient to assume probability densities [...]  and transition probability densities [...] all depending on the parameters $\\theta$.\" \n\n- Why? You mean from a notational perspective?\n---\n\n**Typos**\n\n- \"algorithm. enabling\"\n- \"an weighted\"\n- \"with parametrized by\"\n- \"particle filterunbiasedness\"\n\n\n\n\n", "limitations": "Mostly yes, although the empirical evaluation could be improved in this sense (see above).", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The problem considered in this work is a relevant one.\nWhile the paper is generally well-written, I found it hardly accessible for nonexperts in AD, like myself. \nI wish that the paper included some more examples, for instance a figure showing how these stochastic triples are propagated in a simple program. In this sense, Figures 3 and 4 are not that useful in my opinion. I can't really assess the soundness of the proposed approach, but I didn't spot any major issues. Similarly, I am not familiar with the existing literature and it is possible that I missed some relevant work. I think that the experimental evaluation could be improved (see questions below).\n\n**Update after rebuttal** The authors addressed my points effectively and submitted a significantly improved version of the manuscript.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "MLUmCBz1fgp", "reviewer": "Reviewer_kqcs", "summary": "The paper provides a new method to compute the gradient of the expected output of a stochastic computation graph (SCG) w.r.t one of the input parameters of the graph even when the graph contains discrete randomness. \n\nIn a well known method for computing gradients of stochastic computation graphs when they contain only continuous valued variables all the stochastic variables are reparameterized which allows for the gradient of the expectation to be sampled for any parameter by simply propagating the primal value of the graph variables as well as the forward gradients. This paper extends this well-known method by additionally propagating an alternate value for one of the discrete stochastic nodes in the graph (plus the alternate values of all the children of that flipped node) as well as the gradient of the probability of that alternate value.\n\nThis method has been shown to greatly reduce the variance of the estimated gradient of the expectation as opposed to standard methods such as the score function estimator which is typically used to estimate such gradients for graphs with discrete randomness.\n", "questions": "In lines 311-312 the authors say that the method is only applicable when control flow is independent of random parameters. What does this mean exactly? Does this mean that if the SCG has if-statements and it can take completely different paths based on random values then the method doesn't work. Why?\n\nIn Figure 1 why does it show 3 alternate paths when pruning has been defined to only allow one alternate path?\n\nIn definition 2.2, it seems to me that `w` is the gradient of the probability of the output of the SCG flipping from X to Y (where the gradient is w.r.t. p). Is this an accurate charaterization?\n\nIn lines 151-159 there seems to be some confusion as to whether w is a vector or a scalar based on the description. Could you please give a clearer description?\n\nIn line 181 where does the $\\delta$ come from? The statement of the proposition has $\\delta_1$ and $\\delta_2$? This proof is also hard to follow in the supplement. Could you provide the key points of the proof here? I would give a higher rating and confidence if I can convince myself that these composition propositions are correct.\n\nThe description of smoothed stochastic duals is very vague. Could you give a more mathematical definition?\n\nIn Figure 5 what does `rand(Categorical(probs[X]))` mean?\n\n# Minor points\n\n- Line 142 \"program with parameterized\" seems to be missing something?\n \n- Figure 6 left figure takes very long to render. Please consider alternate image formats.\n\n- Line 242 typo: \"filterunbiasedness\" \n\n- In equation 4.1 can x be negative? If yes then the probabilities don't make sense.\n\n", "limitations": "No known negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The method is clearly original. There are a number of papers in the literature which attempt to solve this particular problem (estimate gradients of SCG with discrete random variables) such as reinforce, concrete, score function estimator, path wise gradients etc. But they are all rather different from the current paper and those methods all suffer from some form of high variance issues that are explicitly addressed in this paper.\n\nThis paper's key contribution is the realization that simply estimating gradients along multiple independent paths may introduce excessive variance in the estimate and it is better to look for correlated paths.\n\nThe work is very signification since estimating gradients in SCGs are an important problem in various branches of AI in particular Reinforcement Learning and Bayesian Inference. There are a number of papers routinely published every year in NeurIPS which make some reference or the other to this problem and use one of the aforementioned methods. So, an improved solution to this problem is hugely impactful.\n\nThe weakness of the approach is that it is only defined for forward mode gradients and the authors only spend any effort to describe the gradient computation for a single input. In most machine learning applications reverse mode gradients are preferable to compute since we have one loss function which we want to differentiate w.r.t. multiple parameters. \n\nThe quality of the technical writing is somewhat patchy though and that is why I hesitate to give a higher rating than would be deserved for work that has provided such an interesting breakthrough.\n\nThe experiments have been selectively picked for examples with a single parameter, but that seems a bit contrived. There are a number of reinforcement learning papers in the literature which rely on SCG differentiation. Using one of them in this paper would have been more impactful than the game of life example. The resampling example on the other hand was interesting as that is a common use case.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "8XH-I0WYr9N", "reviewer": "Reviewer_bKdb", "summary": "This paper proposes a method to compute a derivative of stochastic programs which have discrete dependencies on their parameters. The paper introduces stochastic triples generalizing dual numbers to programs with discrete randomness. The resulting AD method is shown to produce unbiased derivative estimates with lower variance compared to score function methods. A working implementation is developed in Julia. Forward mode and reverse mode AD (using the smoothing technique) is demonstrated on 3 examples. The experiments demonstrate the utility of stochastic triples for unbiased, lower-variance derivates in forward and reverse mode AD on 3 examples.\n\n---\n\nUPDATE: I thank the authors for their detailed responses. After reading the other reviews and author responses, I'm more positively inclined towards the paper and have upgraded my score to \"Accept\".", "questions": "1. Can you provide more detail about how the proposed implementation compares with existing popular reverse-mode AD libraries? What classes of popular applications of reverse-mode AD can StochasticAD.jl be used / not used?\n\n\n\n", "limitations": "Yes", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n  + The paper tackles an important and challenging problem of AD in programs with discrete randomness. \n  + The paper makes novel algorithmic contributions (stochastic triples, smoothing) as far as I can tell.\n  + The paper includes working code and demonstrates the unbiased and low-variance properties empirically.\n  + The paper is intuitively clear and well written.\n\nWeaknesses\n  - The novelty of the technical contributions compared to existing AD frameworks could be made a bit clearer. Also, a deeper discussion of which program classes or constructs are not differentiable by popular reverse-mode AD frameworks might help make the contributions clearer.\n  - A deeper empirical investigation would have been good to see. In particular, I'd have liked to see experiments demonstrating the limits of the proposed approach as well as the use of \"larger\" applications instead of the toy examples.\n\nOverall, this seems like a nice paper making progress on a hard problem. I didn't detect any technical issues with the paper but I don't have a deep background in this research area. I'm open to adjusting my scores based on the other reviews.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept