Title: IMPROVED GRADIENT ESTIMATORS FOR STOCHASTIC DISCRETE VARIABLES

Abstract: In many applications we seek to optimize an expectation with respect to a distribution over discrete variables. Estimating gradients of such objectives with respect to the distribution parameters is a challenging problem. We analyze existing solutions including finite-difference (FD) estimators and continuous relaxation (CR) estimators in terms of bias and variance. We show that the commonly used Gumbel-Softmax estimator is biased and propose a simple method to reduce it. We also derive a simpler piece-wise linear continuous relaxation that also possesses reduced bias. We demonstrate empirically that reduced bias leads to a better performance in variational inference and on binary optimization tasks.

1. INTRODUCTION:

Discrete stochastic variables arise naturally in many applications including topic modeling, semisupervised learning, clustering, variational memory addressing and reinforcement learning. In many cases, the objective is to minimize an expectation of a function of discrete random variables with respect to the distribution parameters:

min φ L[φ],

where L[φ] = E q φ f (z z z) = z z z q φ (z z z)f (z z z).

(1)

Here, z z z is a vector of discrete variables under a φ-parameterized distribution q φ (z z z). For example, in variational inference, f (z z z) is the variational lower bound and q φ (z z z) is the approximating posterior.

Eq. ( 1) is commonly minimized by gradient-based methods, which require estimating the gradient ∂ φ L[φ] (see Schulman et al. (2015) for an overview). For continuous random variables whose sampling can be reparameterized as a function of other parameter-independent random variables, the reparameterization trick (Kingma & Welling (2013) , Rezende et al. ( 2014)) gives low variance gradient estimates. However, this trick is not directly applicable to discrete variables due to the discontinuous cumulative distribution function (CDF).

Extensions of the reparameterization trick to discrete distributions can be grouped into finitedifference (FD) estimators, and continuous relaxation (CR) estimators. We present a summary of three recent FD estimators (Tokui & Sato (2017) ; Yin & Zhou (2018) ; Lorberbom et al. (2018) to guide practitioners to the salient aspects of their bias-variance tradeoff and computational complexity. Most importantly, we propose a scalable (but still unbiased) variant of Tokui & Sato (2017) that trades decreased computation for slightly increased variance.

We examine CR estimators including the popular Gumbel-Softmax (Jang et al. (2016) ; Maddison et al. (2016) ) in terms of the gradient bias and propose a new method to reduce the bias of all CR estimators for both binary and categorical variables. Based upon this understanding, we develop a piecewise-linear estimator for binary and categorical variables that is simpler and less biased than Gumbel-Softmax.

Lastly, we provide empirical evidence that these improved estimators allow for faster optimization when training variational autoencoders and when optimizing a continuous relaxation of a combinatorial optimization problem.

1.1. RELATED WORK:

The most generic approximator of ∂ φ L[φ] is the score function (SF) estimator (a.k.a. REINFORCE Williams (1992) , Glynn (1990) ). SF suffers from high variance and many remedies have been proposed to reduce this variance (Mnih & Gregor (2014) ; Gregor et al. (2013) ; Gu et al. (2015) ; Mnih & Rezende (2016) ; Tucker et al. (2017) ; Grathwohl et al. (2017) ). Unbiased estimators that require multiple function evaluations have also been proposed Tokui & Sato (2017) , Titsias & Lázaro-Gredilla (2015) , Lorberbom et al. (2018) , Yin & Zhou (2018) . These estimators have lower variance but can be computationally demanding. CR estimators often trade bias for variance and previous CR proposals include straight-through estimators (Bengio et al. (2013) ; Raiko et al. (2014) ), the Gumbel-Softmax estimator (Jang et al. (2016) ; Maddison et al. (2016) ), and overlapping smoothing (Vahdat et al. (2018b; a) ; Rolfe (2016)).

2. THE REPARAMETERIZATION AND MARGINALIZATION ESTIMATOR:

We begin a summary of FD estimators with the reparameterization and marginalization (RAM) method of Tokui & Sato (2017) . For a single binary variable the expectation in Eq. ( 1) is enumerated as L = z q φ (z)f (z) = q φ f (1) + (1 -q φ )f (0), where q φ ≡ q φ (z = 1). The gradient is:

EQUATION

where q φ = σ(l φ ) = (1 + e -l φ ) -1 and l φ = logit(q φ ). This derivative involves two function evaluations and contains a finite-difference of f (z). Eq. ( 2) is an unbiased zero-variance estimate since the summation over z is done explicitly. The generalization to M factorially-distributed random variables q φ (z z z)

EQUATION

where again the summation over z i is performed exactly and q φ,i ≡ q φ,i (z i = 1). The summation over z z z \i can be estimated with a single sample but the derivative requires M +1 function evaluations. This limits the applicability of RAM to moderately-sized models. Note that both f (z i = 0, z z z \i ) and f (z i = 1, z z z \i ) are evaluated at the same z z z \i which leads to a lower-variance estimator. For hierarchical q φ (z z z) = M i=1 q φ,i (z i |z z z <i ) the derivative takes the form:

∂ φ L = i z z z<i q φ (z z z <i )∂ φ q φ,i z z z>i q φ (z z z >i |1, z z z <i )f (z z z >i , 1, z z z <i ) - z z z>i q φ (z z z >i |0, z z z <i )f (z z z >i , 0, z z z <i ) .

(4) where q φ,i ≡ q φ,i (1|z z z <i ). The key insight of Tokui & Sato (2017) is to use common random variates z z z >i when sampling from q φ (z z z >i |1, z z z <i ) and q φ (z z z >i |0, z z z <i ). This reduces both the variance and the computational cost. Tokui & Sato (2017) show that this estimator is optimal because it exactly sums over the binary variables whose probability distribution is being differentiated.

A RAM estimator can also be constructed for categorical variables. For a single one-hot encoded categorical variable y = (y 0 , ...y A-1 ), y a ∈ {0, 1}, a y a = 1, the derivative of Eq. ( 1) is

EQUATION

where f a = f (y a = 1), and q a φ = e l a φ / b e l b φ . The generalization to many categorical variables proceeds as for binary variables. For example, the derivative of a factorial categorical distribution over y y y = {y a

EQUATION

This derivative can again be estimated with a single sample but requires M A function evaluations. Since RAM is unbiased and has the minimal variance (due to the explicit summation over the differentiated variable), we use it as a baseline to evaluate computationally cheaper alternatives.

2.1. SAMPLED REPARAMETERIZATION AND MARGINALIZATION:

We propose a modification to RAM that allows us to trade decreased computational cost for increased variance. For binary variables, ∂ φ q φ,i = q φ,i (1 -q φ,i )∂ φ l φ,i where q φ,i = σ(l φ,i ), so that each term in Eq. ( 3) or ( 4) is proportional to q φ,i (1 -q φ,i ). In many applications, we observe that the distribution of a large number of variables (q φ,i ) are drawn to 0 or 1 early during optimization. Such variables have negligible contribution to the full derivative. We exploit this observation to reduce the computational cost by including variable z i in the full gradient with probability p i = 4q φ,i (1-q φ,i )/β, where β is an adjustable hyperparameter. This means that we replace the derivative in Eq. ( 3) with

EQUATION

where ξ i ∈ {0, 1} are Bernoulli variables with probabilities p i indicating whether z i is included or not. We evaluate Eq. ( 7) by sampling ξ and only then evaluating non-zero terms. In Section 4 we show that in the context of variational inference on MNIST the number of function evaluations is reduced by an order of magnitude while still allowing for effective optimization. As always, this computational saving comes at the cost of increased gradient variance which slows training.

In the categorical case, each term in Eq. ( 6) is accompanied by q a φ,i q b φ,i that can be used to assign importance to the (a, b) edge of the simplex for variable y i . The computational cost M A can be reduced by keeping each term with probability p a,b i = 4q a φ,i q b φ,i /β. Appendices A and B summarize two other FD estimators, ARGMAX Lorberbom et al. (2018) and ARM Yin & Zhou (2018) , by noting their bias-variance tradeoff and computational complexity.

3. CONTINUOUS RELAXATION ESTIMATORS:

Unlike FD estimators that use multiple function evaluations to approximate the gradient, continuous relaxation estimators extend the reparameterization trick to discrete variables by approximating them with continuous variables: z → ζ = ζ β (ρ, q φ ), where β is a parameter that controls the approximation, ρ ∈ U[0, 1] is uniform random variable and q φ are parameterized probabilities. 1 We use the same symbol ζ to denote the random variable and its reparameterization by ρ. The objective function L

[φ] = z z z q φ (z z z)f (z z z) is replaced with L[φ] = E ρ ρ ρ f ζ ζ ζ(ρ, q φ )

and its gradients are computed using the chain rule

∂ φ L[φ] = E ρ ρ ρ i ∂ ζi f (ζ ζ ζ)(∂ qi ζ i )(∂ φ q i ) ζi(ρi,q φ,i (1|ζ ζ ζ <i ) (8)

These gradients can be computed efficiently by automatic differentiation libraries. However, because the objective function is changed, CR estimators in Eq. ( 8) are biased. Nevertheless, in practice the bias is often small enough to allow for effective optimization.

3.1. THE GUMBEL-SOFTMAX ESTIMATOR:

The most popular CR estimator is Gumbel-Softmax (GSM) (Jang et al. (2016) ; Maddison et al. (2016)) 2 . For binary variables, this relaxation takes the simple form

EQUATION

1 Continuous relaxations have been instrumental in constructing advanced control variate for SF estimates in Tucker et al. (2017) and Grathwohl et al. (2017) . We discuss the REBAR estimator from Tucker et al. (2017) in Appendix E.

2 More precisely, Maddison et al. (2016) considered problems in variational inference where f (z z z) is relaxed by replacing a generative distribution with its continuous relaxation. In contrast, Jang et al. (2016) directly relaxed discrete variables z → ζ without changing the objective, thus replacing f (z z z) → f (ζ ζ ζ). This does not produce a consistent objective for a probabilistic model, but as Jang et al. (2016) have shown, works well in practice. We work with generic f (z z z) and thus only consider the approach of Jang et al. (2016) . (c) shows the effect of a modification to GSM that reduces its bias.

where q i ≡ q φ,i (1|ζ ζ ζ <i ). This relaxation and its derivative ∂ q ζ are shown in Fig. 1 (a). β controls the sharpness of the relaxation and tunes the trade-off between the bias (closeness of ζ to z) and variance of ∂ q ζ. We note that the slope of the relaxed ζ(ρ

) at ζ = 1/2 is α ≡ β/[4q(1 -q)

] and thus becomes large when q approaches 0 or 1.

3.2. IMPROVED CONTINUOUS RELAXATION ESTIMATORS:

In this section, we analyze the bias introduced by Eq. ( 8) and propose a simple method to reduce it. The bias of

E ρ ρ ρ [∂ ζi f (ζ ζ ζ)∂ φ ζ i ]

has two sources: (a) the relaxation of ζ j for j = i and (b) the relaxation of ζ i . To characterize the latter bias, we start with a single binary variable and write the gradient as the following integral:

EQUATION

Here, ξ(ρ) is any continuous function satisfying ξ(0) = 0 and ξ(1) = 1. For a non-decreasing function ξ(ρ), we can view ξ as a random variable with inverse CDF ξ(ρ) and ρ as a uniform random variable ρ ∈ U[0, 1]. With this interpretation the gradient can be written as an expectation

EQUATION

which can be estimated by sampling. If ∂ ξ f (ξ) does not vary significantly in the interval ξ ∈ [0, 1] then the variance of this estimate is controlled by var(∂ξ/∂ρ) = 1 0 dρ (∂ξ/∂ρ) 2 -1. Thus, the more non-linear ξ(ρ) is, the higher will be the variance of estimate Eq. ( 11). This idea can be extended to factorial q φ (z z z) = M i=1 q φ,i (z i ) as in Eq. ( 3):

EQUATION

At this point there is no relationship between ξ i and z i , and Eq. ( 12) is just a higher variance version of Eq. ( 3). However, if we relax

EQUATION

where the bias comes from the deviation of ζ ζ ζ \i from z z z \i in the function evaluations. In other words, Eq. ( 13) uses an unbiased form for the differentiated variable ζ i and the only bias comes from relaxing the remaining variables ζ ζ ζ \i .

The variance of each term is controlled by var ∂ζi ∂ρi and is reduced by making ζ i (ρ i , q φ,i ) more linear. The bias is reduced by making ζ ζ ζ closer z z z. Varying ζ i (ρ i , q φ,i ) between linear (low variance) and step function (low bias) allows for a controllable trade-off. We call Eq. ( 13) the improved continuous relaxation (ICR) estimator. The original CR estimator of Eq. ( 8) for factorial q φ has the form

EQUATION

and comparing this to the ICR estimate in Eq. ( 13), we see that CR can be transformed into ICR by the replacing ∂ qi ζ i with ∂ ρi ζ i . This change can be simply implemented using TensorFlow's stop gradient ≡ sg notation by replacing

ζ i (ρ i , q i ) with ζ i ρ i + q i -sg(q i ), sg(q i ) (15)

in Eq. ( 8). We emphasize that the ICR estimator in Eq. ( 13) is less biased than the direct CR estimator because, in the case of a single variable, ICR is unbiased while CR is not. Further, this decrease in bias is not accompanied by an increase in variance. We show in Appendix C that similar benefits are obtained for hierarchical q φ . Finally, we note a conceptual relationship between sampled RAM and ICR: both estimators evaluate the gradient only through a subset of variables z i . Sampled RAM chooses this subset explicitly and evaluates the gradients with FD while ICR samples relaxed variables where only a subset of them will possess non-negligible gradients.

3.3. PIECE-WISE LINEAR RELAXATION:

Inspired by ICR and a better understanding of bias-variance trade-off, we propose a piece-wise linear relaxation (PWL) depicted in Fig. 1(b ). The linear part is centered at ρ = 1 -q so that the corresponding binary variable is obtained by z = round(ζ). The slope is given by α = β/[4q(1-q)] similar to the Gumbel-Softmax slope3 . The explicit expression for PWL smoothing is

EQUATION

where

[x] 1 0 ≡ min (1, max [0, x]

) is the hard sigmoid function (note that ∂ q α = 0). This relaxation has several attractive properties. Firstly, we have

∂ qi ζ i = ∂ ρi ζ i ,

which means that the CR and ICR estimators coincide for PWL. Secondly, PWL has easily interpretable expressions for bias and variance. In the case of a function of a single variable the variance is given by var

(∂ ρ ζ) = α -1 while the bias in computing expectation z f (z) over the relaxed distribution is dρf (ζ(ρ)) - z f (z) = { 1 0 dxf (x) -[f (0) + f (1)]/2}/α.

This clearly shows that α trades bias for variance. Finally, the PWL relaxation defined in Eq. ( 16) can be considered as the inverse CDF of the PDF

EQUATION

Eq. ( 17) is a mixture of two delta distributions centered at zero and one, and a uniform distribution defined in the interval [0, 1]. Samples from q(ζ) are in the continuous interval with probability [4q(1 -q)]/β, and ζ = 0/1 have probability proportional to the probability of the binary states.

3.4. CATEGORICAL PIECE-WISE LINEAR RELAXATION:

We now extend ICR estimators to categorical variables. For a single categorical variable we apply the integral representation to each edge (a, b) of the simplex in Eq. ( 5) and relax this pair of variables using PWL:

y → y a,b = {y a = 0.5 + α a,b ρ a,b -q b /(q a + q b ) 1 0 , y b = 1 -y a , y c =a,b = 0} where ρ a,b ∼ U[0, 1]

and α a,b is the slope. We replace the summation over the edges of the simplex by sampling one edge at a time with probability p a,b = (q a + q b )/(A -1). Details are found in Appendix D and we provide the final result:

EQUATION

where ỹa,b has the same value as y a,b but has the gradient scaled by γ a,b = (A -1)(q a + q b ). 4The probabilities of edge selection and the scale factor are chosen to give correct values for the objective and its gradient. Extension of this categorical PWL estimator to multivariate distributions is straightforward; For example, for factorial distributions one relaxes each categorical variable y i by sampling an edge with probability p a,b i = (q a i + q b i )(A -1) and uses a single function evaluation. The resulting gradient is unbiased for a single variable and introduces a bias in the multivariate case which makes it an ICR estimator. 

3.5. IMPROVED CATEGORICAL GUMBEL-SOFTMAX ESTIMATOR:

The sampling done in the categorical PWL estimator leads to increased variance of the gradients.

As an alternative, we suggest an improved version of the categorical Gumbel-Softmax estimator.

Recall that the Gumbel-Softmax estimator for the categorical case has the form:

EQUATION

In Eq. ( 20), we propose a new version of this estimator by applying the same trick as in Eq. ( 15). This estimator remains biased but we empirically demonstrate that its bias is reduced.

ζ(ρ, q q q) → ζ(ρ + q q q -sg(q), sg(q)).

(20)

4. EXPERIMENTS:

In this section we compare the FD and CR estimators and their improved variants on a number of examples. We start with one-variable toy examples that illustrate the bias of GSM estimator and then move to the training of variational auto-encoders and a combinatorial optimization problem.

4.1. TOY EXAMPLE:

We begin with an illustrative single-variable example (Tucker et al. (2017) ) with objective L = z q φ (z)f (z) where f (z) = (z -0.45) 2 . The relaxed convex function f (ζ) depicted in Fig. 2 (a) has two local maxima, one of which is the minimum over the discrete domain. We compare five gradient methods: RAM, Eq. (2); ARM (see Appendix B); PWL, Eq. ( 16); GSM, Eq. ( 9); and improved Gumbel-Softmax (IGSM), Eq. ( 15) (see Appendix F for experimental details). Fig. 2(b) shows the evolution of q φ (z = 1) during training which demonstrates the bias associated with GSM. To quantify this bias, we plot the value of the gradient of all estimators for different values of q φ (z = 1) in Fig. 2(c ). We observe that the GSM gradient has the wrong sign for a large interval in q φ (z = 1) which prevents GSM from converging to the true minimum.

To understand the nature of the bias that GSM introduces, we plot the derivatives ∂ q ζ and ∂ ρ ζ corresponding to GSM and IGSM respectively in Fig. 1(c ). We see that ∂ q ζ is biased towards the value of z = round(ζ) that has the highest probability. This means that GSM in Eq. ( 14) will oversample the derivative ∂ ζ f (ζ) from the most probable mode. In example Fig. 2 (a) this results in oversampling the derivative from z = 0 mode which creates a gradient that pushes optimization away from the true minimum z = 0. The bias is reduced as β is increased.

In Appendix I.1 we consider an example for a concave function over a binary variable and observe similar effects. We also consider both concave and convex functions over a categorical variable in Appendix I.2. In all scenarios, the bias of GSM prevents its convergence to the correct minimum. 

EQUATION

Here, p θ θ θ (z z z) is the prior, p θ θ θ (x x x|z z z) is the decoder, and q φ (z z z|x x x) is the approximating posterior. L is minimized with respect to the θ θ θ parameters of the generative model and the φ parameters of the approximate posterior. The latter minimization corresponds to Eq. ( 1) and we can apply the various estimators to propagate φ-derivatives through discrete samples z z z. For CR estimators we replace z z z

with ζ ζ ζ to compute ∂ φ L[φ, θ θ θ] ≈ x x x∼data E ρ [∂ φ f θ θ θ,φ (ζ ζ ζ, x x x)],

where the expectation is evaluated with a single relaxed sample per data point x x x. The θ θ θ-derivative can be calculated directly from Eq. ( 21) using the discrete z z z variables. Thus, the φ and θ θ θ derivatives are evaluated separately requiring two passes through the computation graph with either relaxed or discrete samples. 7). Many units become deterministic early in training leading to significant computational savings with the sampled RAM method.

In Fig. 4 (a), we include several CR estimators on the same architecture and observe that the GSM estimator performs significantly worse than other estimators. With linear decoders, the objective function Eq. ( 21) is convex in z i . Thus, similar to the example in Fig. 2 , GSM learns a distribution with higher entropy leading to poorer performance. The entropy of all estimators during training is shown in Fig. 4(b) . We note that two-pass training performs better then one-pass training for all estimators. Although two-pass training requires twice the computation in the worst case, this overhead is negligible for these models due to GPU parallelization. We observe that ARM performs poorly confirming that its high variance impedes training. The PWL estimator with two-pass training performs on par with RAM, while being more computationally efficient. Lastly, the REBAR estimator (see Appendix E), using the settings of Tucker et al. (2017) , performs on par with PWL thereby indirectly confirming that the relaxed objective is a good control variate for the discrete one.

We used β = 2 in the above experiments, similar to Jang et al. (2016); Maddison et al. (2016) . To understand the dependence on β, we plot final train NELBO in Fig. 4(c ). We find that improved estimators are less sensitive to the choice of β.

In Fig. 5 , we repeat the experiments for the non-linear architecture 200H ∼ 784V. The one-pass GSM estimator exhibits instability which is remedied by two-pass training. However, two-pass GSM still performs worse than IGSM and PWL due to its bias. Interestingly, one-pass training works better for IGSM/PWL. We observe this repeatedly in the nonlinear models. Unlike the linear case, the RAM estimator converges faster initially but later in training is outperformed by the highervariance IGSM, PWL and ARM. It is likely that additional noise prevents the latent units from being 

5. CONCLUSION:

We have reviewed several finite-difference (FD) and continuous relaxation (CR) estimators of the gradients of the objective Eq. ( 1). FD estimators like RAM Tokui & Sato (2017) can give unbiased low variance estimates but often require multiple function evaluations. We proposed a less expensive version of RAM that requires an order of magnitude fewer function evaluations with a controllable decrease in performance. In contrast, CR estimators, like Gumbel-Softmax (GSM), require a single pass through the objective function and can be computed efficiently giving low variance but biased gradients. We analyzed the nature of the bias introduced by CR estimators and proposed a way to reduce it. This gives rise to improved CR (ICR) estimators, like improved GSM and piece-wise linear.

These ICR estimators are unbiased for a single variable and less biased for many variables. Experiments on VAE training and discrete optimization confirm the theoretical predictions and illustrate the advantages of lower-bias estimators.

A ARGMAX Lorberbom et al. (2018) proposed a FD estimator that we refer to as ARGMAX. For a single variable, ARGMAX relies on an identity that approximates Lorberbom et al. (2018) show that

∂ φ L = ∂ φ q φ [f (1) -f (0)] = (∂ φ l φ )q φ [1 - q φ ][f (1) -f (0)]. With q φ = σ(l φ ),

σ(l φ )[1 -σ(l φ )][f (1) -f (0)] = lim →0 E ρ Θ f (1) + l φ + σ -1 (ρ) -Θ f (0) + l φ + σ -1 (ρ) (22)

where Θ is the Heaviside step function and ρ ∼ U[0, 1]. 6 Lorberbom et al. (2018) approximate the right side of Eq. ( 22) by sampling ρ and evaluating at finite which introduces bias and variance. Decreasing decreases bias but increases variance.

In the categorical case the gradient of Eq. ( 5) can correspondingly be written as

EQUATION

)

where γ b = -log(-log(ρ b )) are Gumbel variables and [pred] is the indicator function (1 if pred is true and 0 otherwise). This derivative requires A function evaluations, similar to RAM. Lorberbom et al. (2018) extend the single variable result to multivariate distributions similar to Eqs. ( 3) and ( 4). ARGMAX has the same computational complexity as RAM, but is biased and has higher variance than RAM. Thus, it is suboptimal to Eqs. ( 3) and ( 4) and for this reason we do not perform experiments with ARGMAX.

B THE AUGMENT-REINFORCE-MERGE ESTIMATOR:

FD estimators are computationally expensive and require multiple function evaluations per gradient. A notable exception is the Augment-REINFORCE-Merge (ARM) estimator introduced in Yin & Zhou (2018) . ARM provides an unbiased estimate using only two function evaluations for the factorized multivariate distribution, regardless of the number of variables. For a single binary variable, the ARM derivative is given by

∂ φ L = ∂ φ l q φ (1 -q φ )(f (1) -f (0)) = ∂ φ l E ρ∼U [0,1] [(f (z (2) ) -f (z (1) ))(ρ -0.5)],

z (1) = Θ(q φ -ρ) z (2) = Θ(ρ -1 + q φ ).

(24) The expectation is approximated with a single sample ρ. This estimator has a significantly lower variance than REINFORCE since the expectation contains the difference f (z (2) ) -f (z (1) ) rather then the function itself. For multivariate f and factorial q φ (z z z) the derivative is

EQUATION

where z

(1) i = Θ(q φ,i -ρ i ) and z

(2) i = Θ(ρ i -1 + q φ,i ). Yin & Zhou (2018) observed that one can replace f (z (2) i , z z z (1) /i ) → f (z (2) i , z z z (2)

/i ) without changing the expectation which allows evaluation by a single sample and two function evaluations f (z z z (1) ), f (z z z (2) ) regardless of M :

EQUATION

)

However, this change comes at the cost of higher variance of the expectation 7 . Thus, while ARM estimator provides a low-variance gradient estimate for a single variable, it introduces high variance 6 We motivate this identity by noting that the non-zero contribution comes from the region σ(-f (1)l φ ) ≤ ρ ≤ σ(-f (0) -l φ ) assuming f (1) > f (0). For small , samples land within this region with probability p ∼ σ(l φ )(1 -σ(l φ )) [f (1) -f (0)] giving rise to the identity. The variance of Eq. ( 22) is

EQUATION

If the functions g (a) (z z z /i ) are highly correlated the ARM estimator will have a much higher variance than RAM.

for multivariate functions compared to RAM. The ARM estimator has straightforward extensions to hierarchical q(z z z) (similar to Eq. ( 4)) and to categorical variables Yin & Zhou (2018) .

C IMPROVED CR ESTIMATOR FOR BAYESIAN NETWORKS:

In this appendix we derive the improved CR estimator for a hierarchical q φ (z z z) = i q φ,i (z i |z z z <i ).

For simplicity, we do this for two variables q φ (z 1 , z 2 ) = q φ,2 (z 2 |z 1 )q φ,1 (z 1 ) with the gradient given in Eq. ( 4):

∂ φ L = ∂ φ q φ,1 z2 q φ,2 (z 2 |1)f (1, z 2 ) - z2 q φ,2 (z 2 |0)f (0, z 2 ) + z1 q φ,1 (z 1 )∂ φ q φ,2 (1|z 1 ) [f (z 1 , 1) -f (z 1 , 0)]

where q φ,1 = q φ,1 (z 1 = 1). We denote the two contributions to ∂ φ L by ∂ φ L (1,2) and determine them separately. Using the integral trick (10) the second contribution can be written as

∂ φ L (2) = z1 q φ,1 (z 1 )∂ φ q φ,2 (1|z 1 ) dρ 1 ∂ζ 2 ∂ρ 2 ∂ ζ2 f (z 1 , ζ 2 ) ≈ E ρ ρ ρ ∂ φ q φ,2 (1|ζ 1 ) ∂ζ 2 ∂ρ 2 ∂ ζ2 f (ζ 1 , ζ 2 )

(28) where ζ 2 = g(ρ 2 , q 2 (1|ζ 1 )) and ζ 1 = g(ρ 1 , q 1 ). The first term can be handled similarly:

∂ φ L (1) = ∂ φ q φ,1 dρ 1 ∂ζ 1 ∂ρ 1 ∂ ζ1 z2 q φ,2 (z 2 |ζ 1 )f (ζ 1 , z 2 ) ≈ E ρ ρ ρ ∂ φ q φ,1 ∂ζ 1 ∂ρ 1 ∂ ζ1 f (ζ 1 , ζ 2 )

(29) Combining these contributions we arrive at an expression very similar to the reparameterization trick with the replacement

EQUATION

D CATEGORICAL PWL ESTIMATOR:

Here, we derive the PWL estimator for categorical variables. The derivative for a single categorical variable y = (y 0 , ...y A-1 ) with y a ∈ {0, 1} and a y a = 1 is given in Eq. ( 5) as

EQUATION

We again apply the integral trick (10) to represent the difference f a -f b . To do that we relax the variable y so that it interpolates between (y a , y b ) = (1, 0) and (y a , y b ) = (0, 1)

as y → y a,b = y a = 0.5 + α a,b (ρ a,b -q b /(q a + q b )) 1 0 , y b = 1 -y a , y c =a,b = 0 where ρ a,b ∈ U[0, 1] and α a,b

is the slope. The relaxed objective then takes the form

EQUATION

where w a,b are weights to be determined. The gradient of this relaxed objective with respect to the logit l a is

EQUATION

Comparison with Eq. ( 31) gives w a,b = (q a + q b ) 2 . However, computing the sum over the edges of the simplex is prohibitively expensive, so we choose to replace it with sampling from the set of edges with probability p a,b = (q a + q b )/( a<b q a + q b ) = (q a + q b )/(A -1). The reason for choosing this distribution is that Eq. ( 32) must give correct value of the objective (not just the derivative) as the relaxation parameter β → ∞, which requires the probability of each state y a = 1 to be equal to q a . For relaxed edge (a, b) the probability of y a = 1 is equal to q a /(q a + q b ), and thus the total probability of y a = 1 is:

b =a p a,b q a q a + q b = b =a q a + q b A -1 q a q a + q b = q a .

(34)

Finally, to get the correct weights w a,b we must rescale each term by a factor γ a,b = (A-1)(q a +q b ) so that w a,b = γ a,b p a,b . In summary, the relaxed objective has the following form:

EQUATION

where ỹa,b has the same value as y a,b but has the gradient scaled by γ a,b = (A -1)(q a + q b ).

E REBAR:

In this section, we derive a simple expression for REBAR gradients Tucker et al. (2017) and describe its connection to ICR estimators. We also present additional experimental results for VAE training with REBAR for different continuous relaxations.

REBAR gradients of the objective (1) can be written as:

∂ φ L = E q φ (z z z)q φ (ζ ζ ζ|z z z) [∂ φ log q φ (z z z) (f (z z z) -f (ζ ζ ζ))] + ∂ φ E q φ (ζ ζ ζ) [f (ζ ζ ζ)] -∂ φ E q sg(φ) (z z z)q φ (ζ ζ ζ|z z z) [f (ζ ζ ζ)] ,

where sg(φ) denotes "stop gradient", and distributions q φ 

∂ φ L = E q φ (z z z)q φ (ζ ζ ζ|z z z) [∂ φ log q φ (z z z) (f (z z z) -f (ζ ζ ζ))]+E ρ ρ ρ [∂ ζ ζ ζ f (ζ ζ ζ)∂ φ ζ ζ ζ(ρ ρ ρ, q φ )]-E q φ (z z z),u [∂ ζ ζ ζ f (ζ ζ ζ)∂ φ ζ ζ ζ(u, q φ |z z z)] .

(36) Focusing on binary variables, the reparameterization for the conditional distribution has the form ζ ζ ζ(u, q φ |z z z) = ζ ζ ζ(ρ ρ ρ, q φ ), where ρ ρ ρ = 1 -q φ + u (z z zq φ -(1 -z z z)(1 -q φ )). In order to minimize the variance, it is preferable to tie the parameters ρ ρ ρ and (z z z, u). This can be achieved by setting z z z = Θ(ρ ρ ρ -1 + q φ ), and choosing

EQUATION

)

This is equivalent to sampling

(z z z, ζ ζ ζ) ∼ q φ (z z z)q φ (ζ ζ ζ|z z z) as z z z = Θ(ρ ρ ρ -1 + q φ ), ζ ζ ζ = ζ ζ ζ(ρ ρ ρ, q φ ) where ρ ρ ρ = 1 -q φ + sg(ρ ρ ρ -1 + q φ ) z z z q φ sg(q φ ) + (1 -z z z) 1 -q φ 1 -sg(q φ )

.

This way ρ ρ ρ = ρ ρ ρ by value but ρ ρ ρ does depend on φ such that ζ ζ ζ(ρ ρ ρ, q φ ) has the correct gradient w.r.t. φ that follows from the conditional distribution q φ (ζ ζ ζ|z z z). We can rewrite (36) as

∂ φ L = E q φ (z z z)q φ (ζ ζ ζ|z z z) [∂ φ log q φ (z z z) (f (z z z) -f (ζ ζ ζ))]+E ρ ρ ρ [∂ ζ ζ ζ f (ζ ζ ζ) (∂ φ ζ ζ ζ(ρ ρ ρ, q φ ) -∂ φ ζ ζ ζ(ρ ρ ρ, q φ ))] . (38) Since the explicit dependence of ζ ζ ζ(ρ ρ ρ, q φ ) and ζ ζ ζ(ρ ρ ρ, q φ ) on q φ is the same, we have ∂ φ ζ ζ ζ(ρ ρ ρ, q φ ) - ∂ φ ζ ζ ζ(ρ ρ ρ, q φ ) = -∂ φ ζ ζ ζ(ρ ρ ρ, sg(q φ ))

and we arrive at our final expression for the REBAR gradient:

EQUATION

)

The advantage of Eq. ( 39) is two-fold: first, it allows for a simple implementation valid for any continuous relaxation and uses only two function evaluations. Second, it gives a suggestive relation to the ICR estimator: indeed using

∂ φ ρ ρ ρ = ∂ φ q φ -1 + z z z ρ ρ ρ-1+q φ q φ -(1 -z z z) ρ ρ ρ-1+q φ 1-q φ

we can write REBAR gradient as Here R 1,2 are correction terms. The term R 2 can have high variance due to the fact that var(∂ q φ log q φ (z z z)) ∼

EQUATION

1 q φ (1-q φ ) while R 1 has variance comparable to ∂ φ L ICR since var(z z z ρ ρ ρ-1+q φ q φ - (1 -z z z) ρ ρ ρ-1+q φ 1-q φ ) ∼ O(1)

. In summary REBAR estimator can be viewed as ICR plus corrections that remove the bias of ICR estimator. In practice the correction term R 2 can increase the variance of REBAR estimator compared to ICR and slow down the training.

We now perform the comparison of two types of continuous relaxations used by REBAR estimator. We train VAE as in Section 4.2 with 200H -784V and 200H ∼ 784V architectures using REBAR with GSM and PWL relaxations. Fig. 6 shows the dependence of the final NELBO on β. In the case of the linear mode 200H -784V we find that PWL relaxation is slightly more stable while for the non-linear model the two relaxations perform similarly.

F EXPERIMENTAL DETAILS:

For the toy example of Section 4.1 we optimize L with respect to q φ (z = 1) using Adam (Kingma & Ba (2014) ) for 2000 iterations with learning rate r = 0.01 using minibatches of size 100 to reduce the variance of gradients. We initialize q φ (z = 1) to the wrong maximum of the relaxed function by setting q φ (z = 1) = σ(5). Heredenotes a linear layer, while ∼ denotes two layers of 200 hidden units with tanh nonlinearity and batch normalization. In our experiments we use the Adam optimizer with default parameters and a fixed learning rate 0.0003, run for 4 • 10 6 steps with minibatches of size 100. We repeat each experiment 5 times with random seed and plot the mean.

G VAE ENCODER TRAINING:

During training of VAEs, the encoder and decoder models interact in complex ways. To eliminate these effects to more directly identify the impact of better gradients, we use a fixed pre-trained decoder, and minimize Eq. ( 21) with respect to encoder parameters φ only. As shown in Fig. 7 , the conclusions from joint training remain unchanged: GSM underperforms due to its bias. Interestingly however, RAM still overfits early in training for the 200H ∼ 784V model, showing that better gradients can negatively affect optimization.

H OPTIMIZATION:

We apply gradient estimators for discrete optimization. Following Patish & Ullman (2018) we study the NP-hard task of finding a maximal clique in a graph. To model this problem, the binary variable z i = 1/0 indicates the presence/absence of vertex i in a maximal clique. If A i,j is the adjacency matrix of the graph and d = i z i is the size of the clique, then the objective function considered 

EQUATION

where κ ∈ [0, 1] is a hyperparameter. Patish & Ullman (2018) use factorial q φ (z) to minimize L[φ] = z z z q φ (z z z)f (z z z). At the end of optimization q φ (z) typically collapses to a single state corresponding to a clique. Most often this clique is a local minimum and not the maximal clique.

We illustrate the performance of CR estimators on a particular graph (1000 nodes, 450000 edges), labeled C1000.9 from the DIMACS data set Johnson & Trick (1996) . We minimize f (z z z) using Adam with default settings, learning rate 0.01 for 40000 iterations. We perform 1000 minimizations in parallel and choose the best clique found at each iteration. Fig. 8 (a),(b) show the size of clique found by each of the estimators for two values of the hyperparameter κ = 0.1, 0.9. At κ = 0.1 the bias of GSM causes trapping in the wrong minima but ICR converges to a good local minimum.

In contrast, at κ = 0.9 the GSM bias accelerates convergence to a good local minimum. We consider a categorical example with convex and concave functions of a single categorical variable y having 10 values. We take f (y) = ± a (g a -y a ) 2 such that g 0 = 0.9, g 1 = 1.1, and g i>1 = 1. The convex function has a minimum at y 1 = 1, while the concave function is minimized at y 0 = 1. We compare 4 estimators for minimizing this function: RAM of Eq. ( 5), PWL of Eq. ( 18), GSM of Eq. ( 19) and IGSM of Eq. ( 20). The probability of the true minimum is shown in Fig. 10 . In both the convex and concave cases the GSM estimator exhibits a bias preventing it from finding the minimum. IGSM is less biased than GSM which allows it to find the true minimum. The PWL estimator is unbiased but has higher variance then IGSM which slows down its optimization. 

:

More generally, the slope of the linear part α can be chosen arbitrarily as long as α ≥ 0.5/min(q, 1 -q) so that both ζ = 0 and ζ = 1 have non-zero probability.

In Tensorflow notation ỹa,b ≡ sg(y a,b ) + sg(γ a,b )(y a,bsg(y a,b )).

In experiments not reported here we observed that the entropy regularizing benefits of the GSM bias can be achieved with explicit entropy regularization of the objective.

Figures and Tables:

Figure1: Continuous relaxations at β = 2 to a Bernoulli random variable Ber(q φ = 0.8): (a) shows the GSM approximation and (b) shows a piece-wise linear approximation to the discontinuous CDF. (c) shows the effect of a modification to GSM that reduces its bias.

Figure 2: Convex, single binary-variable toy example: (a) The relaxed objectivef (ζ) = (ζ -0.45) 2 . (b) The probability q φ (z = 1) during optimization. (c) The gradient ∂ l L of all estimators; the bias of GSM prevents proper minimization.

Figure 3: NELBO of the RAM and sampled RAM estimators on MNIST trained using the 200H -784V architecture having a linear decoder: (a) plots the decrease in training NELBO for RAM and two variants of sampled RAM. (b) shows the computational savings of sampled RAM.

Jang et al. (2016) evaluate both derivatives in one pass (using ζ ζ ζ) thereby introducing bias in the θ θ θ derivatives. We refer to these two possibilities as one/two-pass training and compare their performance. Following Maddison et al. (2016); Jang et al. (2016); Tucker et al. (2017), we consider four architectures with binary variables denoted by 200H -784V, 200H -200H -784V, 200H ∼ 784V, and 200H ∼ 200H ∼ 784V (see Appendix F for details). First, we compare the sampled RAM estimator of Eq. (7) on 200H -784V model for different values of β, where the probability of updating variable z i is p i = 4q i (1 -q i )/β. Fig. 3(a) shows the NELBO on the training set. As expected, increased β leads to increased gradient variance which slows training. Fig. 3(b) shows the average number of function evaluations performed in Eq. (

Figure 4: MNIST training on the linear architecture 200H -784V: (a) compares training NELBO of CR estimators (all estimators use β = 2). Solid/dashed lines correspond to one/two-pass training. (b) shows the entropy of the learned posterior approximations; the GSM bias induces much higher entropy. (c) dependence of the final trained NELBO on β; higher β corresponds to lower bias.

(ζ ζ ζ) and q φ (ζ ζ ζ|z z z) are reparameterizable. Using explicit reparameterizations for ζ ζ ζ(ρ ρ ρ, q φ ) and ζ ζ ζ(u, q φ |z z z) with u ∈ U[0, 1], we can rewrite the gradient as

Figure 6: MNIST training with REBAR using GSM and PWL relaxations, dependence of the final trained NELBO on β: (a) the linear architecture. (b) the non-linear architecture.

VAE training experiments, following Maddison et al. (2016); Jang et al. (2016); Tucker et al. (2017), we consider four architectures with binary variables denoted by 200H -784V, 200H -200H -784V, 200H ∼ 784V, and 200H ∼ 200H ∼ 784V.

Figure 7: Encoder training the with fixed pre-trained decoder on MNIST: (a) the linear architecture. (b) the non-linear architecture.

Figure 9: Single binary variable concave toy example. (a) Relaxed function f (ζ) = -(ζ -0.45) 2 (b) Probability q φ (z = 1) during optimization

Figure 10: The probability of the true minimum for a single categorical-variable toy example with (a) concave function f (y) =a (g a -y a ) 2 and (b) convex function f (y) = a (g a -y a ) 2

Peer review (iclr2019):
 The paper proposed a modification to RAM that allows us to trade decreased computational cost for increased variance. It also proposes an improved continuous relaxation (ICR) estimator to reduce the bias of CR, which is extended to categorical variables. The proposed piece-wise linear relaxation (PWL) can be considered as the inverse CDF of the random variable is very interesting. The ICR estimators can also be extended to categorical variables. The paper is well written. I have some questions: 1. How does the dimension of the variables affect the bias and variance of the proposed estimator? 2. Dose the proposed estimators applicable to hierarchical models with multi-discrete latent variables? 3. What’s the performance of the proposed method compared with the others in terms of running time?