Stochastic Deep Learning for Pytorch

Documentation Status

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


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.


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]


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