# Storchastic: A Framework for General Stochastic Automatic Differentiation

arXiv: https://arxiv.org/abs/2104.00428 · HTML: https://mlsanity.com/paper/2104.00428/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 62 references that are in the map.
- First posted on arXiv 1 Apr 2021.

## Builds on (most cited first, 30 of 34)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Attention, Learn to Solve Routing Problems!](https://mlsanity.com/paper/1803.08475/) (2018)
- [Self-critical Sequence Training for Image Captioning](https://mlsanity.com/paper/1612.00563/) (2016)
- [Pyro: Deep Universal Probabilistic Programming](https://mlsanity.com/paper/1810.09538/) (2018)
- [Meta-SGD: Learning to Learn Quickly for Few-Shot Learning](https://mlsanity.com/paper/1707.09835/) (2017)
- [Black Box Variational Inference](https://mlsanity.com/paper/1401.0118/) (2013)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [Learning with Opponent-Learning Awareness](https://mlsanity.com/paper/1709.04326/) (2017)
- [Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement](https://mlsanity.com/paper/1903.06059/) (2019)
- [Gradient Estimation Using Stochastic Computation Graphs](https://mlsanity.com/paper/1506.05254/) (2015)
- [Neural Variational Inference and Learning in Belief Networks](https://mlsanity.com/paper/1402.0030/) (2014)
- [Backpropagation through the Void: Optimizing control variates for black-box gradient estimation](https://mlsanity.com/paper/1711.00123/) (2017)
- [Implicit Reparameterization Gradients](https://mlsanity.com/paper/1805.08498/) (2018)
- [Variational inference for Monte Carlo objectives](https://mlsanity.com/paper/1602.06725/) (2016)
- [REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models](https://mlsanity.com/paper/1703.07370/) (2017)
- [SparseMAP: Differentiable Sparse Structured Inference](https://mlsanity.com/paper/1802.04223/) (2018)
- [Memory Augmented Policy Optimization for Program Synthesis and Semantic Parsing](https://mlsanity.com/paper/1807.02322/) (2018)
- [DiCE: The Infinitely Differentiable Monte-Carlo Estimator](https://mlsanity.com/paper/1802.05098/) (2018)
- [Estimating Gradients for Discrete Random Variables by Sampling without Replacement](https://mlsanity.com/paper/2002.06043/) (2020)
- [DisARM: An Antithetic Gradient Estimator for Binary Latent Variables](https://mlsanity.com/paper/2006.10680/) (2020)
- [Rao-Blackwellized Stochastic Gradients for Discrete Distributions](https://mlsanity.com/paper/1810.04777/) (2018)
- [Efficient Marginalization of Discrete and Structured Latent Variables via Sparsity](https://mlsanity.com/paper/2007.01919/) (2020)
- [Credit Assignment Techniques in Stochastic Computation Graphs](https://mlsanity.com/paper/1901.01761/) (2019)
- [Loaded DiCE: Trading off Bias and Variance in Any-Order Score Function Estimators for Reinforcement Learning](https://mlsanity.com/paper/1909.10549/) (2019)
- [ARSM: Augment-REINFORCE-Swap-Merge Estimator for Gradient Backpropagation Through Categorical Variables](https://mlsanity.com/paper/1905.01413/) (2019)

## Built on by (most cited first, 11 of 11)

- [A-NeSI: A Scalable Approximate Method for Probabilistic Neurosymbolic Inference](https://mlsanity.com/paper/2212.12393/) (2022)
- [Automatic Differentiation of Programs with Discrete Randomness](https://mlsanity.com/paper/2210.08572/) (2022)
- [Sparse Graph Learning from Spatiotemporal Time Series](https://mlsanity.com/paper/2205.13492/) (2022)
- [Neurosymbolic Diffusion Models](https://mlsanity.com/paper/2505.13138/) (2025)
- [ULLER: A Unified Language for Learning and Reasoning](https://mlsanity.com/paper/2405.00532/) (2024)
- [Differentiating Metropolis-Hastings to Optimize Intractable Densities](https://mlsanity.com/paper/2306.07961/) (2023)
- [Branches of a Tree: Taking Derivatives of Programs with Discrete and Branching Randomness in High Energy Physics](https://mlsanity.com/paper/2308.16680/) (2023)
- [Probabilistic Programming with Programmable Variational Inference](https://mlsanity.com/paper/2406.15742/) (2024)
- [Generalizing Stochastic Smoothing for Differentiation and Gradient Estimation](https://mlsanity.com/paper/2410.08125/) (2024)
- [Auto-Differentiation of Relational Computations for Very Large Scale Machine Learning](https://mlsanity.com/paper/2306.00088/) (2023)
- [GradInf: Gradient Estimation as Probabilistic Inference](https://mlsanity.com/paper/2607.07840/) (2026)

## Abstract (arXiv)

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 https://github.com/HEmile/storchastic.