Original: I MPROVED G RADIENT E STIMATORS FOR S TOCHASTIC D ISCRETE V ARIABLES
Revised: I MPROVED G RADIENT E STIMATORS FOR S TOCHASTIC D ISCRETE V ARIABLES

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: In many applications we seek to optimize an expectation with respect to a distribution over discrete variables.
Revised: In many applications we seek to optimize an expectation with respect to a distribution over discrete variables.

Original: Estimating gradients of such objectives with respect to the distribution parameters is a challenging problem.
Revised: Estimating gradients of such objectives with respect to the distribution parameters is a challenging problem.

Original: We analyze existing solutions including ﬁnite-difference (FD) estimators and continuous relaxation (CR) estimators in terms of bias and variance.
Revised: We analyze existing solutions including ﬁnite-difference (FD) estimators and continuous relaxation (CR) estimators in terms of bias and variance.

Original: We show that the commonly used Gumbel-Softmax estimator is biased and propose a simple method to reduce it.
Revised: We show that the commonly used Gumbel-Softmax estimator is biased and propose a simple method to reduce it.

Original: We also derive a simpler piece-wise linear continuous relaxation that also possesses reduced bias.
Revised: We also derive a simpler piece-wise linear continuous relaxation that also possesses reduced bias.

Original: We demonstrate empirically that reduced bias leads to a better performance in variational inference and on binary optimization tasks.
Revised: We demonstrate empirically that reduced bias leads to a better performance in variational inference and on binary optimization tasks.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Discrete stochastic variables arise naturally in many applications including topic modeling, semisupervised learning, clustering, variational memory addressing and reinforcement learning.
Revised: Discrete stochastic variables arise naturally in many applications including topic modeling, semisupervised learning, clustering, variational memory addressing and reinforcement learning.

Original: In many cases, the objective is to minimize an expectation of a function of discrete random variables with respect to the distribution parameters:
Revised: In many cases, the objective is to minimize an expectation of a function of discrete random variables with respect to the distribution parameters:

Original: [Equation]
Revised: [Equation]

Original: Here, zzz is a vector of discrete variables under a φ -parameterized distribution q φ ( zzz ) .
Revised: Here, zzz is a vector of discrete variables under a φ -parameterized distribution q φ ( zzz ) .

Original: For example, in variational inference, f ( zzz ) is the variational lower bound and q φ ( zzz ) is the approximating posterior.
Revised: For example, in variational inference, f ( zzz ) is the variational lower bound and q φ ( zzz ) is the approximating posterior.

Original: Eq.
Revised: Eq.

Original: (1) is commonly minimized by gradient-based methods, which require estimating the gradient ∂ φ L [ φ ] (see Schulman et al.
Revised: (1) is commonly minimized by gradient-based methods, which require estimating the gradient ∂ φ L [ φ ] (see Schulman et al.

Original: for an overview).
Revised: for an overview).

Original: 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.
Revised: 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.

Original: (2014)) gives low variance gradient estimates.
Revised: (2014)) gives low variance gradient estimates.

Original: However, this trick is infeasible for discrete variables due to the discontinuous cumulative distribution function (CDF).
Revised: However, this trick is not directly applicable to discrete variables due to the discontinuous cumulative distribution function (CDF).

Original: Extensions of the reparameterization trick to discrete distributions can be grouped into ﬁnitedifference (FD) estimators, and continuous relaxation (CR) estimators.
Revised: Extensions of the reparameterization trick to discrete distributions can be grouped into ﬁnitedifference (FD) estimators, and continuous relaxation (CR) estimators.

Original: We present a summary of three recent FD estimators (Tokui & Sato (2017); Yin & Zhou (2018); Lorberbom et al.
Revised: We present a summary of three recent FD estimators (Tokui & Sato (2017); Yin & Zhou (2018); Lorberbom et al.

Original: (2018) to guide practitioners to the salient aspects of their bias-variance tradeoff and computational complexity.
Revised: (2018) to guide practitioners to the salient aspects of their bias-variance tradeoff and computational complexity.

Original: Most importantly, we propose a scalable (but still unbiased) variant of Tokui & Sato (2017) that trades decreased computation for slightly increased variance.
Revised: Most importantly, we propose a scalable (but still unbiased) variant of Tokui & Sato (2017) that trades decreased computation for slightly increased variance.

Original: We examine CR estimators including the popular Gumbel-Softmax (Jang et al.
Revised: We examine CR estimators including the popular Gumbel-Softmax (Jang et al.

Original: ; Maddison et al.
Revised: ; Maddison et al.

Original: 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.
Revised: 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.

Original: Based upon this understanding, we develop a piecewise-linear estimator for binary and categorical variables that is simpler and less biased than
Revised: Based upon this understanding, we develop a piecewise-linear estimator for binary and categorical variables that is simpler and less biased than

Original: Gumbel-Softmax.
Revised: Gumbel-Softmax.

Original: 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.
Revised: 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.

Original: 1. R ELATED WORK
Revised: 1. R ELATED WORK

Original: The most generic approximator of ∂ φ L [ φ ] is the score function (SF) estimator (a.k.a.
Revised: The most generic approximator of ∂ φ L [ φ ] is the score function (SF) estimator (a.k.a.

Original: REINFORCE
Revised: REINFORCE

Original: Williams (1992), Glynn (1990)).
Revised: Williams (1992), Glynn (1990)).

Original: SF suffers from high variance and many remedies have been proposed to reduce this variance (Mnih & Gregor (2014); Gregor et al.
Revised: SF suffers from high variance and many remedies have been proposed to reduce this variance (Mnih & Gregor (2014); Gregor et al.

Original: ; Gu et al.
Revised: ; Gu et al.

Original: ; Mnih & Rezende (2016); Tucker et al.
Revised: ; Mnih & Rezende (2016); Tucker et al.

Original: ; Grathwohl et al.
Revised: ; Grathwohl et al.

Original: Unbiased estimators that require multiple function evaluations have also been proposed Tokui & Sato (2017), Titsias & L´azaroGredilla (2015), Lorberbom et al.
Revised: Unbiased estimators that require multiple function evaluations have also been proposed Tokui & Sato (2017), Titsias & L´azaroGredilla (2015), Lorberbom et al.

Original: (2018), Yin & Zhou (2018).
Revised: (2018), Yin & Zhou (2018).

Original: These estimators have lower variance but can be computationally demanding.
Revised: These estimators have lower variance but can be computationally demanding.

Original: CR estimators often trade bias for variance and previous
Revised: CR estimators often trade bias for variance and previous

Original: CR proposals include straight-through estimators (Bengio et al.
Revised: CR proposals include straight-through estimators (Bengio et al.

Original: ; Raiko et al.
Revised: ; Raiko et al.

Original: the
Revised: the

Original: Gumbel-Softmax estimator (Jang et al.
Revised: Gumbel-Softmax estimator (Jang et al.

Original: ; Maddison et al.
Revised: ; Maddison et al.

Original: and overlapping smoothing (Vahdat et al.
Revised: and overlapping smoothing (Vahdat et al.

Original: (2018b;a); Rolfe (2016)).
Revised: (2018b;a); Rolfe (2016)).

Original:  T HE R EPARAMETERIZATION AND M ARGINALIZATION E STIMATOR
Revised:  T HE R EPARAMETERIZATION AND M ARGINALIZATION E STIMATOR

Original: We begin a summary of FD estimators with the reparameterization and marginalization (RAM) method of Tokui & Sato (2017).
Revised: We begin a summary of FD estimators with the reparameterization and marginalization (RAM) method of Tokui & Sato (2017).

Original: For a single binary variable the expectation in Eq.
Revised: For a single binary variable the expectation in Eq.

Original: (1) is enumerated as L = (cid:80)
Revised: (1) is enumerated as L = (cid:80)

Original: z q φ ( z ) f
Revised: z q φ ( z ) f

Original: [Equation]
Revised: [Equation]

Original: z ) =
Revised: z ) =

Original: q
Revised: q

Original: [Equation]
Revised: [Equation]

Original: f
Revised: f

Original: [Equation]
Revised: [Equation]

Original: 1) +
Revised: +

Original: [Equation]
Revised: [Equation]

Original: φ ) f (0) , where q φ ≡ q φ ( z = 1) .
Revised: ) f (0) , where q φ ≡ q φ ( z = 1) .

Original: The gradient is:
Revised: The gradient is:

Original: [Equation]
Revised: [Equation]

Original: where q φ = σ ( l φ ) = (1 + e − l φ ) − 1 and l φ = logit ( q φ ) .
Revised: where q φ = σ ( l φ ) = (1 + e − l φ ) − 1 and l φ = logit ( q φ ) .

Original: This derivative involves two function evaluations and contains a ﬁnite-difference of f ( z ) .
Revised: This derivative involves two function evaluations and contains a ﬁnite-difference of f ( z ) .

Original: Eq.
Revised: Eq.

Original: (2) is an unbiased zero-variance estimate since the summation over z is done explicitly.
Revised: (2) is an unbiased zero-variance estimate since the summation over z is done explicitly.

Original: The generalization to M factorially-distributed random variables q φ ( zzz ) = (cid:81) Mi
Revised: The generalization to M factorially-distributed random variables q φ ( zzz ) = (cid:81) Mi

Original: =1 q φ ,i ( z i ) is
Revised: =1 q φ ,i ( z i ) is

Original: [Equation]
Revised: [Equation]

Original: where again the summation over z i is performed exactly and q φ ,i ≡ q φ ,i ( z i = 1) .
Revised: where again the summation over z i is performed exactly and q φ ,i ≡ q φ ,i ( z i = 1) .

Original: The summation over zzz \ i can be estimated with a single sample but the derivative requires M +1 function evaluations.
Revised: The summation over zzz \ i can be estimated with a single sample but the derivative requires M +1 function evaluations.

Original: This limits the applicability of RAM to moderately-sized models.
Revised: This limits the applicability of RAM to moderately-sized models.

Original: Note that both f ( z i = 0 ,zzz \ i ) and f ( z i = 1 ,zzz \ i ) are evaluated at the same zzz \ i which leads to a lower-variance estimator.
Revised: Note that both f ( z i = 0 ,zzz \ i ) and f ( z i = 1 ,zzz \ i ) are evaluated at the same zzz \ i which leads to a lower-variance estimator.

Original: For hierarchical q φ ( zzz ) =
Revised: For hierarchical q φ ( zzz ) =

Original: (cid:81) Mi
Revised: (cid:81) Mi

Original: =1 q φ ,i ( z i | zzz <i ) the derivative takes the form:
Revised: =1 q φ ,i ( z i | zzz <i ) the derivative takes the form:

Original: [Equation]
Revised: [Equation]

Original: where q φ ,i ≡ q φ ,i (1 | zzz <i ) .
Revised: where q φ ,i ≡ q φ ,i (1 | zzz <i ) .

Original: The key insight of Tokui & Sato (2017) is to use common random variates zzz >i when sampling from q φ ( zzz >i | 1 ,zzz <i ) and q φ ( zzz >i | 0 ,zzz <i ) .
Revised: The key insight of Tokui & Sato (2017) is to use common random variates zzz >i when sampling from q φ ( zzz >i | 1 ,zzz <i ) and q φ ( zzz >i | 0 ,zzz <i ) .

Original: This reduces both the variance and the computational cost.
Revised: This reduces both the variance and the computational cost.

Original: Tokui & Sato (2017) show that this estimator is optimal because it exactly sums over the binary variables whose probability distribution is being differentiated.
Revised: Tokui & Sato (2017) show that this estimator is optimal because it exactly sums over the binary variables whose probability distribution is being differentiated.

Original: A RAM estimator can also be constructed for categorical variables.
Revised: A RAM estimator can also be constructed for categorical variables.

Original: For a single one-hot encoded categorical variable y
Revised: For a single one-hot encoded categorical variable y

Original: =
Revised: =

Original: ( y 0 , ...
Revised: ( y 0 , ...

Original: y A − 1 ) , y a ∈ { 0 , 1 } , (cid:80) a y a = 1 , the derivative of Eq.
Revised: y A − 1 ) , y a ∈ { 0 , 1 } , (cid:80) a y a = 1 , the derivative of Eq.

Original: (1) is
Revised: (1) is

Original: [Equation]
Revised: [Equation]

Original: where f a = f ( y a = 1) , and q a φ = e l a φ / (cid:80)
Revised: where f a = f ( y a = 1) , and q a φ = e l a φ / (cid:80)

Original: b e l b φ .
Revised: b e l b φ .

Original: The generalization to many categorical variables proceeds as for binary variables.
Revised: The generalization to many categorical variables proceeds as for binary variables.

Original: For example, the derivative of a factorial categorical distribution over yyy =
Revised: For example, the derivative of a factorial categorical distribution over yyy =

Original: { y ai | 0 ≤ a < A, 1 ≤ i ≤ M } is
Revised: { y ai | 0 ≤ a < A, 1 ≤ i ≤ M } is

Original: [Equation]
Revised: [Equation]

Original: This derivative can again be estimated with a single sample but requires MA function evaluations.
Revised: This derivative can again be estimated with a single sample but requires MA function evaluations.

Original: 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.
Revised: 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.

Original: 2. S AMPLED R EPARAMETERIZATION AND M ARGINALIZATION
Revised: 2. S AMPLED R EPARAMETERIZATION AND M ARGINALIZATION

Original: We propose a modiﬁcation to RAM that allows us to trade decreased computational cost for increased variance.
Revised: We propose a modiﬁcation to RAM that allows us to trade decreased computational cost for increased variance.

Original: For binary variables, ∂ φ q φ ,i =
Revised: For binary variables, ∂ φ q φ ,i =

Original: q φ ,i (1 − q φ ,i ) ∂
Revised: q φ ,i (1 − q φ ,i ) ∂

Original: φ l φ ,i where q φ ,i =
Revised: φ l φ ,i where q φ ,i =

Original: σ ( l φ ,i ) , so that each term in Eq.
Revised: σ ( l φ ,i ) , so that each term in Eq.

Original: (3) or (4) is proportional to q φ ,i (1 − q φ ,i ) .
Revised: (3) or (4) is proportional to q φ ,i (1 − q φ ,i ) .

Original: 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.
Revised: 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.

Original: Such variables have negligible contribution to the full derivative.
Revised: Such variables have negligible contribution to the full derivative.

Original: We exploit this observation to reduce the computational cost by including variable z i in the full gradient with probability p
Revised: We exploit this observation to reduce the computational cost by including variable z i in the full gradient with probability p

Original: i
Revised: i

Original: =
Revised: =

Original: q φ ,i (1 − q φ ,i ) /β , where β is an adjustable hyperparameter.
Revised: q φ ,i (1 − q φ ,i ) /β , where β is an adjustable hyperparameter.

Original: This means that we replace the derivative in Eq.
Revised: This means that we replace the derivative in Eq.

Original: (3) with
Revised: (3) with

Original: [Equation]
Revised: [Equation]

Original: where ξ i ∈ { 0 , 1 } are Bernoulli variables with probabilities p
Revised: where ξ i ∈ { 0 , 1 } are Bernoulli variables with probabilities p

Original: i indicating whether z i is included or not.
Revised: i indicating whether z i is included or not.

Original: We evaluate Eq.
Revised: We evaluate Eq.

Original: (7) by sampling ξ and only then evaluating non-zero terms.
Revised: (7) by sampling ξ and only then evaluating non-zero terms.

Original: 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.
Revised: 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.

Original: As always, this computational saving comes at the cost of increased gradient variance which slows training.
Revised: As always, this computational saving comes at the cost of increased gradient variance which slows training.

Original: In the categorical case, each term in Eq.
Revised: In the categorical case, each term in Eq.

Original: (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 .
Revised: (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 .

Original: The computational cost MA can be reduced by keeping each term with probability p a,bi = 4 q a φ ,i q b φ ,i /β .
Revised: The computational cost MA can be reduced by keeping each term with probability p a,bi = 4 q a φ ,i q b φ ,i /β .

Original: Appendices A and B summarize two other FD estimators, ARGMAX Lorberbom et al.
Revised: Appendices A and B summarize two other FD estimators, ARGMAX Lorberbom et al.

Original: and
Revised: and

Original: ARM Yin & Zhou (2018), by noting their bias-variance tradeoff and computational complexity.
Revised: ARM Yin & Zhou (2018), by noting their bias-variance tradeoff and computational complexity.

Original:  C ONTINUOUS R ELAXATION E STIMATORS
Revised:  C ONTINUOUS R ELAXATION E STIMATORS

Original: 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 → ζ = g βφ ( ρ )for ρ ∈ U [0 , 1] , where β is a parameter that controls the approximation.
Revised: 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 ρ .

Original: The objective function L [ φ ] = (cid:80) zzz q φ ( zzz ) f ( zzz ) is replaced with ˜ L [ φ ] =E [Equation]
Revised: The objective function L [ φ ] = (cid:80) zzz q φ ( zzz ) f ( zzz ) is replaced with ˜ L [ φ ] =

Original: = g βφ [Equation] [Equation] and the gradients are computed using  chain rule
Revised: E ρρρ (cid:2) f (cid:0) ζζζ ( ρ, q φ ) (cid:1)(cid:3) and its gradients are computed using the chain rule

Original: [Equation]
Revised: [Equation]

Original: These gradients can be computed efﬁciently by automatic differentiation libraries.
Revised: These gradients can be computed efﬁciently by automatic differentiation libraries.

Original: However, because the objective function is changed, CR estimators in Eq.
Revised: However, because the objective function is changed, CR estimators in Eq.

Original: (8) are biased.
Revised: (8) are biased.

Original: Nevertheless, in practice the bias is often small enough to allow for effective optimization.
Revised: Nevertheless, in practice the bias is often small enough to allow for effective optimization.

Original: 3. T HE G UMBEL -S OFTMAX E STIMATOR
Revised: 3. T HE G UMBEL -S OFTMAX E STIMATOR

Original: The most popular CR estimator is Gumbel-Softmax (GSM)
Revised: The most popular CR estimator is Gumbel-Softmax (GSM)

Original: (Jang et al.
Revised: (Jang et al.

Original: Maddison et al.
Revised: Maddison et al.

Original: For binary variables, this relaxation takes the simple form
Revised: For binary variables, this relaxation takes the simple form

Original: [Equation]
Revised: [Equation]

Original:  where q i ≡ q φ,i (1 | ζζζ <i ) .
Revised: = 4 q (1 q ) z q (a) Gumbel-Softmax 0. z q (b) Piecewise Linear 0. q (c) ∂ q ζ where q i ≡ q φ,i (1 | ζζζ <i ) .

Original: This relaxation and its derivative ∂ q ζ are shown in Fig. 1(a).
Revised: This relaxation and its derivative ∂ q ζ are shown in Fig. 1(a).

Original: β controls the sharpness of the relaxation and tunes the trade-off between the bias (closeness of ζ to z ) and variance of ∂ q ζ .
Revised: β controls the sharpness of the relaxation and tunes the trade-off between the bias (closeness of ζ to z ) and variance of ∂ q ζ .

Original: We note that the slope of the relaxed ζ ( ρ ) at ζ = 1 / 2 is α ≡ β/ [4 q (1 − q )] and thus becomes large when q approaches 0 or 1 .= 4 q (1 q ) z q (a) Gumbel-Softmax 0. z q (b) Piecewise Linear 0.
Revised: We note that the slope of the relaxed ζ ( ρ ) at ζ = 1 / 2 is α ≡ β/ [4 q (1 − q )] and thus becomes large when q approaches 0 or 1 .

Original: q (c) ∂

Original: q ζ

Original: 3. I MPROVED C ONTINUOUS R ELAXATION E STIMATORS
Revised: 3. I MPROVED C ONTINUOUS R ELAXATION E STIMATORS

Original: In this section, we analyze the bias introduced by Eq.
Revised: In this section, we analyze the bias introduced by Eq.

Original: (8) and propose a simple method to reduce it.
Revised: (8) and propose a simple method to reduce it.

Original: The bias of E ρρρ [ ∂ ζ i f ( ζζζ ) ∂ φ ζ i ] has two sources: (a) the relaxation of ζ j for j (cid:54) =
Revised: The bias of E ρρρ [ ∂ ζ i f ( ζζζ ) ∂ φ ζ i ] has two sources: (a) the relaxation of ζ j for j (cid:54) =

Original: i and (b) the relaxation of ζ i .
Revised: i and (b) the relaxation of ζ i .

Original: To characterize the latter bias, we start with a single binary variable and write the gradient as the following integral:
Revised: To characterize the latter bias, we start with a single binary variable and write the gradient as the following integral:

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: Here, ζ ( ρ ) is any continuous function satisfying ζ (0) = [-
-] and ζ (1) = 1 .
Revised: Here, ξ ( ρ ) is any continuous function satisfying ξ (0) = [-
-]0 and ξ (1) = 1 .

Original: For a non-decreasing function ζ ( ρ ) , we can view ζ as a random variable with inverse CDF ζ ( ρ ) and ρ as a uniform random variable ρ ∈ U [0 , 1] .
Revised: For a non-decreasing function ξ ( ρ ) , we can view ξ as a random variable with inverse CDF ξ ( ρ ) and ρ as a uniform random variable ρ ∈ U [0 , 1] .

Original: Given that the corresponding probability density function (PDF) is q ( ζ ) = ( ∂ζ/∂ρ ) − 1 , we rewrite the derivative as 
Revised: With this interpretation the gradient can be written as an expectation

Original: [Equation]
Revised: [Equation]

Original: This expectation can be estimated by sampling ζ ∼ q ( ζ ) . For smooth enough f ( ζ )  ,  the variance of this estimate is controlled by var (cid:0) q ( ζ )− 1 (cid:1) =
Revised: which can be estimated by sampling. If ∂ ξ f ( ξ ) does not vary signiﬁcantly in the interval ξ ∈ [0 , 1] then the variance of this estimate is controlled by var( ∂ξ/∂ρ )

Original:  (cid:82) 1 0 dζ q{+
 dρ+} ( ζ ) − 1 − 1 .
Revised: = (cid:82) 1 {+
 dρ+} ( ∂ξ/∂ρ ) 2 − 1 .

Original: Thus, the more non-linear ζ ( ρ ) is, the higher will be the variance of estimate Eq.
Revised: Thus, the more non-linear ξ ( ρ ) is, the higher will be the variance of estimate Eq.

Original: This idea can be extended to factorial q φ ( zzz ) =
Revised: This idea can be extended to factorial q φ ( zzz ) =

Original: (cid:81) Mi
Revised: (cid:81) Mi

Original: =1 q φ ,i ( z i ) as in Eq.
Revised: =1 q φ ,i ( z i ) as in Eq.

Original: [Equation]
Revised: [Equation]

Original: At this point there is no relationship between q ( ζ i) andq φ ,i ( z i) , and Eq.
Revised: At this point there is no relationship between ξ i and z i , and Eq.

Original: (12) is just a higher variance version of Eq.
Revised: (12) is just a higher variance version of Eq.

Original: However, if we relax zzz \ i → ζζζ \ i  and choose ζ j = g ( ρ j ; q j ) we obtain a biased estimator
Revised: However, if we relax zzz \ i → ζζζ \ i ( ρ, q φ ) and choose ξ i ( ρ i ) = ζ i ( ρ i , q φ ,i ) we obtain a biased estimator

Original: [Equation]
Revised: [Equation]

Original: where the bias comes from the deviation of ζζζ \ i from zzz \ i in the function evaluations.
Revised: where the bias comes from the deviation of ζζζ \ i from zzz \ i in the function evaluations.

Original: In other words,
Revised: In other words,

Original: Eq.
Revised: Eq.

Original: (13) uses an unbiased form for the differentiated variable ζ i and the only bias comes from relaxing the remaining variables ζζζ \ i .
Revised: (13) uses an unbiased form for the differentiated variable ζ i and the only bias comes from relaxing the remaining variables ζζζ \ i .

Original: The variance of each term is controlled by var (cid:0) q ( ζ i ) − 1 (cid:1) and is reduced by making ζ i ( ρ i  ) more linear.
Revised: The variance of each term is controlled by var (cid:16) ∂ζ i ∂ρ i (cid:17) and is reduced by making ζ i ( ρ i , q φ ,i ) more linear.

Original: The bias is reduced by making ζζζ closer zzz .
Revised: The bias is reduced by making ζζζ closer zzz .

Original: Varying ζ i ( ρ i  ) between linear (low variance) and step function (low bias) allows for a controllable trade-off.
Revised: Varying ζ i ( ρ i , q φ ,i ) between linear (low variance) and step function (low bias) allows for a controllable trade-off.

Original: We call Eq.
Revised: We call Eq.

Original: (13) the improved continuous relaxation (ICR) estimator.
Revised: (13) the improved continuous relaxation (ICR) estimator.

Original: The original CR estimator of Eq.
Revised: The original CR estimator of Eq.

Original: (8) for factorial q φ has the form
Revised: (8) for factorial q φ has the form

Original: [Equation]
Revised: [Equation]

Original: and comparing this to the ICR estimate in Eq.
Revised: and comparing this to the ICR estimate in Eq.

Original: (13), we see that CR can be transformed into ICR by the replacing ∂ q i ζ i with ∂ ρ i ζ i .
Revised: (13), we see that CR can be transformed into ICR by the replacing ∂ q i ζ i with ∂ ρ i ζ i .

Original: This change can be simply implemented using TensorFlow’s stop gradient ≡ sg notation by replacingζ i ( ρ i , q i ) with ζ i (cid:0)
Revised: This change can be simply implemented using TensorFlow’s stop gradient ≡ sg notation by replacing

Original: [Equation]
Revised: [Equation]

Original: in Eq.
Revised: in Eq.

Original: We emphasize that the ICR estimator in Eq.
Revised: We emphasize that the ICR estimator in Eq.

Original: (13) is less biased than the direct CR estimator because, in the case of a single variable, ICR is unbiased while CR is not.
Revised: (13) is less biased than the direct CR estimator because, in the case of a single variable, ICR is unbiased while CR is not.

Original: Further, this decrease in bias is not accompanied by an increase in variance.
Revised: Further, this decrease in bias is not accompanied by an increase in variance.

Original: We show in Appendix C that similar beneﬁts are obtained for hierarchical q φ .
Revised: We show in Appendix C that similar beneﬁts are obtained for hierarchical q φ .

Revised: Finally, we note a conceptual relationship between sampled

Revised: RAM and ICR: both estimators evaluate the gradient only through a subset of variables z i .

Revised: Sampled

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

Original: 3. P IECE - WISE L INEAR R ELAXATION
Revised: 3. P IECE - WISE L INEAR R ELAXATION

Original: Inspired by ICR and a better understanding of variance-bias trade-off, we propose a piece-wise linear relaxation (PWL) depicted in Fig. 1(b).
Revised: 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).

Original: The linear part is centered at ρ = 1 − q so that the corresponding binary variable is obtained by z = round( ζ ) .
Revised: The linear part is centered at ρ = 1 − q so that the corresponding binary variable is obtained by z = round( ζ ) .

Original: The slope is given by α = β/ [4 q (1 − q )] similar to the Gumbel-Softmax slope 2 .
Revised: The slope is given by α = β/ [4 q (1 − q )] similar to the Gumbel-Softmax slope 3 .

Original: The explicit expression for PWL smoothing is
Revised: The explicit expression for PWL smoothing is

Original: [Equation]
Revised: [Equation]

Original: (16) where [ x ] 10 ≡ min (1 , max [0 , x ]) is the hard sigmoid function.
Revised: (16) where [ x ] 10 ≡ min (1 , max [0 , x ]) is the hard sigmoid function (note that ∂ q α = 0 ).

Original: This relaxation has several attractive properties.
Revised: This relaxation has several attractive properties.

Original: Firstly, we have ∂ q i ζ i =
Revised: Firstly, we have ∂ q i ζ i =

Original: ∂
Revised: ∂

Original: ρ i ζ i , which means that the CR and ICR estimators coincide for PWL.
Revised: ρ i ζ i , which means that the CR and ICR estimators coincide for PWL.

Original: Secondly, PWL has easily interpretable expressions for bias and variance.
Revised: Secondly, PWL has easily interpretable expressions for bias and variance.

Original: In the case of a function of a single variable the variance is given by var( ∂ ρ ζ ) =
Revised: In the case of a function of a single variable the variance is given by var( ∂ ρ ζ ) =

Original: α − 1 while the bias in computing expectation (cid:80)
Revised: α −while the bias in computing expectation (cid:80)

Original: z f ( z ) over the relaxed distribution is (cid:82) dρf ( ζ ( ρ ))
Revised: z f ( z ) over the relaxed distribution is (cid:82) dρf ( ζ ( ρ ))

Original: − (cid:80) z
Revised: − (cid:80) z

Original: f ( z ) = (cid:16)(cid:82) 1 [-
-] dxf ( x ) − [ f (0) + f (1)] / 2 (cid:17) /α .
Revised: f ( z ) = { (cid:82) 1 [-
-]0 dxf ( x ) − [ f (0) + f (1)] / 2 } /α .

Original: This clearly shows that α trades bias for variance.
Revised: This clearly shows that α trades bias for variance.

Original: Thirdly, we provide a qualitative motivation for the slope to have the form α ≈ 1 /

Original: [ q (1 − q )] .

Original: We recall that the gradient in Eq.

Original: (13) has the form ∂ φ L ≈ E q ( ζζζ ) [ (cid:80) i ∂ ζ i f ( ζζζ )(

Original: ∂ ρ i ζ i ) q i (1 − q i ) ∂

Original: φ l φ ,i ] .

Original: For smooth enough f ( ζζζ ) , the variance of this estimate arises mainly from the variance of ∂ ρ i ζ i and the relative size of the terms is controlled by q i (1 − q i ) .

Original: Since the variance of ∂ ρ i ζ i is given by var( ∂ ρ i ζ i ) =

Original: α

Original: i − 1 , to minimize the variance of the sum, we can choose var( ∂ ρ i ζ i ) ≈ 1 /

Original: [ q i (1 − q i )] , leading to α i ≈ 1 / [ q i (1 − q i )] .

Original: Finally, the PWL relaxation deﬁned in Eq.
Revised: Finally, the PWL relaxation deﬁned in Eq.

Original: (16) can be considered as the inverse CDF of the PDF:
Revised: (16) can be considered as the inverse CDF of the PDF

Original: [Equation]
Revised: [Equation]

Original: [Equation] 1] , where (cid:15) =
Revised: ] , where (cid:15) [Equation] [Equation]

Original: [4 q (1 − q )]/β.
Revised: [Equation] q )]

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (17) is a mixture of two delta distributions centered at zero and one, and a uniform distribution deﬁned in the interval [0 , 1] .
Revised: (17) is a mixture of two delta distributions centered at zero and one, and a uniform distribution deﬁned in the interval [0 , 1] .

Original: Samples from q ( ζ ) are in the continuous interval with probability [4 q (1 − q )] /β , and ζ = 0 / 1 have probability proportional to the probability of the binary states.
Revised: Samples from q ( ζ ) are in the continuous interval with probability [4 q (1 − q )] /β , and ζ = 0 / 1 have probability proportional to the probability of the binary states.

Original: 3. C ATEGORICAL P IECE - WISE L INEAR R ELAXATION
Revised: 3. C ATEGORICAL P IECE - WISE L INEAR R ELAXATION

Original: We now extend ICR estimators to categorical variables.
Revised: We now extend ICR estimators to categorical variables.

Original: For a single categorical variable we apply the integral representation to each edge ( a, b ) of the simplex in Eq.
Revised: For a single categorical variable we apply the integral representation to each edge ( a, b ) of the simplex in Eq.

Original: (5) and relax this pair of variables using PWL: y → y a,b = { y a = (cid:2) 0 .
Revised: (5) and relax this pair of variables using PWL: y → y a,b = { y a = (cid:2) 0 .

Original: 5 + α a, [Equation] (cid:0) ρ a,
Revised: 5 + α a,b (cid:0) ρ a,b − q b

Original: [Equation] − q [Equation] / ( q a + q 
Revised: ( q a + q b ) (cid:1)(cid:3) 1 0 , y b = 1 − y a , y c (cid:54) =

Original: [Equation] ) (cid:1)(cid:3) 1 0 , y b = 1 − y a , y c (cid:54) = a,b = 0 } where ρ a,b ∼ U [0 , 1] and α a,b is the slope.
Revised: a,b = 0 } where ρ a,b ∼ U [0 , 1] and α a,b is the slope.

Original: We replace the summation over the edges of the simplex by sampling one edge at a time with probability p a,b
Revised: We replace the summation over the edges of the simplex by sampling one edge at a time with probability p a,b

Original: = ( q a + q b ) /
Revised: = ( q a + q b ) /

Original: ( A − 1) .
Revised: ( A − 1) .

Original: Details are found in
Revised: Details are found in

Original: Appendix D and we provide the ﬁnal result:
Revised: Appendix D and we provide the ﬁnal result:

Original: [Equation]
Revised: [Equation]

Original: (18) where ˜ y a,b has the same value as y a,b but has the gradient scaled by γ a,b
Revised: 18) where ˜ y a,b has the same value as y a,b but has the gradient scaled by γ a,b

Original: = ( A − 1)( q a + q b ) .
Revised: = ( A − 1)( q a + q b ) .

Original: The probabilities of edge selection and the scale factor are chosen to give correct values for the objective and its gradient.
Revised: The probabilities of edge selection and the scale factor are chosen to give correct values for the objective and its gradient.

Original: Extension of this categorical PWL estimator to multivariate distributions is straightforward; For example, for factorial distributions one relaxes each categorical variable
Revised: Extension of this categorical PWL estimator to multivariate distributions is straightforward; For example, for factorial distributions one relaxes each categorical variable

Original: y i by sampling an edge with probabilityp a,b i =
Revised: y i by sampling an edge with probability

Original:  ( q ai + q b i ) ( A − 1) .
Revised: p a,bi = ( q ai + q bi )( A − 1) and uses a single function evaluation.

Original: The resulting gradient is unbiased for a single variable and introduces a bias in the multivariate case which makes it an ICR estimator.
Revised: The resulting gradient is unbiased for a single variable and introduces a bias in the multivariate case which makes it an ICR estimator.

Original: f ( ) (a) Relaxed function    2000 Steps 0.
Revised: f ( ) (a) Relaxed function    2000 Steps 0.

Original: q ( z = 1 )
Revised: q ( z = 1 )

Original: GSM IGSM
Revised: GSM IGSM

Original: PWL RAM
Revised: PWL RAM

Original: ARM (b) Probability q ( z = 1) 0.
Revised: ARM (b) Probability q ( z = 1) 0.

Original: 1.0 q(z=1) 0.02 0.
Revised: 1.0 q(z=1) 0.02 0.

Original: l
Revised: l

Original: GSMIGSMPWLRAMARM (c) Gradient ∂ l L
Revised: GSMIGSMPWLRAMARM (c) Gradient ∂ l L

Original: 3. I MPROVED C ATEGORICAL G UMBEL -S OFTMAX E STIMATOR
Revised: 3. I MPROVED C ATEGORICAL G UMBEL -S OFTMAX E STIMATOR

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

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

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

Original: [Equation]
Revised: [Equation]

Original: a )) , where
Revised: a )) , where

Original: [Equation]
Revised: [Equation]

Original: In Eq.
Revised: In Eq.

Original: (20), we propose a new version of this estimator by applying the same trick as in Eq.
Revised: (20), we propose a new version of this estimator by applying the same trick as in Eq.

Original: This estimator remains biased but we empirically demonstrate that its bias is reduced.
Revised: This estimator remains biased but we empirically demonstrate that its bias is reduced.

Original: [Equation]
Revised: [Equation]

Original: sg ( q )) .
Revised: sg ( q )) .

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section we compare the FD and CR estimators and their improved variants on a number of examples.
Revised: In this section we compare the FD and CR estimators and their improved variants on a number of examples.

Original: 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.
Revised: 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.

Original: 4. T OY EXAMPLE
Revised: 4. T OY EXAMPLE

Original: We begin with an illustrative single-variable example (Tucker et al.
Revised: We begin with an illustrative single-variable example (Tucker et al.

Original: with objective L = (cid:80)
Revised: with objective L = (cid:80)

Original: z q φ ( z ) f
Revised: z q φ ( z ) f

Original: ( z ) where f ( z ) =
Revised: ( z ) where f ( z ) =

Original: ( z − 0 . 45) 2 .
Revised: ( z − 0 . 45) 2 .

Original: The relaxed convex function f
Revised: The relaxed convex function f

Original: ( ζ ) depicted in Fig.
Revised: ( ζ ) depicted in Fig.

Original: 2(a) has two local maxima, one of which is the minimum over the discrete domain.
Revised: 2(a) has two local maxima, one of which is the minimum over the discrete domain.

Original: We compare ﬁve gradient methods: RAM, Eq.
Revised: We compare ﬁve gradient methods: RAM, Eq.

Original: ARM (see Appendix B); PWL, Eq.
Revised: ARM (see Appendix B); PWL, Eq.

Original: (16); GSM, Eq.
Revised: (16); GSM, Eq.

Original: (9); and improved Gumbel-Softmax (IGSM), Eq.
Revised: (9); and improved Gumbel-Softmax (IGSM), Eq.

Original: (15) (see Appendix E for experimental details).
Revised: (15) (see Appendix F for experimental details).

Original: Fig. 2(b) shows the evolution of q φ ( z = 1) during training which demonstrates the bias associated with
Revised: Fig. 2(b) shows the evolution of q φ ( z = 1) during training which demonstrates the bias associated with

Original: GSM.
Revised: GSM.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: To understand the nature of the bias that GSM introduces, we plot the derivatives ∂ q ζ and ∂ ρ ζ corresponding to GSM and IGSM respectively in Fig.
Revised: To understand the nature of the bias that GSM introduces, we plot the derivatives ∂ q ζ and ∂ ρ ζ corresponding to GSM and IGSM respectively in Fig.

Original: 1(c).
Revised: 1(c).

Original: We see that ∂ q ζ is biased towards the value of z = round( ζ ) that has the highest probability.
Revised: We see that ∂ q ζ is biased towards the value of z = round( ζ ) that has the highest probability.

Original: This means that GSM in Eq.
Revised: This means that GSM in Eq.

Original: (14) will oversample the derivative ∂ ζ f ( ζ ) from the most probable mode.
Revised: (14) will oversample the derivative ∂ ζ f ( ζ ) from the most probable mode.

Original: In example Fig.
Revised: In example Fig.

Original: 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 .
Revised: 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 .

Original: The bias is reduced as β is increased.
Revised: The bias is reduced as β is increased.

Original: In Appendix H.1 we consider an example for a concave function over a binary variable and observe similar effects.
Revised: In Appendix I.1 we consider an example for a concave function over a binary variable and observe similar effects.

Original: We also consider both concave and convex functions over a categorical variable in
Revised: We also consider both concave and convex functions over a categorical variable in

Original: Appendix H.2.
Revised: Appendix I.2.

Original: In all scenarios, the bias of GSM prevents its convergence to the correct minimum.
Revised: In all scenarios, the bias of GSM prevents its convergence to the correct minimum.

Original: [Figure]
Revised: [Figure]

Original: N E L B O
Revised: N E L B O

Original: RAM
Revised: RAM

Original: [Figure]
Revised: [Figure]

Original: a) NELBO    4000 Steps(thousands)
Revised: a) NELBO    4000 Steps(thousands)

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: % 40% 60% 80% 100% e v a l u a t i o n s ( %
Revised: % 40% 60% 80% 100% e v a l u a t i o n s ( %

Original: M )
Revised: M )

Original: [Figure]
Revised: [Figure]

Original: 4. D ISCRETE V ARIATIONAL A UTOENCODERS
Revised: 4. D ISCRETE V ARIATIONAL A UTOENCODERS

Original: Next, we test the estimators by training variational autoencoders Kingma & Welling (2013) with discrete priors.
Revised: Next, we test the estimators by training variational autoencoders Kingma & Welling (2013) with discrete priors.

Original: The objective is the negative expectation lower bound on the log-likelihood (NELBO):
Revised: The objective is the negative expectation lower bound on the log-likelihood (NELBO):

Original: L [ φ ,θ θ θ ] = (cid:88)
Revised: L [ φ ,θ θ θ ] = (cid:88)

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: Here, p θθθ ( zzz ) is the prior, p θθθ ( xxx | zzz ) is the decoder, and q φ ( zzz | xxx ) is the approximating posterior.
Revised: Here, p θθθ ( zzz ) is the prior, p θθθ ( xxx | zzz ) is the decoder, and q φ ( zzz | xxx ) is the approximating posterior.

Original: L is minimized with respect to the θθθ parameters of the generative model and the φ parameters of the approximate posterior.
Revised: L is minimized with respect to the θθθ parameters of the generative model and the φ parameters of the approximate posterior.

Original: The latter minimization corresponds to Eq.
Revised: The latter minimization corresponds to Eq.

Original: (1) and we can apply the various estimators to propagate φ -derivatives through discrete samples zzz .
Revised: (1) and we can apply the various estimators to propagate φ -derivatives through discrete samples zzz .

Original: For CR estimators we replace zzz with ζζζ to compute ∂ φ L [ φ ,θθθ ] ≈ (cid:80) xxx ∼ data E ρ [ ∂ φ f θθθ, φ ( ζζζ,xxx )] , where the expectation is evaluated with a single relaxed sample per data point xxx .
Revised: For CR estimators we replace zzz with ζζζ to compute ∂ φ L [ φ ,θθθ ] ≈ (cid:80) xxx ∼ data E ρ [ ∂ φ f θθθ, φ ( ζζζ,xxx )] , where the expectation is evaluated with a single relaxed sample per data point xxx .

Original: The θθθ -derivative can be calculated directly from Eq.
Revised: The θθθ -derivative can be calculated directly from Eq.

Original: (21) using the discrete zzz variables.
Revised: (21) using the discrete zzz variables.

Original: Thus, the φ and θθθ derivatives are evaluated separately requiring two passes through the computation graph with either relaxed or discrete samples.
Revised: Thus, the φ and θθθ derivatives are evaluated separately requiring two passes through the computation graph with either relaxed or discrete samples.

Original: Jang et al.
Revised: Jang et al.

Original: (2016) evaluate both derivatives in one pass (using ζζζ ) thereby introducing bias in the θθθ derivatives.
Revised: (2016) evaluate both derivatives in one pass (using ζζζ ) thereby introducing bias in the θθθ derivatives.

Original: We refer to these two possibilities as one/two-pass training and compare their performance.
Revised: We refer to these two possibilities as one/two-pass training and compare their performance.

Original: Following Maddison et al.
Revised: Following Maddison et al.

Original: ; Jang et al.
Revised: ; Jang et al.

Original: ; Tucker et al.
Revised: ; Tucker et al.

Original: (2017), we consider four architectures with binary variables denoted by 200H − 784V , 200H − 200H − 784V , 200H ∼ 784V , and 200H ∼ 200H ∼ 784V (see Appendix E for details).
Revised: (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).

Original: First, we compare the sampled RAM estimator of Eq.
Revised: First, we compare the sampled RAM estimator of Eq.

Original: (7) on 200H − 784V model for different values of β , where  probability of updating variable z
Revised: (7) on 200H − 784V model for different values of β , where the probability of updating variable z

Original: i is p i =
Revised: i is p i =

Original: 4 q i (1 − q i ) /β .
Revised: 4 q i (1 − q i ) /β .

Original: Fig.
Revised: Fig.

Original: 3(a) shows the NELBO on the training set.
Revised: 3(a) shows the NELBO on the training set.

Original: As expected, increased β leads to increased gradient variance which slows training.
Revised: As expected, increased β leads to increased gradient variance which slows training.

Original: Fig. 3(b) shows the average number of function evaluations performed in Eq.
Revised: Fig. 3(b) shows the average number of function evaluations performed in Eq.

Original: Many units become deterministic early in training leading to signiﬁcant computational savings with the sampled RAM method.
Revised: Many units become deterministic early in training leading to signiﬁcant computational savings with the sampled RAM method.

Original: In Fig. 4(a), we include several CR estimators on the same architecture and observe that the GSM estimator performs signiﬁcantly worse than other estimators.
Revised: In Fig. 4(a), we include several CR estimators on the same architecture and observe that the GSM estimator performs signiﬁcantly worse than other estimators.

Original: With linear decoders, the objective function Eq.
Revised: With linear decoders, the objective function Eq.

Original: (21) is convex in z i .
Revised: (21) is convex in z i .

Original: Thus, similar to the example in Fig.
Revised: Thus, similar to the example in Fig.

Original: 2, GSM learns a distribution with higher entropy which leads to poorer performance.
Revised: 2, GSM learns a distribution with higher entropy leading to poorer performance.

Original: The entropy of all estimators during training is shown in Fig. 4(b).
Revised: The entropy of all estimators during training is shown in Fig. 4(b).

Original: Wealso note that two-pass training performs better then one-pass training for all estimators.
Revised: We note that two-pass training performs better then one-pass training for all estimators.

Original: Although two-pass training requires twice the computation in the worst case, this overhead is negligible for these models due to GPU parallelization.
Revised: Although two-pass training requires twice the computation in the worst case, this overhead is negligible for these models due to GPU parallelization.

Original: We observe that ARM performs poorly conﬁrming that its high variance impedes training.
Revised: We observe that ARM performs poorly conﬁrming that its high variance impedes training.

Original: The PWL estimator with two-pass training performs on par with RAM, while being more computationally efﬁcient.
Revised: The PWL estimator with two-pass training performs on par with RAM, while being more computationally efﬁcient.

Revised: Lastly, the REBAR estimator (see Appendix E), using the settings of Tucker et al.

Revised: (2017), performs on par with PWL thereby indirectly conﬁrming that the relaxed objective is a good control variate for the discrete one.

Original: We used β = 2 in the above experiments, similar to Jang et al.
Revised: We used β = 2 in the above experiments, similar to Jang et al.

Original: ; Maddison et al.
Revised: ; Maddison et al.

Original: To understand the dependence on β , we plot ﬁnal train NELBO in Fig. 4(c).
Revised: To understand the dependence on β , we plot ﬁnal train NELBO in Fig. 4(c).

Original: We ﬁnd that improved estimators are less sensitive to the choice of β .
Revised: We ﬁnd that improved estimators are less sensitive to the choice of β .

Original: In Fig.
Revised: In Fig.

Original: 5, we repeat the experiments for the non-linear architecture 200H ∼ 784V .
Revised: 5, we repeat the experiments for the non-linear architecture 200H ∼ 784V .

Original: The one-pass
Revised: The one-pass

Original: GSM estimator exhibits instability which is remedied by two-pass training.
Revised: GSM estimator exhibits instability which is remedied by two-pass training.

Original: However, two-pass
Revised: However, two-pass

Original: GSM still performs worse than IGSM and PWL due to its bias.
Revised: GSM still performs worse than IGSM and PWL due to its bias.

Original: Interestingly, one-pass training works better for IGSM/PWL.
Revised: Interestingly, one-pass training works better for IGSM/PWL.

Original: We observe this repeatedly in the nonlinear models.
Revised: We observe this repeatedly in the nonlinear models.

Original: Unlike the linear case, the RAM estimator converges faster initially but later in training is outperformed by the highervariance IGSM, PWL and ARM.
Revised: Unlike the linear case, the RAM estimator converges faster initially but later in training is outperformed by the highervariance IGSM, PWL and ARM.

Original: It is likely that additional noise prevents the latent units from being  turned off early in training which is known to cause poor performance in VAEs.
Revised: It is likely that additional noise prevents the latent units from being [Figure] [Figure] turned off early in training which is known to cause poor performance in VAEs.

Original: As in the linear case we plot the dependence of the ﬁnal train NELBO on β in Fig. 5(b).
Revised: As in the linear case we plot the dependence of the ﬁnal train NELBO on β in Fig. 5(b).

Original: The GSM estimator outperforms
Revised: The GSM estimator outperforms

Original: [Figure] [Figure] GSM. IGSM and PWL for β ≥ 4 because its bias favors higher entropy approximating posteriors which inhibit latent units from turning off early in training.
Revised: IGSM and PWL for β ≥ 4 because its bias favors higher entropy approximating posteriors which inhibit latent units from turning off early in training.

Original: A well known resolution for inactive latent units is KL-annealing Bowman et al.
Revised: A well known resolution for inactive latent units is KL-annealing Bowman et al.

Original: Fig. 5(c) shows that KL annealing indeed improves
Revised: Fig. 5(c) shows that KL annealing indeed improves

Original: IGSM and PWL by inhibiting over-pruning of latent units 4 .
Revised: IGSM and PWL by inhibiting over-pruning of latent units 5 .

Original: As with the linear case, the IGSM and PWL estimators are more stable with respect to the choice of β .
Revised: As with the linear case, the IGSM and PWL estimators are more stable under variations in β .

Revised: Here, REBAR underperforms the ICR estimators likely due to its higher variance.

Original: Additional results for other VAE models and for categorical variables on both MNIST and OMNIGLOT are presented in Appendix H.3 with similar conclusions.
Revised: Additional results for other VAE models and for categorical variables on both MNIST and OMNIGLOT are presented in Appendix I.3 with similar conclusions.

Original: Interested reader is referred to Appendix F for experimental results on training encoder part of a VAE with pretrained generative model.
Revised: The interested reader is referred to Appendix G for experimental results on training encoder part of a VAE with pretrained generative model.

Original: Appendix G compares estimators for solving the discrete maximum clique optimization problem analyzed in Patish & Ullman (2018).
Revised: Appendix H compares estimators for solving the discrete maximum clique optimization problem analyzed in Patish & Ullman (2018).

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We have reviewed several ﬁnite-difference (FD) and continuous relaxation (CR) estimators of the gradients of the objective Eq.
Revised: We have reviewed several ﬁnite-difference (FD) and continuous relaxation (CR) estimators of the gradients of the objective Eq.

Original: FD estimators like RAM Tokui & Sato (2017) can give unbiased low variance estimates but often require multiple function evaluations.
Revised: FD estimators like RAM Tokui & Sato (2017) can give unbiased low variance estimates but often require multiple function evaluations.

Original: We proposed a less expensive version of RAMestimator that requires  order of magnitude fewer function evaluations with a controllable decrease in performance.
Revised: We proposed a less expensive version of RAM that requires an order of magnitude fewer function evaluations with a controllable decrease in performance.

Original: On the other hand, CR estimators, like Gumbel-Softmax (GSM), require a single pass through the objective function and can be computed efﬁciently giving low variance but biased gradients.
Revised: In contrast, CR estimators, like Gumbel-Softmax (GSM), require a single pass through the objective function and can be computed efﬁciently giving low variance but biased gradients.

Original: We analyzed the nature of the bias introduced by CR estimators and proposed a way to reduce it.
Revised: We analyzed the nature of the bias introduced by CR estimators and proposed a way to reduce it.

Original: This gives rise to improved CR (ICR) estimators, like improved GSM and piece-wise linear.
Revised: This gives rise to improved CR (ICR) estimators, like improved GSM and piece-wise linear.

Original: These ICR estimators are unbiased for a single variable and less biased for many variables.
Revised: These ICR estimators are unbiased for a single variable and less biased for many variables.

Original: Experiments on VAE training and discrete optimization conﬁrm the theoretical predictions and illustrate the advantage of less-biased estimators.
Revised: Experiments on VAE training and discrete optimization conﬁrm the theoretical predictions and illustrate the advantages of lower-bias estimators.

Original: A ARGMAX
Revised: A ARGMAX

Original: Lorberbom et al.
Revised: Lorberbom et al.

Original: (2018) proposed a FD estimator that we refer to as ARGMAX.
Revised: (2018) proposed a FD estimator that we refer to as ARGMAX.

Original: For a single variable,
Revised: For a single variable,

Original: ARGMAX relies on an identity that approximates ∂ φ L = ∂
Revised: ARGMAX relies on an identity that approximates ∂ φ L = ∂

Original: φ q φ [ f (1) − f (0)]
Revised: φ q φ [ f (1) − f (0)]

Original: = ( ∂ φ l φ ) q φ
Revised: = ( ∂ φ l φ ) q φ

Original: [1 − q φ ][ f (1) − f (0)] .
Revised: [1 − q φ ][ f (1) − f (0)] .

Original: With q φ = σ ( l φ ) , Lorberbom et al.
Revised: With q φ = σ ( l φ ) , Lorberbom et al.

Original: (2018) show that
Revised: (2018) show that

Original: [Equation]
Revised: [Equation]

Original: where Θ is the Heaviside step function and ρ ∼ U [0 , 1] .
Revised: where Θ is the Heaviside step function and ρ ∼ U [0 , 1] .

Original: 5 Lorberbom et al.
Revised: 6 Lorberbom et al.

Original: (2018) approximate the right side of Eq.
Revised: (2018) approximate the right side of Eq.

Original: (22) by sampling ρ and evaluating at ﬁnite (cid:15) which introduces bias and variance.
Revised: (22) by sampling ρ and evaluating at ﬁnite (cid:15) which introduces bias and variance.

Original: Decreasing (cid:15) decreases bias but increases variance.
Revised: Decreasing (cid:15) decreases bias but increases variance.

Original: In the categorical case the gradient of Eq.
Revised: In the categorical case the gradient of Eq.

Original: (5) can correspondingly be written as ∂ φ L = (cid:88)
Revised: (5) can correspondingly be written as ∂ φ L = (cid:88)

Original: [Equation]
Revised: [Equation]

Original: where γ b = − log( − log( ρ b )) are Gumbel variables and [ pred ] is the indicator function (1 if pred is true and 0 otherwise).
Revised: where γ b = − log( − log( ρ b )) are Gumbel variables and [ pred ] is the indicator function (1 if pred is true and 0 otherwise).

Original: This derivative requires A function evaluations, similar to RAM.
Revised: This derivative requires A function evaluations, similar to RAM.

Original: Lorberbom et al.
Revised: Lorberbom et al.

Original: (2018) extend the single variable result to multivariate distributions similar to Eqs.
Revised: (2018) extend the single variable result to multivariate distributions similar to Eqs.

Original: (3) and (4).
Revised: (3) and (4).

Original: ARGMAX has the same computational complexity as RAM, but is biased and has higher variance than RAM.
Revised: ARGMAX has the same computational complexity as RAM, but is biased and has higher variance than RAM.

Original: Thus, it is suboptimal to Eqs.
Revised: Thus, it is suboptimal to Eqs.

Original: (3) and (4) and for this reason we do not perform experiments with ARGMAX.
Revised: (3) and (4) and for this reason we do not perform experiments with ARGMAX.

Original: B T HE A UGMENT -REINFORCE-M ERGE E STIMATOR
Revised: B T HE A UGMENT -REINFORCE-M ERGE E STIMATOR

Original: FD estimators are computationally expensive and require multiple function evaluations per gradient.
Revised: FD estimators are computationally expensive and require multiple function evaluations per gradient.

Original: A notable exception is the Augment-REINFORCE-Merge (ARM) estimator introduced in Yin &
Revised: A notable exception is the Augment-REINFORCE-Merge (ARM) estimator introduced in Yin &

Original: Zhou (2018).
Revised: Zhou (2018).

Original: ARM provides an unbiased estimate using only two function evaluations for the factorized multivariate distribution, regardless of the number of variables.
Revised: ARM provides an unbiased estimate using only two function evaluations for the factorized multivariate distribution, regardless of the number of variables.

Original: For a single binary variable, the ARM derivative is given by
Revised: For a single binary variable, the ARM derivative is given by

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: The expectation is approximated with a single sample ρ .
Revised: The expectation is approximated with a single sample ρ .

Original: This estimator has a signiﬁcantly lower variance than REINFORCE since the expectation contains the difference f
Revised: This estimator has a signiﬁcantly lower variance than REINFORCE since the expectation contains the difference f

Original: ( z (2) )
Revised: ( z (2) )

Original: − f ( z (1) ) rather then the function itself.
Revised: − f ( z (1) ) rather then the function itself.

Original: For multivariate f and factorial q φ ( zzz ) the derivative is
Revised: For multivariate f and factorial q φ ( zzz ) the derivative is

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where z (1) i =
Revised: where z (1) i =

Original: Θ( q φ ,i − ρ i ) and z (2) i
Revised: Θ( q φ ,i − ρ i ) and z (2) i

Original: =
Revised: =

Original: Θ( ρ i − 1 + q φ ,i ) .
Revised: Θ( ρ i − 1 + q φ ,i ) .

Original: Yin & Zhou (2018) observed that one can replace f ( z (2) i ,zzz (1) /i )
Revised: Yin & Zhou (2018) observed that one can replace f ( z (2) i ,zzz (1) /i )

Original: → f ( z (2) i ,zzz (2) /i ) without changing the expectation which allows evaluation by a single sample and two function evaluations f ( zzz (1) ) , f ( zzz (2) )
Revised: → f ( z (2) i ,zzz (2) /i ) without changing the expectation which allows evaluation by a single sample and two function evaluations f ( zzz (1) ) , f ( zzz (2) )

Original: regardless of M :
Revised: regardless of M :

Original: [Equation]
Revised: [Equation]

Original: However, this change comes at the cost of higher variance of the expectation 6 .
Revised: However, this change comes at the cost of higher variance of the expectation 7 .

Original: Thus, while ARM estimator provides a low-variance gradient estimate for a single variable it introduces high variance
Revised: Thus, while ARM estimator provides a low-variance gradient estimate for a single variable, it introduces high variance

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: VAR
Revised: VAR

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: /i )
Revised: /i )]

Original: [Equation]
Revised: [Equation]

Original: /i
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: (1) ( zzz /i ) g (2) ( zzz /i )]
Revised: /i ) g (2) ( zzz /i )]

Original: zzz /i ,zzz
Revised: zzz /i ,zzz

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

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

Original: for multivariate functions compared to RAM.
Revised: for multivariate functions compared to RAM.

Original: The ARM estimator has straightforward extensions to hierarchical q ( zzz ) (similar to Eq.
Revised: The ARM estimator has straightforward extensions to hierarchical q ( zzz ) (similar to Eq.

Original: (4)) and to categorical variables
Revised: (4)) and to categorical variables

Original: Yin & Zhou (2018).
Revised: Yin & Zhou (2018).

Original: C I MPROVED CR ESTIMATOR FOR B AYESIAN NETWORKS
Revised: C I MPROVED CR ESTIMATOR FOR B AYESIAN NETWORKS

Original: In this appendix we derive the improved CR estimator for a hierarchical q φ ( zzz ) =
Revised: In this appendix we derive the improved CR estimator for a hierarchical q φ ( zzz ) =

Original: (cid:81)
Revised: (cid:81)

Original: i q φ ,i ( z i | zzz <i ) .
Revised: i q φ ,i ( z i | zzz <i ) .

Original: For simplicity, we do this for two variables q φ ( z 1 , z 2 ) =
Revised: For simplicity, we do this for two variables q φ ( z 1 , z 2 ) =

Original: q φ , 2 ( z 2 | z 1 ) q φ , 1 ( z 1 ) with the gradient given in Eq.
Revised: q φ , 2 ( z 2 | z 1 ) q φ , 1 ( z 1 ) with the gradient given in Eq.

Original: [Equation]
Revised: [Equation]

Original: where q φ , 1 = q φ , 1 ( z 1 = 1) .
Revised: where q φ , 1 = q φ , 1 ( z 1 = 1) .

Original: We denote the two contributions to ∂ φ L by ∂ φ L (1 , 2) and determine them separately.
Revised: We denote the two contributions to ∂ φ L by ∂ φ L (1 , 2) and determine them separately.

Original: Using the integral trick (10) the second contribution can be written as ∂ φ L (2) =
Revised: Using the integral trick (10) the second contribution can be written as ∂ φ L (2) =

Original: (cid:88)
Revised: (cid:88)

Original: [Equation]
Revised: [Equation]

Original: where ζ 2 = g ( ρ 2 ; q 2 (1 | ζ 1 )) and ζ 1 = g ( ρ 1 ; q 1 ) .
Revised: where ζ 2 = g ( ρ 2 , q 2 (1 | ζ 1 )) and ζ 1 = g ( ρ 1 , q 1 ) .

Original: The ﬁrst term can be handled similarly:
Revised: The ﬁrst term can be handled similarly:

Original: [Equation]
Revised: [Equation]

Original: Combining these contributions we arrive at an expression very similar to the reparameterization trick with the replacement ∂ q i ζ i → ∂ ρ i ζ i :
Revised: Combining these contributions we arrive at an expression very similar to the reparameterization trick with the replacement ∂ q i ζ i → ∂ ρ i ζ i :

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: D C ATEGORICAL PWL ESTIMATOR
Revised: D C ATEGORICAL PWL ESTIMATOR

Original: Here, we derive the PWL estimator for categorical variables.
Revised: Here, we derive the PWL estimator for categorical variables.

Original: The derivative for a single categorical variable y
Revised: The derivative for a single categorical variable y

Original: =
Revised: =

Original: ( y 0 , ...
Revised: ( y 0 , ...

Original: y A − 1 ) with y a ∈ { 0 , 1 } and (cid:80) a y a = 1 is given in Eq.
Revised: y A − 1 ) with y a ∈ { 0 , 1 } and (cid:80) a y a = 1 is given in Eq.

Original: (5) as
Revised: (5) as

Original: [Equation]
Revised: [Equation]

Original: We again apply the integral trick (10) to represent the difference f a − f b .
Revised: We again apply the integral trick (10) to represent the difference f a − f b .

Original: To do that we relax the variable y
Revised: To do that we relax the variable y

Original: so that it interpolates between ( y a , y b )
Revised: so that it interpolates between ( y a , y b )

Original: =
Revised: =

Original: (1 , 0) and ( y a , y b )
Revised: (1 , 0) and ( y a , y b )

Original: = (0 , 1) as y
Revised: = (0 , 1) as y

Original: → y a,b = (cid:8) y a = (cid:2) 0 .
Revised: → y a,b = (cid:8) y a = (cid:2) 0 .

Original: 5 + α a,b ( ρ a,b − q b / ( q a + q b )) (cid:3) 1 0 , y b = 1 − y a , y c (cid:54) = a,b = 0 (cid:9) where ρ a,b ∈ U [0 , 1] and α a,b is the slope.
Revised: 5 + α a,b ( ρ a,b − q b / ( q a + q b )) (cid:3) 1 0 , y b = 1 − y a , y c (cid:54) = a,b = 0 (cid:9) where ρ a,b ∈ U [0 , 1] and α a,b is the slope.

Original: The relaxed objective then takes the form
Revised: The relaxed objective then takes the form

Original: [Equation]
Revised: [Equation]

Original: 32) where w a,b are weights to be determined.
Revised: 32) where w a,b are weights to be determined.

Original: The gradient of this relaxed objective with respect to the logit l a is
Revised: The gradient of this relaxed objective with respect to the logit l a is

Original: [Equation]
Revised: [Equation]

Original: Comparison with Eq.
Revised: Comparison with Eq.

Original: (31) gives w a,b
Revised: (31) gives w a,b

Original: = ( q a + q b ) 2 .
Revised: = ( q a + q b ) 2 .

Original: 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
Revised: 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

Original: =
Revised: =

Original: ( q a + q b ) /
Revised: ( q a + q b ) /

Original: ( (cid:80) a<b q a + q b ) =
Revised: ( (cid:80) a<b q a + q b ) =

Original: ( q a + q b ) /
Revised: ( q a + q b ) /

Original: ( A − 1) .
Revised: ( A − 1) .

Original: The reason for choosing this distribution is that Eq.
Revised: The reason for choosing this distribution is that Eq.

Original: (32) must give correct value of the objective (not just the derivative) as[Figure] the relaxation parameter β → ∞ , which requires the probability of each state y a = 1 to be equal to q a .
Revised: (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 .

Original: 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:
Revised: 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:

Original: [Equation]
Revised: [Equation]

Original: Finally, to get the correct weights w a,b we must rescale each term by a factor γ a,b
Revised: Finally, to get the correct weights w a,b we must rescale each term by a factor γ a,b

Original: = ( A − 1)( q a + q b ) so that w a,b = γ a,b p a,b .
Revised: = ( A − 1)( q a + q b ) so that w a,b = γ a,b p a,b .

Original: In summary, the relaxed objective has the following form:
Revised: In summary, the relaxed objective has the following form:

Original: [Equation]
Revised: [Equation]

Original: (35) where ˜ y a,b has the same value as y a,b but has the gradient scaled by γ a,b
Revised: where ˜ y a,b has the same value as y a,b but has the gradient scaled by γ a,b

Original: = ( A − 1)( q a + q b ) .
Revised: = ( A − 1)( q a + q b ) .

Original: E E XPERIMENTAL DETAILS
Revised: E REBAR

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

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

Revised: [Equation]

Revised: where sg ( φ ) denotes “stop gradient”, and distributions q φ ( ζζζ ) and q φ ( ζζζ | zzz ) are reparameterizable.

Revised: Using explicit reparameterizations for ζζζ ( ρρρ, q φ ) and ζζζ ( u , q φ | zzz ) with u ∈ U [0 , 1] , we can rewrite the gradient as

Revised: [Equation]

Revised: [Equation]

Revised: Focusing on binary variables, the reparameterization for the conditional distribution has the form ζζζ ( u , q φ | zzz ) =

Revised: ζζζ (˜ ρρρ, q φ ) , where ˜ ρρρ = 1 − q φ + u ( zzzq φ − (1 − zzz )(1 − q φ )) .

Revised: In order to minimize the variance, it is preferable to tie the parameters ρρρ and ( zzz, u ) .

Revised: This can be achieved by setting zzz =

Revised: Θ( ρρρ − 1 + q φ ) , and choosing

Revised: [Equation]

Revised: This is equivalent to sampling ( zzz,ζζζ ) ∼ q φ ( zzz ) q φ ( ζζζ | zzz ) as zzz =

Revised: Θ( ρρρ − 1 + q φ ) ,ζζζ = ζζζ (˜ ρρρ, q φ )

Revised: where

Revised: [Equation]

Revised: This way ˜ ρρρ = ρρρ by value

Revised: but ˜ ρρρ does depend on φ such that ζζζ (˜ ρρρ, q φ ) has the correct gradient w.r.t.

Revised: φ that follows from the conditional distribution q φ ( ζζζ | zzz ) .

Revised: We can rewrite (36) as

Revised: [Equation]

Revised: Since the explicit dependence of ζζζ ( ρρρ, q φ ) and ζζζ (˜ ρρρ, q φ ) on q φ is the same, we have ∂ φ ζζζ ( ρρρ, q φ ) − ∂ φ ζζζ (˜

Revised: [Equation]

Revised: , q φ ) = − ∂ φ ζζζ (˜ ρρρ, sg ( q φ )) and we arrive at our ﬁnal expression for the REBAR gradient:

Revised: [Equation]

Revised: [Equation]

Revised: The advantage of Eq.

Revised: (39) is two-fold: ﬁrst, it allows for a simple implementation valid for any continuous relaxation and uses only two function evaluations.

Revised: Second, it gives a suggestive relation to the ICR estimator: indeed using ∂ φ ˜ ρρρ = ∂

Revised: φ q φ (cid:16) − 1 + zzz ρρρ − 1+ q φ q φ

Revised: [Equation]

Revised: φ

Revised: [Equation]

Revised: we can write

Revised: REBAR gradient as

Revised: [Equation]

Revised: N E L B O

Revised: REBAR-GSM REBAR-PWL

Revised: [Figure]

Original: {+200H − 784V    
+}
Revised: {+200H − 784V    
+}

Revised: N E L B O

Revised: [Figure]

Revised: GSM

Revised: [Figure]

Revised: PWL

Revised: [Figure]

Revised: The term R 2 can have high variance due to the fact that var( ∂ q φ log q φ ( zzz )) ∼ 1 q φ (1 − q φ ) while R 1 has variance comparable to ∂ φ L ICR since var( zzz ρρρ − 1+ q φ q φ − (1 − zzz )

Revised: ρρρ − 1+ q φ 1 − q φ ) ∼ O (1) .

Revised: In summary REBAR estimator can be viewed as ICR plus corrections that remove the bias of ICR estimator.

Revised: In practice the correction term R 2 can increase the variance of

Revised: REBAR estimator compared to ICR and slow down the training.

Revised: We now perform the comparison of two types of continuous relaxations used by REBAR estimator.

Revised: We train VAE as in Section 4.2 with 200H − 784V and 200H ∼ 784V architectures using REBAR with GSM and PWL relaxations.

Revised: Fig.

Revised: 6 shows the dependence of the ﬁnal NELBO on β .

Revised: In the case of the linear mode 200H − 784V

Revised: we ﬁnd that PWL relaxation is slightly more stable while for the non-linear model the two relaxations perform similarly.

Original: 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 .
Revised: 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 .

Original: 01 using minibatches of size 100 to reduce the variance of gradients.
Revised: 01 using minibatches of size 100 to reduce the variance of gradients.

Original: We initialize q φ ( z = 1) to the wrong maximum of the relaxed function by setting q φ ( z = 1) =
Revised: We initialize q φ ( z = 1) to the wrong maximum of the relaxed function by setting q φ ( z = 1) =

Original: σ (5) .
Revised: σ (5) .

Original: For VAE training experiments, following Maddison et al.
Revised: For VAE training experiments, following Maddison et al.

Original: ; Jang et al.
Revised: ; Jang et al.

Original: ; Tucker et al.
Revised: ; Tucker et al.

Original: (2017), we consider four architectures with binary variables denoted by 200H − 784V , 200H − 200H − 784V , 200H ∼ 784V , and 200H ∼ 200H ∼ 784V .
Revised: (2017), we consider four architectures with binary variables denoted by 200H − 784V , 200H − 200H − 784V , 200H ∼ 784V , and 200H ∼ 200H ∼ 784V .

Original: Here − denotes a linear layer, while ∼ denotes two layers of 200 hidden units with tanh nonlinearity and batch normalization.
Revised: Here − denotes a linear layer, while ∼ denotes two layers of 200 hidden units with tanh nonlinearity and batch normalization.

Original: In our experiments we use the Adam optimizer with default parameters and a ﬁxed learning rate 0 .
Revised: In our experiments we use the Adam optimizer with default parameters and a ﬁxed learning rate 0 .

Original: 0003 , run for 4 · 10 6 steps with minibatches of size 100.
Revised: 0003 , run for 4 · 10 6 steps with minibatches of size 100.

Original: We repeat each experiment 5 times with random seed and plot the mean.
Revised: We repeat each experiment 5 times with random seed and plot the mean.

Original: F VAE ENCODER TRAINING
Revised: G VAE ENCODER TRAINING

Original: During training of VAEs, the encoder and decoder models interact in complex ways.
Revised: During training of VAEs, the encoder and decoder models interact in complex ways.

Original: To eliminate these effects to more directly identify the impact of better gradients, we use a ﬁxed pre-trained decoder, and minimize Eq.
Revised: To eliminate these effects to more directly identify the impact of better gradients, we use a ﬁxed pre-trained decoder, and minimize Eq.

Original: (21) with respect to encoder parameters φ only.
Revised: (21) with respect to encoder parameters φ only.

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 6, the conclusions from joint training remain unchanged: GSM underperforms due to its bias.
Revised: 7, the conclusions from joint training remain unchanged: GSM underperforms due to its bias.

Original: Interestingly however, RAM still overﬁts early in training for the 200H ∼ 784V model, showing that better gradients can negatively affect optimization.
Revised: Interestingly however, RAM still overﬁts early in training for the 200H ∼ 784V model, showing that better gradients can negatively affect optimization.

Original: G O PTIMIZATION
Revised: H O PTIMIZATION

Original: We apply gradient estimators for discrete optimization.
Revised: We apply gradient estimators for discrete optimization.

Original: Following Patish & Ullman (2018)
Revised: Following Patish & Ullman (2018)

Original: we study the NP-hard task of ﬁnding a maximal clique in a graph.
Revised: we study the NP-hard task of ﬁnding a maximal clique in a graph.

Original: To model this problem, the binary variable z i = 1 / 0 indicates the presence/absence of vertex i in a maximal clique.
Revised: To model this problem, the binary variable z i = 1 / 0 indicates the presence/absence of vertex i in a maximal clique.

Original: If A i,j is the adjacency[Figure] matrix of the graph and d = (cid:80) 
Revised: If A i,j is the adjacency matrix of the graph and d = (cid:80) i

Original: i z i is the size of the clique, then the objective function considered  in Patish & Ullman (2018) is
Revised: z i is the size of the clique, then the objective function considered [Figure] [Figure] in Patish & Ullman (2018) is

Original: [Equation]
Revised: [Equation]

Original: (36) where κ ∈ [0 , 1] is a hyperparameter.
Revised: (41) where κ ∈ [0 , 1] is a hyperparameter.

Original: Patish & Ullman (2018) use factorial q φ ( z ) to minimize
Revised: Patish & Ullman (2018) use factorial q φ ( z ) to minimize

Original: L [ φ ] = (cid:80) zzz
Revised: L [ φ ] = (cid:80) zzz

Original: q φ ( zzz ) f ( zzz ) .
Revised: q φ ( zzz ) f ( zzz ) .

Original: At the end of optimization q φ ( z ) typically collapses to a single state corresponding to a clique.
Revised: At the end of optimization q φ ( z ) typically collapses to a single state corresponding to a clique.

Original: Most often this clique is a local minimum and not the maximal clique.
Revised: Most often this clique is a local minimum and not the maximal clique.

Original: We illustrate the performance of CR estimators on a particular graph (1000 nodes, 450000 edges), labeled C 1000 .
Revised: We illustrate the performance of CR estimators on a particular graph (1000 nodes, 450000 edges), labeled C 1000 .

Original: 9 from the DIMACS data set Johnson & Trick (1996).
Revised: 9 from the DIMACS data set Johnson & Trick (1996).

Original: We minimize f ( zzz ) using
Revised: We minimize f ( zzz ) using

Original: Adam with default settings, learning rate 0 .
Revised: Adam with default settings, learning rate 0 .

Original: 01 for 40000 iterations.
Revised: 01 for 40000 iterations.

Original: We perform 1000 minimizations in parallel and choose the best clique found at each iteration.
Revised: We perform 1000 minimizations in parallel and choose the best clique found at each iteration.

Original: Fig.
Revised: Fig.

Original: 7(a),(b) show the size of clique found by each of the estimators for two values of the hyperparameter κ = 0 .
Revised: 8(a),(b) show the size of clique found by each of the estimators for two values of the hyperparameter κ = 0 .

Original: At κ = 0 .the bias of GSM causes trapping in the wrong minima but ICR converges to a good local minimum.
Revised: At κ = 0 .the bias of GSM causes trapping in the wrong minima but ICR converges to a good local minimum.

Original: In contrast, at κ = 0 . 

Revised: In contrast, at κ = 0 .

Original: the GSM bias accelerates convergence to a good local minimum.
Revised: the GSM bias accelerates convergence to a good local minimum.

Original: Fig. 7(c) shows the dependence of the ﬁnal clique size on κ ; unbiased estimators provide more stable results.
Revised: Fig. 8(c) shows the dependence of the ﬁnal clique size on κ ; unbiased estimators provide more stable results.

Original: H A DDITIONAL E XPERIMENTS H. B INARY C ONCAVE T OY E XAMPLE
Revised: I A DDITIONAL E XPERIMENTS I. B INARY C ONCAVE T OY E XAMPLE

Original: We evaluate performance on the concave function f ( ζ ) = − ( ζ − 0 . 45) 2 shown in Fig. 8(a).
Revised: We evaluate performance on the concave function f ( ζ ) = − ( ζ − 0 . 45) 2 shown in Fig. 9(a).

Original: The training setup is identical to Section 4.1 but with q φ ( z = 1) initialized to σ ( − 5) .Fig.
Revised: The training setup is identical to Section 4.1 but with q φ ( z = 1) initialized to σ ( − 5) .

Original: 8(b) demonstrates that the GSM optimization gets trapped in the wrong minimum due to its bias towards the dominant mode ζ = 0 . 
Revised: Fig. 9(b) demonstrates that the GSM optimization gets trapped in the wrong minimum due to its bias towards the dominant mode ζ = 0 . f ( ) (a) Concave relaxed function    2000 Steps 0. q ( z =) GSM IGSM PWL

Original: [Figure]
Revised: [Figure]

Revised: ARM

Revised: [Figure]

Revised: [Figure]

Original: H. C ATEGORICAL T OY E XAMPLE
Revised: C ATEGORICAL T OY E XAMPLE

Original: We consider a categorical example with convex and concave functions of a single categorical variable
Revised: We consider a categorical example with convex and concave functions of a single categorical variable

Original: y having 10 values.
Revised: y having 10 values.

Original: We take f ( y ) = ± (cid:80) a ( g a − y a ) 2 such that g 0 = 0 .
Revised: We take f ( y ) = ± (cid:80) a ( g a − y a ) 2 such that g 0 = 0 .

Original: 9 , g 1 = 1 . 1 , and g i> 1 = 1 .
Revised: 9 , g 1 = 1 . 1 , and g i> 1 = 1 .

Original: The convex function has a minimum at y 1 = 1 , while the concave function is minimized at y 0 = 1 .
Revised: The convex function has a minimum at y 1 = 1 , while the concave function is minimized at y 0 = 1 .

Original: We compare 4 estimators for minimizing this function: RAM of Eq.
Revised: We compare 4 estimators for minimizing this function: RAM of Eq.

Original: (5), PWL of Eq.
Revised: (5), PWL of Eq.

Original: GSM of Eq.
Revised: GSM of Eq.

Original: (19) and IGSM of Eq.
Revised: (19) and IGSM of Eq.

Original: The probability of the true minimum is shown in Fig.
Revised: The probability of the true minimum is shown in Fig.

Original: In both the convex and concave cases the GSM estimator exhibits a bias preventing it from ﬁnding the minimum.
Revised: In both the convex and concave cases the GSM estimator exhibits a bias preventing it from ﬁnding the minimum.

Original: IGSM is less biased than GSM which allows it to ﬁnd the true minimum.
Revised: IGSM is less biased than GSM which allows it to ﬁnd the true minimum.

Original: The PWL estimator is unbiased but has higher variance then IGSM which slows down its optimization.
Revised: The PWL estimator is unbiased but has higher variance then IGSM which slows down its optimization.

Original: [Figure]
Revised: [Figure]

Original:  DISCRETE V ARIATIONAL A UTOENCODERS
Revised:  D

Revised: V A

Original: Fig.
Revised: Fig.

Original: 10 shows the results for all 4 architectures 200H − 784V , 200H − 200H − 784V , 200H ∼ 784V , and 200H ∼ 200H ∼ 784V at β = 2 .
Revised: 11 shows the results for all 4 architectures 200H − 784V , 200H − 200H − 784V , 200H ∼ 784V , and 200H ∼ 200H ∼ 784V at β = 2 .

Original: Hierarchical models with two layers of latent units in
Revised: Hierarchical models with two layers of latent units in

Original: Fig.
Revised: Fig.

Original: 10(b),(d) exhibit similar trends to the factorial case considered in the main text.
Revised: 11(b),(d) exhibit similar trends to the factorial case considered in the main text.

Original: The GSM estimator converges to a higher NELBO in the linear case and becomes unstable in the non-linear case.
Revised: The GSM estimator converges to a higher NELBO in the linear case and becomes unstable in the non-linear case.

Original: [Figure]
Revised: [Figure]

Original: We also compare the performance of GSM and IGSM estimators in the categorical case 20 × 10H − 784V with latent space being 20 categorical variables with 10 classes Fig.
Revised: We also compare the performance of GSM and IGSM estimators in the categorical case 20 × 10H − 784V with latent space being 20 categorical variables with 10 classes Fig.

Original: We see that IGSM outperforms GSM estimator when using two-pass training.
Revised: We see that IGSM outperforms GSM estimator when using two-pass training.

Original: This indirectly conﬁrms that IGSM is less biased.
Revised: This indirectly conﬁrms that IGSM is less biased.

Original: The PWL estimator performs best indicating the advantages ICR estimators.
Revised: The PWL estimator performs best indicating the advantages ICR estimators.

Original: [Figure]
Revised: [Figure]

Original: For completeness we also show the training curves for OMNIGLOT dataset at β = 2 in Fig 12.
Revised: For completeness we also show the training curves for OMNIGLOT dataset at β = 2 in Fig 13.

Original: Interestingly, one-pass training performs better in all cases.
Revised: Interestingly, one-pass training performs better in all cases.

Original: PWL estimator performs best among considered CR estimators
Revised: PWL estimator performs best among considered CR estimators.

Original: [Figure]
Revised: [Figure]