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

1 INTRODUCTION . Given a function f ( z ) of a random variable z = ( z1 , . . . , zV ) 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 E ( φ ) = ∫ f ( z ) qφ ( z ) dz = Ez∼qφ ( z ) [ f ( z ) ] . ( 1 ) In particular , this expectation objective appears in both maximizing the evidence lower bound ( ELBO ) for variational inference ( Jordan et al. , 1999 ; Blei et al. , 2017 ) and approximately maximizing the log marginal likelihood of a hierarchal Bayesian model ( Bishop , 1995 ) , two fundamental problems in statistical inference . To maximize ( 1 ) , if ∇zf ( 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 ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) to compute the gradient as ∇φE ( φ ) = ∇φE ∼p ( ) [ f ( Tφ ( ) ) ] = E ∼p ( ) [ ∇φf ( Tφ ( ) ) ] . ( 2 ) This trick , however , is often inapplicable to discrete random variables , as widely used to construct discrete latent variable models such as sigmoid belief networks ( Neal , 1992 ; Saul et al. , 1996 ) . To maximize ( 1 ) for discrete z , using the score function∇φ log qφ ( z ) = ∇φqφ ( z ) /qφ ( z ) , one may compute ∇φE ( φ ) via REINFORCE ( Williams , 1992 ) as ∇φE ( φ ) = Ez∼qφ ( z ) [ f ( z ) ∇φ log qφ ( z ) ] ≈ 1 K ∑K k=1 f ( z ( k ) ) ∇φ log qφ ( z ( k ) ) , where z ( k ) iid∼ qφ ( z ) are independent , and identically distributed ( iid ) . This unbiased estimator is also known as ( a.k.a . ) the score-function ( Fu , 2006 ) or likelihood-ratio estimator ( Glynn , 1990 ) . While it is unbiased and only requires drawing iid random samples from qφ ( z ) and computing∇φ log qφ ( z ( k ) ) , ∗Corresponding author . 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 Ez∼qφ ( z ) [ ∇φf ( z ) ] = 0 . For example , in variational inference , we need to maximize the ELBO as Ez∼qφ ( z ) [ f ( z ) ] , where f ( z ) = log [ p ( x | z ) p ( z ) /qφ ( z ) ] . In this case , although f ( z ) depends on φ , as Ez∼qφ ( z ) [ ∇φ log qφ ( z ) ] = ∫ ∇φqφ ( z ) dz = ∇φ ∫ qφ ( z ) dz = 0 , we have Ez∼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 ( Paisley et al. , 2012 ; Ranganath et al. , 2014 ; Mnih & Gregor , 2014 ; Gu et al. , 2016 ; Mnih & Rezende , 2016 ; Ruiz et al. , 2016 ; Kucukelbir et al. , 2017 ; Naesseth et al. , 2017 ) . 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 ( Maddison et al. , 2017 ; Jang et al. , 2017 ) . Combining both REINFORCE and the continuous relaxation of discrete random variables , REBAR of Tucker et al . ( 2017 ) and RELAX of Grathwohl et al . ( 2018 ) 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 ( Titsias & Lázaro-Gredilla , 2015 ) . 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 Owen ( 2013 ) , 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 . 2 ARM : AUGMENT-REINFORCE-MERGE ESTIMATOR . In this section , we first present the key theorem of the paper , and then provide its derivation . With this theorem , we summarize ARM gradient ascent for multivariate binary latent variables in Algorithm 1 , as shown in Appendix A . Let us denote σ ( φ ) = eφ/ ( 1 + eφ ) as the sigmoid function and 1 [ · ] as an indicator function that equals to one if the argument is true and zero otherwise . Theorem 1 ( ARM ) . For a vector of V binary random variables z = ( z1 , . . . , zV ) T , the gradient of E ( φ ) = Ez∼∏Vv=1 Bernoulli ( zv ; σ ( φv ) ) [ f ( z ) ] ( 3 ) with respect to φ = ( φ1 , . . . , φV ) T , the logits of the Bernoulli probability parameters , can be expressed as ∇φE ( φ ) = Eu∼∏Vv=1 Uniform ( uv ; 0,1 ) [ ( f ( 1 [ u > σ ( −φ ) ] ) − f ( 1 [ u < σ ( φ ) ] ) ) ( u− 1/2 ) ] , ( 4 ) where 1 [ u > σ ( −φ ) ] : = ( 1 [ u1 > σ ( −φ1 ) ] , . . . ,1 [ uV > σ ( −φV ) ] ) T . 2.1 UNIVARIATE ARM ESTIMATOR . Below we will first present the ARM estimator for a univariate binary latent variable ( i.e. , V = 1 ) , and then generalize it to a multivariate one ( i.e. , V > 1 ) . In the univariate case , we need to evaluate the gradient of E ( φ ) = Ez∼Bernoulli ( σ ( φ ) ) [ f ( z ) ] with respect to φ , which has an analytic expression as ∇φE ( φ ) = ∇φ [ σ ( φ ) f ( 1 ) + σ ( −φ ) f ( 0 ) ] = σ ( φ ) σ ( −φ ) [ f ( 1 ) − f ( 0 ) ] . ( 5 ) Since ∫ σ ( φ ) 0 ( 1− 2u ) du = σ ( φ ) σ ( −φ ) and ∫ 1 σ ( φ ) ( 1− 2u ) du = −σ ( φ ) σ ( −φ ) , we can rewrite ( 5 ) as ∇φE ( φ ) = ∫ σ ( φ ) 0 f ( 1 ) ( 1− 2u ) du+ ∫ 1 σ ( φ ) f ( 0 ) ( 1− 2u ) du = ∫ 1 0 f ( 1 [ u < σ ( φ ) ] ) ( 1− 2u ) du = Eu∼Uniform ( 0,1 ) [ f ( 1 [ u < σ ( φ ) ] ) ( 1− 2u ) ] . ( 6 ) We refer to ( 6 ) as the univariate augment-REINFORCE ( AR ) estimator , as we initially derived it by combining variable augmentation ( Tanner & Wong , 1987 ; Van Dyk & Meng , 2001 ) , REINFORCE , and Rao Blackwellization ( Casella & Robert , 1996 ) ; we defer the details to Appendix B . Applying antithetic sampling ( Owen , 2013 ) to the AR estimator in ( 6 ) , with ũ = 1− u , we have ∇φE ( φ ) = Eu∼Uniform ( 0,1 ) [ f ( 1 [ u < σ ( φ ) ] ) ( 1/2− u ) ] + Eũ∼Uniform ( 0,1 ) [ f ( 1 [ ũ < σ ( φ ) ] ) ( 1/2− ũ ) ] = Eu∼Uniform ( 0,1 ) [ f ( 1 [ u < σ ( φ ) ] ) ( 1/2− u ) + f ( 1 [ ũ < σ ( φ ) ] ) ( 1/2− ũ ) ] = Eu∼Uniform ( 0,1 ) [ ( f ( 1 [ u > σ ( −φ ) ] ) − f ( 1 [ u < σ ( φ ) ] ) ) ( u− 1/2 ) ] , ( 7 ) which provides the proof for Theorem 1 for V = 1 . We refer to ( 7 ) as the univariate ARM estimator , as we initially derived it by improving the AR estimator with an additional merge step , which shares random numbers between two expectations to reduce the variance of Monte Carlos integration ; we defer the details to Appendix B . Note that since Eu∼Uniform ( 0,1 ) [ f ( 1 [ u < σ ( φ ) ] ) ( 1/2−u ) ] = −Eu∼Uniform ( 0,1 ) [ f ( 1 [ u > σ ( −φ ) ] ) ( 1/2−u ) ] , we have Eu∼Uniform ( 0,1 ) [ ( f ( 1 [ u < σ ( φ ) ] ) + f ( 1 [ u > σ ( −φ ) ] ) ) ( 1/2− u ) ] = 0 , subtracted which from the AR estimator in ( 6 ) leads to the ARM estimator in ( 7 ) . For this reason , we can also consider the ARM estimator as the AR estimator subtracted by a zero-mean baseline function as b ( u ) = ( f ( 1 [ u < σ ( φ ) ] ) + f ( 1 [ u > σ ( −φ ) ] ) ) ( 1/2− u ) . This baseline function is distinct from previously proposed ones in being parameterized by the function f itself over two correlated binary latent variables and satisfying b ( u ) = −b ( 1− u ) . From this point of view , the ARM estimator can be considered as a “ self-control ” gradient estimator that exploits the function f itself to control the variance of Monte Carlo integration .

Summary: In this paper the authors propose a new variance-reduction technique to use when computing an expected loss gradient where the expectation is with respect to independent binary random variables, e.g. for training VAEs with a discrete latent space. The paper is interesting, highly relevant, simple to implement, suggests many possible extensions, and shows good results on the experiments performed. However the exposition leaves a lot to be desired.