Stochastic Automatic Differentiation library for PyTorch.
GPL-3.0 License
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.
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},
booktitle = {Advances in Neural Information Processing Systems},
editor = {M. Ranzato and A. Beygelzimer and Y. Dauphin and P.S. Liang and J. Wortman Vaughan},
pages = {7574--7587},
url = {https://proceedings.neurips.cc/paper_files/paper/2021/file/3dfe2f633108d604df160cd1b01710db-Paper.pdf},
volume = {34},
year={2021}
}