ARM: Augment-REINFORCE-Merge Gradient for Stochastic Binary Networks

Abstract

To backpropagate the gradients through discrete stochastic layers, we encode the true gradients into a multiplication between random noises and the difference of the same function of two different sets of discrete latent variables, which are correlated with these random noises. The expectations of that multiplication over iterations are zeros combined with spikes from time to time. To modulate the frequencies, amplitudes, and signs of the spikes to capture the temporal evolution of the true gradients, we propose the augment-REINFORCE-merge (ARM) estimator that combines data augmentation, the score-function estimator, permutation of the indices of latent variables, and variance reduction for Monte Carlo integration using common random numbers. The ARM estimator provides low-variance and unbiased gradient estimates for the parameters of discrete distributions, leading to state-of-the-art performance in both auto-encoding variational Bayes and maximum likelihood inference, for discrete latent variable models with one or multiple discrete stochastic layers.

Low-Variance Black-Box Gradient Estimates for the Plackett-Luce Distribution

“…Linderman et al (2018) relaxes the discrete set of permutations to Birkhoff polytope, the set of doubly-stochastic matrices, and extend stick-breaking approach (Sethuraman 2017) and derive control variate for black-box function optimization combining the REINFORCE estimator and reparametrization trick. Yin and Zhou (2018) propose gradient estimator that estimates the gradients of discrete distribution parameters in an augmented space.…”

“…Despite the recent breakthroughs in gradient estimation for continuous latent variables (Kingma and Welling 2013;Rezende, Mohamed, and Wierstra 2014;Mohamed et al 2019), gradient estimation for discrete latent variables remains a challenge. Currently, general-purpose estimators (Williams 1992;Mnih and Gregor 2014) remain unreliable and the state-of-the-art methods (Tucker et al 2017;Grathwohl et al 2018;Yin and Zhou 2018) exclusively consider the categorical distribution. Although the reduction to the categorical case allows benefiting from gradient estimators for continuous relaxations, such solutions are hard to translate to discrete distributions with large support.…”

GO Gradient for Expectation-Based Objectives

“…Attracted by the generalization of REINFORCE, many works try to improve its performance via efficient variance-reduction techniques, like control variants (Mnih & Gregor, 2014;Titsias & Lázaro-Gredilla, 2015;Gu et al, 2015;Mnih & Rezende, 2016;Tucker et al, 2017;Grathwohl et al, 2017) or via data augmentation and permutation techniques (Yin & Zhou, 2018). Most of this research focuses on discrete random variables, likely because Rep (if it exists) works well for continuous random variables but it may not exist for discrete random variables.…”

Improved Gradient-Based Optimization Over Discrete Distributions

“…SF suffers from high variance and many remedies have been proposed to reduce this variance (Mnih & Gregor (2014); Gregor et al (2013); Gu et al (2015); Mnih & Rezende (2016); Tucker et al (2017); Grathwohl et al (2017)). Unbiased estimators that require multiple function evaluations have also been proposed Tokui & Sato (2017), Titsias & Lázaro-Gredilla (2015), Lorberbom et al (2018), Yin & Zhou (2018). These estimators have lower variance but can be computationally demanding.…”

“…FD estimators are computationally expensive and require multiple function evaluations per gradient. A notable exception is the Augment-REINFORCE-Merge (ARM) estimator introduced in Yin & Zhou (2018). ARM provides an unbiased estimate using only two function evaluations for the factorized multivariate distribution, regardless of the number of variables.…”