ARM: A UGMENT -REINFORCE-M ERGE G RADIENT FOR S TOCHASTIC B INARY N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 available at https://github.com/ABC-anonymous-1.
I NTRODUCTION
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 signiﬁcant 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 φ ( (cid:15) ) , (cid:15) ∼ p ( (cid:15) ) , where (cid:15) 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
[Equation]
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 )]
= (cid:82) ∇ φ q φ ( z )
d z = ∇ φ (cid:82) q
φ ( z ) d z = 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 ﬁrst 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.
and RELAX of Grathwohl et al.
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 conﬂict, 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 overﬁtting.
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´azaro-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 signiﬁcant 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.
ARM: A UGMENT -REINFORCE- MERGE ESTIMATOR
In this section, we ﬁrst 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 >σ ( − φ )] := (cid:0) 1 [ u 1 >σ ( − φ 1 )] , . .
, 1 [ u V >σ ( − φ V )] (cid:1) T .
2. U NIVARIATE ARM ESTIMATOR
Below we will ﬁrst 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 (cid:82) σ ( φ )(1 − 2 u ) du
=
σ ( φ ) σ ( − φ ) and (cid:82) 1 σ ( φ ) (1 − 2 u ) du = − σ ( φ ) σ ( − φ ) , we can rewrite (5) as
[Equation]
We refer to (6) as the univariate augment-REINFORCE (AR) estimator, as originally derived 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.
Apply antithetic sampling (Owen, 2013) to the AR estimator in (6), with ˜ u
= 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 originally derived 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) (cid:2)(cid:0) f
( 1 [ u<σ ( φ )] )
+ f ( 1 [ u>σ ( − φ )] ) (cid:1) (1 / 2 − u ) (cid:3) = 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
[Equation]
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. M ULTIVARIATE 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 a multivariate setting.
Let us denote ( · ) \ v as a vector whose v th element is removed.
For the expectation in (3), applying the univariate ARM estimator in (36) together with the law of total expectation, we have
[Equation]
[Equation]
Since z \ v ∼ (cid:81) ν
(cid:54) = v Bernoulli ( z ν ; σ ( φ ν )) can be equivalently generated as z \ v = 1 [ u \ v <σ ( φ \ v )] or as z \ v =
[ u \ v >σ ( − φ \ v )] , where u \ v ∼ (cid:81) ν
(cid:54) = v Uniform ( u ν ; 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 ﬁrst 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, satisﬁes 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. E FFECTIVENESS 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
[Equation]
( z )] , with a single Monte-Carlo sample u ∼ Uniform (0 , 1) or z ∼ Bernoulli ( σ ( φ )) , the ARM gradient is expressed as g ARM ( u, φ ) =
(cid:0) f ( 1 [ u>σ ( − φ )] ) − f ( 1 [ u<σ ( φ )] ) (cid:1) ( 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
[Equation]
In the general setting, with u (1) , . . .
, u ( K ) iid ∼ (cid:81)
Vv =1 Uniform (0 , 1) , we deﬁne the ARM estimate of ∇ φ v E ( φ ) with K Monte Carlo samples, denoted as g ARM K ,v , and the AR estimate with 2 K Monte
Carlo samples, denoted as g AR 2 K ,v , using
[Equation]
where g v ( u ( k ) )
=
f ( 1 [ u ( k ) <σ ( φ )] )(1 − 2 u ( k ) v ) .
Similar to the analysis in Owen (2013), the amount of variance reduction brought by the ARM estimator can be reﬂected by the ratio as var [ g
[Equation]
K ,v ]
[Equation]
u ) , g v (1 − u )) .
Note − g v ( u ) =
f ( 1 [ u <σ ( φ )] )(2 u v − 1) , g v (1 − u ) =
f ( 1 [ u >σ ( − φ )] )(2 u v − 1) , and P ( 1 [ u v <σ ( φ v )]
= 1 [ u v >σ ( − φ 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 2 K ,v shown in (12) , the variance of g ARM
K ,v is guaranteed to be lower than that of g AR K ,v , and if f ≥ 0 (or f ≤ 0 ), then the variance of g ARM K ,v is guaranteed to be lower than that of g AR 2 K ,v .
We show below that under the anti-symmetric constraint b ( u ) = − b (1 − u ) , which implies that E
[Equation]
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 − 2 u ) , which minimizes the variance of Monte Carlo integration, can be expressed as
[Equation]
(13) 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, 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-speciﬁc constant whose value can be optimized for variance reduction.
B ACKPROPAGATION THROUGH DISCRETE STOCHASTIC LAYERS
A latent variable model with multiple stochastic hidden layers can be constructed as
[Equation] [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 overﬁtting (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. ARM VARIATIONAL AUTO - ENCODER
For auto-encoding variational inference (Kingma & Welling, 2013; Rezende et al., 2014), we construct a variational distribution as
[Equation]
(16) with which the ELBO can be expressed as
[Equation]
Proposition 6 (ARM backpropagation) .
For a stochastic binary network with T binary stochastic hidden layers, constructing a variational auto-encoder (VAE) deﬁned with b 0 = x and
[Equation]
for t = 1 , . . .
, T , the gradient of the ELBO with respect to w t can be expressed as
[Equation]
)] ∇ w
[Equation]
T
[Equation]
t
[Equation]
b t
[Equation] [Equation]
where f ∆ ( u t , T
[Equation]
b
[Equation] [Equation]
>σ
[Equation] [Equation]
The gradient presented in (19) can be estimated with a single Monte Carlo sample as
[Equation]
[Equation]
where b (1) t = 1 [
[Equation]
>
[Equation]
w t ( b t − 1 ))] , b (1) t +1: T ∼ q ( b t +1: T | b (1) t ) , b (2) t = 1 [ u t <σ ( T w t ( b t − 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 ∇ w t T w t ( 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 (2 T ) .
3. 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.
[Figure]
Proposition 7.
For a stochastic binary network deﬁned as
[Equation]
22) 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 u t ∼ 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 b 1: t − 1 ∼ p ( b 1: t − 1 | b t ) , b t = 1 [ u t>σ ( −T θ t ( b t +1))] )
[log p θ
( x | b 1 )] − E b 1: t − 1 ∼ p ( b 1: t − 1 | b t ) , b t = 1 [ u t<σ ( T θ t ( b t +1))] )
[log p θ
( x | b 1 )] .
E XPERIMENTAL 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
[Equation]
[( z
[Equation]
p [Equation] ) 2 ] , where p 0 ∈ { 0 . 49 , 0 .
The optimal solution is σ ( φ ) = {+
+}
[ p 0 < 0 .
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]
(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 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
[Equation]
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 sufﬁciently 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
[Table]
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 true gradients, despite that the trace plots of the ARM gradients are distinct from these of the true gradients.
More speciﬁcally, 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 (37) 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 speciﬁc column, we plot against the value of φ the sample mean ¯ g , sample standard deviation s g , and the gradient signal-to-noise ratio deﬁned as SNR g = ¯ g/s g ; for each φ value, we use
K
= 1000 single-Monte-Carlo-sample gradient estimates to calculate ¯ g , s g , and SNR g .
Both ﬁgures 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 function.
In Figure 5 of Appendix D, it is also interesting to notice that the absolute value of 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 veriﬁed to be true using (23) and (38) in Appendix C, as the ratio of the true gradient g φ to (cid:112) var [ g φ, ARM ] , the standard deviation of the ARM estimate in (24), can be expressed as
[Equation]
We ﬁnd that the values of the ratio shown above are almost exactly matched by the values of
SNR g = ¯ 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 to noise ratios regardless of the value of p 0 .
4. D ISCRETE 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, efﬁcient 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´azaro-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 deﬁned in (15) and decoder deﬁned in (16), we consider three different network architectures, as summarized in Table 1, including “Nonlinear” that has one stochastic but two LeakyReLU
(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
[Table]
(Salakhutdinov & Murray, 2008; Larochelle & Murray, 2011), 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 predeﬁned 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.
Note as shown in Figures 2(a,d) and 7(a,d), both REBAR and RELAX exhibit clear signs of overﬁtting 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 overﬁtting 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 overﬁt 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
[Figure]
1.0 Steps 1e6 7.
l o g ( v a r )
ARMREBARRELAX 0.
0.8 Steps 1e 8.
l o g ( v a r )
ARMREBARRELAX
Nonlinear
[Figure]
the vanilla REINFORCE on MNIST-static, as shown in Table 2 (a), ARM achieves signiﬁcantly lower test 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 ﬁnish 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 signiﬁcantly better, suggesting ARM is more robust, better resists overﬁtting, and has better generalization ability.
On both MNIST-static and OMNIGLOT, with the “Nonlinear” network architecture, both REBAR and RELAX exhibit severe overﬁtting, 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 overﬁtting linear and two-stochastic-layer networks, ARM overall performs better than both REBAR and RELAX and runs signiﬁcantly 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 overﬁt 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
[Table]
variance; in other words, they may have difﬁculty 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 overﬁtting; 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).
4. M AXIMUM 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 x l ∼ Bernoulli ( σ ( T θ 0 ( b 1 ))) , b 1 ∼ Bernoulli ( σ ( T θ 1 ( b 2 ))) , b 2 ∼ Bernoulli ( σ ( T θ 2 ( x u ))) .
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.
We approximate log p θ 0:
( x l | x u ) with log 1 K (cid:80)
Kk =1 Bernoulli ( x l ; σ ( T θ 0 ( b ( k ) 1 ))) , where b ( k )∼ Bernoulli ( σ ( T θ 1 ( b ( k ) 2 ))) , b ( k )∼ 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 efﬁcient, 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).
C ONCLUSIONS
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 learnable parameters for variance reduction, it maintains efﬁcient computation and avoids increasing the risk of overﬁtting.
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 localexpectation based variance reduction method, and applying it to reinforcement learning whose action space is discrete.
Appendix A T HE ARM GRADIENT ASCENT A LGORITHM
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.
[Table] [Table]
− 1 t =1 q w t +1 ( b t +1 | b t ) (cid:105) where q w t ( b t | b t − 1 ) =
Bernoulli ( b t ; σ ( T w t ( b t − 1 ))) , target f
( b 1: T ; ψ ) output: w 1: T and ψ that maximize E ( w 1: T , ψ ) =
E b 1: T ∼ q w 1: T [ f ( b 1: T ; ψ )]
Initialize w 1: T , ψ randomly; while not converged do
Sample a mini-batch of x from data; for t = 1:T do
If t ≥ 2 , sample b t − 1 ∼ q ( b t − 1 | b t − 2 ) , if t = 2 , b 1: t − 1 = b 1 , else b 1: t − 1 = [ b 1: t − 2 , b t − 1 ] ; sample u t ∼ (cid:81) Uniform (0 , 1) ; b 1 t = 1 [ u t >σ ( −T w t ( b t − 1 ))] ; b 2 t = 1 [ u t <σ ( T w t ( b t − 1 ))] ; if b 1 t = b 2 t then f ∆ ( b 1: t − 1 , b 1 t : T , b 2 t : T ) = 0 ; else b 1 t +1: T ∼ q ( b t +1: T | b 1 t ) ; b 2 t +1: T ∼ q ( b t +1: T | b 2 t ) ; f ∆ ( b 1: t − 1 , b 1 t : T , b 2 t : T ) =
f ( b 1: t − 1 , b 1 t : T ) − f ( b 1: t − 1 , b 2 t : T ) ; g w t = f ∆ ( b 1: t − 1 , b 1 t : T , b 2 t : T )( u t − 12 ) T ∇ w t T w t ( b t − 1 ) ; end w t =
w t + ρ t g w t with step-size ρ t end ψ
= ψ + η t ∇ ψ f ( b 1: T ; ψ ) with step-size η t end
O
RIGINAL DERIVATION OF THE AR AND ARM ESTIMATORS
Let us denote t ∼ Exp ( λ ) as an exponential distribution, whose probability density function is deﬁned 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 =
(cid:15)/λ, (cid:15) ∼ 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 (cid:15) 1 /λ 1 , (cid:15)
1 ∼ Exp (1) and t 2 ∼
Exp ( λ 2 ) is equal in distribution to (cid:15) 2 /λ 2 , (cid:15)
2 ∼ Exp (1) , we have
P ( t 1 < t 2 ) = P ( (cid:15) 1 /λ 1 < (cid:15) 2 /λ 2 ) = P ( (cid:15) 1 < (cid:15) 2 λ 1 /λ 2 ) =
λ 1 / ( λ 1 + λ 2 ) .
B. A UGMENTATION OF A B ERNOULLI RANDOM VARIABLE AND R EPARAMETERIZATION
From (26) 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. REINFORCE ESTIMATOR IN THE AUGMENTED SPACE
Since the indicator function 1 [ (cid:15) 1 e − φ <(cid:15) 2 ] is not differentiable, the reparameterization trick in (2) is not directly applicable to computing the gradient of (28).
Fortunately, as t 1 = (cid:15) 1 e − φ , (cid:15) 1 ∼ Exp (1) is equal in distribution to t 1 ∼
Exp ( e φ ) , the expectation in (28) can be further reparameterized as
[Equation]
[Equation]
and hence, via REINFORCE and then another reparameterization, we can express the gradient as
[Equation]
Similarly, we have
[Equation]
φ
[Equation]
and hence can also express the gradient as
[Equation]
Note that letting (cid:15) 1 , (cid:15) 2 iid ∼ Exp (1) is the same in distribution as letting
[Equation]
, where u ∼ Uniform
[Equation] [Equation]
∼ Gamma (2 ,
[Equation] [Equation]
[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) }
[Equation]
and Lemma IV.3 of Zhou & Carin (2012).
Thus, (B.2) can be reparameterized as
[Equation]
Applying Rao Blackwellization (Casella & Robert, 1996)
[Equation]
we can further express the gradient as
[Equation]
Therefore, the gradient estimator shown above, the same as (6), is referred to as the AugmentREINFORCE (AR) estimator.
B. M ERGE OF REINFORCE GRADIENTS
A key observation of the paper is that by swapping the indices of the two iid standard exponential random variables in (31) , the gradient ∇ φ E ( φ ) can be equivalently expressed as
[Equation]
As the term inside the expectation in (30) and that in (34) could be highly positively correlated
[Equation]
we are motivated to merge (30) and (34) by sharing the same set of standard exponential random variables for Monte Carlo integration
[Equation]
which provides a new opportunity to well control the estimation variance (Owen, 2013).
More speciﬁcally, simply taking the average of (30) and (34) leads to
[Equation]
Note one may also take a weighted average of (30) and (34), and optimize the combination weight to potentially further reduce the variance of the estimator.
We leave that for future study.
Note that (35) can be reparameterized as
[Equation]
Applying Rao Blackwellization (Casella & Robert, 1996)
[Equation]
we can further express the gradient as
[Equation]
Therefore, the gradient estimator shown above, the same as (7), is referred to as the AugmentREINFORCE-merge (ARM) estimator.
C P ROOFS
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 g true ( φ )
=
σ ( φ )(1 − σ ( φ ))
[ f (1) − f (0)] .
Denote f ∆ ( u, φ ) =
f
( 1 [ u>σ ( − φ )] ) − f ( 1 [ u<σ ( φ )] ) .
Since
[Equation]
The second moment of g ARM ( u, φ ) can be expressed as
[Equation]
[Equation]
Denoting t = σ ( | φ | ) − σ ( −| φ | ) , we can re-express g true ( φ ) = 14 (1 − t 2 )
[ f (1) − f (0)] .
Thus, the variance of g ARM ( u, φ ) can be expressed as
[Equation]
which reaches its maximum at 0 .
[ f (1) − f (0)] 2 when t = √ 5 − 1 2 .
For the REINFORCE gradient, we have
[Equation]
Therefore the variance can be expressed as
[Equation]
The largest variance satisﬁes
[Equation]
and hence when f is always positive or negative, we have
[Equation]
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 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 ARM 1 ,v ) − var ( g AR 1 ,v ) =
− 3 E u [ f ( 1 [ u <σ ( φ )] ) 2 ( u v − 1 / 2) 2 ] + E u
[ f ( 1 [ u >σ ( − φ )] ) 2 ( u v − 1 / 2) 2 ] − 2 E u
[ f ( 1 [ u >σ ( − φ )] ) f ( 1 [ u <σ ( φ )] )( u v − 1 /
[Equation]
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 ARM 1 ,v ) − var ( g AR 2 ,v ) =
E u
[Equation]
( 1 [ u <σ ( φ )] ) − f (
[Equation]
[ u >σ
[Equation] [Equation]
[Equation] [Equation]
when f is always positive or negative, the variance of g ARM K ,v is lower than that of g AR 2 K ,v .
Proof of Proposition 4.
Denoting g ( u ) =
g AR ( u ) − b ( u ) , we have var
[Equation]
g
[Equation]
u
[Equation] [Equation]
var [ g AR ,
[Equation]
u
[Equation] [Equation]
u
[Equation]
To maximize the variance reduction, it is equivalant to consider the constrained optimization problem
[Equation]
subject to:
[Equation]
u
[Equation]
which is the same as a Lagrangian problem as
[Equation]
Setting δ L δλ v = 0 gives b v ( u ) + b v (1 − u ) = 0 .
By writing (cid:82) λ v ( u )( b v ( u ) + b v (1 − u ))
d u =
[Equation]
( λ v ( u ) + λ v (1 − u ))
b v ( u ) d
u
and setting δ L δb v = 0 , we have
[Equation]
g
[Equation]
Interchange u and 1 − u gives
[Equation]
Solving (39) and (40) 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 ) =
12 ( 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 − 2 u ) satisﬁes 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 − 2 u v ) , we have var ( g C,
[Equation]
Proof of Proposition 6.
First
[Equation]
to compute the gradient with respect to w 1 , since
[Equation]
41) we have
[Equation]
42) where
[Equation]
Second, to compute the gradient with respect to w t , where 2 ≤ t ≤ T − 1 , since
E (
[Equation]
we have
[Equation]
45) where
[Equation] [Equation]
Finally
[Equation]
to compute the gradient with respect to w T
[Equation]
we have
[Equation]
where
[Equation]
D A DDITIONAL EXPERIMENTAL RESULTS
[Figure] [Figure]
[Table]
[Figure] [Figure]