Storchastic: A Framework for General Stochastic Automatic Differentiation
arXiv:2104.00428
Abstract
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.
arXiv comment30 pages, 2 figures, 1 table, accepted in NeurIPS 2021
References in corpus (21)
- PyTorch: An Imperative Style, High-Performance Deep Learning Library
- Auto-Encoding Variational Bayes
- Continuous control with deep reinforcement learning
- Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks
- Categorical Reparameterization with Gumbel-Softmax
- The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
- Meta-SGD: Learning to Learn Quickly for Few-Shot Learning
- Neural Variational Inference and Learning in Belief Networks
- Attention, Learn to Solve Routing Problems!
- Gradient Estimation Using Stochastic Computation Graphs
- REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models
- Variational inference for Monte Carlo objectives
- Implicit Reparameterization Gradients
- ARM: Augment-REINFORCE-Merge Gradient for Stochastic Binary Networks
- Rao-Blackwellized Stochastic Gradients for Discrete Distributions
- ARSM: Augment-REINFORCE-Swap-Merge Estimator for Gradient Backpropagation Through Categorical Variables
- Credit Assignment Techniques in Stochastic Computation Graphs
- Efficient Marginalization of Discrete and Structured Latent Variables via Sparsity
- Total stochastic gradient algorithms and applications in reinforcement learning
- GO Gradient for Expectation-Based Objectives
- Backprop-Q: Generalized Backpropagation for Stochastic Computation Graphs