Storchastic: A Framework for General Stochastic Automatic Differentiation

1 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 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 . 2 Background . 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 . 2.1 Stochastic Computation Graphs . 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 . 2.2 Problem Statement . 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 ] . 2.3 Example : Discrete Variational Autoencoder . 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 can not 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 can not apply the reparameterization method here without introducing bias . 2.4 Formalizing AD libraries and DiCE . 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 .

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