DisARM: An Antithetic Gradient Estimator for Binary Latent Variables
arXiv:2006.10680
Abstract
Training models with discrete latent variables is challenging due to the difficulty of estimating the gradients accurately. Much of the recent progress has been achieved by taking advantage of continuous relaxations of the system, which are not always available or even possible. The Augment-REINFORCE-Merge (ARM) estimator provides an alternative that, instead of relaxation, uses continuous augmentation. Applying antithetic sampling over the augmenting variables yields a relatively low-variance and unbiased estimator applicable to any model with binary latent variables. However, while antithetic sampling reduces variance, the augmentation process increases variance. We show that ARM can be improved by analytically integrating out the randomness introduced by the augmentation process, guaranteeing substantial variance reduction. Our estimator, DisARM, is simple to implement and has the same computational cost as ARM. We evaluate DisARM on several generative modeling benchmarks and show that it consistently outperforms ARM and a strong independent sample baseline in terms of both variance and log-likelihood. Furthermore, we propose a local version of DisARM designed for optimizing the multi-sample variational bound, and show that it outperforms VIMCO, the current state-of-the-art method.
References in corpus (11)
- Auto-Encoding Variational Bayes
- Categorical Reparameterization with Gumbel-Softmax
- Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
- The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
- Neural Variational Inference and Learning in Belief Networks
- REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models
- Variational inference for Monte Carlo objectives
- Variational Bayesian Inference with Stochastic Search
- Automated Variational Inference in Probabilistic Programming
- ARM: Augment-REINFORCE-Merge Gradient for Stochastic Binary Networks
- ARSM: Augment-REINFORCE-Swap-Merge Estimator for Gradient Backpropagation Through Categorical Variables