Title: ARM: AUGMENT-REINFORCE-MERGE GRADIENT FOR STOCHASTIC BINARY NETWORKS

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.

1. INTRODUCTION:

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

EQUATION

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 ∇ 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 (Kingma & Welling, 2013; Rezende et al., 2014) to compute the gradient as

EQUATION

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(φ) = E z∼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) ), 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 (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 & Lá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 = (z 1 , . . . , z V ) T , the gradient of

EQUATION

with respect to φ = (φ 1 , . . . , φ V ) T , the logits of the Bernoulli probability parameters, can be expressed as

EQUATION

where

1 [u>σ(-φ)] := 1 [u1>σ(-φ1)] , . . . , 1 [u V >σ(-φ 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(φ) = E z∼Bernoulli(σ(φ)) [f (z)

] with respect to φ, which has an analytic expression as

EQUATION

Since

σ(φ) 0 (1 -2u)du = σ(φ)σ(-φ) and 1 σ(φ) (1 -2u)du = -σ(φ)σ(-φ)

, we can rewrite (5) as

EQUATION

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 ũ = 1 -u, we have

EQUATION

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

E u∼Uniform(0,1) [f (1 [u<σ(φ)] )(1/2-u)] = -E u∼Uniform(0,1) [f (1 [u>σ(-φ)] )(1/2-u)],

we have

E u∼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 .

2.2. MULTIVARIATE GENERALIZATION VIA THE LAW OF TOTAL EXPECTATION:

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

EQUATION

Since z \v ∼ j =v Bernoulli(z j ; σ(φ j )) can be equivalently generated as

z \v = 1 [u \v <σ(φ \v )] or as z \v = 1 [u \v >σ(-φ \v )]

, where u \v ∼ j =v Uniform(u j ; 0, 1), exchanging the order of the two expectations in (8) and applying reparameterization, we conclude the proof for (4) of Theorem 1 with

EQUATION

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

EQUATION

)

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

EQUATION

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.

2.3. EFFECTIVENESS OF ARM FOR VARIANCE REDUCTION:

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

E z∼Bernoulli(σ(φ)) [f (z)],

with a single Monte-Carlo sample u ∼ Uniform(0, 1) or z ∼ Bernoulli(σ(φ)), the ARM gradient is expressed as

g ARM (u, φ) = f (1 [u>σ(-φ)] ) -f (1 [u<σ(φ)] ) (u -1/2), and the REINFORCE gradient as g R (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 ≤ 16 25 .

In the general setting, with u (1) , . . . , u (K) iid ∼ V v=1 Uniform(0, 1), we define the ARM estimate of ∇ φv E(φ) with K Monte Carlo samples, denoted as g ARM K ,v , and the AR estimate with 2K Monte Carlo samples, denoted as g AR 2K ,v , using

EQUATION

where

g v (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[g ARM K ,v ] var[g AR 2K ,v ] = var[g v (u)] -Cov(-g v (u), g v (1 -u)) var[g v (u)] = 1-ρ v , ρ v = Corr(-g v (u), g v (1-u)). Note -g v (u) = f (1 [u<σ(φ)] )(2u v -1), g v (1-u) = f (1 [u>σ(-φ)] )(2u v -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 g ARM K ,v and AR estimate g AR 2K ,v shown in (12), the variance of g ARM K ,v is guaranteed to be lower than that of g AR K ,v ; moreover, if f ≥ 0 (or f ≤ 0), then the variance of g ARM K ,v is guaranteed to be lower than that of g AR 2K ,v .

We show below that under the anti-symmetric constraint b(u) = -b(1 -u), which implies that E u∼ V v=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 E z∼q φ (z) [f (z)], the optimal anti-symmetric baseline function to be subtracted from the AR estimator g AR (u) = f (1 [u<σ(φ)] )(1 -2u), which minimizes the variance of Monte Carlo integration, can be expressed as

EQUATION

where

B = {b(u) : b v (u) = -b v (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

b v (u) = c v (1/2 -u v )

, where c v is a dimension-specific constant whose value can be optimized for variance reduction.

3. BACKPROPAGATION THROUGH DISCRETE STOCHASTIC LAYERS:

A latent variable model with multiple stochastic hidden layers can be constructed as

EQUATION

whose joint likelihood given the distribution parameters θ 0:T = {θ 0 , . . . , θ T } is expressed as

EQUATION

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.

3.1. ARM VARIATIONAL AUTO-ENCODER:

For auto-encoding variational inference (Kingma & Welling, 2013; Rezende et al., 2014) , we construct a variational distribution as

EQUATION

with which the ELBO can be expressed as

E(w 1:T ) = E b 1:T ∼qw 1:T (b 1:T | x) [f (b 1:T )] ,

where

EQUATION

)

Proposition 6 (ARM backpropagation). For a stochastic binary network with T binary stochastic hidden layers, constructing a variational auto-encoder (VAE) defined with b 0 = x and

EQUATION

for t = 1, . . . , T , the gradient of the ELBO with respect to w t can be expressed as

∇ wt E(w 1:T ) = E q(b1:t-1) [E ut∼Uniform(0,1) [f ∆ (u t , T wt (b t-1 ), b 1:t-1 )(u t -1/2)]∇ wt T wt (b t-1 )] ,

where

EQUATION

The gradient presented in ( 19) can be estimated with a single Monte Carlo sample as

f∆ (u t , T wt (b t-1 ), b 1:t-1 ) = 0, if b (1) t = b (2) t , f (b 1:t-1 , b (1) t:T ) -f (b 1:t-1 , b (2) t:T ), otherwise, (20) where b (1) t = 1 [ut>σ(-Tw t (bt-1))] , b (1) t+1:T ∼ q(b t+1:T | b (1) t ), b (2) t = 1 [ut<σ(Tw t (bt-1))] , and b (2) t+1:T ∼ q(b t+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 is O(1), which involves a single evaluation of the function f and gradient backpropagation as ∇ wt T wt (b t-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 ).

3.2. ARM MAXIMUM LIKELIHOOD ESTIMATION:

For maximum likelihood estimation, the log marginal likelihood can be expressed as

EQUATION

Generalizing Proposition 6 leads to the following proposition. Proposition 7. For a stochastic binary network defined as

EQUATION

)

the gradient of the lower bound in (21) with respect to θ t can be expressed as

∇ θt E(θ 1:T ) = E p(b t+1:T ) [E ut∼Uniform(0,1) [f ∆ (u t , T θt (b t+1 ), b t+1:T )(u t -1/2)]∇ θt T θt (b t+1 )] ,

where

f ∆ (u t , T θt (b t+1 ), b t+1:T ) = E b1:t-1∼p(b1:t-1 | bt), bt=1 [u t >σ(-T θ t (b t+1 ))] ) [log p θ0 (x | b 1 )] -E b1:t-1∼p(b1:t-1 | bt), bt=1 [u t <σ(T θ t (b t+1 ))] ) [log p θ0 (x | b 1 )].

4. EXPERIMENTAL RESULTS:

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(φ) = E z∼Bernoulli(σ(φ)) [(z -p 0 ) 2 ],

where p 0 ∈ {0.49, 0.499, 0.501, 0.51}.

The optimal solution is σ(φ) = 1 [p0<0.5] . The closer p 0 is to 0.5, the more challenging the optimization becomes. We compare both the AR and ARM estimators to the true gradient as

EQUATION

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 RE-BAR in our experiments for this toy problem (as also shown in Grathwohl et al. (2018) for p 0 = 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<σ(φ)] -p 0 ) 2 (1 [u<σ(φ)] -σ(φ)), g φ,AR = (1 [u<σ(φ)] -p 0 ) 2 (1 -2u),

while the ARM gradient can be expressed as

EQUATION

)

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 p 0 = 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. 

(φ) = E z∼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 ḡ, sample standard deviation s g , and the gradient signal-to-noise ratio defined as SNR g = |ḡ|/s g ; for each φ value, we use K = 1000 single-Monte-Carlo-sample gradient estimates to calculate ḡ, s g , and SNR g . 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 p 0 ; 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

EQUATION

)

We find that the values of the ratio shown above are almost exactly matched by the values of SNR g = |ḡ|/s g 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 p 0 .

4.1. DISCRETE VARIATIONAL AUTO-ENCODERS:

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 & Lázaro-Gredilla, 2015), MuProp (Gu et al., 2016) , Concrete (Gumbel-Softmax) (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 

4.2. MAXIMUM LIKELIHOOD ESTIMATION FOR A STOCHASTIC BINARY NETWORK:

Denoting x l , x u ∈ R 394 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 (x l | x u ) (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

EQUATION

)

We set the network structure as 392-200-200-392, which means both b 1 and b 2 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

(x l | x u ) with log 1 K K k=1 Bernoulli(x l ; σ(T θ0 (b (k) 1 ))), where b (k) 1 ∼ Bernoulli(σ(T θ1 (b (k) 2 ))), b (k) 2 ∼ Bernoulli(σ(T θ2 (x u ))

). 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 (x l | x u ), averaging over the test set using K = 1000. We show example results of predicting the activation probabilities of the pixels of x l given x u 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) .

5. CONCLUSION:

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. t = /λ, ∼ Exp(1). It is well known, e.g., in Ross (2006) , that if t 1 ∼ Exp(λ 1 ) and t 2 ∼ Exp(λ 2 ) are two independent exponential random variables, then the probability that t 1 is smaller than t 2 can be expressed as P (t 1 < t 2 ) = λ 1 /(λ 1 + λ 2 ); moreover, since t 1 ∼ Exp(λ 1 ) is equal in distribution to 1 /λ 1 , 1 ∼ Exp(1) and t 2 ∼ Exp(λ 2 ) is equal in distribution to 2 /λ 2 , 2 ∼ Exp(1), we have

EQUATION

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

EQUATION

Consequently, the expectation with respect to the Bernoulli random variable can be reparameterized as one with respect to two augmented exponential random variables as

EQUATION

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 t 1 = 1 e -φ , 1 ∼ Exp( 1) is equal in distribution to t 1 ∼ Exp(e φ ), the expectation in ( 29) can be further reparameterized as

EQUATION

and hence, via REINFORCE and then another reparameterization, we can express the gradient as

EQUATION

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

EQUATION

Note that letting 1 , 2 iid ∼ Exp(1) is the same in distribution as letting

EQUATION

which can be proved using Exp(1)

d = Gamma(1, 1), (u, 1 -u) T d = Dirichlet (1 2 )

, 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(φ) = E u∼Uniform(0,1), ∼Gamma(2,1) f (1 [u<σ(φ)] )(1 -u) ,

Applying Rao-Blackwellization (Casella & Robert, 1996) , we can further express the gradient as

EQUATION

Therefore, the gradient estimator shown above, the same as ( 6), is referred to as the Augment-REINFORCE (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

EQUATION

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

EQUATION

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(φ) = E u∼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

EQUATION

Therefore, the gradient estimator shown above, the same as ( 7), is referred to as the Augment-REINFORCE-merge (ARM) estimator.

C PROOFS:

Proof of Proposition 2. Since the gradients g ARM (u, φ), g AR (u, φ), and g R (z, φ) are all unbiased, their expectations are the same as the true gradient

EQUATION

The second moment of g ARM (u, φ) can be expressed as

E u∼Uniform(0,1) [g 2 ARM (u, φ)] = E u∼Uniform(0,1) [f 2 ∆ (u, φ)(u -1/2) 2 ] = 1 σ(|φ|) [f (1) -f (0)] 2 (u -1/2) 2 du + σ(-|φ|) 0 [f (0) -f (1)] 2 (u -1/2) 2 du = 1 12 [1 -(σ(|φ|) -σ(-|φ|)) 3 ][f (1) -f (0)] 2 Denoting t = σ(|φ|) -σ(-|φ|), we can re-express g true (φ) = 1 4 (1 -t 2 )[f (1) -f (0)].

Thus, the variance of g ARM (u, φ) can be expressed as

var[g ARM (u, φ)] = 1 4 1 3 (1 -t 3 ) - 1 4 (1 -t 2 ) 2 [f (1) -f (0)] 2 = 1 16 (1 -t)(t 3 + 7 3 t 2 + 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

E z∼Bernoulli(σ(φ)) [g 2 R (z, φ)] =E z∼Bernoulli(σ(φ)) f 2 (z)(z(1 -σ(φ)) -σ(φ)(1 -z)) 2 =σ(φ)(1 -σ(φ))[(1 -σ(φ))f 2 (1) + σ(φ)f 2 (0)].

Therefore the variance can be expressed as

var[g R (u, φ)] =σ(φ)(1 -σ(φ)) (1 -σ(φ))f 2 (1) + σ(φ)f 2 (0) -σ(φ)(1 -σ(φ))[f (1) -f (0)] 2 =σ(φ)(1 -σ(φ))[(1 -σ(φ))f (1) + σ(φ)f (0)] 2 .

The largest variance satisfies

sup φ var[g R (z, φ)] ≥ var[g R (z, 0)] = 1 16 [f (1) + f (0)] 2 ,

and hence when f is always positive or negative, we have

sup φ var[g ARM (z, φ)] sup φ var[g R (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 1 25 (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 E u [f (1 [u<σ(φ)] ) 2 (u v -1/2) 2 ] = E u [f (1 [u>σ(-φ)] ) 2 (u v -1/2) 2 ], we have var(g ARM1,v ) -var(g AR1,v ) = -3E u [f (1 [u<σ(φ)] ) 2 (u v -1/2) 2 ] + E u [f (1 [u>σ(-φ)] ) 2 (u v -1/2) 2 ] -2E u [f (1 [u>σ(-φ)] )f (1 [u<σ(φ)] )(u v -1/2) 2 ] = -E u [f (1 [u<σ(φ)] ) 2 (u v -1/2) 2 ] -E u [f (1 [u>σ(-φ)] ) 2 (u v -1/2) 2 ] -2E u [f (1 [u>σ(-φ)] )f (1 [u<σ(φ)] )(u v -1/2) 2 ] = -E u f (1 [u>σ(-φ)] ) + f (1 [u<σ(φ)] ) 2 (u v -1/2) 2 ≤ 0,

which shows that the estimation variance of g ARM K ,v is guaranteed to be lower than that of the

g AR K ,v , unless f (1 [u>σ(-φ)] ) + f (1 [u<σ(φ)] ) = 0 almost surely. Furthermore, since var(g ARM1,v ) -var(g AR2,v ) =E u [(f (1 [u<σ(φ)] ) -f (1 [u>σ(-φ)] )) 2 (u v -1/2) 2 ] -E u (1) ,u (2) [(f (1 [u (1) <σ(φ)] )(u (1) v -1/2) + f (1 [u (2) <σ(φ)] )(u (2) v -1/2)) 2 ] = -2E u (1) [f (1 [u (1) <σ(φ)] )(u (1) v -1/2)]E u (2) [f (1 [u (2) <σ(φ)] )(u (2) v -1/2))] -2E u [f (1 [u<σ(φ)] )f (1 [u>σ(-φ)] ))(u v -1/2) 2 ] = -2 E u [f (1 [u<σ(φ)] )(u v -1/2)] 2 -2E u [f (1 [u<σ(φ)] )f (1 [u>σ(-φ)] ))(u v -1/2) 2 ],

when f is always positive or negative, the variance of g ARM K ,v is lower than that of g AR 2K ,v .

Proof of Proposition:

4. Denoting g(u) = g AR (u) -b(u), we have var[g v (u)] -var[g AR,v (u)] = -2E u [g AR,v (u)b v (u)] + E u [b 2 v (u)].

To maximize the variance reduction, it is equivalant to consider the constrained optimization problem

min bv(u) -2E u [g AR,v (u)b v (u)] + E u [b 2 v (u)] subject to: b v (u) = -b v (1 -u),

which is the same as a Lagrangian problem as

EQUATION

Interchange u and 1 -u gives

EQUATION

Solving ( 40) and ( 41) with b v (u) + b v (1 -u) = 0 and p(u) = p(1 -u), we have the optimal baseline function as b

* v (u) = 1 2 (g AR,v (u) -g AR,v (1 -u)).

The proof is completed by noticing that g AR (u) -b * (u) is the same as the single sample gradient estimate under the ARM estimator.

Proof of Corollary 5. Since b v (u) = c v (1-2u) satisfies the anti-symmetric property, we can directly arrive at Corollary 5 using Proposition 4. Alternatively, since

E u [f (1 [u<σ(φ)] ) 2 (u v -1/2) 2 ] = E u [f (1 [u>σ(-φ)] ) 2 (u v -1/2) 2 ] and E u [f (1 [u<σ(φ)] )(u v -1/2) 2 ] = E u [f (1 [u>σ(-φ)] )(u v - 1/2) 2 ], for g C,v = (f (1 [u<σ(φ)] ) -c v )(1 -2u v ), we have var(g C,v ) -var(g ARM,v ) = E u [(f (1 [u<σ(φ)] ) -c v ) 2 (1 -2u v ) 2 ] -E u [(f (1 [u<σ(φ)] ) -f (1 [u>σ(-φ)] )) 2 (u v -1/2) 2 ] = E u 4c 2 v -8c v f (1 [u<σ(φ)] ) + 2f 2 (1 [u<σ(φ)] ) + 2f (1 [u<σ(φ)] )f (1 [u>σ(-φ)] ) (u v -1/2) 2 = E u f (1 [u>σ(-φ)] ) + f (1 [u<σ(φ)] ) -2c v 2 (u v -1/2) 2 ≥ 0.

Proof of Proposition 6. First, to compute the gradient with respect to w 1 , since

EQUATION

we have

EQUATION

where

EQUATION

Second, to compute the gradient with respect to w t , where

EQUATION

we have

EQUATION

where

EQUATION

Finally, to compute the gradient with respect to w T , we have

EQUATION

where

EQUATION

ARM(K=5000) 

(φ) = E z∼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 ḡ, sample standard deviation sg, and gradient signal-tonoise ratio SNRg = |ḡ|/sg. In each estimator specific column, we plot ḡ, 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). 

:



Figures and Tables:

Figure 1: Comparison of a variety of gradient estimators in maximizing E(φ) = E z∼Bernoulli(σ(φ)) [(z -p0) 2 ] via

Figure 2: Training and validation negative ELBOs on MNIST-static with respect to the training iterations, shown in the top row, and with respect to the wall clock times on Tesla-K40 GPU, shown in the bottom row, for three differently structured Bernoulli VAEs.

Figure 3: Trace plots of the log variance of the gradient estimators on the MNIST-static data for "Nonlinear"and "Linear" network architectures, whose corresponding trace plots of both the training and validation -ELBO are shown in Figures2(a) and 2(b), respectively. The variance of the gradient of each element is estimated by performing exponential smoothing, with the smoothing factor as 0.999, on its first two moments; the logarithm of the average over all elements' gradient variances is shown for each stochastic gradient ascent step.

(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 Figure2(b).

Figure 5: Comparison of a variety of gradient estimators for estimating the gradient of E(φ) =

Figure 6: Training and validation negative ELBOs on MNIST-threshold with respect to the training iterations, shown in the top row, and with respect to the wall clock times on Tesla-K40 GPU, shown in the bottom row, for three differently structured Bernoulli VAEs.

Figure 7: Training and validation negative ELBOs on OMNIGLOT with respect to the training iterations, shown in the top row, and with respect to the wall clock times on Tesla-K40 GPU, shown in the bottom row, for three differently structured Bernoulli VAEs.

The constructions of three differently structured discrete variational auto-encoders. The following symbols "→", "]", )", and " " represent deterministic linear transform, leaky rectified linear units (LeakyReLU)(Maas et al., 2013) nonlinear activation, sigmoid nonlinear activation, and random sampling respectively, in the encoder (a.k.a. recognition network); their reversed versions are used in the decoder (a.k.a. generator).

Table content: <table><tr><td>Nonlinear</td><td>Linear</td><td>Linear two layers</td></tr><tr><td colspan="3">Encoder 784→200]→200]→200) 200 784→200) 200 784→200) 200→200) 200</td></tr><tr><td colspan="3">Decoder 784 (784←[200←[200←200 784 (784←200 784 (784←200 (200←200</td></tr></table>

Test negative log-likelihoods of discrete VAEs trained with a variety of stochastic gradient estimators on MNIST-static and OMNIGLOT, where * , , †, ‡ represent the results reported inMnih & Gregor (2014),Tucker et al. (2017),Gu et al. (2016), andGrathwohl et al. (2018), respectively. The results for LeGrad(Titsias & Lázaro-Gredilla, 2015) are obtained by running the code provided by the authors. We report the results of ARM using the sample mean and standard deviation over five independent trials with random initializations.

Table content: <table><tr><td/><td/><td colspan="2">(a) MNIST-static</td><td/><td/></tr><tr><td colspan="2">Linear</td><td colspan="2">Nonlinear</td><td colspan="3">Two layers</td></tr><tr><td>Algorithm</td><td>-log p(x)</td><td>Algorithm</td><td>-log p(x)</td><td colspan="2">Algorithm</td><td>-log p(x)</td></tr><tr><td>AR</td><td>= 164.1</td><td>AR</td><td>= 114.6</td><td>AR</td><td/><td>= 162.2</td></tr><tr><td>REINFORCE</td><td>= 170.1</td><td>REINFORCE</td><td>= 114.1</td><td colspan="2">REINFORCE</td><td>= 159.2</td></tr><tr><td>Wake-Sleep  *</td><td>= 120.8</td><td>Wake-Sleep  *</td><td>-</td><td colspan="2">Wake-Sleep  *</td><td>= 107.7</td></tr><tr><td>NVIL  *</td><td>= 113.1</td><td>NVIL  *</td><td>= 102.2</td><td>NVIL  *</td><td/><td>= 99.8</td></tr><tr><td>LeGrad</td><td>≤ 117.5</td><td>LeGrad</td><td>-</td><td>LeGrad</td><td/><td>-</td></tr><tr><td>MuProp  †</td><td>≤ 113.0</td><td>MuProp</td><td>= 99.1</td><td>MuProp  †</td><td/><td>≤ 100.4</td></tr><tr><td>Concrete</td><td>= 107.2</td><td>Concrete</td><td>= 99.6</td><td>Concrete</td><td/><td>= 95.6</td></tr><tr><td>REBAR</td><td>= 107.7</td><td>REBAR</td><td>= 100.7</td><td>REBAR</td><td/><td>= 95.7</td></tr><tr><td>RELAX  ‡</td><td>≤ 113.6</td><td>RELAX  ‡</td><td>≤ 119.2</td><td>RELAX  ‡</td><td/><td>≤ 100.9</td></tr><tr><td>ARM</td><td>= 107.2 ± 0.1</td><td>ARM</td><td>= 98.4 ± 0.3</td><td>ARM</td><td colspan="2">= 96.7 ± 0.3</td></tr><tr><td/><td/><td colspan="2">(b) OMNIGLOT</td><td/><td/></tr><tr><td/><td>Linear</td><td colspan="2">Nonlinear</td><td colspan="2">Two layers</td></tr><tr><td>Algorithm</td><td>-log p(x)</td><td>Algorithm</td><td>-log p(x)</td><td>Algorithm</td><td colspan="2">-log p(x)</td></tr><tr><td>NVIL  *</td><td>= 117.6</td><td>NVIL  *</td><td>= 116.6</td><td>NVIL  *</td><td colspan="2">= 111.4</td></tr><tr><td>MuProp</td><td>= 117.6</td><td>MuProp</td><td>= 117.5</td><td>MuProp</td><td colspan="2">= 111.2</td></tr><tr><td>Concrete</td><td>= 117.7</td><td>Concrete</td><td>= 116.7</td><td>Concrete</td><td colspan="2">= 111.3</td></tr><tr><td>REBAR</td><td>= 117.7</td><td>REBAR</td><td>= 118.0</td><td>REBAR</td><td colspan="2">= 110.8</td></tr><tr><td>RELAX  ‡</td><td>≤ 122.1</td><td>RELAX  ‡</td><td>≤ 128.2</td><td>RELAX  ‡</td><td colspan="2">≤ 115.4</td></tr><tr><td>ARM</td><td>= 115.8 ± 0.2</td><td>ARM</td><td>= 117.6 ± 0.4</td><td>ARM</td><td colspan="2">= 109.8 ± 0.3</td></tr></table>

Comparison of the test negative log-likelihoods between ARM and various gradient estimators inJang et al. (2017), for the MNIST conditional distribution estimation benchmark task.

Table content: <table><tr><td>Gradient estimator</td><td>ARM</td><td colspan="4">ST DARN Annealed ST ST Gumbel-S. SF</td><td>MuProp</td></tr><tr><td colspan="3">-log p(x l | xu) 57.9 ± 0.1 58.9 59.7</td><td>58.7</td><td>59.3</td><td>72.0</td><td>58.9</td></tr></table>

-p0) 2 ], where p0 ∈ {0.49, 0.499, 0.501, 0.51} and the values of φ are updated via gradient ascent with the true gradients. Shown in Rows 1, 2, and 3 are the trace plots of the estimated gradients using K = 10, 100, and 5000 Monte Carlo samples, respectively.

Table content: <table><tr><td/><td/><td/><td/><td/><td/><td/><td>0.49 0.499 0.501 0.51 True</td></tr><tr><td/><td/><td/><td>Iteration</td><td/><td/><td/><td/></tr><tr><td colspan="8">Figure 4: Comparison of a variety of gradient estimators for estimating the gradient of E(φ) =</td></tr><tr><td>E z∼Bernoulli(σ(φ)) [(z 0.01 0.00 0.01 mean REINFORCE 0.06 0.08 0.10 0.12 std</td><td>0.01 0.00 0.01 0.140 0.145</td><td>AR</td><td>0.01 0.00 0.01 0.00 0.05 0.10</td><td>RELAX</td><td>0.00 0.01 0.01 0.002 0.004</td><td>ARM</td><td>0.49 0.499 0.501 0.51 True 0.49 0.499 0.501 0.51 True</td></tr></table>

Test negative ELBOs of discrete VAEs trained with four different stochastic gradient estimators.MNIST-threshold is the binarized MNIST thresholded at 0.5 and MNIST-static is the binarized MNIST used inSalakhutdinov & Murray (2008);Larochelle & Murray (2011).

Table content: <table><tr><td>ARM RELAX REBAR ST Gumbel-Softmax</td></tr></table>

Peer review (iclr2019):
 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. Major comments: The authors devote several pages of fairly dense mathematics to deriving the ARM estimate in section 2 (up to section 2.5). However I found it relatively easy to derive (15) directly, using elementary results such as the law of total expectation and a single 1-dimensional integral, in about 10 lines of equations. As the authors note, deriving (4) from (15) requires an extra line or two. In my opinion it would greatly improve the clarity of the paper to use a more direct and straightforward derivation (perhaps with the interesting historical account of how the authors first derived this result given in an appendix). I could understand the more lengthy derivation being helpful if it gave insight into the source of variance reduction, but I don*t see this personally, and the current discussion of variance reduction does not refer to the derivation of (15) at all. The analysis of variance in section 2.6 leaves a lot to be desired. The central claim of the paper is that this method reduces variance, so it is an important section! Firstly, the variance of ARM vs AR is interesting, but the variance of ARM vs REINFORCE seems also highly relevant. Secondly, it seems like it would be very informative to look at the ratio of stdev to the mean for the ARM gradient estimate, since the true gradient is multiplied by sigmoid(phi) sigmoid(-phi) and so is very small if the probability of z = 1 is close to 0 or 1, exactly in the same regime where ARM has an advantage in variance reduction over AR. For example, it may be that learning in this regime is very difficult due to the weak gradient even if the estimate is extremely low variance. Thirdly and somewhat relatedly, in this same regime (z = 1 close to 0 or 1) the ARM gradient estimate is very often 0, meaning no learning takes place, so it seems a bit strange to argue that the new method is fantastic in the regime where it*s almost always not learning! Of course, not learning is better than adding lots of spurious variance as reinforce would, but perhaps this could be made clearer. Finally, the theoretical analysis involving correlation gives very little insight and is extremely hand-wavy. A short worked example in the 1D or 2D case explicitly computing the variance of REINFORCE, AR and ARM seems like it would be highly informative. Minor comments: In the introduction, "*approximately* maximizing the marginal likelihood" might be more accurate, since as given in (28) the exact marginal likelihood is not optimized in practice, and the exact marginal likelihood is not of the form (1) but is rather the logarithm of something of the form (1). I wasn*t clear why "equal in distribution" was used a few things for things that are simply equal, such as just above (5). In section 2.3, I don*t see any real reason the estimates in (9) and (11) "could be highly positively correlated", other than an argument along the lines of the simple one given in section 2.6 that they*re often equal and so zero. As an aside, in section 3.1, it is great not to assume conditional independence of the binary latent variables across layers, but assuming conditional independence within each layer is still very restrictive. It is reasonable for the generative distribution to have this property, since the resulting net can still be essentially "universal" by stacking enough layers, but assuming this factorization in the variational distribution is highly restrictive with hard-to-reason-about consequences for the learned generative model. I realize this is a commonly used assumption and the authors are interested in the variance reduction properties of their approach rather than the training itself, but I just mention that it would be great to see extensions of the current work that can cope tractably with correlated latent variables within each layer. In section 3.2, according to my understanding of standard terminology, "maximum likelihood inference" is a misnomer and would normally be "maximum likelihood estimation", since maximum likelihood is a method for estimating parameters whereas inference is about inferring latent variable values given parameters. In section 4, it would be great to see some plots of explicit variance estimates of the different methods, given the overall goal of the paper (unless I just missed this?), even though figure 1 gives some insight into the variance characteristics. In section 4.2, the expression log 1/K .sum_k Bernoulli... differs in the placement of log from Jang et al (2017). Which is the standard convention for this task?