Storchastic is a PyTorch library for stochastic gradient estimation in Deep Learning

Stochastic Deep Learning for Pytorch

Documentation on Read the Docs. 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

Installation

pip install storchastic

Requires Pytorch 1.5 (older versions will not do!) and Pyro. The code is build on Python 3.7. The master branch works with PyTorch 1.7, but the version on pip is not compatible. Binaries will be updated soon.

Algorithms

Feel free to create an issue if an estimator is missing here.

• Reparameterization [1, 3]
• Score Function (REINFORCE) with Moving Average baseline [1, 4]
• Score Function with Batch Average Baseline [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]

In development

• Memory Augmented Policy Optimization [11]
• Rao-Blackwellized REINFORCE [12]

Planned

• Measure valued derivatives [1, 14]
• ARM [15]
• Automatic Credit Assignment [16]
• ...

References

• [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