ARM: AUGMENT-REINFORCE-MERGE GRADIENT FOR STOCHASTIC BINARY NETWORKS
Year: 2019

Abstract:
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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 ARM: AUGMENT-REINFORCE-MERGE ESTIMATOR
text: 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 .

heading: 2.1 UNIVARIATE ARM ESTIMATOR
text: 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 .

heading: 2.2 MULTIVARIATE GENERALIZATION VIA THE LAW OF TOTAL EXPECTATION
text: Although the ARM estimator for univariate binary is of little use in practice, as the true gradient, shown in (5), has an analytic expression, it serves as the foundation for generalizing it to multivariate settings. Let us denote (·)\v as a vector whose vth element is removed. For the expectation in (3), applying the univariate ARM estimator in (37) together with the law of total expectation, we have
∇φvE(φ) = Ez\v∼∏j 6=v Bernoulli(zj ;σ(φj)){∇φvEzv∼Bernoulli(σ(φv)[f(z)]} = Ez\v∼∏j 6=v Bernoulli(zj ;σ(φj)){Euv∼Uniform(0,1)[(uv − 1/2) × ( f(z\v, zv = 1[uv>σ(−φv)])− f(z\v, zv = 1[uv<σ(φv)]) )]} . (8)
Since z\v ∼ ∏ j 6=v Bernoulli(zj ;σ(φj)) can be equivalently generated as z\v = 1[u\v<σ(φ\v)] or
as z\v = 1[u\v>σ(−φ\v)], where u\v ∼ ∏ j 6=v Uniform(uj ; 0, 1), exchanging the order of the two
expectations in (8) and applying reparameterization, we conclude the proof for (4) of Theorem 1 with ∇φvE(φ) = Euv∼Uniform(0,1) { (uv − 1/2) Ez\v∼∏j 6=v Bernoulli(zj ;σ(φj))[ f(z\v, zv = 1[uv>σ(−φv)])− f(z\v, zv = 1[uv<σ(φv)])
]} = Eu∼∏Vv=1 Uniform(uv ;0,1) [ (uv − 1/2)f(z\v = 1[u\v>σ(−φ\v)], zv = 1[uv>σ(−φv)])
] − Eu∼∏Vv=1 Uniform(uv;0,1) [ (uv − 1/2)f(z\v = 1[u\v<σ(φ\v)], zv = 1[uv<σ(φv)])
] = Eu∼∏Vv=1 Uniform(uv ;0,1) [ (uv − 1/2) ( f(1[u>σ(−φ)])− f(1[u<σ(φ)]) )] . (9)
Alternatively, instead of generalizing the univariate ARM gradient as in (8) and (9), we can first do a multivariate generalization of the univariate AR gradient in (6) as
∇φvE(φ) = Ez\v∼∏j 6=v Bernoulli(zj ;σ(φj)){∇φvEzv∼Bernoulli(σ(φv)[f(z)]} = Ez\v∼∏j 6=v Bernoulli(zj ;σ(φj)){Euv∼Uniform(0,1)[(1− 2uv)f(z\v, zv = 1[uv<σ(φv)])]} = Eu∼∏Vv=1 Uniform(uv ;0,1) [ (1− 2uv)f(1[u<σ(φ)]) ] . (10)
The same as the derivation of the univariate ARM estimator, here we can arrive at (4) from (10) by either adding an antithetic sampling step, or subtracting the AR estimator by a baseline function as
b(u) = ( f(1[u<σ(φ)]) + f(1[u>σ(−φ)]) ) (1/2− u), (11)
which has zero mean, satisfies b(u) = −b(1−u), and is distinct from previously proposed baselines in taking advantage of “self-control” for variance reduction and adding no extra parameters to learn.

heading: 2.3 EFFECTIVENESS OF ARM FOR VARIANCE REDUCTION
text: For the univariate case, we show below that the ARM estimator has smaller worst-case variance than REINFORCE does. The proof is deferred to Appendix C. Proposition 2 (Univariate gradient variance). For the objective function Ez∼Bernoulli(σ(φ))[f(z)], with a single Monte-Carlo sample u ∼ Uniform(0, 1) or z ∼ Bernoulli(σ(φ)), the ARM gradient is expressed as gARM(u, φ) = ( f(1[u>σ(−φ)]) − f(1[u<σ(φ)]) ) (u − 1/2), and the REINFORCE gradient as gR(z, φ) = f(z)∇φ log Bernoulli(z;σ(φ)) = f(z)(z − σ(φ)). Assuming f ≥ 0 (or f ≤ 0), then supφvar[gARM(u,φ)]supφvar[gR(u,φ)] ≤ 16 25 (1− 2 f(0) f(0)+f(1) ) 2 ≤ 1625 . In the general setting, with u(1), . . . ,u(K) iid∼ ∏V v=1 Uniform(0, 1), we define the ARM estimate of ∇φvE(φ) with K Monte Carlo samples, denoted as gARMK ,v, and the AR estimate with 2K Monte Carlo samples, denoted as gAR2K ,v , using
gARMK ,v = 1
2K ∑K k=1(gv(u (k)) + gv(1− u(k))), gAR2K ,v = 12K ∑2K k=1 gv(u (k)), (12)
where gv(u(k)) = f(1[u(k)<σ(φ)])(1− 2u (k) v ). Similar to the analysis in Owen (2013), the amount of variance reduction brought by the ARM estimator can be reflected by the ratio as
var[gARMK ,v] var[gAR2K ,v] = var[gv(u)]− Cov(−gv(u), gv(1− u)) var[gv(u)] = 1−ρv, ρv = Corr(−gv(u), gv(1−u)).
Note−gv(u) = f(1[u<σ(φ)])(2uv−1), gv(1−u) = f(1[u>σ(−φ)])(2uv−1), and P (1[uv<σ(φv)] = 1[uv>σ(−φv)]) = σ(|φv|) − σ(−|φv|). Therefore a strong positive correlation (i.e., ρv → 1) and hence noticeable variance reduction are likely, especially if φv moves far away from zero during training. Concretely, we have the following proposition. Proposition 3 (Variance reduction). For the ARM estimate gARMK ,v and AR estimate gAR2K ,v shown in (12), the variance of gARMK ,v is guaranteed to be lower than that of gARK ,v; moreover, if f ≥ 0 (or f ≤ 0), then the variance of gARMK ,v is guaranteed to be lower than that of gAR2K ,v .
We show below that under the anti-symmetric constraint b(u) = −b(1 − u), which implies that Eu∼∏Vv=1 Uniform(uv ;0,1)[b(u)] is a vector of zeros, Equation (11) is the optimal baseline function to be subtracted from the AR estimator for variance reduction. The proof is deferred to Appendix C.
Proposition 4 (Optimal anti-symmetric baseline). For the gradient of Ez∼qφ(z)[f(z)], the optimal anti-symmetric baseline function to be subtracted from the AR estimator gAR(u) = f(1[u<σ(φ)])(1− 2u), which minimizes the variance of Monte Carlo integration, can be expressed as
arg min bv(u)∈B
var[gAR,v(u)− bv(u)] = 1
2 (gAR,v(u)− gAR,v(1− u)), (13)
where B = {b(u) : bv(u) = −bv(1 − u) for all v} is the set of all zero-mean anti-symmetric baseline functions. Note the optimal baseline function shown in (13) is exactly the same as (11), which is subtracted from the AR estimator in (10) to arrive at the ARM estimator in (4). Corollary 5 (Lower variance than constant baseline). The optimal anti-symmetric baseline function for the AR estimator, as shown in (13) and also in (11), leads to lower estimation variance than any constant based baseline function as bv(u) = cv(1/2−uv), where cv is a dimension-specific constant whose value can be optimized for variance reduction.

heading: 3 BACKPROPAGATION THROUGH DISCRETE STOCHASTIC LAYERS
text: A latent variable model with multiple stochastic hidden layers can be constructed as
x ∼ pθ0(x | b1), b1 ∼ pθ1(b1 | b2), . . . , bt ∼ pθt(bt | bt+1), . . . , bT ∼ pθT (bT ), (14)
whose joint likelihood given the distribution parameters θ0:T = {θ0, . . . ,θT } is expressed as p(x, b1:T |θ0:T ) = pθ0(x | b1) [∏T−1
t=1 pθt(bt | bt+1)
] pθT (bT ). (15)
In comparison to deterministic feedforward neural networks, stochastic ones can represent complex distributions and show natural resistance to overfitting (Neal, 1992; Saul et al., 1996; Tang & Salakhutdinov, 2013; Raiko et al., 2014; Gu et al., 2016; Tang & Salakhutdinov, 2013). However, the training of the network, especially if there are stochastic discrete layers, is often much more challenging. Below we show for both auto-encoding variational inference and maximum likelihood estimation, how to apply the ARM estimator for gradient backpropagation in stochastic binary networks.

heading: 3.1 ARM VARIATIONAL AUTO-ENCODER
text: For auto-encoding variational inference (Kingma & Welling, 2013; Rezende et al., 2014), we construct a variational distribution as
qw1:T (b1:T |x) = qw1(b1 |x) [∏T−1
t=1 qwt+1(bt+1 | bt)
] , (16)
with which the ELBO can be expressed as
E(w1:T ) = Eb1:T∼qw1:T (b1:T |x) [f(b1:T )] , where f(b1:T ) = log pθ0(x | b1) + log pθ1:T (b1:T )− log qw1:T (b1:T |x).
(17)
Proposition 6 (ARM backpropagation). For a stochastic binary network with T binary stochastic hidden layers, constructing a variational auto-encoder (VAE) defined with b0 = x and
qwt(bt | bt−1) = Bernoulli(bt;σ(Twt(bt−1))) (18)
for t = 1, . . . , T , the gradient of the ELBO with respect to wt can be expressed as
∇wtE(w1:T ) = Eq(b1:t−1) [Eut∼Uniform(0,1)[f∆(ut, Twt(bt−1), b1:t−1)(ut − 1/2)]∇wtTwt(bt−1)] , where f∆(ut, Twt(bt−1), b1:t−1) = Ebt+1:T∼q(bt+1:T | bt), bt=1[ut>σ(−Twt (bt−1))] [f(b1:T )]
− Ebt+1:T∼q(bt+1:T | bt), bt=1[ut<σ(Twt (bt−1))] [f(b1:T )]. (19)
The gradient presented in (19) can be estimated with a single Monte Carlo sample as
f̂∆(ut, Twt(bt−1), b1:t−1) =
{ 0, if b(1)t = b (2) t ,
f(b1:t−1, b (1) t:T )− f(b1:t−1, b (2) t:T ), otherwise,
(20)
where b(1)t = 1[ut>σ(−Twt (bt−1))], b (1) t+1:T ∼ q(bt+1:T | b (1) t ), b (2) t = 1[ut<σ(Twt (bt−1))], and b (2) t+1:T ∼ q(bt+1:T | b (2) t ). The proof of Proposition 6 is provided in Appendix C. Suppose the computation complexity of vanilla REINFORCE for a stochastic hidden layer isO(1), which involves a single evaluation of the function f and gradient backpropagation as ∇wtTwt(bt−1), then for a T -stochastic-hidden-layer network, the computation complexity of vanilla REINFORCE is O(T ). By contrast, if evaluating f is much less expensive in computation than gradient backpropagation, then the ARM estimator also has O(T ) complexity, whereas if evaluating f dominates gradient backpropagation in computation, then its worst-case complexity is O(2T ).

heading: 3.2 ARM MAXIMUM LIKELIHOOD ESTIMATION
text: For maximum likelihood estimation, the log marginal likelihood can be expressed as
log pθ0:T (x) = logEb1:T∼pθ1:T (b1:T )[pθ0(x | b1)] ≥ E(θ1:T ) = Eb1:T∼pθ1:T (b1:T )[log pθ0(x | b1)]. (21)
Generalizing Proposition 6 leads to the following proposition. Proposition 7. For a stochastic binary network defined as
pθt(bt | bt+1) = Bernoulli(bt;σ(Tθt(bt+1))), (22)
the gradient of the lower bound in (21) with respect to θt can be expressed as
∇θtE(θ1:T ) = Ep(bt+1:T ) [Eut∼Uniform(0,1)[f∆(ut, Tθt(bt+1), bt+1:T )(ut − 1/2)]∇θtTθt(bt+1)] , where f∆(ut, Tθt(bt+1), bt+1:T ) = Eb1:t−1∼p(b1:t−1 | bt), bt=1[ut>σ(−Tθt (bt+1))])[log pθ0(x | b1)]
− Eb1:t−1∼p(b1:t−1 | bt), bt=1[ut<σ(Tθt (bt+1))])[log pθ0(x | b1)].

heading: 4 EXPERIMENTAL RESULTS
text: To illustrate the working mechanism of the ARM estimator, related to Tucker et al. (2017) and Grathwohl et al. (2018), we consider learning φ to maximize
E(φ) = Ez∼Bernoulli(σ(φ))[(z − p0)2], where p0 ∈ {0.49, 0.499, 0.501, 0.51}.
The optimal solution is σ(φ) = 1[p0<0.5]. The closer p0 is to 0.5, the more challenging the optimization becomes. We compare both the AR and ARM estimators to the true gradient as
gφ = (1− 2p0)σ(φ)(1− σ(φ)) (23)
and three previously proposed unbiased estimators, including REINFORCE, REBAR (Tucker et al., 2017), and RELAX (Grathwohl et al., 2018). Since RELAX is closely related to REBAR in introducing stochastically estimated control variates to improve REINFORCE, and clearly outperforms REBAR in our experiments for this toy problem (as also shown in Grathwohl et al. (2018) for p0 = 0.49), we omit the results of REBAR for brevity. With a single random sample u ∼ Uniform(0, 1) for Monte Carlo integration, the REINFORCE and AR gradients can be expressed as
gφ,REINFORCE = (1[u<σ(φ)] − p0)2(1[u<σ(φ)] − σ(φ)), gφ,AR = (1[u<σ(φ)] − p0)2(1− 2u),
while the ARM gradient can be expressed as gφ,ARM = [ (1[u>σ(−φ)] − p0)2 − (1[u<σ(φ)] − p0)2 ] (u− 1/2). (24)
See Grathwohl et al. (2018) for the details on RELAX.
As shown in Figure 1, the REINFORCE gradients have large variances. Consequently, a REINFORCE based gradient ascent algorithm may diverge if the gradient ascent stepsize is not sufficiently small. For example, when p0 = 0.501, the optimal value for the Bernoulli probability σ(φ) is 0, but the algorithm with 0.1 as the stepsize infers it to be close to 1 at the end of 2000 iterations of a random trial. The AR estimator behaves similarly as REINFORCE does. By contrast, both RELAX and ARM exhibit clearly lower estimation variance. It is interesting to note that the trace plots of the estimated probability σ(φ) with the univariate ARM estimator almost exactly match these with the
-0.01
0.01
Gr ad
ien t
True
0.2
0.2 REINFORCE
0.2
0.2 AR
0.2
0.2 RELAX
0.01
0.01 ARM
0.49 0.499 0.501 0.51
0.00
0.50
1.00
Pro ba
bil ity
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5 1.0 0.49 0.499 0.501 0.51
0 2000 -0.05
0.00
0.05
G_ va
r
0 2000 1.6
1.8
2.0
2.2 1e 2
0 2000 1.6
1.8
2.0
2.2 1e 2
0 2000 0.0
0.5
1.0 1e 3
0 2000
1.00
1.25
1.50
1e 5
0.49 True
Iteration
Figure 1: Comparison of a variety of gradient estimators in maximizing E(φ) = Ez∼Bernoulli(σ(φ))[(z−p0)2] via gradient ascent, where p0 ∈ {0.49, 0.499, 0.501, 0.51}; the optimal solution is σ(φ) = 1(p0 < 0.5). Shown in Rows 1 and 2 are the trace plots of the true/estimated gradients∇φE(φ) and estimated Bernoulli probability parameters σ(φ), with φ updated via gradient ascent. Shown in Row 3 are the gradient variances for p0 = 0.49, estimated using K = 5000 Monte Carlo samples at each iteration; the theoretical gradient variances are also shown if they can be analytically calculated (see Appendices C and D and for related analytic expressions).
true gradients, despite that the trace plots of the ARM gradients are distinct from these of the true gradients. More specifically, while the true gradients smoothly evolve over iterations, the univariate ARM gradients are characterized by zeros and random spikes; this distinct behavior is expected by examining (38) in Appendix C, which suggests that at any given iteration, the univariate ARM gradient based on a single Monte Carlo sample is either exactly zero, which happens with probability σ(|φ|) − σ(−|φ|), or taking |[f(1)− f(0)](1/2− u)| as its absolute value. These observations suggest that by adjusting the frequencies and amplitudes of spike gradients, the univariate ARM estimator very well approximates the behavior of the true gradient for learning with gradient ascent.
In Figure 4 of Appendix D, we plot the gradient estimated with multiple Monte Carlo samples against the true gradient at each iteration, further showing the ARM estimator has the lowest estimation variance given the same number of Monte Carlo samples. Moreover, in Figure 5 of Appendix D, for each estimator specific column, we plot against the value of φ the sample mean ḡ, sample standard deviation sg, and the gradient signal-to-noise ratio defined as SNRg = |ḡ|/sg; for each φ value, we use K = 1000 single-Monte-Carlo-sample gradient estimates to calculate ḡ, sg, and SNRg. Both figures further show that the ARM estimator outperforms not only REINFORCE, which has large variance, but also RELAX, which improves REINFORCE with an adaptively estimated baseline.
In Figure 5 of Appendix D, it is also interesting to notice that the gradient signal-to-noise ratio for the ARM estimator appears to be only a function of φ but not a function of p0; this can be verified to be true using (23) and (39) in Appendix C, as the ratio of the absolute value of the true gradient |gφ| to√
var[gφ,ARM], the standard deviation of the ARM estimate in (24), can be expressed as σ(φ)(1− σ(φ)) /√
1 16 (1− t)(t3 + 7 3 t 2 + 13 t+ 1 3 ), t = σ(|φ|)− σ(−|φ|). (25)
We find that the values of the ratio shown above are almost exactly matched by the values of SNRg = |ḡ|/sg under the ARM estimator, shown in the bottom right subplot of Figure 5. Therefore, for this example optimization problem, the ARM estimator exhibits a desirable property in providing high gradient signal-to-noise ratios regardless of the value of p0.

heading: 4.1 DISCRETE VARIATIONAL AUTO-ENCODERS
text: To optimize a variational auto-encoder (VAE) for a discrete latent variable model, existing solutions often rely on biased but low-variance stochastic gradient estimators (Bengio et al., 2013; Jang et al., 2017), unbiased but high-variance ones (Mnih & Gregor, 2014), or unbiased REINFORCE combined with computationally expensive baselines, whose parameters are estimated by minimizing the sample variance of the estimator with SGD (Tucker et al., 2017; Grathwohl et al., 2018). By contrast, the ARM estimator exhibits low variance and is unbiased, efficient to compute, and simple to implement.
For discrete VAEs, we compare ARM with a variety of representative stochastic gradient estimators for discrete latent variables, including Wake-Sleep (Hinton et al., 1995), NVIL (Mnih & Gregor, 2014), LeGrad (Titsias & Lázaro-Gredilla, 2015), MuProp (Gu et al., 2016), Concrete (GumbelSoftmax) (Jang et al., 2017; Maddison et al., 2017), REBAR (Grathwohl et al., 2018), and RELAX (Tucker et al., 2017). Following the settings in Tucker et al. (2017) and Grathwohl et al. (2018), for the encoder defined in (15) and decoder defined in (16), we consider three different network architectures, as summarized in Table 1, including “Nonlinear” that has one stochastic but two Leaky-ReLU (Maas et al., 2013) deterministic hidden layers, “Linear” that has one stochastic hidden layer, and “Linear two layers” that has two stochastic hidden layers. We consider a widely used binarization (Salakhutdinov & Murray, 2008; Larochelle & Murray, 2011; Yin & Zhou, 2018), referred to as MNIST-static and available at http://www.dmi.usherb.ca/∼larocheh/mlpython/ modules/datasets/binarized mnist.html, making our numerical results directly comparable to those reported in the literature. In addition to MNIST-static, we also consider MNIST-threshold (van den Oord et al., 2017), which binarizes MNIST by thresholding each pixel value at 0.5, and the binarized OMNIGLOT dataset.
We train discrete VAEs with 200 conditionally iid Bernoulli random variables as the hidden units of each stochastic binary layer. We maximize a single-Monte-Carlo-sample ELBO using Adam (Kingma & Ba, 2014), with the learning rate selected from {5, 1, 0.5} × 10−4 by the validation set. We set the batch size as 50 for MNIST and 25 for OMNIGLOT. For each dataset, using its default
training/validation/testing partition, we train all methods on the training set, calculate the validation log-likelihood for every epoch, and report the test negative log-likelihood when the validation negative log-likelihood reaches its minimum within a predefined maximum number of iterations.
We summarize the test negative log-likelihoods in Table 2 for MNIST-static. We also summarize the test negative ELBOs in Table 4 of the Appendix, and provide related trace plots of the training and validation negative ELBOs on MNIST-static in Figure 2, and these on MNIST-threshold and OMNIGLOT in Figures 6 and 7 of the Appendix, respectively. For these trace plots, for a fair comparison of convergence speed between different algorithms, we use publicly available code and setting the learning rate of ARM the same as that selected by RELAX in Grathwohl et al. (2018). Note as shown in Figures 2(a,d) and 7(a,d), both REBAR and RELAX exhibit clear signs of overfitting on both MNIST-static and Omniglot using the “Nonlinear” architecture; as ARM runs much faster per iteration than both of them and do not exhibit overfitting given the same number of iterations, we allow ARM to run more stochastic gradient ascent steps under these two scenarios to check whether it will eventually overfit the training set.
These results show that ARM provides state-of-the-art performance in delivering not only fast convergence, but also low negative log-likelihoods and negative ELBOs on both the validation and test sets, with low computational cost, for all three different network architectures. In comparison to the vanilla REINFORCE on MNIST-static, as shown in Table 2 (a), ARM achieves significantly lower test negative log-likelihoods, which can be explained by having much lower variance in its gradient
estimation, while only costing 20% to 30% more computation time to finish the same number of iterations.
The trace plots in Figures 2, 6, and 7 show that ARM achieves its objective better or on a par with the competing methods in all three different network architectures. In particular, the performance of ARM on MNIST-threshold is significantly better, suggesting ARM is more robust, better resists overfitting, and has better generalization ability. On both MNIST-static and OMNIGLOT, with the “Nonlinear” network architecture, both REBAR and RELAX exhibit severe overfitting, which could be caused by their training procedure, which updates the parameters of the baseline function by minimizing the sample variance of the gradient estimator using SGD. For less overfitting linear and two-stochastic-layer networks, ARM overall performs better than both REBAR and RELAX and runs significantly faster (about 6-8 times faster) in terms of the computation time per iteration.
To understand why ARM has the best overall performance, we examine the trace plots of the logarithm of the estimated variance of gradient estimates in Figure 3. On the MNIST-static dataset with the “Nonlinear” network, the left subplot of Figure 3 shows that both REBAR and RELAX exhibit lower variance than ARM does for their single-Monte-Carlo-sample based gradient estimates; however, the corresponding trace plots of the validation negative ELBOs, shown in Figure 2(a), suggest they both severely overfit the training data as the learning progresses; our hypothesis for this phenomenon is that REBAR and RELAX may favor suboptimal solutions that are associated with lower gradient variance; in other words, they may have difficulty in converging to local optimal solutions that are associated with high gradient variance. For the “Linear” network architecture, the right subplot of Figure 3 shows that ARM exhibits lower variance for its gradient estimate than both REBAR and RELAX do, and Figure 2(b) shows that none of them exhibit clear signs of overfitting; this observation could be used to explain why ARM results in the best convergence for both the training and validation negative ELBOs, as shown in Figure 2(b).

heading: 4.2 MAXIMUM LIKELIHOOD ESTIMATION FOR A STOCHASTIC BINARY NETWORK
text: Denoting xl,xu ∈ R394 as the lower and upper halves of an MNIST digit, respectively, we consider a standard benchmark task of estimating the conditional distribution pθ0:2(xl |xu) (Raiko et al., 2014; Bengio et al., 2013; Gu et al., 2016; Jang et al., 2017; Tucker et al., 2017), using a stochastic binary network with two stochastic binary hidden layers, expressed as
xl ∼ Bernoulli(σ(Tθ0(b1))), b1 ∼ Bernoulli(σ(Tθ1(b2))), b2 ∼ Bernoulli(σ(Tθ2(xu))). (26)
We set the network structure as 392-200-200-392, which means both b1 and b2 are 200 dimensional binary vectors and the transformation Tθ are linear so the results are directly comparable with those in Jang et al. (2017). We approximate log pθ0:2(xl |xu) with log 1K ∑K k=1 Bernoulli(xl;σ(Tθ0(b (k) 1 ))), where b(k)1 ∼ Bernoulli(σ(Tθ1(b (k) 2 ))), b (k) 2 ∼ Bernoulli(σ(Tθ2(xu))). We perform training with K = 1, which can also be considered as optimizing on a single-Monte-Carlo-sample estimate of the lower bound of the log marginal likelihood shown in (21). We use Adam (Kingma & Ba, 2014), with the learning rate set as 10−4, mini-batch size as 100, and number of epochs for training as 2000. Given the inferred point estimate of θ0:2 after training, we evaluate the accuracy of conditional density estimation by estimating the negative log-likelihood as − log pθ0:2(xl |xu), averaging over the test set using K = 1000. We show example results of predicting the activation probabilities of the pixels of xl given xu in Figure 8 of the Appendix.
As shown in Table 3, optimizing a stochastic binary network with the ARM estimator, which is unbiased and computationally efficient, achieves the lowest test negative log-likelihood, outperforming previously proposed biased stochastic gradient estimators on similarly structured stochastic networks, including DARN (Gregor et al., 2013), straight through (ST) (Bengio et al., 2013), slope-annealed ST (Chung et al., 2016), and ST Gumbel-softmax (Jang et al., 2017), and unbiased ones, including score-function (SF) and MuProp (Gu et al., 2016).

heading: 5 CONCLUSION
text: 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.

heading: ACKNOWLEDGMENTS
text: M. Zhou acknowledges the support of Award IIS-1812699 from the U.S. National Science Foundation, and support from the McCombs Research Excellence Grant. The authors acknowledge the support of NVIDIA Corporation with the donation of the Titan Xp GPU used for this research, and the computational support of Texas Advanced Computing Center.

heading: A THE ARM GRADIENT ASCENT ALGORITHM
text: We summarize the algorithm to compute ARM gradient for binary latent variables. Here we show the gradient with respect to the logits associated with the probability of Bernoulli random variables. If the logits are further generated by deterministic transform such as neural networks, the gradient with respect to the transform parameters can be directly computed by the chain rule. For stochastic transforms, the implementation of ARM gradient is discussed in detail in Section 3 and summarized in Algorithm 2.
Algorithm 1: ARM gradient for a V -dimensional binary latent vector input : Bernoulli distribution {qφv (zv)}v=1:V with probability {σ(φv)}v=1:V , target f(z); z = (z1, · · · , zV ), φ = (φ1, · · · , φV ) output :φ and ψ that maximize E(φ,ψ) = Ez∼∏Vv=1 qφv (zv)[f(z;ψ)] Initialize φ, ψ randomly; while not converged do
Sample a mini-batch of x from the data; Sample zv ∼ Bernoulli(σ(φv)) for v = 1, · · · , V ; sample uv ∼ Uniform(0, 1) for v = 1, · · · , V , u = (u1, · · · , uV ) ; gψ = ∇ψf(z;ψ) ; f∆(u,φ) = f ( 1[u>σ(−φ)] ) − f ( 1[u<σ(φ)] ) ; gφ = f∆(u,φ)(u− 0.5) φ = φ+ ρtgφ, ψ = ψ + ηtgψ, with stepsizes ρt, ηt
end
Algorithm 2: ARM gradient for a T -stochastic-hidden-layer binary network input : Inference network qw1:T (b1:T |x) = qw1(b1 |x) [∏T−1 t=1 qwt+1(bt+1 | bt) ] where qwt(bt | bt−1) = Bernoulli(bt;σ(Twt(bt−1))), target f(b1:T ;ψ) output : w1:T and ψ that maximize E(w1:T ,ψ) = Eb1:T∼qw1:T [f(b1:T ;ψ)]
Initializew1:T , ψ randomly; while not converged do
Sample a mini-batch of x from data; for t = 1:T do
If t ≥ 2, sample bt−1 ∼ q(bt−1|bt−2), if t = 2, b1:t−1 = b1, else b1:t−1 = [b1:t−2, bt−1] ; sample ut ∼ ∏ Uniform(0, 1) ; b1t = 1[ut>σ(−Twt (bt−1))] ; b2t = 1[ut<σ(Twt (bt−1))] ; if b1t = b2t then
f∆(b1:t−1, b 1 t:T , b 2 t:T ) = 0 ;
else b1t+1:T ∼ q(bt+1:T |b1t ) ; b2t+1:T ∼ q(bt+1:T |b2t ) ; f∆(b1:t−1, b 1 t:T , b 2 t:T ) = f(b1:t−1, b 1 t:T )− f(b1:t−1, b2t:T ) ;
gwt = f∆(b1:t−1, b 1 t:T , b 2 t:T )(ut − 12 ) T∇wtTwt(bt−1) ; end wt = wt + ρtgwt with step-size ρt
end ψ = ψ + ηt∇ψf(b1:T ;ψ) with step-size ηt
end

heading: B ORIGINAL DERIVATION OF THE AR AND ARM ESTIMATORS
text: Let us denote t ∼ Exp(λ) as an exponential distribution, whose probability density function is defined as p(t |λ) = λe−λt, where λ > 0 and t > 0. The mean and variance are E[t] = λ−1 and var[t] = λ−2, respectively. The exponential random variable t ∼ Exp(λ) can be reparameterized as
t = /λ, ∼ Exp(1). It is well known, e.g., in Ross (2006), that if t1 ∼ Exp(λ1) and t2 ∼ Exp(λ2) are two independent exponential random variables, then the probability that t1 is smaller than t2 can be expressed as P (t1 < t2) = λ1/(λ1 + λ2); moreover, since t1 ∼ Exp(λ1) is equal in distribution to 1/λ1, 1 ∼ Exp(1) and t2 ∼ Exp(λ2) is equal in distribution to 2/λ2, 2 ∼ Exp(1), we have
P (t1 < t2) = P ( 1/λ1 < 2/λ2) = P ( 1 < 2λ1/λ2) = λ1/(λ1 + λ2). (27)
B.1 AUGMENTATION OF A BERNOULLI RANDOM VARIABLE AND REPARAMETERIZATION
From (27) it becomes clear that the Bernoulli random variable z ∼ Bernoulli(σ(φ)) can be reparameterized by comparing two augmented exponential random variables as
z = 1[ 1< 2eφ], 1 ∼ Exp(1), 2 ∼ Exp(1). (28)
Consequently, the expectation with respect to the Bernoulli random variable can be reparameterized as one with respect to two augmented exponential random variables as
E(φ) = Ez∼Bernoulli(σ(φ))[f(z)] = E 1, 2iid∼ Exp(1)[f(1[ 1e−φ< 2])]. (29)
B.2 REINFORCE ESTIMATOR IN THE AUGMENTED SPACE
Since the indicator function 1[ 1e−φ< 2] is not differentiable, the reparameterization trick in (2) is not directly applicable to computing the gradient of (29). Fortunately, as t1 = 1e−φ, 1 ∼ Exp(1) is equal in distribution to t1 ∼ Exp(eφ), the expectation in (29) can be further reparameterized as
E(φ) = E 1, 2 iid∼ Exp(1) [f(1[ 1e−φ< 2])] = Et1∼Exp(eφ), 2∼Exp(1)[f(1[t1< 2])], (30)
and hence, via REINFORCE and then another reparameterization, we can express the gradient as
∇φE(φ) = Et1∼Exp(eφ), 2∼Exp(1)[f(1[t1< 2])∇φ log Exp(t1; e φ)]
= Et1∼Exp(eφ), 2∼Exp(1)[f(1[t1< 2])(1− t1e φ)]
= E 1, 2 iid∼ Exp(1) [f(1[ 1e−φ< 2])(1− 1)]. (31)
Similarly, we have E(φ) = E 1, 2 iid∼ Exp(1) [f(1[ 1< 2eφ])] = E 1∼Exp(1), t2∼Exp(−eφ)[f(1[ 1<t2])], and hence can also express the gradient as
∇φE(φ) = E 1∼Exp(1), t2∼Exp(e−φ)[f(1[ 1<t2])∇φ log Exp(t2; e −φ)]
= −E 1∼Exp(1), t2∼Exp(e−φ)[f(1[ 1<t2])(1− t2e −φ)]
= −E 1, 2 iid∼ Exp(1) [f(1[ 1e−φ< 2])(1− 2)]. (32)
Note that letting 1, 2 iid∼ Exp(1) is the same in distribution as letting
1 = u, 2 = (1− u), where u ∼ Uniform(0, 1), ∼ Gamma(2, 1), (33)
which can be proved using Exp(1) d= Gamma(1, 1), (u, 1 − u)T d= Dirichlet (12), where u ∼ Uniform(0, 1)}, together with Lemma IV.3 of Zhou & Carin (2012); we use “ d=” to denote “equal in distribution.” Thus, (B.2) can be reparameterized as
∇φE(φ) = Eu∼Uniform(0,1), ∼Gamma(2,1) [ f(1[u<σ(φ)])(1− u) ] ,
Applying Rao-Blackwellization (Casella & Robert, 1996), we can further express the gradient as ∇φE(φ) = Eu∼Uniform(0,1) [ f(1[u<σ(φ)])(1− 2u) ] . (34)
Therefore, the gradient estimator shown above, the same as (6), is referred to as the AugmentREINFORCE (AR) estimator.
B.3 MERGE OF REINFORCE GRADIENTS
A key observation of the paper is that by swapping the indices of the two iid standard exponential random variables in (32) , the gradient∇φE(φ) can be equivalently expressed as
∇φE(φ) = −E 1, 2 iid∼ Exp(1) [f(1[ 2e−φ< 1])(1− 1)]. (35)
As the term inside the expectation in (31) and that in (35) could be highly positively correlated, we are motivated to merge (31) and (35) by sharing the same set of standard exponential random variables for Monte Carlo integration, which provides a new opportunity to well control the estimation variance (Owen, 2013). More specifically, simply taking the average of (31) and (35) leads to
∇φE(φ) = E 1, 2 iid∼ Exp(1)
[( f(1[ 1e−φ< 2])− f(1[ 2e−φ< 1]) ) (1/2− 1/2) ] . (36)
Note one may also take a weighted average of (31) and (35), and optimize the combination weight to potentially further reduce the variance of the estimator. We leave that for future study.
Note that (36) can be reparameterized as ∇φE(φ) = Eu∼Uniform(0,1), ∼Gamma(2,1) [( f(1[u>σ(−φ)])− f(1[u<σ(φ)]) ) ( u/2− 1/2) ] ,
Applying Rao-Blackwellization (Casella & Robert, 1996), we can further express the gradient as ∇φE(φ) = Eu∼Uniform(0,1) [( f(1[u>σ(−φ)])− f(1[u<σ(φ)]) ) (u− 1/2) ] . (37)
Therefore, the gradient estimator shown above, the same as (7), is referred to as the AugmentREINFORCE-merge (ARM) estimator.

heading: C PROOFS
text: Proof of Proposition 2. Since the gradients gARM(u, φ), gAR(u, φ), and gR(z, φ) are all unbiased, their expectations are the same as the true gradient gtrue(φ) = σ(φ)(1−σ(φ))[f(1)−f(0)]. Denote f∆(u, φ) = f(1[u>σ(−φ)])− f(1[u<σ(φ)]). Since
f∆(u, φ) =  0, if σ(−|φ|) < u < σ(|φ|), f(1)− f(0), if u > σ(|φ|), f(0)− f(1), if u < σ(−|φ|),
(38)
The second moment of gARM(u, φ) can be expressed as
Eu∼Uniform(0,1)[g2ARM(u, φ)] = Eu∼Uniform(0,1)[f2∆(u, φ)(u− 1/2)2]
= ∫ 1 σ(|φ|) [f(1)− f(0)]2(u− 1/2)2du+ ∫ σ(−|φ|) 0 [f(0)− f(1)]2(u− 1/2)2du
= 1
12 [1− (σ(|φ|)− σ(−|φ|))3][f(1)− f(0)]2
Denoting t = σ(|φ|)− σ(−|φ|), we can re-express gtrue(φ) = 14 (1− t 2)[f(1)− f(0)]. Thus, the variance of gARM(u, φ) can be expressed as
var[gARM(u, φ)] = 1
4
[ 1
3 (1− t3)− 1 4 (1− t2)2
] [f(1)− f(0)]2
= 1 16 (1− t)(t3 + 7 3 t2 + 1 3 t+ 1 3 )[f(1)− f(0)]2 (39) ≤ 1 25 [f(1)− f(0)]2,
which reaches its maximum at 0.039788[f(1)− f(0)]2 when t = √
5−1 2 .
For the REINFORCE gradient, we have Ez∼Bernoulli(σ(φ))[g2R(z, φ)] =Ez∼Bernoulli(σ(φ)) [ f2(z)(z(1− σ(φ))− σ(φ)(1− z))2 ] =σ(φ)(1− σ(φ))[(1− σ(φ))f2(1) + σ(φ)f2(0)].
Therefore the variance can be expressed as
var[gR(u, φ)] =σ(φ)(1− σ(φ)) [ (1− σ(φ))f2(1) + σ(φ)f2(0)− σ(φ)(1− σ(φ))[f(1)− f(0)]2 ] =σ(φ)(1− σ(φ))[(1− σ(φ))f(1) + σ(φ)f(0)]2.
The largest variance satisfies
sup φ
var[gR(z, φ)] ≥ var[gR(z, 0)] = 1
16 [f(1) + f(0)]2,
and hence when f is always positive or negative, we have
supφ var[gARM(z, φ)] supφ var[gR(z, φ)] ≤ 16 25
( 1− 2 f(0)
f(0) + f(1) )2 ≤ 16 25 .
In summary, the ARM gradient has a variance that is bounded by 125 (f(1)−f(0)) 2, and its worst-case variance is smaller than that of REINFORCE.
Proof of Proposition 3. We only need to prove for K = 1 and the proof for K > 1 automatically follows. Since Eu[f(1[u<σ(φ)])2(uv − 1/2)2] = Eu[f(1[u>σ(−φ)])2(uv − 1/2)2], we have
var(gARM1,v)− var(gAR1,v) = −3Eu[f(1[u<σ(φ)])2(uv − 1/2)2] + Eu[f(1[u>σ(−φ)])2(uv − 1/2)2] − 2Eu[f(1[u>σ(−φ)])f(1[u<σ(φ)])(uv − 1/2)2]
= −Eu[f(1[u<σ(φ)])2(uv − 1/2)2]− Eu[f(1[u>σ(−φ)])2(uv − 1/2)2] − 2Eu[f(1[u>σ(−φ)])f(1[u<σ(φ)])(uv − 1/2)2]
= −Eu [( f(1[u>σ(−φ)]) + f(1[u<σ(φ)]) )2 (uv − 1/2)2 ] ≤ 0,
which shows that the estimation variance of gARMK ,v is guaranteed to be lower than that of the gARK ,v , unless f(1[u>σ(−φ)]) + f(1[u<σ(φ)]) = 0 almost surely. Furthermore, since
var(gARM1,v)− var(gAR2,v) =Eu[(f(1[u<σ(φ)])− f(1[u>σ(−φ)]))2(uv − 1/2)2] − Eu(1),u(2) [(f(1[u(1)<σ(φ)])(u(1)v − 1/2) + f(1[u(2)<σ(φ)])(u(2)v − 1/2))2]
=− 2Eu(1) [f(1[u(1)<σ(φ)])(u(1)v − 1/2)]Eu(2) [f(1[u(2)<σ(φ)])(u(2)v − 1/2))] − 2Eu[f(1[u<σ(φ)])f(1[u>σ(−φ)]))(uv − 1/2)2]
=− 2 ( Eu[f(1[u<σ(φ)])(uv − 1/2)] )2 − 2Eu[f(1[u<σ(φ)])f(1[u>σ(−φ)]))(uv − 1/2)2],
when f is always positive or negative, the variance of gARMK ,v is lower than that of gAR2K ,v .
Proof of Proposition 4. Denoting g(u) = gAR(u)− b(u), we have
var[gv(u)]− var[gAR,v(u)] = −2Eu[gAR,v(u)bv(u)] + Eu[b2v(u)].
To maximize the variance reduction, it is equivalant to consider the constrained optimization problem
min bv(u)
− 2Eu[gAR,v(u)bv(u)] + Eu[b2v(u)]
subject to: bv(u) = −bv(1− u),
which is the same as a Lagrangian problem as
min bv(u),λv(u)
L(bv(u), λv(u)) = −2Eu[gAR,v(u)bv(u)] + Eu[b2v(u)] + ∫ λv(u)(bv(u) + bv(1− u))du.
Setting δLδλv = 0 gives bv(u) + bv(1 − u) = 0. By writing ∫ λv(u)(bv(u) + bv(1 − u))du =∫
(λv(u) + λv(1− u))bv(u)du and setting δLδbv = 0, we have
[2gAR,v(u)− 2bv(u)]p(u) = λv(u) + λv(1− u). (40) Interchange u and 1− u gives
[2gAR,v(1− u)− 2bv(1− u)]p(1− u) = λv(1− u) + λv(u). (41) Solving (40) and (41) with bv(u) + bv(1 − u) = 0 and p(u) = p(1 − u), we have the optimal baseline function as b∗v(u) = 1 2 (gAR,v(u)− gAR,v(1− u)). The proof is completed by noticing that gAR(u)− b∗(u) is the same as the single sample gradient estimate under the ARM estimator.
Proof of Corollary 5. Since bv(u) = cv(1−2u) satisfies the anti-symmetric property, we can directly arrive at Corollary 5 using Proposition 4. Alternatively, since Eu[f(1[u<σ(φ)])2(uv − 1/2)2] = Eu[f(1[u>σ(−φ)])2(uv − 1/2)2] and Eu[f(1[u<σ(φ)])(uv − 1/2)2] = Eu[f(1[u>σ(−φ)])(uv − 1/2)2], for gC,v = (f(1[u<σ(φ)])− cv)(1− 2uv), we have
var(gC,v)− var(gARM,v) = Eu[(f(1[u<σ(φ)])− cv)2(1− 2uv)2]− Eu[(f(1[u<σ(φ)])− f(1[u>σ(−φ)]))2(uv − 1/2)2] = Eu [( 4c2v − 8cvf(1[u<σ(φ)]) + 2f2(1[u<σ(φ)]) + 2f(1[u<σ(φ)])f(1[u>σ(−φ)]) ) (uv − 1/2)2
] = Eu [( f(1[u>σ(−φ)]) + f(1[u<σ(φ)])− 2cv )2 (uv − 1/2)2
] ≥ 0.
Proof of Proposition 6. First, to compute the gradient with respect to w1, since
E(w1:T ) = Eq(b1)Eq(b2:T | b1)[f(b1:T )], (42) we have
∇w1E(w1:T ) = Eu1∼Uniform(0,1)[f∆(u1, Tw1(x))(u1 − 1/2)]∇w1Tw1(x), (43) where
f∆(u1, Tw1(x)) = Eb2:T∼q(b2:T | b1), b1=1[u1>σ(−Tw1 (x))])[f(b1:T )]
− Eb2:T∼q(b2:T | b1), b1=1[u1<σ(Tw1 (x))])[f(b1:T )]. (44)
Second, to compute the gradient with respect to wt, where 2 ≤ t ≤ T − 1, since E(w1:T ) = Eq(b1:t−1)Eq(bt | bt−1)Eq(bt+1:T | bt)[f(b1:T )], (45)
we have ∇wtE(w1:T )=Eq(b1:t−1) [ Eut∼Uniform(0,1)[f∆(ut, Twt(bt−1), b1:t−1)(ut − 1/2)]∇wtTwt(bt−1) ] ,
(46)
where
f∆(ut, Twt(bt−1), b1:t−1) = Ebt+1:T∼q(bt+1:T | bt), bt=1[ut>σ(−Twt (bt−1))])[f(b1:T )]
− Ebt+1:T∼q(bt+1:T | bt), bt=1[ut<σ(Twt (bt−1))])[f(b1:T )]. (47)
Finally, to compute the gradient with respect to wT , we have ∇wT E(w1:T ) = Eq(b1:T−1) [ EuT∼Uniform(0,1)[f∆(uT , TwT (bT−1), b1:T−1)(uT − 1/2)] × ∇wT TwT (bT−1)] , (48)
where
f∆(uT , TwT (bT−1), b1:T−1) = f(b1:T−1, bT = 1[uT>σ(−TwT (bT−1))]) − f(b1:T−1, bT = 1[uT<σ(TwT (bT−1))]). (49)

heading: D ADDITIONAL EXPERIMENTAL RESULTS
text: For the univariate AR gradient, we have Eu∼Uniform(0,1)[g2AR(u, φ)] = ∫ σ(φ)
0
f2(1)(1− 2u)2du+ ∫ 1 σ(φ) f2(0)(1− 2u)2du
= 1
6 (f2(0) + f2(1)) +
1 6 (1− 2σ(φ))3(f2(0)− f2(1)).
Thus its variance can be expressed as
var[gAR(u, φ)] = 1
6 (f2(0)+f2(1))+
1 6 (1−2σ(φ))3(f2(0)−f2(1))−σ2(φ)(1−σ(φ))2[f(1)−f(0)]2,
which is used for related plots in Figures 1 and 5.
0.01
0.00
0.01
me an
REINFORCE
0.01
0.00
0.01 AR
0.01
0.00
0.01 RELAX
0.01
0.00
0.01 ARM
0.49 0.499 0.501 0.51 True
0.06
0.08
0.10
0.12
std
0.140
0.145
0.00
0.05
0.10
0.002 0.004 0.49 0.499 0.501 0.51 True
2 0 2 0.00
0.05
0.10
SN R
2 0 2 0.00
0.05
0.10
2 0 2 0.0
0.5
1.0
1.5
2.0
2 0 2 0.0
0.5
1.0
1.5 2.0 0.49 0.499 0.501 0.51 True
Phi
Figure 5: Comparison of a variety of gradient estimators for estimating the gradient of E(φ) = Ez∼Bernoulli(σ(φ))[(z − p0)2], where p0 ∈ {0.49, 0.499, 0.501, 0.51} and the values of φ range from −2.5 to 2.5. For each φ value, we compute for each estimator K = 1000 single-Monte-Carlo-sample gradient estimates, and use them to calculate their sample mean ḡ, sample standard deviation sg , and gradient signal-tonoise ratio SNRg = |ḡ|/sg . In each estimator specific column, we plot ḡ, sg , and SNRg in Rows 1, 2, and 3, respectively. The theoretical gradient standard deviations and gradient signal-to-noise ratios are also shown if they can be analytically calculated (see Eq. 25 and Appendices C and D and for related analytic expressions).

Review summaries:
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.

For binary layers, how to calculate and backpropagate gradients is a big problem, particularly for the binary neural networks. To solve the problem, this paper proposes an unbiased and low variance augment-REINFORCE-merge (ARM) estimator. With the help of an appropriate reparameterization, the antithetic sampling in an augmented space can be used to drive a variance-reduction mechanism. The experimental results show that ARM estimator converges fast, has low computational complexity, and provides advanced prediction performance.