Storchastic: A Framework for General Stochastic Automatic Differentiation
Year: 2021

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 any-order derivative estimates. Finally, we implement Storchastic as a PyTorch library at github.com/HEmile/storchastic.

heading: 1 Introduction
text: 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 reparameterization [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 uses the score function estimator to create a surrogate loss that provides gradient estimates when differentiated. While this approach is more general as 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 a new framework called Storchastic to 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 gradient 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 (NeurIPS 2021).
ing function, gradient function and control variate. We can use this decomposition to 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 satisfy certain conditions, performing n-th order differentiation on the Storchastic surrogate loss gives unbiased estimates of the n-th order derivative of the stochastic computation graph. We show these conditions hold for a wide variety of gradient estimation methods for first order differentiation. For many score function-based methods like RELAX [16], MAPO [28] and the unordered set estimator [25], the conditions also hold for any-order differentiation. In Storchastic, we only have to prove these conditions locally. This means that modelers are free to choose the gradient estimation method that best suits each sampling step, while guaranteeing that the gradient remains unbiased. Storchastic 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 are the following:
1. We introduce Storchastic, a new framework for general stochastic AD that uses four gradient 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, we introduce a mathematical formalization of forward-mode evaluation in AD libraries in Section 2.4.
3. We derive a technique for extending variance reduction using control variates to any-order derivative estimation in Section 3.5.
4. We implement Storchastic as an open source library for PyTorch, Section 3.7.

heading: 2 Background
text: We use capital letters N,F, S1, ..., Sk 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 xi.

heading: 2.1 Stochastic Computation Graphs
text: We start by introducing Stochastic Computation Graphs (SCGs) [45], which is a formalism for stochastic AD. A Stochastic Computation Graph (SCG) is a directed acyclic graph (DAG) G = (N , E) where nodes N are partitioned in stochastic nodes S and deterministic nodes F . We define the set of parameters Θ ⊆ F such that all θ ∈ Θ have no incoming edges, and the set of cost nodes C ⊆ F such that all c ∈ C have no outgoing edges. The set of parents pa(N) is the set of incoming nodes of a node N ∈ N , that is pa(N) = {M ∈ N|(M,N) ∈ E}. Each stochastic node S ∈ S represents a random variable with sample space ΩS and probability distribution pS conditioned on its parents. Each deterministic node F represents a (deterministic) function fF of its parents.
M influences N , denoted M ≺ N , if there is a directed path from M to N . We denote with N≺N = {M ∈ N|M ≺ N} the set of nodes that influence N . The joint probability of all random variables xS ∈ ∏ S∈S ΩS is defined as p(xS) = ∏ S∈S pS(xS |xpa(S)), where xpa(S) is the set of values of the nodes pa(S). The expected value of a deterministic node F ∈ F is its expected value over sampling stochastic nodes that influence that node, that is,
E[F ] = ES≺F [fF (pa(F ))] = ∫
ΩS≺F
p(x≺S)fF (xpa(F ))dxS≺F .

heading: 2.2 Problem Statement
text: In this paper, we aim to define a surrogate loss that, when differentiated using an AD library, gives an unbiased estimate of the n-th order derivative of a parameter θ with respect to the expected total cost∇(n)θ E[ ∑ C∈C C]. This gradient can be written as ∑ C∈C ∇ (n) θ E[C], and we focus on estimating the gradient of a single cost node∇(n)θ E[C].

heading: 2.3 Example: Discrete Variational Autoencoder
text: Next, we introduce a running example: A variational autoencoder (VAE) with a discrete latent space [20, 19]. First, we represent the model as an SCG: The deterministic nodes are N = {φ, θ, qz, px, `KLD, `Rec} and the stochastic nodes are S = {x, z}. These are connected as shown in Figure 2. The parameters are Θ = {θ, φ} which respectively are the parameters of the variational posterior q and the model likelihood p, and the cost nodes C = {`KLD, `Rec} that represent the KL-divergence between the posterior and the prior, and the ‘reconstruction loss’, or the model log-likelihood after decoding the sample z. Finally, qz represents the parameters of the multivariate categorical distribution of the amortized variational posterior qφ(z|x). This SCG represents the equation Ex,z[`KLD + `Rec] = Ex[`KLD] + Ex,z∼qφ(z|x)[`Rec]. The problem we are interested in is estimating the gradients of these expectations with respect to the parameters. Since x is not influenced by the parameters, we have ∇(n)θ Ex[`KLD] = 0
and ∇(n)φ Ex[`KLD] = Ex[∇ (n) φ `KLD]. The second term is more challenging. We can move the gradient with respect to θ in, since z is not influenced by θ: ∇(n)θ Ex,z∼qφ(z|x)[`Rec] = Ex,z∼qφ(z|x)[∇ (n) θ `Rec]. However, we cannot compute ∇ (n) φ Ex,z∼qφ(z|x)[`Rec] without gradient estimation methods. This is because sampling from qφ(z|x) is dependent on φ. Furthermore, since we are dealing with a discrete stochastic node, we cannot apply the reparameterization method here without introducing bias.

heading: 2.4 Formalizing AD libraries and DiCE
text: To be able to properly formalize and prove the propositions in this paper, we introduce the ‘forwardmode’ operator that simulates forward-mode evaluation using AD libraries. This operator properly handles the common ‘stop-grad’ operator, which ensures that its argument is only evaluated during forward-mode evaluations of the computation graph. It is implemented in Tensorflow and Jax with the name stop gradient [1, 4] and in PyTorch as detach or no grad [39]. ‘stop-grad’ is necessary 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⊥(x) = 0. The forward-mode operator −→, which is denoted as an arrow above the argument it evaluates, acts as an identity function, except that −−→ ⊥(a) = −→a . Additionally, we define the MagicBox operator as
(x) = exp(x−⊥(x)).
Importantly, the definition of −→ implies that −−−−−→ ∇xf(x) does not equal ∇x −−→ f(x) if f contains a stopgrad 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 −−−→ (x) = 1 and∇x (f(x)) = (f(x))∇xf(x). 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-order derivative estimation that defines a surrogate loss using the score function:
∇(n)θ E[ ∑ C∈C C] = E [−−−−−−−−−−−−−−−−−−−−−−−−−−−−→ ∇(n)θ ∑ C∈C ( ∑ S∈S≺C log p(xS |xpa(S)) ) C ] . (1)
DiCE correctly handles the credit assignment problem: The score function is only applied to the stochastic nodes that influence a cost node. It also handles pathwise dependencies of the parameter through cost functions. However, it has high variance since it is based on a straightforward application of the score function.
3 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 a suitable gradient estimation method for each stochastic node. First, we present 5 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 in 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 PyTorch 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 we used to develop Storchastic.
1. Modelers should be able to choose a different gradient estimation method for each stochastic 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 gradient estimation methods, including score function-based methods with complex sampling techniques [52, 25, 28] or control variates [16, 47], and other methods such as measurevalued derivatives [18, 40] and SPSA [46] which are missing AD implementations [37].
3. Storchastic should define a 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. Differentiating the surrogate loss n times should give estimates of the nth-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.

heading: 3.2 Gradient Estimators in Storchastic
text: Next, we introduce each of the four components and motivate why each is needed to ensure Requirement 2 is satisfied. First, we note that several recent gradient estimators, like MAPO [28], unordered set estimator [25] and self-critical baselines [21, 42] act on sequences of stochastic nodes instead of on a single stochastic node. Therefore, we create a partition S1, ...,Sk of S≺C topologically ordered by the influence relation, and define the shorthand xi := xSi . For each partition Si, we choose a gradient estimator, which is a 4-tuple 〈qi, wi, li, ai〉. Here, q(Xi|x<i) is the proposal distribution, wi(xi) is the weighting function, li(xi) is the gradient function and ai is the control variate.

heading: 3.2.1 Proposal distribution
text: Many gradient estimation methods in the literature do not sample a single value xi ∼ p(xi|x<i), but sample, often multiple, values from possibly a different distribution. Some instances of sampling schemes are taking multiple i.i.d. samples, importance sampling [32] which is very common in offpolicy 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 q(Xi|x<i) samples a set of values Xi = {xi,1, ...,xi,m} where each xi,j ∈ ΩSi . The sample is conditioned on x<i = ∪S∈Sixpa(S), the values of the parent nodes of the stochastic nodes in Si. This is illustrated in Figure 1.b.

heading: 3.2.2 Weighting function
text: 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 wi : ΩSi → R+. 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].

heading: 3.2.3 Gradient function
text: The gradient function is an unbiased gradient estimator together with the weighting function. It 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.

heading: 3.2.4 Control variate
text: Modelers can use control variates to reduce the variance of gradient estimates [17, 37]. It is a function that has zero-mean when differentiated. Within the context of score functions, a common control variate is a baseline, which is a function that is independent of the sampled value. We also 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.

heading: 3.2.5 Example: Leave-one-out baseline
text: As an example, we show how to formulate the score function with the leave-one-out baseline [35, 22] in Storchastic. This method samples m values with replacement and uses the average of the other values as a baseline.
• Proposal distribution: We use m samples with replacement, which can be formulated as q(Xi|x<i) = ∏m j=1 p(xi,j |x<i).
• Weighting function: Since samples are independent, we use wi(xi) = 1m . • Gradient function: The score-function uses the log-probability li(xi) = log p(xi|x<i). • Control variate: We use ai,j(x<i,j ,Xi) = (1 − (li(xi))) 1m−1 ∑ j′ 6=j fC(x<i,xi,j),
where 1m−1 ∑ j′ 6=j fC(x<i,xi,j) is the leave-one-out baseline. (1 − (li(xi))) is used to ensure the baseline will be subtracted from the cost before multiplication with the gradient function. It will not affect the forward evaluation since −−−−−−−−−→ 1− (li(xi)) evaluates to 0.
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 n times, and then evaluating the result using an AD library, will give unbiased estimates of the n-th order derivative of the parameter θ with respect to the cost C. Furthermore, according to Requirement 1, we assume the modeler has chosen a gradient estimator 〈qi, wi, li, ai〉 for each partition Si, which can all be different. Then the Storchastic surrogate loss is
SLStorch = ∑
x1∈X1
w1(x1) [ a1(x1,Xi) + ∑ x2∈X2 w2(x2) [ (l1(x1))a2(x<2,Xi) + . . .
+ ∑
xk∈Xk
wk(xk) [ ( k−1∑ j=1 lj(xj))ak(x<k,Xi) + ( k∑ i=1 li(xi) ) C ] . . . ]] , (2)
where X1 ∼ q(X1),X2 ∼ q(X2|x1), ...,Xk ∼ q(Xk|x<k).
When this loss is differentiated n times using AD libraries, it will produce unbiased estimates of the n-th derivative, as we will show later. To help understand the Storchastic surrogate loss and why
Algorithm 1 The Storchastic framework: Compute a Monte Carlo estimate of the n-th order gradient given k gradient estimators 〈qi, wi, li, ai〉.
1: function ESTIMATE GRADIENT(n, θ) 2: SLStorch ← SURROGATE LOSS(1, {}, 0) . Compute surrogate loss 3: return −−−−−−−−−→ ∇(n)θ SLStorch . Differentiate and use AD library to evaluate surrogate loss
4: 5: function SURROGATE LOSS(i, x<i, L) 6: if i = k + 1 then 7: return ( L)fC(x≤k) . Use MagicBox to distribute cost
8: Xi ∼ q(Xi|x<i) . Sample from proposal distribution 9: sum← 0 10: for xi ∈ Xi do . Iterate over options in sampled set 11: A← (L)ai(x<i,Xi) . Compute control variate 12: Li ← L+ li(xi) . Compute gradient function 13: Ĉ ← SURROGATE LOSS(i+ 1, x≤i, Li) . Compute surrogate loss for xi 14: sum← sum + wi(xi)(Ĉ +A) . Weight and add 15: return sum
it satisfies the requirements, we will break it down using Algorithm 1. The ESTIMATE GRADIENT function computes the surrogate loss for the SCG, and then differentiates it n ≥ 0 times using the AD library to return an estimate of the n-th order gradient, which should be unbiased according to Requirement 4. If n is set to zero, this returns an estimate of the expected cost.
The SURROGATE LOSS function computes the equation using a 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<i is the set of previously sampled values for partitions S<i, and L is the sum of gradient functions of those previously sampled values. In line 8, we sample a set of values Xi for partition i from q(Xi|x<i). In lines 9 to 14, we compute the sum over values xi in Xi, 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 xi. 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 x<i together with the newly sampled value xi. 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 xi ∈ Xi, meaning that this computation is exponential in the size of the sampled sets of values Xi. For example, the surrogate loss samples |X1| times from q2, one for each value x1 ∈ X1. 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 k 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.

heading: 3.4 Conditions for Unbiased Estimation
text: 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 are 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 is an unbiased estimate of ∇(n)θ E[C] under the following conditions. First, all functions fF corresponding to deterministic nodes F and all probability measures pS corresponding to stochastic nodes S are identical under evaluation. Secondly, for each gradient estimator 〈qi, wi, li, ai〉, i = 1, ..., k, all the following hold for m = 0, ..., n:
1. Eq(Xi|x<i)[ ∑
xi∈Xi −−−−−−−−−−−−−−−−−−−→ ∇(m)θ wi(xi) (li(xi))f(xi)] = −−−−−−−−−−−→ ∇(m)θ ESi [f(xi)] for any deterministic
function f ;
2. Eq(Xi|x<i)[ ∑
xi∈Xi −−−−−−−−−−−−−−−−→ ∇(m)θ wi(xi)ai(x<i,Xi)] = 0;
3. for n ≥ m > 0, Eq(Xi|x<i)[ ∑
xi∈Xi −−−−−−−−→ ∇(m)θ wi(xi)] = 0;
4. −−−−−−→ q(Xi|x<i) = q(Xi|x<i), for all permissible Xi.
The first condition defines a 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 estimate of the gradient for all orders of differentiation m = 0, ..., n. 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
function to be 0 in expectation for orders of differentiation larger than 0. 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 be 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.
A full formalization and the proof of Theorem 1 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 expectation i over the proposal distribution qi is an unbiased estimate of the nth-order derivative over Si conditional on the previous samples. To reduce the multiple sums over gradient functions inside MagicBox, we make use of a property of MagicBox proven in Appendix A: Proposition 4. Summation inside a MagicBox is equivalent under evaluation to multiplication of the arguments in individual MagicBoxes, ie:
(l1(x) + l2(x))f(x) −→≡ (l1(x)) (l2(x))f(x).
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+b = eaeb.

heading: 3.5 Any-order variance reduction using control variates
text: To satisfy Requirement 5, we investigate implementing control variates such that the variance of anyorder 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 a slightly different formulation that we show misses some terms. Theorem 2 (informal). Let Li = ∑i j=1 li. The Storchastic surrogate loss of (2) can equivalently be computed as
SLStorch −→≡ ∑ x1∈X1 ∑ x2∈X2 · · · ∑ xk∈Xk k∏ i=1 wi(xi) k∑ i=1 (Li−1) ( ai(x<i,Xi) + ( (li)− 1)C ) + C.
This gives insight into how control variates are used in Storchastic. They are added to the gradient function, but only during differentiation since −−−−−−→ (Li)− 1 = 0. Furthermore, since both terms are multiplied with (Li−1) (see line 11 of Algorithm 1), both terms correctly distribute over the same any-order derivative terms. By choosing a control variate of the form ai(x<i,Xi) = (1− (li)) · bi, we recover baselines which are common in the context of score functions [13, 37]. For the proof, we use the following proposition also proven in Appendix C: Proposition 7. For orders of differentiation n > 0,
−−−−−−−→ ∇(n)N (Lk) =
−−−−−−−−−−−−−−−−−−−−−→ ∇(n)N k∑ i=1 ( (li)− 1 ) (Li−1).

heading: 3.6 Gradient Estimation Methods
text: In Appendix D we show how several prominent examples of gradient estimation methods in the literature can be formulated using Storchastic, and prove for what orders of 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 baseline [35, 22].
Furthermore, we show how importance sampling, sum-and-sample 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 n ≤ 1 since it relies on a 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-Valued Derivative [18], and show that the last two will only be unbiased for n ≤ 1. Finally, we note that reparameterization [20, 43] can be implemented by transforming the SCG such that the sampling step is outside the path from the parameter to the cost [45].

heading: 3.7 Implementation
text: We implemented Storchastic as an open source PyTorch [39] library 1. To ensure modelers can easily use this library, it automatically handles sets of samples as extra dimensions to PyTorch tensors which allows running multiple sample evaluations in parallel. This approach is illustrated in Figure 1. By making use of PyTorch broadcasting semantics, this allows defining models for simple single-sample computations that are automatically parallelized using Storchastic when using multiple 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.

heading: 3.7.1 Example: Leave-one-out baseline in Discrete Variational Autoencoder
text: As a small case study, we show how to implement the score function with the leave-one-out baseline introduced in Section 3.2.5 for the discrete variational autoencoder introduced in Section 2.3 in PyTorch using Storchastic. While the code listed is simplified, it 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 a 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 ScoreFunctionLOO method defined in Figure 3. Then, we run the training loop: In line 6 we create the stochastic node x by denoting the minibatch dimension as an independent dimension. In line 8 we run the encoder with parameters φ to find the variational posterior qz . We call the gradient estimation method in line 9 to get a sample of z. Note that this interface is independent of gradient estimation method chosen, meaning that if we wanted to compare our implemented method with a 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.
1Code is available at github.com/HEmile/storchastic.
We run this model on our currently implemented set of gradient estimation methods for discrete variables in Appendix E and report the results, which are meant purely to illustrate the case study.

heading: 4 Related Work
text: 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], see Appendix D, and works focused on Reinforcement Learning [49, 29, 36] or Variational Inference [35]. For a recent overview, see [37].
The literature focused on 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 on 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 a 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-like techniques and bootstrapping for SCGs to incorporate reparameterization 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.

heading: 5 Conclusion
text: 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 a large variety of gradient estimation methods from the literature. We also investigated 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 compounds when using different gradient estimation methods. Furthermore, Storchastic could be extended to allow for biased methods. We are also interested in closely analyzing the different components of gradient estimators, both from a theoretical and empirical point of view, to develop new estimators that combine the strengths of estimators in the literature.

Review summaries:
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. 

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.  The 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.  In 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.  The paper shows that under certain reasonable local conditions the obtained estimators are unbiased under any order of differentiation. These conditions are as follows: - for all orders of differentiation the surrogate should be unbiased for any deterministic function under the proposal distribution + weighting - for all orders of differentiation the control variate should be unbiased under the proposal distribution + weighting - the expectation of the weighting function needs to be zero when differentiated - the proposal distribution should be id under forward mode evaluation  The 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. 

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.  The 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.  The 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.