Storchastic: A Framework for General Stochastic Automatic Differentiation
Authors: Emile van Krieken
Jakub Mikolaj Tomczak
Annette Ten Teije
Venue: NeurIPS 2021 Poster (2021)
Keywords: gradient estimation
automatic differentiation
optimization
stochastic computation graphs

Abstract: Modelers use automatic differentiation (AD) of computation graphs to implement complex Deep Learning models without defining gradient computations. Stochastic AD extends AD to stochastic computation graphs with sampling steps, which arise when modelers handle the intractable expectations common in Reinforcement Learning and Variational Inference. However, current methods for stochastic AD are limited: They are either only applicable to continuous random variables and differentiable functions, or can only use simple but high variance score-function estimators. To overcome these limitations, we introduce Storchastic, a new framework for AD of stochastic computation graphs. Storchastic allows the modeler to choose from a wide variety of gradient estimation methods at each sampling step, to optimally reduce the variance of the gradient estimates. Furthermore, Storchastic is provably unbiased for estimation of any-order gradients, and generalizes variance reduction techniques to higher-order gradient estimates. Finally, we implement Storchastic as a PyTorch library at github.com/HEmile/storchastic.

Storchastic: A Framework for General Stochastic Automatic Differentiation
Emile van Krieken Vrije Universiteit Amsterdam e. van. krieken@vu.nl
Jakub M. Tomezak Vrije Universiteit Amsterdam m. tomczak@uu.n..
Annette ten Teije Vrije Universiteit Amsterdam annette. ten. tei je@vu.
Abstract
Modelers use automatic differentiation (AD) of computation graphs to imple- ment complex deep learning models without defining gradient computations. Stochastic AD extends AD to stochastic computation graphs with sampling steps, which arise when modelers handle the intractable expectations common in re- inforcement learning and variational inference. However, current methods for stochastic AD are limited: They are either only applicable to continuous ran- dom variables and differentiable functions, or can only use simple but high vari- introduce ance score-function estimators. To overcome these limitations, Storchastic, a new framework for AD of stochastic computation graphs. Stor- chastic allows the modeler to choose from wide variety of gradient estimation methods at each sampling step, to optimally reduce the variance of the gradi- ent estimates. Furthermore, Storchastic is provably unbiased for estimation of any-order gradients, and generalizes variance reduction techniques to any-order derivative estimates. Finally, implement Storchastic PyTorch library github. com www.ccooommmmeleee /ttm//iiitttoocaattii.cc
Introduction
One of the driving forces behind deep learning is automatic differentiation (AD) libraries of complex computation graphs. Deep learning modelers are relieved by accessible AD of the need to implement complex derivation expressions of the computation graph. However, modelers are currently limited in settings where the modeler uses intractable expectations over random variables [37, 8). Two common examples are reinforcement learning methods using policy gradient optimization [49, 29, 36] and latent variable models, especially when inferred using amortized variational inference [34, 20, 41, 43]. Typically, modelers estimate these expectations using Monte Carlo methods, that is, sampling, and resort to gradient estimation techniques [37] to differentiate through the expectation.
A popular approach for stochastic AD is reparameterizatioo [20], which is both unbiased and has low variance, but is limited to continuous random variables and differentiable functions. The other popular approach [49, 45, 13] analyzes the computation graph and then the score function estimator to create a surrogate loss that provides gradient estimates when differentiated. While this approach is more general it can also be applied to discrete random variables and non-differentiable functions, naive applications of the score function will have high variance, which leads to unstable and slow convergence. Furthermore, this approach is often implemented incorrectly [13], which can introduce bias in gradients.
We therefore develop new framework called Storchastic support deep learning modelers. They can use Storchastic to focus on defining stochastic deep learning models without having to worry about complex gradient estimation implementations. Storchastic extends DiCE [13] to other gradi- ent estimation techniques than basic applications of the score function. It defines a surrogate loss by decomposing gradient estimation methods into four components: The proposal distribution, weight-
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Stochastic compun computation graph
Samples from proposal distribution
Weighting function
Storchastic
surrogate loss
Σ
001
122
Control variate
122
Ag de xx
Cradient function
+
Cost evaluation
Compute surrogate loss
Analyze graph
Sample and evaluate
Compute components
An illustration of the (parallelized) Storchastic loss computation. a. Assign the stochastic nodes of the input stochastic computation graph (SCG) into topologically sorted partitions. b. Evaluate the SCG. We first sample the set of values X; from the proposal distribution. For each of the samples x, Xi, we then sample set of samples X2. The rows in the figure indicate different samples in Xi, while the columns indicate samples in X2. The different samples are used to evaluate the cost function, [|2]] [222| times. c. Compute the weighting function, control variate and gradient function for all samples. d. Using those components and the cost function evaluation, compute the storchastic surrogate loss, mimicking Algorithm 1. C refers to element- wise multiplication, to element-wise summation and > for summing the entries of a matrix.
Figure I:
ing function, gradient function and control variate. We can use this decomposition get insight into how gradient estimators differ, and use them to further reduce variance by adapting components of different gradient estimators.
Our main contribution is a framework with a formalization and a proof that, if the components sat- isfy certain conditions, performing 7-th order differentiation the Storchastic surrogate loss gives unbiased estimates of the nith order derivative of the stochastic computation graph. We show these conditions hold for wide variety of gradient estimation methods for first order differentiation. For many score sunction-bssed methods like RELAX [16], MAPO [28] and the unordered set estimator Storchastic, we only have to prove
[25], the conditions also hold for any-order differentiation. these conditions locally. This means that modelers are free to choose the gradient estimation method that best suits each sampling step, while guarantceing that the gradient remains unbiased. Storchas- tie is the first stochastic AD framework to incorporate the measure- valued derivative [40, 18, 37] and SPSA [46, 2], and the first to guarantee variance reduction of any-order derivative estimates through control variates.
In short, our contributions the following:
1. We introduce Storchastic, new framework for general stochastic AD that uses four gradi- ent estimation components, in Section 3.1.3.3.
2. We prove Theorem 1, which provides conditions under which Storchastic gives unbiased any-order derivative estimates in Section 3.4. To this end, introduce mathematical formalization forward-mode evaluation in AD libraries in Section 2.4.
3. We derive technique for extending variance reduction using control variates to any-order derivative estimation in Section 3.5.
4. We implement Storchastic as open source library for PyTorch, Section 3.7.

---

Z
Figure 2: A Stochastic Computation Graph representing the computation of the losses of an VAE with discrete latent space.
2
Background
We capital letters N, .1,..... St. for nodes in a graph, calligraphic capital letters S,F for sets and non-capital letters for concrete computable objects such as functions f and values x,
2.1 Stochastic Computation Graphs
We start by introducing Stochastic Computation Graphs (SCGs) [45], which is formalism for stochastic AD. A Stochustic Computation Gruph (SCG) is a directed acyelic graph (DAG) G = (N,E) where nodes N are partitioned in stochustic nodes S and deterministic nodes F. We define the set of parameters C F such that all 0 e θ have no incoming edges, and the set of cost nodes C F such that all C have outgoing edges.
The of parents pa(N) is the of incoming nodes of node N N, that pa(N) {M (M,, E}. Each stochastic node S S represents a random variable with sample space Us and probability distribution ps conditioned on its parents. Each deterministic node F represents (deterministic) function fF of its parents.
M influences N, denoted M N = [ E M - } the set of nodes that influence N. The joint probability of all random variables Xg E IIses Pg is defined as p(xs) IIses Pss Xp((S)) where xpa(S) is the set of values of the nodes pa(S). The expected value of deterministic node F E F is its expected value over sampling stochastic nodes that influence that node, that is,
N, if there is directed path from M to N. We denote with
E F = Es (F(pa((F))
2.2 Problem Statement
In this paper, we aim to define a surrogate loss that, when differentiated using an AD library, gives unbiased estimate of the n-th order derivative of parameter 0 with respect the expected total cost v EEEEc CJ. This gradient can written as Ecee "("ECC" and we focus on estimating the gradient of single cost node ((")
2.3 Example: Discrete Variational Autoencoder
Next, we introduce a running example: A variational autoencoder (VAE) with a discrete la- tent space [20, 19]. First, we represent the model N = [Ø. Par. LKLD, (te)) and the stochastic nodes are S = {x.z). These are connected shown in Figure! 2. The parameters are θ = (0, Φ) which respectively the variational posterior and the model likelihood p. and the cost nodes = {LKLD. (te) that represent the KL-divergence between the posterior and the prior, and the 'reconstruction loss", the model log-likelihood after decoding the sample z. Finally, şs represents the parameters of the mul- tivariate categorical distribution of the amortized variational posterior (z(r)) This SCG represents the equation
SCG: The deterministic nodes
the parameters of
EL.LDLLD + Reel [.[(<L]] (xn,=ngg(
The problem we are interested in is estimating the gradients these expectations with respect to the parameters. Since is not influenced by the parameters, have v."" E, 0

---

and V EglleeLD " (LL)) The second term is more challenging. We can move the gradient with respect to 8 in, since is not influenced by O: "x Ree) without gradient es-
( (z(a)) ( Ree] However, cannot compute V (ag Bagaa timation methods. This is because sampling from (=(2(=) is dependent ø. Furthermore, since we are dealing with a discrete stochastic node, we cannot apply the reparameterization method here without introducing bias.
2.4 Formalizing AD libraries and DiCE
To be able properly formalize and prove the propositions in this paper, we introduce the 'forward- mode' operator that simulates forward-mode evaluation using AD libraries. This operator properly handles the common 'stop-grad' operator, which ensures that its argument only evaluated during forward-man mode evaluations of the computation graph. It is implemented in Tensorflow and Jax with the name stop gradiens ent [l, 4] and in PyTorch as detach or no grad [39]. stop-grad' is neces- sary to define surrogate losses for gradient estimation, which is why it is essential to properly define it. For formal definitions of the following operators and proofs we refer the reader to Appendix A.
Definition 1 (informal). The stop-grad operator is a function such that (x) 0. The forward-mode operator -y, which is denoted arrow above the argument it evaluates, acts as an identity function, except that L(a) = Z. Additionally, we define the MagicBox operator as [()) exp(x (z:)..
Importantly. the definition of + implies that (z) does not equal V (zf(z) if f contains stop- grad operator. MagicBox, which was first introduced in [13], is particularly useful for creating surrogate losses that remain unbiased for any-order differentiation. It is defined such that (zz) 1 and V (z)) Of (z)) fxr).. This allows injecting multiplicative factors to the computation graph only when computing gradients.
Making use of MagicBox, DiCE [13] is an estimator for automatic nth-orrdr derivative estimation that defines surrogate loss using the score function:
E CEC
-
log 1(xs (pppxxs(ppss)))CC)))))pp))))))))
(I)
CEC
SES.c
DiCE correctly handles the credit assignment problem: The score function only applied to the stochastic nodes that influence cost node. It also handles pathwise dependencies of the param- eter through cost functions. However, it has high variance since it is based on a straightforward application of the score function.
The Storchastic Framework
In this section, we introduce Storchastic, a framework for general any-order gradient estimation in SCGs that gives modelers the freedom to choose suitable gradient estimation method for each stochastic node. First, we present requirements that we used to develop the framework in Section 3.1. Storchastic deconstructs gradient estimators into four components that we present in Section
3.2. We use these components to introduce the Storchastic surrogate loss Section 3.3, and give conditions that need to hold for unbiased estimation in Section 3.4. In Section 3.5 we discuss variance reduction, in Section 3.6 we discuss several estimators that fit in Storchastic, and in Section 3.7 we discuss our Py Torch implementation. An overview of our approach is outlined in Figure 1.
3.1 Requirements of the Storchastic Framework
First, we introduce the set of requirements used to develop Storchastic.
1. Modelers should be able to choose a different gradient estimation method for each stochas- tie node. This allows for choosing the method best suited for that stochastic node, or adding background knowledge in the estimator.

---

2. Storchastic should be flexible enough to allow implementing a wide range of reviewed gra- dient estimation methods, including score function-bases methods with complex sampling techniques [52, 25, 28] or control variates [16, 47]. and other methods such as measure- valued derivatives [18, 40] and SPSA [46] which are missing AD implementations [37].
3. Storchastic should define surrogate loss [45], which gives gradients of the SCG when differentiated using an AD library. This makes it easier to implement gradient estimation methods as modelers get the computation of derivatives for free.
4. Differenniating the surrogate loss n times should give estimates of the rah-order derivative, which are used in for example reinforcement learning [15, 14] and meta-learning [12, 27).
5. Variance reduction methods through better sampling and control variates should generalize in higher-order derivative estimation.
6. Storchastic should be provably unbiased. To reduce the effort of developing new methods, researchers should only have to prove a set of local conditions that generalize to any SCG.
3.2 Gradient Estimators in Storchastic
Next, we introduce each of the four components and motivate why each is needed to ensure Require- ment is satisfied. First, we note that several recent gradient estimators, like MAPO [28], unordered set estimator [25] and self-critical baselines [21, 42] act sequences of stochastic nodes instead of on a single stochastic node. Therefore, we create a partition St, ..., Sg of S S<c topologically ordered by the influence relation, and define the shorthand x; xs.- For each partition S,, we choose gradient estimator, which is a 4-tuple (qs, Wi, 44, a.). Here, ( xxs) is the proposal distribution, (x,) is the weighting function, (xx) is the gradient function and a; is the control variate.
3.2.1 Proposal distribution
Many gradient estimation methods in the literature do not sample a single value x; ~ p(x; (x=)). but sample, often multiple, values from possibly different distribution. Some instances of sampling schemes are taking multiple i.i.d. samples, importance sampling [32] which is very common off-
policy reinforcement learning, sampling without replacement [25]. memory-augmented sampling [28] and antithetic sampling [51]. Furthermore, measure-valued derivatives [37, 18] and SPSA [46] also sample from different distributions by comparing the performance of two related distributions. To capture this, the proposal distribution (X,((xx)) samples a set of values X, = {xx,1. m,,m) where each x,, E Uts.- The sample is conditioned on x USS,, the values of the parent nodes of the stochastic nodes in S,. This is illustrated in Figure 1.b.
3.2.2 Weighting function
When a gradient estimator uses a different sampling scheme, we have to weight each individual sample to ensure it remains a valid estimate of the expectation. For this, we use a nonnegative weighting function w, Ωs, + Rt. Usually, this function is going to be detached from the computation graph, but we allow it to receive gradients as well to support implementing expectations and gradient estimation methods that compute the expectation over (a subset of) values [25, 28, 30].
3.2.3 Gradient function
The gradient function is an unbiased gradient estimator together with the weighting function. distributes the empirical cost evaluation to the parameters of the distribution. In the case of score function methods, this is the log-probability. For measure-valued derivatives and SPSA we can use the parameters of the distribution itself.
3.2.4 Control variate
Modelers can use control variates to reduce the variance of gradient estimates [17, 37]. It function that has zero-mean when differentiated. Within the context of score functions, common control variate is a baseline, which a function that independent found that LAX, RELAX, and REBAR (Appendix D.2.4), and the GO gradient [7] (Appendix D.2.6) have natural implementations using a control variate. We discuss how we implement control variates in Storchastic in Section 3.5.
the sampled value. We also
5

---

3.2.5 Example: Leave-one-out baseline
As an example, we show how to formulate the score function with the leave-one-our baseline [35, 22] in Storchastic. This method samples values with replacement and uses the average of the other values baseline.
Proposal distribution: use m samples with replacement, which can be formulated as gXii(x))) P(x,gl(
Weighting function: Since samples are independent, we use wy (x;) = 1
Gradient function: The score-function uses the log-probability 1 (x,)) = log p(x, (xe)).
Control variate: We use at. (X
(..ii.j.. X.) (1 (x(()))) - Enze fo(x<s Xt,s):
where m-I Egzj fo(xci, Xx,)) is the leave-one-out baseline. (1 (x.())))) is used to
ensure the baseline will be subtracted from the cost before multiplication with the gradient
function. will not affect the forward evaluation since 1 (x:()) evaluates to O.
3.3 The Storchastic Surrogate Loss
As mentioned in Requirement 3, we would like to define a surrogate loss, which we will introduce next. Differentiating this loss times, and then evaluating the result using AD library, will give unbiased estimates of the n-th order derivative of the parameter 0 with respect to the cost C. Furthermore, according to Requirement 1, assume the modeler has chosen gradient estimator (4,, Wi, li, as) for each partition S; which can all be different. Then the Storchastic surrogate loss is
Slsoorch
Σ X1EX
wi (x1) ai (x), +
Σ X2
w2(x2)
+
Σ wx((Xk) X
1 (xz)))x((xea)) N.) + + 4(xx))
where Xi ~ {(Ki),, ~
Xx ~ ((K))x=)).
When this loss is differentiated n times using libraries, will produce unbiased estimates the n-th derivative, as we will show later. To help understand the Storchastic surrogate loss and why
Algorithm The Storchastic framework: Compute a Monte Carlo estimate of the n-th order gradi- ent given & gradient estimators (q,, wi, 4, a,).
1: function ESTIMATE. CEETTTE.GRADIENN GRADIENT(Z,, 0)
2: SLtoorc <- SURROGATE (..0)
Compute surrogate loss
3: return (1) ST.torch
D Differentiate and use AD library to evaluate surrogate loss
4:
5: function SURROGATE X
6: i then
7: return F (L)
> Use MagicBox to distribute cost Sample from proposal distribution
8:
X; ~ ((((xxe))
9:
<- 0
for
x; E X; do
Iterate over options in sampled set Compute control variate b Compute gradient function
11:
A ExL)a X))
12:
L, + L 4(x,)
13:
Ò + SURROGATE SURROGATE.LOSSE + 1,x L;)
D. Compute surrogate loss for x; Weight and add
14:
sum +- sum + (x; + A)
return sum
satisfies the requirements, we will break down using Algorithm The ESTIMATE GRADIENT function computes the surrogate loss for the SCG, and then differentiates it 2 0 times using the AD library to return an estimate the n-th order gradient, which should unbiased according to Requirement 4. If n is set zero, this returns estimate of the expected cost.

---

The SURROGATE LOSS function computes the equation using recursive computation, which is illustrated in Figure 1.b-d. It iterates through the partitions and uses the gradient estimator to sample and compute the output of each component. It receives three inputs: The first input i indexes the partitions and gradient estimators, the second input x the set of previously sampled values for partitions Szi. and L is the sum of gradient functions of those previously sampled values. In line 8, we sample a set of values X. for partition i from ((X,(x<)) In lines to 14, we compute the sum over values x; in X, which reflects the i-th sum of the equation. Within this summation, in lines 11 and 12, we compute the gradient function and control variate for each value x, We will explain in Section 3.5 why we multiply the control variate with the MagicBox of the sum of the previous gradient function.
In line 13, we go into recursion by moving to the next partition. We condition the surrogate loss on the previous samples Xz; together with the newly sampled value x, We pass the sum of gradient functions for later usage in the recursion. Finally, in line 14, the sample performance and the control variate are added in a weighted sum. The recursion call happens for each x; X, meaning that this computation is exponential in the size of the sampled sets of values X, For example, the surrogate loss samples | times from 92. one for each value X1 E XT However, this computation can be trivially parallelized by using tensor operations in AD libraries. An illustration of this parallelized computation is given in Figure 1.
Finally, in line 7 after having sampled values for all l partitions, we compute the cost, and multiply it with the MagicBox of the sum of gradient functions. This is similar to what happens in the DiCE estimator in Equation (1). Storchastic can be extended to multiple cost nodes by computing surrogate losses for each cost node, and adding these together before differentiation. For stochastic nodes that influence multiple cost nodes, the algorithm can share samples and gradient estimation methods to reduce overhead.
3.4 Conditions for Unbiased Estimation
We next introduce our main result that shows Storchastic satisfies Requirements 4 and 6, namely the conditions the gradient estimators should satisfy such that the Storchastic surrogate loss gives estimates of the n-th order gradient of the SCG. A useful part of our result is that, in line with Requirement 6, only local conditions of gradient estimators have to be proven to ensure estimates unbiased. Our result gives immediate generalization of these local proofs to any SCG.
Theorem 1. Evaluating the n-th order derivative of the Storchastic surrogate loss in Equation (2) using an AD library an unbiased estimate of V [["E under the following conditions. First, all functions fr corresponding to deterministic nodes F and all probability measures ps correspond- ing to stochastic nodes S are identical under evaluation. Secondly, for each gradient estimator (4, Wi, li, az), i= 1,...,,, k, all the following hold for m O....,
1. Ex(x((x))) ExIEX, V (x.(( = V (m)) Es, [f for any deterministic function f:
2. (x(()x ExEx,
V (m) u.
Xz)1 = 0:
3. for 2 m > 0, E (x Ex,e
V ((m) (i: (x:)] 0;
4. ( = (((()xx)) for all permissible X,
The first condition defines local surrogate loss for single expectations of any function under the proposal distribution. The condition then says that this surrogate loss should give an unbiased esti- mate of the gradient for all orders of differentiation m 0, ..., Note that since 0 is included, the forward evaluation should also be unbiased. This is the main condition used to prove unbiasedness of the Storchastic framework, and can be proven for the score function and expectation, and for measure-valued derivatives and SPSA for zeroth and first-order differentiation.
The second condition says that the control variate should be 0 in expectation under the proposal distribution for all orders of differentiation. This is how control variates are defined in previous work [37], and should usually not restrict the choice. The third condition constrains the weighting
7

---

function to be 0 in expectation for orders of differentiation larger than O. Usually, this is satisfied by the fact that weighting functions are detached from the computation graph, but when enumerating expectations, this can be shown by using that the sum of weights is constant. The final condition is a regularity condition that says proposal distributions should not different under forward mode. We also assume that the SCG is identical under evaluation. This means that all functions and probability densities evaluate to the same value with and without the forward-mode operator, even when differentiated. This concept is formally introduced in Appendix A.
full formalization and the proof of Theorem are given in Appendix B.1. The general idea is to rewrite each sampling step as an expectation. and then inductively show that the inner expecta- tion i over the proposal distribution q: is an unbiased estimate of the rihh-rddr derivative over S conditional on the previous samples. To reduce the multiple sums over gradient functions inside MagicBox, we make use of property of MagicBox proven in Appendix A:
Proposition 4. Summation inside MagicBox is equivalent under evaluation multiplication of the arguments in individual MagicBoxes, ie:
(I1 +
Equivalence under evaluation, denoted ž. informally means that, under evaluation of -; the two expressions and their derivatives are equal. This equivalence is closely related to ea E" c
3.5 Any-order variance reduction using control variates
To satisfy Requirement 5, we investigate implementing control variates such that the variance of any- order derivatives is properly decreased. This is challenging in general SCG's [33], since in higher orders of differentiation, derivatives of gradient functions will interact, but naive implementations of control variates only reduce the variance of the gradient function corresponding to a single stochastic node. Storchastic solves this problem similarly to the method introduced in [33]. In line 11 of the algorithm, we multiply the control variate with the sum of preceding gradient functions (L)). We prove that this ensures every term of the any-order derivative will be affected by a control variate in Appendix C. This proof is new, since [33] only showed this for first and second order differentiation, not for general control variates, and uses slightly different formulation that we show misses some terms.
Theorem (informal). Let L, = Ej= 4. The Storchastic surrogate loss of (2) can equivalently be computed as
I (x.
LStorch E Σ EII
E ( (x)
+ +C.
XIEA X2EX,
x1EXi 1=1
This gives insight into how control variates are used in Storchastic. They are added to the gradient function, but only during differentiation since (L.) - O. Furthermore, since both terms multiplied with EX (11) (see line 11 of Algorithm 1), both terms correctly distribute over the same any-order derivative terms. By choosing control variate of the form ax(x;í, X)) (1 ([(()) by we recover baselines which are common in the context of score functions [13, 37J. For the proof, use the following proposition also proven in Appendix
Proposition 7. For orders of differentiation > 0,
(()) Σ (1)) ()--))
3.6 Gradient Estimation Methods
In Appendix D we show how several prominent examples of gradient estimation methods in the liter- ature can be formulated using Storchustic, and prove for what orders differentiation the conditions hold. Starting off, we show that for finite discrete random variables, we can formulate enumerating over all possible options using Storchastic. The score function fits by mimicking DiCE [13]. We extend it to multiple samples with replacement to allow using the leave-one-out bascline [35, 22].
8

---

class ScoreFunctionL00sstooe method. Kethod): def proposal_dist (self, distribution, amt_samples): return distr. sample (amt_samples,,,
def functing (self, distribution, amt_samples). return torch full (amt (amt_samples. campleem
def estimator (self, sample, cost):
# Compute gradient function (log-probability) log-prob sample. distribution log-prob (tensor) sum_sosts storch. sum (costs detach(), sample.. name) # Compute control variate
12
13
baseline (sum_costs costs) (sample.n 1) return log_prob, (1.0 magic box box
14
(log-prob)) baseline
Figure 3: Implementing the score function with the leave-one-out baseline the Storchastic library.
Furthermore, we show how importance sampling, sum-and aed-aandddlllla estimators such as MAPO [28], the unordered set estimator [25] and RELAX and REBAR [16, 47] fit in Storchastic. We also discuss the antithetic sampling estimator ARM [51]. Unfortunately. condition 2 only holds for this estimator for since it relies particular property of the score function that holds only for first-order gradient estimation. In addition to score function based methods, we discuss the GO gradient, SPSA [44] and Measure-Valucee Derivative [18], and show that the last will only be unbiased for S 1. Finally, we note that reparameterization [20, 43] can be implemented by transforming the SCG such that the sampling step outside the path from the parameter the cost [45].
3.7 Implementation
We implemented Storchastic open source PyTorch [39] library 1. To ensure modelers ily use this library, automatically handles sets of samples as extra dimensions to Py PyTorch tensors which allows running multiple sample evaluations in parallel. This approach is illustrated in Fig- ure 1. By making use of PyTorch broadcasting semantics, this allows defining models for simple single-sample computations that automatically parallelized using Storchastic when using mul- tiple samples. The Storchastic library has implemented most of the gradient estimation methods mentioned in Section 3.6. Furthermore, new gradient estimation methods can seamlessly be added.
3.7.1 Example: Leave-one-out baseline in Discrete Variational Autoencoder
As small case study, we show how implement the score function with the leave-one-out baseline introduced Section 3.2.5 for the discrete variational autoencoder introduced in Section 2.3 in PyTorch using Storchastic. While the code listed is simplified, shows the flexibility with which one can compute gradients in SCGs.
We list in Figure 3 how to implement the score function with the leave-one-out baseline. Line 3 implements the proposal distribution, line 6 the weighting function, line 10 the gradient function and line 13 and 14 the control variate. Gradient estimation methods in Storchastic all extend common base class storch. method. Method to allow easy interoperability between different methods.
In Figure 4, we show how to implement the discrete VAE. The implementation directly follows the SCG shown in Figure 2. In line 2, we create the EcoreFunctionia Then, we run the training loop: In line 6 we create the stochastic node x by denoting the minibatch dimension as an independent dimension. variational posterior 4:- We call the gradient estimation method in line 9 to get sample of z. Note that this interface is independent of gradient estimation method chosen, meaning that if wanted to compare our implemented method with baseline, all that is needed is to change line 2. After the decoder. we compute the two costs in lines 12 and 13. Finally, we call Storchastic main algorithm in line 15 and run the optimizer.
method defined in Figure 3.
line 8 we run the encoder with parameters
find the
Code available github. www////aalle/soooceeeestteeeeeee

---

from vae import minibatches, encode, decode, KLD, binary binary_cross_atrrgy method ScoreFunction100 ("2", ("z"", for data minibatches():
optimizer pptimizer.zero.graat zero_grad
# Denote minibatch dimension as independent plate dimension
data = storch denote. benote_.ndeeendente (datta view (-1, 784), 0, "data")
Compute variational distribution given data, sample 2
9 torch distributions OneHotCategoricco (lgiitseencode (data))) (data))
2 method(q)
# Compute costs, form the ELBO reconstruction decode (2)
12
storch add_ (KL)((QQ)))
13
storch, add_cost wwwnarr_ccrossseeeeeeee
(reconstruction, data))
14
Storchastic backvard pass, optimize
ELBO = storch backward
16
optimizer ()
Figure 4: Simplified implementation of the discrete VAE using Storchastic.
We run this model on our currently implemented of gradient estimation methods for discrete variables in Appendix E and report the results, which are meant purely to illustrate the case study.
4
Related Work
The literature on gradient estimation is rich, with papers focusing on general methods that can be implemented in Storchastic [46, 18, 16, 51, 28, 30, 7), Appendix D, and works focused Reinforcement Learning [49, 29, 36] or Variational Inference [35]. For a recent overview, see [37].
The literature focused SCGs is split into methods using reparameterization [43, 20, 11, 31, 19] and those using the score function [45]. Of those, DiCE [13] is most similar to Storchastic, and can do any-order estimation general SCGs. DiCE is used in the probabilistic programming library Pyro [3]. We extend DiCE to allow for incorporating many other gradient estimation methods than just basic score function. We also derive and prove correctness of general implementation for control variates for any-order estimation which is similar to the one conjectured for DiCE in [33].
[38, 50) and [48] study actor-critic-lik techniques and bootstrapping for SCGs to incorporate repa- rameterization using methods inspired by deterministic policy gradients [29]. By using models to differentiate through, these methods are biased through model inaccuracies and thus do not directly fit into Storchastic. However, combining these ideas with the automatic nature of Storchastic could be interesting future work.
5 Conclusion
We investigated general automatic differentiation for stochastic computation graphs. We developed the Storchastic framework, and introduced an algorithm for unbiased any-order gradient estimation that allows using large variety of gradient estimation methods from the literature. We also investi- gated variance reduction and showed how to properly implement control variates such that it affects any-order gradient estimates. The framework satisfies the requirements introduced in Section 3.1.
For future work, we are interested in extending the analysis of Storchastic to how variance com- pounds when using different gradient estimation methods. Furthermore, Storchastic could be ex- tended allow for biased methods. We also interested in closely analyzing the different com- ponents of gradient estimators, both from theoretical and empirical point of view, to develop new estimators that combine the strengths of estimators in the literature.
References
[1] M. Abadi, Agarwal, Barham, E. Brevdo, Chen, C. Citro, G. S. Corrado, Davis,
J. Dean, M. Devin, S. Ghemawat, I. Goodfellow, A. Harp, G. Irving, M. Isard, Y. Jia, R. Joze-
10

---

fowicz, L. Kaiser, M. Kudlur, J. Levenberg, D. Mané, R. Monga, S. Moore, D. Murray, C. Olah, M. Schuster, J. Shlens, B. Steiner, Sutskever, K. Talwar. P. Tucker, V. Vanhoucke, V. Va- sudevan, F. Viégas, O. Vinyals, P. Warden, M. Wattenberg, M. Wicke, Y. Yu, and X. Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems. 2015.
[2] S. Bhatnagar, H. Prasad, and L. Prashanth. Stochastic Recursive Algorithms for Optimization, volume 434 of Lecture Notes in Control and Information Sciences. Springer London, London, doi: 10.1007/978- 111779978-14471-444 1-4471-428-000
2013. ISBN 978-1-4477 188-1-4471-428433
978-1-4477 14471-4285-0.
[3] E. Bingham, J. Chen, M. Jankowiak, F. Obermeyer, N. Pradhan, T. Karaletsos, R. Singh, P. A. Szerlip, Horsfall, and N. D. Goodman. Pyro: Deep universal probabilistic programming. Journal of Machine Learning Research, 20:28:1-28:6, 2019.
[4] J. Bradbury. R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. SaaderrmnnMiile and Q. Zhang. JAX: Composable transfor- mations of Python+NumPy programs, 2018.
[5] I.. Buesing, T. Weber, and S. Mohamed. Stochastic gradient estimation with finite differences. page 4, 2016.
[6] G. Casella and C. P. Robert. Rao-blackwellisationn of sampling schemes. Biometrika, 83(1): 81 81-94, 1996. ISSN 00063444.
[7] Y. Cong, M. Zhao, K. Bai, and L. Carin. GO gradient for expectation- basedd objectives. page 30, 2019.
[8] G. Correia, V. Niculae, W. Aziz, and A. Martins. Efficient marginalization of discrete and structured latent variables via sparsity. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 11789-11802 Curran Associates, Inc., 2020.
[9] 7. Dong, A. Mnih, and G. Tucker. DisARM: An antithetic gradient estimator for binary la- tent variables. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 18637-18647. Curran Associates, Inc., 2020.
[10] G. Farquhar, S. Whiteson, and J. N. Foerster. Loaded DiCE:
Trading off bias and variance in any-order score function gradient estimators for reinforcement learning. In H. M. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché-Buc, E. B. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Pro- cessing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 1149-81601 2019.
[11] M. Figurnov. S. Mohamed, and A. Mnih. Implicit reparameterization gradients. In S. Ben- gio, H. M. Wallach, H. Larochelle, K. Grauman, N. Cesa- and R. Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural In- formation Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pages 439-450, 2018.
[12] C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Confer- ence on Machine Learning. ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 1 1/2--135. PMLR, 2017.
[13] Foerster, G. Farquhar, M. Al-Shedivat, Rocktäschel, F. Xing, and Whiteson. DiCE: The infinitely differentiable monte carlo estimator. In International Conference Machine Learning, pages 1529-1538, 2018
[14] J. N. Foerster, R. Y. Chen, M. Al-Shedivat, S. Whiteson, Abbeel, and I. Mordatch. Learning with oeannnnn- awareness. In E. André, S. Koenig, M. Dastani, and G. Sukthankar, editors, Proceedings of the 17th International Conference on Autonomous Agents and Mul- tiAgent Systems, AAMAS 2018, Stockholm, Sweden, July 10-15, 2018, pages 122-130. Inter- national Foundation for Autonomous Agents and Multiagent Systems Richland, SC, USA / ACM, 2018
[15] Furmston, G. Lever, and D. Barber. Approximate newton methods for policy search in markov decision processes. Journal of Machine Learning Research, 17:227:1-227:51, 2016.
11

---

[16] W. Grathwohl, D. Choi, Y. Wu, G. Roeder, and D. Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In 6th International Con- ference on Learning Representations, ICLR 2018 - Conference Track Proceedings, 2018.
[17] E. Greensmith, P. I.. Bartlett, and J. Baxter. Variance reduction techniques for gradient esti- mates in reinforcement learning. Journal of Machine Learning Research, 5(9), 2004.
[18] B. Heidergott and F. Vázquez-Abad. Measure-valued differentiation for Markov chains. Jour- nal of Optimization Theory and Applications, 136(2):187-209, 2008.
[19] E. Jang, S. Gu, and B. Poole. Categorical reparameterization with gumbel-softmax. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24- 26, 2017, Conference Track Proceedings, 2017.
[20] D. P. Kingma and M. Welling. Auto-encoding variational bayes. In 2nd International Confer- ence on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Confer- ence Track Proceedings, 2014.
[21] W. Kool, H. van Hoof, and M. Welling. Attention, learn to solve routing problems! In 7th International Conference Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
[22] W. Kool, H. van Hoof, and M. Welling. Buy 4 REINFORCE samples, get a baseline for free! page 14, 2019.
[23] W. Kool, H. van Hoof, and M. Welling. Stochastic beams and where to find them: The gumbel- top-k trick for sampling sequences without replacement. In K. Chaudhuri and R. Salakhutdi- nov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 3499-3508. PMLR, 2019.
[24] W. Kool, H. van Hoof, and M. Welling. Ancestral gumbel-top-k sampling for sampling without replacement. Journal of Machine Learning Research, 21(47):1-36, 2020.
[25] W. Kool, H. van Hoof, and M. Welling. Estimating gradients for discrete random variables by sampling without replacement. page 28, 2020.
[26] Y. LeCun and C. Cortes. MNIST handwritten digit database. 2010.
[27] Z. Li, F. Zhou, F. Chen, and H. Li. Meta-SGD: Learning to Learn Quickly for Few-Shot Learning. arXiv: [cs], Sept. 2017.
[28] C. Liang, M. Norouzi, J. Berant, Q. Le, and N. Lao. Memory Augmented Policy Optimization for Program Synthesis and Semantic Parsing. 1rXiv::807. 007222 [cs, staf), Jan. 2019.
[29] T. P. Lillicrap. J. J. Hunt, A. Pritzel, N. Heess, T. Erez, Y. Tassa, D. Silver, and D. Wierstra. Continuous control with deep reinforeement learning. In 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016.
[30] R. Liu, Regier, N. Tripuraneni, M. I. Jordan, and J. D. McAuliffe. Rao-blackwellizze stochastic gradients for discrete distributions. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 4023-4033 PMLR, 2019.
[31] C. J. Maddiso, A. Mnih, and Y. W. Teh. The concrete distribution: A continuous relaxation of discrete random variables. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings, 2017.
[32] A. R. Mahmood, H. Van Hasselt, and R. S. Sutton. Weighted importance sampling for off- policy learning with linear function approximation. In NIPS, pages 3014-3022, 2014.
[33] J. Mao, J. Foerster, T. Rocktäschel, M. Al-Shedivat, G. Farquhar, and S. Whiteson. A base- line for any order gradient estimation in stochastic computation graphs. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 4343-4351 4351, Long Beach, California, USA, June 2019. PMLR.
12

---

[34] A. Mnih and K. Gregor. Neural variational inference and learning in belief networks. In E. P. Xing and T. Jebara, editors, Proceedings of the 31st International Conference on Ma- chine Learning, volume 32 of Proceedings of Machine Learning Research, pages 1791-1799, Bejing, China, June 2014. PMIR.
[35] A. Mnih and D. J. Rezende. Variational inference for monte carlo objectives. In 33rd Interna- tional Conference on Machine Learning, ICMI. 20/6, 2016. ISBN
[36] V. Mnih, A. P. Badia, M. Mirza, A. Graves, T. P. Lillicrap, T. Harley, D. Silver, and K. Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Proceedings of the 33nd International Conference on Machine Learning, ICMI. 2016, New York City, NY, USA, June 19-24, 2016, pages 1928-1937, 2016.
[37] S. Mohamed, M. Rosca, M. Figurnov, and A. Mnih. Monte carlo gradient estimation in ma- chine learning. Journal of Machine Learning Research, 21:132:1-132:62 2020.
[38] P. Parmas. Total stochastic gradient algorithms and applications in reinforcement learning. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Desa-Bianchi, and R. Garnett, edi- tors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[39] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang. Z. DeVito, M. Raison, A. Te- jani, S. Chilamkurthy, B. Steiner, L. Fang. J. Bai, and S. Chintala. PyTorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, d' texetquuoeee Alché-Buu Buc, E. Fox, and R. Garnett, editors, Advances in Neural Informa- tion Processing Systems 32, pages
Curran Associates, Inc., 2019.
[40] G. C. Pflug. Sampling derivatives of probabilities. Computing, 1989. ISSN 0010485X. doi: 10. 10.007//BF0224327.
[41] R. Ranganath, S. Gerrish, and D. M. Blei. Black box variational inference. In Proceedings of the Seventeenth International Conference on Artificial Intelligncce and Statistics, AISTATS 2014, Reykjavik, Iceland, April 22-25, 2014, volume 33 of JMLR Workshop and Conference Proceedings, pages 814-822. JMLR.org. 2014.
[42] S. J. Rennie, E. Marcheret, Y. Mroueh, J. Ross, and V. Goel. Self-Critical Sequence Training for Image Captioning. In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 1 1199-1195, Honolulu, HI, July 2017. IEEE. ISBN 978-1-5386-0457-1 doi:
10.1109/CVPR.2017.11
[43] D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In E. P. Xing and T. Jebara, editors, Proceedings of the 31st International Conference on Machine Learning, volume 32 of Proceedings of Machine Learning Research, pages 1278-1286 Bejing, China, June 2014. PMLR.
[44] R. Y. Rubinstein and D. P. Kroese. Simulation and the Monte Carlo Method, volume 10. John Wiley& & Sons. 2016.
[45] J. Schulman, N. Heess, T. Weber, and P. Abbeel. Gradient estimation using stochastic compu- tation graphs. In Advances in Neural Information Processing Systems, 2015.
[46] C. Spall et al. Multivariate stochastic upproximation using simultaneuus perturbation gradient approximation. IEEE transactions on automatic control, 37(3):332-344 341, 1992.
[47] G. Tucker, A. Mnih, C. J. Maddison.
D. Lawson, and J. Sohl-Dickstein. REBAR: Low- variance, unbiased gradient estimates for discrete latent variable models. Advances in Neural Information Processing Systems, 2017.
[48] T. Weber, N. Heess, L. Buesing, and D. Silver. Credit assignment techniques in stochas- tic computation graphs. In K. Chaudhuri and M. Sugiyama, editors, The 22nd International Conference on Artificial Intelligence and Statistics, AISTATS 2019, 16-18 April 2019, Naha, Okinawa, Japan, volume 89 of Proceedings of Machine Learning Research, pages 2660-2660. PMLR, 2019.
[49] R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Machine Learning, 1992. ISSN 0885-6125, doi: 10. 10.10077b000992266.
[50] X. Xu, Zu, Y. Zhang, H. Zhou, and W. Feng. Backprop-Q Generalized Backpropagation for Stochastic Computation Graphs. arXiv: 1807.0951 [cs, stat), Jan. 2019.
13

---

[51] M. Yin and M. Zhou. ARM: rugment-reinforee meinforce-merge gradient for stochastic binary networks. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
[52] M. Yin, Y. Yue, and M. Zhou. ARSM: Augneen-eeifoore-wwa--rrrgg estimator for gradient backpropagation through categorical variables. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 7095-7104, PMLR, 2019.
14


Reviews:
[{"review_id": "qUNXhzn38G4", "reviewer": "Reviewer_VArr", "summary": "Storchastic presents a general formulation for surrogate losses to be differentiated at any order to evaluate a large array of gradient estimation techniques for expectations of stochastic computational graphs and implements this framework in PyTorch.\n\nThe MagicBox operator idea of the DiCE estimator can be used to make the score function gradient estimator evaluatable by AD methods. The Storchastic framework utilizes this idea for a surrogate loss that works for random variables conditioned on other random variable in the computational graph. The required sample size grows exponentially in the depth of the conditional nesting. In the PyTorch implementation the nested sampling is implemented by extra dimensions in the tensors which allows for good parallelization.\n\nIn Storchastic each random variable can use a suitable individual variance reduction technique, such as proposal distribution with weighting or control variates (or the reparameterization trick). The paper demonstrates that the proposed surrogate loss can be used to implement various well-established techniques from the literature.\n\nThe paper shows that under certain reasonable local conditions the obtained estimators are unbiased under any order of differentiation. These conditions are as follows:\n- for all orders of differentiation the surrogate should be unbiased for any deterministic function under the proposal distribution + weighting\n- for all orders of differentiation the control variate should be unbiased under the proposal distribution + weighting\n- the expectation of the weighting function needs to be zero when differentiated\n- the proposal distribution should be id under forward mode evaluation\n\nThe paper also shows how control variates can be appropriately pulled into the surrogate loss formulation for nested conditional random variables using the MagicBox trick and provide variance reduction at any order of differentiation.\n", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The paper gives a succinct definition of stochastic computational graphs and the essential problem of estimating gradients (and higher order derivatives) of expectations. The definitions, problem statements, assumption and introduced formalisms are stated clearly and explained well.\n\nThe general formulation encapsulating existing gradient estimators and making the connection to modern frameworks for DNN training is an original contribution that has the potential to be of significant use for practitioners that are working on complex stochastic models (such as deep discrete variational auto encoders etc.).\n\nThe combination and generalization of the various different approaches to gradient estimation approaches in a general surrogate loss formalism as presented may also be helpful to researchers trying to derive gradient estimators even when they do not end up using the PyTorch implementation. Nonetheless, I think the paper could be strengthened by including a brief section about the implementation interface so that readers of the paper can already have an idea of how they would interact with the library.\n\nI recommend to accept the paper because it presents a high quality clear methodology of how to elegantly implement state-of-the-art gradient estimators using AD methods.\n\nNit:\n- The sum in Algorithm 1 reads like \"s times u times m\" (the italic letters look like individual variable names), maybe you could change the typesetting here", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "kmCmYotEFve", "reviewer": "Reviewer_3poh", "summary": "This paper introduces a framework for stochastic gradient estimation that is able to capture many existing gradient estimation techniques, and allows a modeler to interchange different gradient estimation techniques in their computation graph. The paper also provides a mathematical framework to prove unbiasedness of any-order derivatives for common gradient estimation techniques using their surrogate loss.\n\nThe paper introduces new notation and a formalization of forward mode AD and gradient computation in stochastic graphs. Using this framework, they introduce an algorithm that, given a partition of stochastic nodes, and a specification of a gradient estimator for each stochastic node, returns a surrogate loss that can be arbitrarily differentiated to give an n-th order derivative of a cost wrt a parameter. Specification of a gradient estimator reduces to specifying 4 components: a gradient function, proposal distribution, weighting function, and control variate.\n\nThe paper then proves that many existing gradient estimation techniques fit in the framework, and for many of these techniques use the framework to prove unbiasedness of arbitrary derivatives.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "I think the unification of several gradient estimation techniques into a framework is a nice contribution. The paper is well polished and clearly written, and the notation introduced seems to help with the proofs significantly.\n\nAlthough similar to DiCE, it seems to extend it in a nontrivial way by incorporating many different gradient estimation functions, and bringing a more formal foundation to it.\n\nThe paper also introduces an implementation of the framework in PyTorch, but unfortunately does not include any case studies or experiments with the code that would help the reader understand whether the framework is particularly useful. I expand on this in the limitations section.\n\nGiven the limitations, I am borderline on the paper, but currently lean towards acceptance because I think the framework is a nice contribution.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The most important missing piece I think is the lack of explicit case studies / experiments with the framework. Although I understand the framework provides a lot of flexibility to the modeler, is this flexibility useful? It would be very helpful for me to see some sort of demonstration of the framework.\n\nHow is the computation cost/overhead for the framework? You mention that the method is exponential in the number of stochastic nodes. I assume this is unavoidable if you want multiple samples per stochastic node, but it would be very helpful to understand if the framework adds any additional overhead over using one of the gradient estimators in a standard way. Having a case study section would also help answering this question.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4 hours"}, "extra_text": {}}, {"review_id": "Ykbscs0QhWD", "reviewer": "Reviewer_6cXR", "summary": "This paper introduces a framework for automatic differentiation for stochastic computation graphs. A major step in stochastic computation graph is the sampling step and this can happen when modelers deal with expectation terms in models such as reinforcement learning or variational inference. This framework provides an unbiased estimator for any order gradient and also generalizes the variance reduction technique for those estimates.  For gradient estimation, it needs a 4-tuple including proposal distribution-to sample from-, weighting function-to weight samples-, gradient function- unbiased gradient estimator- and control variate -to reduce variance-. Based on this 4-tuple it defines a surrogate loss to train the parameters of the graph. On the theoretical side, it shows that its estimator is unbiased and on the practical side this framework is implemented in PyTorch. ", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This work considers an important practical problem that modelers deal with when implementing stochastic models. However, their work is mainly based on previous work such as [13]. \n\nThe major comment I have with regard to this paper is about its clarity. The significance of this problem is not well described for an unfamiliar reader. For example, they could start with a VI model and explain why this is hard to implement in PyTorch. Furthermore, they use the magic box operator from previous work but they don’t explain clearly what it does and why it is important in SCG models. Another part is the section where they explain their surrogate loss. No intuition is given to the reader about how they come up with this loss function. \nAlso, it would be nice to see more detail in the mathematical form about the rule of control variate in reducing the variance in example 3.2.5. \n\nMy questions: \n\n1- Is the proposal distribution fixed through the training or can it be changed? \n\n2-   How can the modeler determine the cardinality of samples for each stochastic node in different partitions? Is it fixed or dynamic? \n\n\n3-  Does this framework allow arbitrary weighting for different cost nodes? Can we train these weights as well? \n\n\n4- Not including the reparameterization trick is a limitation. How much change does your framework need to include this trick? \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "AAxO9EVnMoA", "reviewer": "Reviewer_xpUA", "summary": "The paper presents \"storchastic\", a framework which allows choice and combination of a variety of methods for (potentially higher-order) gradient estimation.  The paper shows how many gradient estimators and control variates can be reconstructed within this framework, and proves that many of the resulting (potentially higher-order) gradient estimators are unbiased given some conditions, and show that control variates can reduce the variance of even higher-order gradient estimates.", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "I am not very confident in my assessment of this paper.  \n\nThe paper describes what sounds like a very useful library, allowing a choice of gradient estimators (proposal distribution, weight function, gradient function, control variate) and facilitating higher-order gradient estimation. The proofs of conditions for unbiased higher-order gradient estimation and any-order variance reduction seem useful and important. \n\nWith the exception of these two proofs, most of the contribution is in the framework / library.  The paper describes the desiderata and approach, but there is not really any qualitative or quantitative evaluation of the library.  It is much more difficult to come up with evaluation/comparison for a paper like this than it is for a method or model which can be measured on some benchmark.\n\nI lean slightly to acceptance as the description of the framework makes the library sound like a useful tool, and as I think we should encourage development of such open-source tools and libraries.  The two proofs are also useful contributions.\n\nHowever I think the paper could be significantly improved if the library/framework could  be demonstrated and could somehow be compared with existing libraries. For example, code snippets could show the ease of implementing different gradient estimators with Storchastic vs in vanilla Pytorch or in Pyro. Experiments showing the efficacy of control variates for higher-order gradient estimation, or showing the efficacy of using different gradient estimators at different stochastic nodes, would also support the premise that this library's functionality is an important improvement.\n\nThe computation being exponential in the size of the sampled sets of values X seems like a major limitation. Do prior libraries also have this limitation? Is it possible to overcome this, at the cost of variance, while still being able to reduce variance with increased sample size? I don't find \"this can be parallelized\" convincing as an argument for exponential computation time not mattering.\n\nNits -\nDon't capitalize \"Deep Learning\", \"Reinforcement Learning\" or \"Variational Inference\", only capitalize proper nouns.\nIs \"Stochastic automatic differentiation\" an accepted term? To me, it seems a bit ambiguous . Stochastic gradient descent, using  AD to compute gradients, sounds like it would fit \"stochastic automatic differentiation\", but it doesn't fit the definition from this paper, which involves gradient estimation in computation graphs where the parameters control the sampling distributions.\nI also think it would be possible to have a debate over whether this is truly still automatic differentiation, given the user needs to specify what estimation methods to use at each node (but this is minor / particularly pedantic).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces a new framework for stochastic gradient estimation that encompasses various existing gradient estimation techniques. The unification of several gradient estimation techniques into a single framework is recognized by the reviewers as a somewhat significant contribution. The reviewers did not raise any severe concern so overall, I think this is a worthwile contribution as a poster.
One reviewer asked for a case study which I would like to authors to include in a revision given that they promised to do so.
Decision: Accept (Poster)