To backpropagate the gradients through stochastic binary layers, we propose the augment-REINFORCE-merge (ARM) estimator that is unbiased, exhibits low variance, and has low computational complexity. Exploiting variable augmentation, REINFORCE, and reparameterization, the ARM estimator achieves adaptive variance reduction for Monte Carlo integration by merging two expectations via common random numbers. The variance-reduction mechanism of the ARM estimator can also be attributed to either antithetic sampling in an augmented space, or the use of an optimal anti-symmetric "self-control" baseline function together with the REINFORCE estimator in that augmented space. Experimental results show the ARM estimator provides state-of-the-art performance in auto-encoding variational inference and maximum likelihood estimation, for discrete latent variable models with one or multiple stochastic binary layers. Python code for reproducible research is publicly available. Given a function f (z) of a random variable z = (z 1 , . . . , z V )T , which follows a distribution q φ (z) parameterized by φ, there has been significant recent interest in estimating φ to maximize (or minimize) the expectation of f (z) with respect to z ∼ q φ (z), expressed as DISPLAYFORM0 In particular, this expectation objective appears in both maximizing the evidence lower bound (ELBO) for variational inference BID12 and approximately maximizing the log marginal likelihood of a hierarchal Bayesian model BID1 , two fundamental problems in statistical inference. To maximize (1), if ∇ z f (z) is tractable to compute and z ∼ q φ (z) can be generated via reparameterization as z = T φ ( ), ∼ p( ), where are random noises and T φ (·) denotes a deterministic transform parameterized by φ, then one may apply the reparameterization trick BID14 BID27 to compute the gradient as DISPLAYFORM1 This trick, however, is often inapplicable to discrete random variables, as widely used to construct discrete latent variable models such as sigmoid belief networks BID22 BID31 .To maximize (1) for discrete z, using the score function ∇ φ log q φ (z) = ∇ φ q φ (z)/q φ (z), one may compute ∇ φ E(φ) via REINFORCE BID38 as its high Monte-Carlo-integration variance often limits its use in practice. Note that if f (z) depends on φ, then we assume it is true that E z∼q φ (z) [∇ φ f (z)] = 0. For example, in variational inference, we need to maximize the ELBO as E z∼q φ (z) [f (z)], where f (z) = log[p(x | z)p(z)/q φ (z)]. In this case, although f (z) depends on φ, as E z∼q φ (z) [∇ φ log q φ (z)] = ∇ φ q φ (z)dz = ∇ φ q φ (z)dz = 0, we have E z∼q φ (z) [∇ φ f (z)] = 0.To address the high-variance issue, one may introduce an appropriate baseline (a.k.a. control variate) to reduce the variance of REINFORCE BID24 BID26 BID19 BID9 BID20 BID29 BID21 . Alternatively , one may first relax the discrete random variables with continuous ones and then apply the reparameterization trick to estimate the gradients, which reduces the variance of Monte Carlo integration at the expense of introducing bias BID11 . Combining both REINFORCE and the continuous relaxation of discrete random variables, REBAR of BID35 and RELAX of BID7 both aim to produce a low-variance and unbiased gradient estimator by introducing a continuous relaxation based baseline function, whose parameters, however, need to be estimated at each mini-batch by minimizing the sample variance of the estimator with stochastic gradient descent (SGD). Estimating the baseline parameters often clearly increases the computation. Moreover, the potential conflict, between minimizing the sample variance of the gradient estimate and maximizing the expectation objective, could slow down or even prevent convergence and increase the risk of overfitting. Another interesting variance-control idea applicable to discrete latent variables is using local expectation gradients, which estimates the gradients based on REINFORCE, by performing Monte Carlo integration using a single global sample together with exact integration of the local variable for each latent dimension BID34 .Distinct from the usual idea of introducing baseline functions and optimizing their parameters to reduce the estimation variance of REINFORCE, we propose the augment-REINFORCE-merge (ARM) estimator, a novel unbiased and low-variance gradient estimator for binary latent variables that is also simple to implement and has low computational complexity. We show by rewriting the expectation with respect to Bernoulli random variables as one with respect to augmented exponential random variables, and then expressing the gradient as an expectation via REINFORCE, one can derive the ARM estimator in the augmented space with the assistance of appropriate reparameterization. In particular, in the augmented space, one can derive the ARM estimator by using either the strategy of sharing common random numbers between two expectations, or the strategy of applying antithetic sampling. Both strategies, as detailedly discussed in BID23 , can be used to explain why the ARM estimator is unbiased and could lead to significant variance reduction. Moreover, we show that the ARM estimator can be considered as improving the REINFORCE estimator in an augmented space by introducing an optimal baseline function subject to an anti-symmetric constraint; this baseline function can be considered as a "self-control" one, as it exploits the function f itself and correlated random noises for variance reduction, and adds no extra parameters to learn. This "self-control" feature makes the ARM estimator distinct from both REBAR and RELAX, which rely on minimizing the sample variance of the gradient estimate to optimize the baseline function.We perform experiments on a representative toy optimization problem and both auto-encoding variational inference and maximum likelihood estimation for discrete latent variable models, with one or multiple binary stochastic layers. Our extensive experiments show that the ARM estimator is unbiased, exhibits low variance, converges fast, has low computation, and provides state-of-the-art out-of-sample prediction performance for discrete latent variable models, suggesting the effectiveness of using the ARM estimator for gradient backpropagation through stochastic binary layers. Python code for reproducible research is available at https://github.com/mingzhang-yin/ARM-gradient. To train a discrete latent variable model with one or multiple stochastic binary layers, we propose the augment-REINFORCE-merge (ARM) estimator to provide unbiased and low-variance gradient estimates of the parameters of Bernoulli distributions. With a single Monte Carlo sample, the estimated gradient is the product of uniform random noises and the difference of a function of two vectors of correlated binary latent variables. Without relying on estimating a baseline function with extra learnable parameters for variance reduction, it maintains efficient computation and avoids increasing the risk of overfitting. Applying the ARM gradient leads to not only fast convergence, but also low test negative log-likelihoods (and low test negative evidence lower bounds for variational inference), on both auto-encoding variational inference and maximum likelihood estimation for stochastic binary feedforward neural networks. Some natural extensions of the proposed ARM estimator include generalizing it to multivariate categorical latent variables, combining it with a baseline or local-expectation based variance reduction method, and applying it to reinforcement learning whose action space is discrete.Initialize w1:T , ψ randomly; while not converged do Sample a mini-batch of x from data; DISPLAYFORM0 ) T ∇w t Tw t (bt−1) ; end wt = wt + ρtgw t with step-size ρt end ψ = ψ + ηt∇ ψ f (b1:T ; ψ) with step-size ηt end
An unbiased and low-variance gradient estimator for discrete latent variable models Proposes a new variance-reduction technique to use when computing an expected loss gradient where the expectation is with respect to independent binary random variables. An algorithm combining Rao-Blackwellization and common random numbers for lowering the variance of the score-function gradient estimator in the special case of stochastic binary networks An unbiased and low variance augment-REINFORCE-merge (ARM) estimator for calculating and backpropagating gradients in binary neural networks