DisARM: An Antithetic Gradient Estimator for Binary Latent Variables
This work addresses a key bottleneck in machine learning for researchers and practitioners dealing with discrete latent variables, offering a more efficient method for gradient estimation.
The paper tackles the challenge of training models with binary latent variables by introducing DisARM, an improved gradient estimator that analytically integrates out randomness from augmentation, reducing variance and outperforming ARM and VIMCO on generative modeling benchmarks with better log-likelihood.
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.