Storchastic is a PyTorch library for stochastic gradient estimation in Deep Learning [1]. Many state of the art deep learning models use gradient estimation, in particular within the fields of Variational Inference and Reinforcement Learning. While PyTorch computes gradients of deterministic computation graphs automatically, it will not estimate gradients on stochastic computation graphs [2].
With Storchastic, you can easily define any stochastic deep learning model and let it estimate the gradients for you. Storchastic provides a large range of gradient estimation methods that you can plug and play, to figure out which one works best for your problem. Storchastic provides automatic broadcasting of sampled batch dimensions, which increases code readability and allows implementing complex models with ease.
When dealing with continuous random variables and differentiable functions, the popular reparameterization method [3] is usually very effective. However, this method is not applicable when dealing with discrete random variables or non-differentiable functions. This is why Storchastic has a focus on gradient estimators for discrete random variables, non-differentiable functions and sequence models.
Documentation on Read the Docs.
Example: Discrete Variational Auto-Encoder
In your virtual Python environment, run
pip install storchastic
Requires Pytorch 1.8 and Pyro. The code is build using Python 3.8.
Feel free to create an issue if an estimator is missing here.
- Reparameterization [1, 3]
- REINFORCE with Moving Average baseline [1, 4]
- REINFORCE with Leave-One-Out baseline (RLOO) [5, 6]
- Expected value for enumerable distributions
- (Straight through) Gumbel Softmax [7, 8]
- LAX, RELAX [9]
- REBAR [10]
- REINFORCE Without Replacement [6]
- Unordered Set Estimator [13]
- ARM [15]
- Rao-Blackwellized REINFORCE [12]
- Memory Augmented Policy Optimization [11]
- Measure valued derivatives [1, 14]
- Automatic Credit Assignment [16]
- ...
- [1] Monte Carlo Gradient Estimation in Machine Learning, Mohamed et al, 2019
- [2] Gradient Estimation Using Stochastic Computation Graphs, Schulman et al, NeurIPS 2015
- [3] Auto-Encoding Variational Bayes, Kingma and Welling, ICLR 2014
- [4] Simple statistical gradient-following algorithms for connectionist reinforcement learning, Williams, Machine Learning 1992
- [5] Variational inference for Monte Carlo objectives, Mnih and Rezende, ICML 2016
- [6] Buy 4 REINFORCE Samples, Get a Baseline for Free!, Kool et al, ICLR Workshop dlStructPred 2019
- [7] Categorical Reparameterization with Gumbel-Softmax, Jang et al, ICLR 2017
- [8] The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables, Maddison et al, ICLR 2017
- [9] Backpropagation through the Void: Optimizing control variates for black-box gradient estimation, Grathwohl et al, ICLR 2018
- [10] REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models, Tucker et al, NeurIPS 2017
- [11] Memory Augmented Policy Optimization for Program Synthesis and Semantic Parsing, Liang et al, NeurIPS 2018
- [12] Rao-Blackwellized Stochastic Gradients for Discrete Distributions, Liu et al, ICML 2019
- [13] Estimating Gradients for Discrete Random Variables by Sampling without Replacement, Kool et al, ICLR 2020
- [14] Measure-Valued Derivatives for Approximate Bayesian Inference, Rosca et al, Workshop on Bayesian Deep Learning (NeurIPS 2019)
- [15] ARM: Augment-REINFORCE-Merge Gradient for Stochastic Binary Networks, Yin and Zhou, ICLR 2019
- [16] Credit Assignment Techniques in Stochastic Computation Graphs, Weber et al, AISTATS 2019
To cite Storchastic, please cite this preprint:
@article{van2021storchastic,
title={Storchastic: A Framework for General Stochastic Automatic Differentiation},
author={van Krieken, Emile and Tomczak, Jakub M and Teije, Annette ten},
journal={arXiv preprint arXiv:2104.00428},
year={2021}
}