Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
main_review: I am not very confident in my assessment of this paper.  

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

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

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

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

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

Nits -
Don't capitalize "Deep Learning", "Reinforcement Learning" or "Variational Inference", only capitalize proper nouns.
Is "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.
I 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).
limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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. 
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]. 

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

My questions: 

1- Is the proposal distribution fixed through the training or can it be changed? 

2-   How can the modeler determine the cardinality of samples for each stochastic node in different partitions? Is it fixed or dynamic? 


3-  Does this framework allow arbitrary weighting for different cost nodes? Can we train these weights as well? 


4- Not including the reparameterization trick is a limitation. How much change does your framework need to include this trick? 

limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: 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.

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.

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

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

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

Nit:
- 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
limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.

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

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

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

Given the limitations, I am borderline on the paper, but currently lean towards acceptance because I think the framework is a nice contribution.
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.

How 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_title: null

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)