Title: GO GRADIENT FOR EXPECTATION-BASED OBJECTIVES

Abstract: Within many machine learning algorithms, a fundamental problem concerns efficient calculation of an unbiased gradient wrt parameters γ for expectation-based objectives E qγ (y) [f (y)]. Most existing methods either (i) suffer from high variance, seeking help from (often) complicated variance-reduction techniques; or (ii) they only apply to reparameterizable continuous random variables and employ a reparameterization trick. To address these limitations, we propose a General and One-sample (GO) gradient that (i) applies to many distributions associated with non-reparameterizable continuous or discrete random variables, and (ii) has the same low-variance as the reparameterization trick. We find that the GO gradient often works well in practice based on only one Monte Carlo sample (although one can of course use more samples if desired). Alongside the GO gradient, we develop a means of propagating the chain rule through distributions, yielding statistical back-propagation, coupling neural networks to common random variables.

1. INTRODUCTION:

Neural networks, typically trained using back-propagation for parameter optimization, have recently demonstrated significant success across a wide range of applications. There has been interest in coupling neural networks with random variables, so as to embrace greater descriptive capacity. Recent examples of this include black-box variational inference (BBVI) (Kingma & Welling, 2014; Rezende et al., 2014; Ranganath et al., 2014; Hernández-Lobato et al., 2016; Ranganath et al., 2016b; Li & Turner, 2016; Ranganath et al., 2016a; Zhang et al., 2018) and generative adversarial networks (GANs) (Goodfellow et al., 2014; Radford et al., 2015; Zhao et al., 2016; Arjovsky et al., 2017; Li et al., 2017; Gan et al., 2017; Li et al., 2018) . Unfortunately, efficiently backpropagating gradients through general distributions (random variables) remains a bottleneck. Most current methodology focuses on distributions with continuous random variables, for which the reparameterization trick may be readily applied (Kingma & Welling, 2014; Grathwohl et al., 2017) .

As an example, the aforementioned bottleneck greatly constrains the applicability of BBVI, by limiting variational approximations to reparameterizable distributions. This limitation excludes discrete random variables and many types of continuous ones. From the perspective of GAN, the need to employ reparameterization has constrained most applications to continuous observations. There are many forms of data that are more-naturally discrete.

The fundamental problem associated with the aforementioned challenges is the need to efficiently calculate an unbiased low-variance gradient wrt parameters γ for an expectation objective of the form E qγ (y) [f (y)] 1 . We are interested in general distributions q γ (y), for which the components of y may be either continuous or discrete. Typically the components of y have a hierarchical structure, and a subset of the components of y play a role in evaluating f (y).

Unfortunately, classical methods for estimating gradients of E qγ (y) [f (y)] wrt γ have limitations. The REINFORCE gradient (Williams, 1992) , although generally applicable (e.g., for continuous and discrete random variables), exhibits high variance with Monte Carlo (MC) estimation of the expectation, forcing one to apply additional variance-reduction techniques. The reparameterization trick (Rep) (Salimans et al., 2013; Kingma & Welling, 2014; Rezende et al., 2014) works well, with as few as only one MC sample, but it is limited to continuous reparameterizable y. Many efforts have been devoted to improving these two formulations, as detailed in Section 6. However, none of these methods is characterized by generalization (applicable to general distributions) and efficiency (working well with as few as one MC sample).

The key contributions of this work are based on the recognition that REINFORCE and Rep are seeking to solve the same objective, but in practice Rep yields lower-variance estimations, albeit for a narrower class of distributions. Recent work (Ranganath et al., 2016b) has made a connection between REINFORCE and Rep, recognizing that the former estimates a term the latter evaluates analytically. The high variance by which REINFORCE approximates this term manifests high variance in the gradient estimation. Extending these ideas, we make the following main contributions. (i) We propose a new General and One-sample (GO) gradient in Section 3, that principally generalizes Rep to many non-reparameterizable distributions and justifies two recent methods (Figurnov et al., 2018; Jankowiak & Obermeyer, 2018) ; the "One sample" motivating the name GO is meant to highlight the low variance of the proposed method, although of course one may use more than one sample if desired. (ii) We find that the core of the GO gradient is something we term a variable-nabla, which can be interpreted as the gradient of a random variable wrt a parameter. (iii) Utilizing variablenablas to propagate the chain rule through distributions, we broaden the applicability of the GO gradient in Sections 4-5 and present statistical back-propagation, a statistical generalization of classic back-propagation (Rumelhart & Hinton, 1986) . Through this generalization, we may couple neural networks to general random variables, and compute needed gradients with low variance.

2. BACKGROUND:

To motivate this paper, we begin by briefly elucidating common machine learning problems for which there is a need to efficiently estimate gradients of γ for functions of the form E qγ (y) [f (y) ]. Assume access to data samples {x i } i=1,N , drawn i.i.d. from the true (and unknown) underlying distribution q(x). We seek to learn a model p θ (x) to approximate q(x). A classic approach to such learning is to maximize the expected log likelihood θ = argmax θ E q(x) [log p θ (x)], perhaps with an added regularization term on θ. Expectation E q(x) (•) is approximated via the available data samples, as θ = argmax θ 1 N N i=1 log p θ (x i ). It is often convenient to employ a model with latent variables z, i.e., p θ (x) = p θ (x, z)dz = p θ (x|z)p(z)dz, with prior p(z) on z. The integral wrt z is typically intractable, motivating introduction of the approximate posterior q φ (z|x), with parameters φ. The well-known evidence lower bound (ELBO) (Jordan et al., 1999; Bishop, 2006) is defined as ELBO(θ, φ; x) = E q φ (z|x) [log p θ (x, z) -log q φ (z|x)]

(1)

EQUATION

where p θ (z|x) is the true posterior, and KL(• •) represents the Kullback-Leibler divergence. Variational learning seeks ( θ, φ) = argmax θ,φ N i=1 ELBO(θ, φ; x i ). While computation of the ELBO has been considered for many years, a problem introduced recently concerns adversarial learning of p θ (x), or, more precisely, learning a model that allows one to efficiently and accurately draw samples x ∼ p θ (x) that are similar to x ∼ q(x). With generative adversarial networks (GANs) (Goodfellow et al., 2014) , one seeks to solve

EQUATION

where D β (x) is a discriminator with parameters β, quantifying the probability x was drawn from q(x), with 1 -D β (x) representing the probability that it was drawn from p θ (x). There have been many recent extensions of GAN (Radford et al., 2015; Zhao et al., 2016; Arjovsky et al., 2017; Li et al., 2017; Gan et al., 2017; Li et al., 2018) , but the basic setup in (3) holds for most.

To optimize (1) and (3), the most challenging gradients that must be computed are of the form ∇ γ E qγ (y) [f (y)]; for (1) y = z and γ = φ, while for (3) y = x and γ = θ. The need to evaluate expressions like ∇ γ E qγ (y) [f (y)] arises in many other machine learning problems, and consequently it has generated much prior attention.

Evaluation of the gradient of the expectation is simplified if the gradient can be moved inside the expectation. REINFORCE (Williams, 1992) is based on the identity

∇ γ E qγ (y) [f (y)] = E qγ (y) f (y)∇ γ log q γ (y) .

(4) While simple in concept, this estimate is known to have high variance when the expectation E qγ (y) (•) is approximated (as needed in practice) by a finite number of samples.

An approach (Salimans et al., 2013; Kingma & Welling, 2014; Rezende et al., 2014) that has attracted recent attention is termed the reparameterization trick, applicable when q γ (y) can be reparametrized as y = τ γ ( ), with ∼ q( ), where τ γ ( ) is a differentiable transformation and q( ) is a simple distribution that may be readily sampled. To keep consistency with the literature, we call a distribution reparameterizable if and only if those conditions are satisfied2 . In this case we have

∇ γ E qγ (y) [f (y)] = E q( ) [∇ γ τ γ ( )][∇ y f (y)]| y=τ γ ( ) .

(5) This gradient, termed Rep, is typically characterized by relatively low variance, when approximating E q( ) (•) with a small number of samples ∼ q( ). This approach has been widely employed for computation of the ELBO and within GAN, but it limits one to models that satisfy the assumptions of Rep.

3. GO GRADIENT:

The reparameterization trick (Rep) is limited to reparameterizable random variables y with continuous components. There are situations for which Rep is not readily applicable, e.g., where the components of y may be discrete or nonnegative Gamma distributed. We seek to gain insights from the relationship between REINFORCE and Rep, and generalize the types of random variables y for which the latter approach may be effected. We term our proposed approach a General and One-sample (GO) gradient. In practice, we find that this approach works well with as few as one sample for evaluating the expectation, and it is applicable to more general settings than Rep.

Recall that Rep was first applied within the context of variational learning (Kingma & Welling, 2014) , as in (1). Specifically, it was assumed q γ (y) = v q γ (y v ), omitting explicit dependence on data x, for notational convenience; y v is component v of y. In Kingma & Welling (2014) 

q γ (y v ) corresponded to a Gaussian distribution q γ (y v ) = N (y v ; µ v (γ), σ 2 v (γ))

, with mean µ v (γ) and variance σ 2 v (γ). In the following we generalize q γ (y v ) such that it need not be Gaussian. Applying integration by parts (Ranganath et al., 2016b )

EQUATION

where y -v denotes y with y v excluded, and Q γ (y v ) is the cumulative distribution function (CDF) of q γ (y v ). The "0" term is readily proven to be zero for any Q γ (y v ), with the assumption that f (y) doesn't tend to infinity faster than ∇ γ Q γ (y v ) tending to zero when y v → ±∞.

The "Key" term exactly recovers the one-dimensional Rep when reparameterization (Ranganath et al., 2016b) . Further, applying ∇ γ q γ (y v ) = q γ (y v )∇ γ log q γ (y v ) in (6) yields REINFORCE. Consequently, it appears that Rep yields low variance by analytically setting to zero the unnecessary but high-variance-injecting "0" term, while in contrast REINFORCE implicitly seeks to numerically implement both terms in (7).

y v = τ γ ( v ), v ∼ q( v ) exists

We generalize q γ (y) for discrete y v , here assuming y v ∈ {0, 1, . . . , ∞}. It is shown in Appendix A.2 that this framework is also applicable to discrete y v with a finite alphabet. It may be shown (see Appendix A.2) that

∇ γ E qγ (y) [f (y)] = v E qγ (y -v ) yv f (y)∇ γ q γ (y v ) = v E qγ (y -v ) [f (y)∇ γ Q γ (y v )]| yv=∞ "0" - yv [∇ γ Q γ (y v )][f (y -v , y v + 1) -f (y)]

"Key" (8) where Q γ (y v ) = yv n=0 q γ (n), and Q γ (∞) = 1 for all γ. Theorem 1 (GO Gradient). For expectation objectives E qγ (y) [f (y)], where q γ (y) satisfies (i) q γ (y) = v q γ (y v ); (ii) the corresponding CDF Q γ (y v ) is differentiable wrt parameters γ; and (iii) one can calculate ∇ γ Q γ (y v ), the General and One-sample (GO) gradient is defined as

EQUATION

where

G qγ (y) γ specifies G q(y) κ • • • , g q(yv) κ , • • • with variable-nabla g q(yv) κ -1 q(yv) ∇ κ Q(y v ), D y [f (y)] = • • • , D yv [f (y)], • • •

T , and

D yv [f (y)] ∇ yv f (y), Continous y v f (y -v , y v + 1) -f (y), Discrete y v

All proofs are provided in Appendix A, where we also list g qγ (yv) γ for a wide selection of possible q γ (y), for both continuous and discrete y. Note for the special case with continuous y, GO reduces to Implicit Rep gradients (Figurnov et al., 2018) and pathwise derivatives (Jankowiak & Obermeyer, 2018) ; in other words, GO provides a principled explanation for their low variance, namely their foundation (implicit differentiation) originates from integration by parts. For high-dimensional discrete y, calculating D y [f (y)] is computationally expensive. Fortunately, for f (y) often used in practice special properties hold that can be exploited for efficient parallel computing. Also for discrete y v with finite support, it is possible that one could analytically evaluate a part of expectations in (9) for lower variance, mimicking the local idea in Titsias & Lázaro-Gredilla (2015) ; Titsias (2015) . Appendix I shows an example illustrating how to handle these two issues in practice.

4. DEEP GO GRADIENT:

The GO gradient in Theorem 1 can only handle single-layer mean-field q γ (y), characterized by an independence assumption on the components of y. One may enlarge the descriptive capability of q γ (y) by modeling it as a marginal distribution of a deep model (Ranganath et al., 2016b; Bishop, 2006) . Hereafter, we focus on this situation, and begin with a 2-layer model for simple demonstration. Specifically, consider

q γ (y) = q γ (y, λ)dλ = q γ y (y|λ)q γ λ (λ)dλ = v q γ y (y v |λ) • k q γ λ (λ k )dλ,

where γ = {γ y , γ λ }, y is the leaf variable, and the internal variable λ is assumed to be continuous. Components of y are assumed to be conditionally independent given λ, but upon marginalizing out λ this independence is removed.

The objective becomes E qγ (y) [f (y)] = E qγ y (y|λ)qγ λ (λ) [f (y)], and via Theorem 1

EQUATION

)

Lemma 1. Equation (10) exactly recovers the Rep gradient in (5), if γ = γ y and q γ (y) has reparameterization y = τ γ (λ), λ ∼ q(λ) for differentiable τ γ (λ) and easily sampled q(λ).

Lemma 1 shows that Rep is a special case of our deep GO gradient in the following Theorem 2. Note neither Implicit Rep gradients nor pathwise derivatives can recover Rep in general, because a neural-network-parameterized y = τ γ (λ) may lead to non-trivial CDF Q γ (y).

For the gradient wrt γ λ , we first apply Theorem 1, yielding

∇ γ λ E qγ (y) [f (y)] = E qγ λ (λ) G qγ λ (λ) γ λ D λ E qγ y (y|λ) [f (y)] .

For continuous internal variable λ one can apply Theorem 1 again, from which

EQUATION

Now extending the same procedure to deeper models with L layers, we generalize the GO gradient in Theorem 2. Random variable y (L) is assumed to be the leaf variable of interest, and may be continuous or discrete; latent/internal random variables {y (1) , . . . , y (L-1) } are assumed continuous (these generalize λ from above).

Theorem 2 (Deep GO Gradient). For expectation objectives E qγ (y (L) ) [f (y (L) )] with q γ (y (L) ) being the marginal distribution of

q γ y (1) , • • • , y (L) = q γ (1) y (1) L l=2

q γ (l) y (l) |y (l-1) , -1) are continuous, the leaf variable y (L) is either continuous or discrete, q γ (l) (y (l) |y (l-1) ) = v q γ (l) (y

where γ = {γ (1) , • • • , γ (L) }, all internal variables y (1) , • • • , y (L

(l)

v |y (l-1) ), and one has access to variable-nablas g q γ (l) (y (l) v |y (l-1) ) γ (l) and g q γ (l) (y (l) v |y (l-1) ) y (l-1) , as defined in Theorem 1, the General and One-sample (GO) gradient is defined as

EQUATION

)

where BP[y

(L) ] = D y (L) [f (y (L) )] and BP[y (l) ] = G q γ (l+1) (y (l+1) |y (l) ) y (l)

BP[y (l+1) ] for l < L. Corollary 1. The deep GO gradient in Theorem 2 exactly recovers back-propagation (Rumelhart & Hinton, 1986) when each element distribution q γ (l) (y

(l) v |y (l-1)

) is specified as the Dirac delta function located at the activated value after activation function. Assume all continuous variables y, and abuse notations for better understanding.

Dy[f (y)] → ∇yf (y), G q(y) κ

→ ∇κy:

Figure 1 : Relating back-propagation (Rumelhart & Hinton, 1986) with the deep GO gradient in Theorem 2. (i) In deterministic deep neural networks, one forward-propagates information using activation functions, like ReLU, to sequentially activate {y (l) } l=1,••• ,L (black solid arrows), and then back-propagates gradients from Loss f (•) to each parameter γ (l) via gradient-flow through {y (k) } k=L,••• ,l (red dashed arrows). (ii) Similarly for the deep GO gradient with 1 MC sample, one forward-propagates information to calculate the expected loss function f (y (L) ) using distributions as statistical activation functions, and then uses variable-nablas to sequentially back-propagate gradients through random variables {y (k) } k=L,••• ,l to each γ (l) , as in (12).

5. STATISTICAL BACK-PROPAGATION AND HIERARCHICAL VARIATIONAL INFERENCE:

Recall the motivating discussion in Section 2, in which we considered generative model p θ (x, z) and inference model q φ (z|x), the former used to model synthesis of the observed data x and the latter used for inference of z given observed x. In recent deep architectures, a hierarchical representation for cumulative latent variables z = (z (1) , . . . , z (L) ) has been considered (Rezende et al., 2014; Ranganath et al., 2015; Zhou et al., 2015; 2016; Ranganath et al., 2016b; Cong et al., 2017; Zhang et al., 2018) . As an example, there are models with p θ (x, z) = p θ (x|z (1) ) L) ). When performing inference for such models, it is intuitive to consider first-order Markov chain structure for q φ (z|x) = q φ (z (1) |x) L l=2 q φ (z (l) |z (l-1) ). The discussion in this section is most relevant for variational inference, for computation of ∇ φ E q φ (z (1) ,...,z (L) |x) [f (z (1) , . . . , z (L) )], and consequently we specialize to that notation in the subsequent discussion (we consider representations in terms of z, rather than the more general y notation employed in Section 4).

EQUATION

Before proceeding, we seek to make clear the distinction between this section and Section 4. In the latter, only the leaf variable z (L) appears in ∇ γ E qγ (z (L) ) [f (z (L) )]; see (12), with y (L) → z (L) . That is because in Section 4 the underlying model is a marginal distribution of z (L) , i.e., q γ (z (L) ), which is relevant to the generators of GANs; see (3), with x → z (L) and p θ (x) → q γ (z (L) ). Random variables z (1) , . . . , z (L-1) were marginalized out of q γ (z (1) , . . . , z (L-1) , z (L) ) to represent q γ (z (L) ). As discussed in Section 4, z (1) , . . . , z (L-1) were added there to enhance the modeling flexibility of q γ (z (L) ). In this section, the deep set of random variables z = (z (1) , . . . , z (L) ) are inherent components of the underlying generative model for x, i.e., p θ (x, z) = p θ (x, z (1) , . . . , z (L) ). Hence, all components of z manifested via inference model q φ (z|x) = q φ (z (1) , . . . , z (L) |x) play a role in f (z). Besides, no specific structure is imposed on p θ (x, z) and q φ (z|x) in this section, moving beyond the aforementioned first-order Markov structure. For a practical application, one may employ domain knowledge to design suitable graphical models for p θ (x, z) and q φ (z|x), and then use the following Theorem 3 for training. Theorem 3 (Statistical Back-Propagation). For expectation objectives

E q φ ({z (i) } L i=1 ) f ({z (i) } L i=1 ) ,

where {z (i) } I i=1 denotes I continuous internal variables with at least one child variable, {z (j) } L j=I+1 represents L -I continuous or discrete leaf variables with no children except f (•), and q φ (•) is constructed as a hierarchical probabilistic graphical model

q φ ({z (i) } L i=1 ) = q φ ({z (i) } I i=1 ) L j=I+1 q φ (z (j) |{z (i) } I i=1 )

with each element distribution q φ (z v |pa(z v )) having accessible variable-nablas as defined in Theorem 1, pa(z v ) denotes the parent variables of z v , the General and One-sample (GO) gradient for φ k ∈ φ is defined as

EQUATION

where ch(φ k ) denotes the children variables of φ k , and with

z v ∈ ch(φ k ), G q φ [ch(φ k )] φ k = [• • • , g q φ (zv|pa(zv)) φ k , • • • ], BP[ch(φ k )] = [• • • , BP[z v ], • • • ] T ,

and

BP[z v ] is iteratively calculated as BP[z v ] = G q φ [ch(zv)] zv BP[ch(z v )], until leaf variables where BP[z v ] = D zv [f ({z (i) } L i=1 )].

Statistical back-propagation in Theorem 3 is relevant to hierarchical variational inference (HVI) (Ranganath et al., 2016b; Hoffman & Blei, 2015; Mnih & Gregor, 2014) (see Appendix G), greatly generalizing GO gradients to the inference of directed acyclic probabilistic graphical models. In HVI variational distributions are specified as hierarchical graphical models constructed by neural networks. Using statistical back-propagation, one may rely on GO gradients to perform HVI with low variance, while greatly broadening modeling flexibility.

6. RELATED WORK:

There are many methods directed toward low-variance gradients for expectation-based objectives.

Attracted by the generalization of REINFORCE, many works try to improve its performance via efficient variance-reduction techniques, like control variants (Mnih & Gregor, 2014; Titsias & Lázaro-Gredilla, 2015; Gu et al., 2015; Mnih & Rezende, 2016; Tucker et al., 2017; Grathwohl et al., 2017) or via data augmentation and permutation techniques (Yin & Zhou, 2018) . Most of this research focuses on discrete random variables, likely because Rep (if it exists) works well for continuous random variables but it may not exist for discrete random variables. Other efforts are devoted to continuously relaxing discrete variables, to combine both REINFORCE and Rep for variance reduction (Jang et al., 2016; Maddison et al., 2016; Tucker et al., 2017; Grathwohl et al., 2017) .

Inspired by the low variance of Rep, there are methods that try to generalize its scope. The Generalized Rep (GRep) gradient (Ruiz et al., 2016) employs an approximate reparameterization whose transformed distribution weakly depends on the parameters of interest. Rejection sampling variational inference (RSVI) (Naesseth et al., 2016) exploits highly-tuned transformations in mature rejection sampling simulation to better approximate Rep for non-reparameterizable distributions. Compared to the aforementioned methods, the proposed GO gradient, containing Rep as a special case for continuous random variables, applies to both continuous and discrete random variables with the same low-variance as the Rep gradient. Implicit Rep gradients (Figurnov et al., 2018) and pathwise derivatives (Jankowiak & Obermeyer, 2018) are recent low-variance methods that exploit the gradient of the expected function; they are special cases of GO in the single-layer continuous settings.

The idea of gradient backpropagation through random variables has been exploited before. RE-LAX (Grathwohl et al., 2017) , employing neural-network-parametrized control variants to assist REINFORCE for that goal, has a variance potentially as low as the Rep gradient. SCG (Schulman et al., 2015) utilizes the generalizability of REINFORCE to construct widely-applicable stochastic computation graphs. However, REINFORCE is known to have high variance, especially for highdimensional problems, where the proposed methods are preferable when applicable (Schulman et al., 2015) . Stochastic back-propagation (Rezende et al., 2014; Fan et al., 2015) , focusing mainly on reparameterizable Gaussian random variables and deep latent Gaussian models, exploits the product rule for an integral to derive gradient backpropagation through several continuous random variables. By comparison, the proposed statistical back-propagation based on the GO gradient is applicable to most distributions for continuous random variables. Further, it also flexibly generalizes to hierarchical probabilistic graphical models with continuous internal variables and continuous/discrete leaf ones.

7. EXPERIMENTS:

We 

7.1. GAMMA AND NB ONE-DIMENSIONAL SIMPLE EXAMPLES:

We first consider illustrative one-dimensional "toy" problems, to examine the GO gradient for both continuous and discrete random variables. The optimization objective is expressed as

max φ ELBO(φ) = E q φ (z) [log p(z|x) -log q φ (z)] + log p(x),

where for continuous z we assume p(z|x) = Gam(z; α 0 , β 0 ) for given set (α 0 , β 0 ), with q φ (z) = Gam(z; α, β) and φ = {α, β}; for discrete z we assume p(z|x) = NB(z; r 0 , p 0 ) for given set (r 0 , p 0 ), with q φ (z) = NB(z; r, p) and φ = {r, p}. Stochastic gradient ascent with onesample-estimated gradients is used to optimize the objective, which is equivalent to minimizing KL(q φ (z) p(z|x)).

Figure 2 shows the results (see Appendix H for additional details). For the nonnegative continuous z associated with the gamma distribution, we compare our GO gradient with GRep (Ruiz et al., 2016) , RSVI (Naesseth et al., 2016) , and their modified version using the "sticking" idea (Roeder et al., 2017) , denoted as GRep-Stick and RSVI-Stick, respectively. For RSVI and RSVI-Stick, the shape augmentation parameter is set as 5 by default. The only difference between GRep and GRep-Stick (also RSVI and RSVI-Stick) is the latter does not analytically express the entropy

E q φ (z) [-log q φ (z)].

Figure 2 (a) clearly shows the utility of employing sticking to reduce variance; without it, GRep and RSVI exhibit high variance, that destabilizes the optimization for small gamma shape parameters, as shown in Figures 2(b ) and 2(c). We adopt the sticking approach hereafter for all the compared methods. Among methods with sticking, GO exhibits the lowest variance in general, as shown in Figures 2(a ) and 2(b). GO empirically provides more stable learning curves, as shown in Figure 2 (c). For the discrete case corresponding to the NB distribution, GO is compared to REINFORCE (Williams, 1992) . To address the concern about comparison with the same number of evaluations of the expected function3 , another curve of REINFORCE using 2 samples is also added, termed REINFORCE2. It is apparent from Figure 2 (d) that, thanks to analytically removing the "0" terms in (8), the GO gradient has much lower variance, even in this simple one-dimensional case.

7.2. DISCRETE VARIATIONAL AUTOENCODER:

To demonstrate the low variance of the proposed GO gradient, we consider the discrete variational autoencoder (VAE) experiment from REBAR (Tucker et al., 2017) and RELAX (Grathwohl et al., 2017) , to make a direct comparison with state-of-the-art variance-reduction methods. Since the statistical back-propagation in Theorem 3 cannot handle discrete internal variables, we focus on the single-latent-layer settings (1 layer of 200 Bernoulli random variables), i.e.,

p θ (x, z) : x ∼ Bern NN P x|z (z) , z ∼ Bern P z q φ (z|x) : z ∼ Bern NN P z|x (x)

where P z is the parameters of the prior p θ (z), NN P x|z (z) means using a neural network to project the latent binary code z to the parameters P x|z of the likelihood p θ (x|z), and NN P z|x (x) is similarly defined for q φ (z|x). The objective is given in (1). See Appendix I for more details. We believe GO's better performance originates from: (i) its inherent low-variance nature; (ii) GO has less parameters compared to REBAR and RELAX (no control variant is adopted for GO); (iii) efficient batch processing methods (see Appendix I) are adopted to benefit from parallel computing. Table 1 presents the best training/validation ELBOs under various experimental settings for the compared methods. GO provides the best performance in all situations. Additional experimental results are given in Appendix I.

Many variance-reduction techniques can be used to further reduce the variance of GO, especially when complicated models are of interest. Compared to RELAX, GO cannot be directly applied when f (y) is not computable or where the interested model has discrete internal variables (like multilayer Sigmoid belief networks (Neal, 1992) ). For the latter issue, we present in Appendix B.4 a procedure to assist GO (or statistical back-propagation in Theorem 3) in handling discrete internal variables.

7.3. MULTINOMIAL GAN:

To demonstrate the deep GO gradient in Theorem 2, we adopt multinomial leaf variables x and construct a new multinomial GAN (denoted as MNGAN-GO) for generating discrete observations with a finite alphabet. The corresponding generator p θ (x) is expressed as

∼ N (0, I), x ∼ Mult(1, NN P ( )).

For brevity, we integrate the generator's parameters θ into the NN notation, and do not explicitly express them. Details for this example are provided in Appendix J.

We compare MNGAN-GO with the recently proposed boundary-seeking GAN (BGAN) (Hjelm et al., 2018) on 1-bit (1-state, Bernoulli leaf variables x), 1-bit (2-state), 2-bit (4-state), 3-bit (8-state) and 4-bit (16-state) discrete image generation tasks, using quantized MNIST datasets (LeCun et al., 1998) .

Table 2 presents inception scores (Salimans et al., 2016) of both methods. MNGAN-GO performs better in general. Further, with GO's assistance, MNGAN-GO shows more potential to benefit from richer information coming from more quantized states. For demonstration, Figure 4 shows the generated samples from the 4-bit experiment, where better image quality and higher diversity are observed for the samples from MNGAN-GO.

7.4. HVI FOR DEEP EXPONENTIAL FAMILIES AND DEEP LATENT DIRICHLET ALLOCATION:

To demonstrate statistical back-propagation in Theorem 3, we design variational inference nets for two nonconjugate hierarchical Bayesian models, i.e., deep exponential families (DEF) (Ranganath et al., 2015) and deep latent Dirichlet allocation (DLDA) (Zhou et al., 2015; 2016; Cong et al., 2017) . l+1) , Φ (l) ∼ Dir(η 0 ). For demonstration, we design the inference nets q φ (z|x) following the first-order Markov chain construction in Section 5, namely

EQUATION

EQUATION

EQUATION

β (z (l-1) ) . Further details are provided in Appendix K. One might also wish to design inference nets that have structure beyond the above first-order Markov chain construction, as in Zhang et al. (2018) ; we do not consider that here, but Theorem 3 is applicable to that case. HVI for a 2-layer DEF is first performed, with the ELBO curves shown in Figure 5 . GO enables faster and more stable convergence. Figure 6 presents the HVI results for a 3-layer DLDA, for which stable ELBOs are again observed. More importantly, with the GO gradient, one can utilize pure gradient-based methods to efficiently train such complicated nonconjugate models for meaningful dictionaries (see Appendix K for more implementary details).

8. CONCLUSIONS:

For expectation-based objectives, we propose a General and One-sample (GO) gradient that applies to continuous and discrete random variables. We further generalize the GO gradient to cases for which the underlying model is deep and has a marginal distribution corresponding to the latent variables of interest, and to cases for which the latent variables are hierarchical. The GO-gradient setup is demonstrated to yield the same low-variance estimation as the reparameterization trick, which is only applicable to reparameterizable continuous random variables. Alongside the GO gradient, we constitute a means of propagating the chain rule through distributions. Accordingly, we present statistical back-propagation, to flexibly integrate deep neural networks with general classes of random variables.

A PROOF OF THEOREM 1:

We first prove (7) in the main manuscript, followed by its discrete counterpart, i.e., (8) in the main manuscript. Then, it is easy to verify Theorem 1.

A.1 PROOF OF EQUATION ( 7) IN THE MAIN MANUSCRIPT Similar proof in one-dimension is also given in the supplemental materials of Ranganath et al. (2016b) .

We want to calculate

∇ γ E qγ (y) [f (y)] = ∇ γ E v qγ (yv) [f (y)] = v E qγ (y -v ) f (y)∇ γ q γ (y v )dy v ,

where y -v denotes y with y v excluded. Without loss of generality, we assume y v ∈ (-∞, ∞).

Let v (y v ) = ∇ γ q γ (y v ), and we have

v(y v ) = yv -∞ ∇ γ q γ (t)dt = ∇ γ yv -∞ q γ (t)dt = ∇ γ Q γ (y v ),

where Q γ (y v ) is the cumulative distribution function (CDF) of q γ (y v ). Further define u(y v ) = f (y v , y -v ), we then apply integration by parts (or partial integration) to get

EQUATION

With Q γ (∞) = 1 and Q γ (-∞) = 0, it's straightforward to verify that the first term is always zero for any Q γ (y v ), thus named the "0" term.

A.2 PROOF OF EQUATION (8) IN THE MAIN MANUSCRIPT:

For discrete variables y, we have

∇ γ E qγ (y) [f (y)] = v E qγ (y -v ) N yv=0 f (y)∇ γ q γ (y v ) ,

where y v ∈ {0, 1, • • • , N } and N is the size of the alphabet.

To handle the summation of products of two sequences and develop discrete counterpart of ( 7 

S N = N n=0 a n b n = a 0 b 0 + N n=1 a n (B n -B n-1 ) = a 0 B 0 + N n=1 a n B n - N -1 n=0 a n+1 B n = a N B N + N -1 n=0 a n B n - N -1 n=0 a n+1 B n = a N B N - N -1 n=0 (a n+1 -a n )B n . Substituting n = y v , a n = f (y), b n = ∇ γ q γ (y v ), and B n = ∇ γ Q γ (y v ) into the above equation, we have ∇ γ E qγ (y) [f (y)] = v E qγ (y -v ) f (y -v , y v = N )∇ γ Q γ (y v = N ) "0" - N -1 yv=0 [f (y -v , y v + 1) -f (y)]∇ γ Q γ (y v ) "Key"

.

Note the first term equals zero for both finite alphabet, i.e., N < ∞, and infinite alphabet, i.e., N = ∞. When N = ∞, we get (8) in the main manuscript.

With the above proofs for Eqs. ( 7) and ( 8), one could straightforwardly verify Theorem 1. 

q γ (y) g qγ (y) γ Delta δ(y -µ) g qγ (y) µ = 1 Bernoulli(p) g qγ (y) p = 1/(1 -p) y = 0 0 y = 1 Normal(µ, σ 2 ) g qγ (y) µ = 1 g qγ (y) σ = y-µ σ Log-Normal(µ, σ 2 ) g qγ (y) µ = y g qγ (y) σ = y log y-µ σ Logit-Normal(µ, σ 2 ) g qγ (y) µ = y(1 -y) g qγ (y) σ = y(1 -y) logit(y)-µ σ Cauchy(µ, γ) g qγ (y) µ = 1 g qγ (y) γ = y-µ γ Gamma α, 1 β g qγ (y) α

= [log(βy)-ψ(α)]Γ(α,βy)+βyT (3,α,βy)

β α y α-1 e -βy g qγ (y) β = -y β Beta α, β g qγ (y) α = - log y -ψ(α) + ψ(α + β) y α-1 (1 -y) β-1 B(y; α, β) + y α 2 (1 -y) β-1 • 3F2(α, α, 1 -β; α + 1, α + 1; y) g qγ (y) β =   log(1 -y) -ψ(β) + ψ(α + β) y α-1 (1 -y) β-1 B(1 -y; β, α) - 1 -y β 2 y α-1 • 3 F2(β, β, 1 -α; β + 1, β + 1; 1 -y)   NB r, p g qγ (y) r =   - log(1 -p) -ψ(r) + ψ(r + y + 1) (1 -p) r p y (y + r)B(1 -p; r, y + 1) + y + r p y r 2 •3F2(r, r, -y; r + 1, r + 1; 1 -p)   g qγ (y) p = y+r 1-p Exponential(λ) g qγ (y) λ = -y λ Student's t(v) g qγ (y) v = y 2 1 + y 2 v v+1 2 •        1 v + ψ v 2 -ψ v + 1 2 • 2F1 1 2 , v + 1 2 ; 3 2 ; - y 2 v + 2y(v + 1) 3v • 2F1 3 2 , v + 3 2 ; 5 2 ; - y 2 v + y 2 3v • F 212 201     3 2 , v + 3 2 ; 1; 1, v + 1 2 ; 2, 5 2 ; ; v + 3 2 ; - y 2 v , - y 2 v            Weibull(λ, k) g qγ (y) λ = y λ g qγ (y) k = y k log λ y Laplace(µ, b) g qγ (y) µ = 1 g qγ (y) b = y-µ b Poisson(λ) g qγ (y) λ = 1 Geometric(p) g qγ (y) p = -y+1 p Categorical(p) g qγ (y) p = -1 py [1 0≤y , 1 1≤y , • • • , 1 (N -1)≤y , 0] T for y = {0, 1, • • • , N }, p = [p 0 , p 1 , • • • , p N ] T ψ(x)

is the digamma function. Γ(x, y) is the upper incomplete gamma function. B(x; α, β) is the incomplete beta function. pFq (a1, • • • , ap; b1, • • • , bp; x) is the generalized hypergeometric function. T (m, s, x) is a special case of Meijer G-function (Geddes et al., 1990) .

B WHY DISCRETE Internal VARIABLES ARE CHALLENGING?:

For simpler demonstration, we first use a 2-layer model to show why discrete internal variables are challenging. Then, we present an importance-sampling proposal that might be useful under specific situations. Finally, we present a strategy to learn discrete internal variables with the statistical back-propagation in Theorem 3 of the main manuscript.

Assume q γ (y) being the marginal distribution of the following 2-layer model

q γ (y) = E qγ λ (λ) [q γ y (y|λ)]

where γ = {γ y , γ λ }, q γ (y, λ) = q γ y (y|λ)q γ λ (λ) = v q γ y (y v |λ) • k q γ λ (λ k ), and both the leaf variable y and the internal variable λ could be either continuous or discrete.

Accordingly, the objective becomes

E qγ (y) [f (y)] = E qγ (y,λ) [f (y)] = E qγ y (y|λ)qγ λ (λ) [f (y)].

For gradient wrt γ y , using Theorem 1, it is straight to show

∇ γ y E qγ (y,λ) [f (y)] = E qγ λ (λ) ∇ γ y E qγ y (y|λ) [f (y)] = E qγ (y,λ) G qγ y (y|λ) γ y D y [f (y)] . (15)

For gradient wrt γ λ , we first have

∇ γ λ E qγ (y,λ) [f (y)] = ∇ γ λ E qγ λ (λ) E qγ y (y|λ) [f (y)] .

With f (λ) = E qγ y (y|λ) [f (y)], we then apply Theorem 1 and get

EQUATION

where

EQUATION

D λ k [ f (λ)] ∇ λ k f (λ), Continous λ k f (λ -k , λ k + 1) -f (λ), Discrete λ k

Next, we separately discuss the situations where λ k is continuous or discrete.

B.1 FOR CONTINUOUS λ k:

One can directly apply Theorem 1 again, namely

EQUATION

Substituting ( 17) into ( 16), we have

EQUATION

B.2 FOR DISCRETE λ k:

In this case, we need to calculate

D λ k [ f (λ)] = f (λ -k , λ k + 1) -f (λ).

The keys are again partial integration and Abel transformation. For simplicity, we first assume one-dimensional y, and separately discuss y being continuous and discrete.

For continuous y, we apply partial integration to f (λ) and get

f (λ) = E qγ y (y|λ) [f (y)] = q γ y (y|λ)f (y)dy = f (y)Q γ y (y|λ) ∞ -∞ -Q γ y (y|λ)∇ y f (y)dy.

Accordingly, we have

D λ k [ f (λ)] = f (λ -k , λ k + 1) -f (λ) = Q γ y (y|λ) -Q γ y (y|λ -k , λ k + 1) ∇ y f (y)dy + f (y)Q γ y (y|λ -k , λ k + 1) ∞ -∞ -f (y)Q γ y (y|λ) ∞ -∞ "0"

.

Removing the "0" term, we have

EQUATION

)

For discrete y, by similarly exploiting Abel transformation, we have

f (λ) = E qγ y (y|λ) [f (y)] = y q γ y (y|λ)f (y) = f (∞)Q γ y (∞|λ) - y Q γ y (y|λ) f (y + 1) -f (y) .

Accordingly, we get

EQUATION

Unifying ( 19) for continuous y and ( 20) for discrete y, we have

D λ k [ f (λ)] = f (λ -k , λ k + 1) -f (λ) = E qγ y (y|λ -k ,λ k +1) [f (y)] -E qγ y (y|λ) [f (y)] = E qγ y (y|λ) D y [f (y)] • ḡqγ y (y|λ) λ k (21)

where we define ḡqγ y (y|λ)

λ k Q γ y (y|λ) -Q γ y (y|λ -k , λ k + 1)

q γ y (y|λ) .

Multi-dimensional y. Based on the above one-dimensional foundation, we next move on to multidimensional situations. With definitions y :i {y 1 , • • • , y i } and y i: {y i , • • • , y V }, where V is the dimensionality of y, we have

EQUATION

Apply ( 21) and we have

D λ k [ f (λ)] = E qγ y (y 2: |λ -k ,λ k +1) E qγ y (y1|λ) D y1 [f (y)] • ḡqγ y (y1|λ) λ k + E qγ y (y 2: |λ -k ,λ k +1) E qγ y (y1|λ) [f (y)] -E qγ y (y|λ) [f (y)] = E qγ y (y 2: |λ -k ,λ k +1)qγ y (y1|λ) D y1 [f (y)] • ḡqγ y (y1|λ) λ k + f (y) -E qγ y (y|λ) [f (y)]

Similarly, we add extra terms to the above equation to enable applying (21) again as

D λ k [ f (λ)] = E qγ y (y 2: |λ -k ,λ k +1)qγ y (y1|λ) D y1 [f (y)] • ḡqγ y (y1|λ) λ k + f (y) -E qγ y (y 3: |λ -k ,λ k +1)qγ y (y :2 |λ) D y1 [f (y)] • ḡqγ y (y1|λ) λ k + f (y) + E qγ y (y 3: |λ -k ,λ k +1)qγ y (y :2 |λ) D y1 [f (y)] • ḡqγ y (y1|λ) λ k + f (y) -E qγ y (y|λ) [f (y)].

Accordingly, we apply (21) to the first two terms and have

D λ k [ f (λ)] = E qγ y (y 3: |λ -k ,λ k +1)qγ y (y :2 |λ) D y2 D y1 [f (y)]ḡ qγ y (y1|λ) λ k + f (y) ḡqγ y (y2|λ) λ k + E qγ y (y 3: |λ -k ,λ k +1)qγ y (y :2 |λ) D y1 [f (y)]ḡ qγ y (y1|λ) λ k + f (y) -E qγ y (y|λ) [f (y)] = E qγ y (y 3: |λ -k ,λ k +1)qγ y (y :2 |λ)   D y2 D y1 [f (y)]ḡ qγ y (y1|λ) λ k + f (y) ḡqγ y (y2|λ) λ k + D y1 [f (y)]ḡ qγ y (y1|λ) λ k + f (y)   -E qγ y (y|λ) [f (y)]

So forth, we summarize the pattern into the following equation as

EQUATION

where A λ k (y, V ) is iteratively calculated as

A λ k (y, 0) = f (y) A λ k (y, v) = D yv A λ k (y, v -1) ḡqγ y (yv|λ) λ k + A λ k (y, v -1) A λ k (y, V ) = D y V A λ k (y, V -1) ḡqγ y (y V |λ) λ k + A λ k (y, V -1).

Despite elegant structures within ( 23), to calculate it, one must iterate over all dimensions of y, which is computational expensive in practice. More importantly, it is straightforward to show that deeper models will have similar but much more complicated expressions.

B.3 AN IMPORTANCE-SAMPLING PROPOSAL TO HANDLE DISCRETE Internal VARIABLES:

Next, we present an intuitive proposal that might be useful under specific situations.

The key idea is to use different extra items to enable "easy-to-use" expression for

D λ k [ f (λ)], namely D λ k [ f (λ)] = f (λ -k , λ k + 1) -f (λ) = E qγ y (y|λ -k ,λ k +1) [f (y)] -E qγ y (y|λ) [f (y)] = E qγ y (y|λ -k ,λ k +1) [f (y)] -E qγ y (y :V -1 |λ -k ,λ k +1)qγ y (y V |λ) [f (y)] + E qγ y (y :V -1 |λ -k ,λ k +1)qγ y (y V |λ) [f (y)] • • • -E qγ y (y :i-1 |λ -k ,λ k +1)qγ y (y i: |λ) [f (y)] + E qγ y (y :i-1 |λ -k ,λ k +1)qγ y (y i: |λ) [f (y)] • • • -E qγ y (y1|λ -k ,λ k +1)qγ y (y 2: |λ) [f (y)] + E qγ y (y1|λ -k ,λ k +1)qγ y (y 2: |λ) [f (y)] -E qγ y (y|λ) [f (y)].

Apply (21) to the adjacent two terms for V times, we have

D λ k [ f (λ)] = E qγ y (y :V -1 |λ -k ,λ k +1)qγ y (y V |λ) D y V [f (y)]ḡ qγ y (y V |λ) λ k • • • + E qγ y (y :i-1 |λ -k ,λ k +1)qγ y (y i: |λ) D yi [f (y)]ḡ qγ y (yi|λ) λ k • • • + E qγ y (y V |λ) D y1 [f (y)]ḡ qγ y (y1|λ) λ k ,

where we can apply the idea of importance sampling and modify the above equation to

EQUATION

Published as a conference paper at ICLR 2019 Note that importance sampling may not always work well in practice (Bishop, 2006) .

We further define the generalized variable-nabla as

EQUATION

With the generalized variable-nabla, we unify (24) for discrete λ k and ( 17) for continuous λ k and get

D λ k [ f (λ)] = E qγ y (y|λ) v D yv [f (y)] • gg qγ (yv|λ) λ k

, which apparently obeys the chain rule. Accordingly, we have the gradient for γ λ in ( 16) as

∇ γ λ E qγ (y,λ) [f (y)] = E qγ (y,λ) k v D yv [f (y)] • gg qγ (yv|λ) λ k • gg qγ λ (λ k ) γ λ .

One can straightforwardly verify that, with the generalized variable-nabla defined in (25), the chain rule applies to

∇ γ E qγ (y (L) ) [f (y (L) )],

where one can freely specify both leaf and internal variables to be either continuous or discrete. The only problem is that, for discrete internal variables, the importance sampling trick used in (24) may not always work as expected.

B.4 STRATEGY TO LEARN DISCRETE INTERNAL VARIABLES WITH STATISTICAL BACK-PROPAGATION:

Practically, if one has to deal with a q γ (y (1) , • • • , y (L) ) with discrete internal variables y (l) , l < L, we suggest the strategy in Figure 7 , with which one should expect a close performance but enjoy much easier implementation with statistical back-propagation in Theorem 3 of the main manuscript.

In fact, one can always add additional continuous internal variables to the graphical models to remedy the performance loss or even boost the performance.

(a) (b)

Figure 7 : A strategy for discrete internal variables. Blue and red circles denote continuous and discrete variables, respectively. The centered dots represent the corresponding distribution parameters. (a) Practically, one uses a neural network (black arrow) to connect the left variable to the parameters of the center discrete one, and then uses another neural network to propagate the sampled value to the next. (b) Instead, we suggest "extracting" the discrete variable as a leaf one and propagate its parameters to the next.

C PROOF OF LEMMA 1:

First, a marginal distribution q γ (y) with reparameterization y = τ γ ( ), ∼ q( ) can be expressed as a joint distribution, namely q γ (y) = q γ (y, ) = q γ (y| )q( ),

where q γ (y| ) = δ(y -τ γ ( )), δ(•) is the Dirac delta function, and τ γ ( ) could be flexibly specified as a injective, surjective, or bijective function.

Next, we align notations and rewrite (10) as

EQUATION

where

G qγ (y| ) γ = • • • , g qγ (yv| ) γ , • • • and g qγ (yv| ) γ -1 qγ (yv| ) ∇ γ Q γ (y v | ). With q γ (y v | ) = δ(y v -[τ γ ( )] v ), we have Q γ (y v | ) = 1 [τ γ ( )] v ≤ y v 0 [τ γ ( )] v > y v

Accordingly, we have

∇ γ Q γ (y v | ) = ∇ [τ γ ( )]v Q γ (y v | ) • ∇ γ [τ γ ( )] v = -δ([τ γ ( )] v -y v ) • ∇ γ [τ γ ( )] v and g qγ (yv| ) γ = -1 q γ (y v | ) ∇ γ Q γ (y v | ) = ∇ γ [τ γ ( )] v .

Substituting the above equations into (26), we get

EQUATION

which is the multi-dimensional Rep gradient in (5) of the main manuscript.

D PROOF OF THEOREM 2:

Firstly, with the internal variable λ being continuous, ( 10) and ( 11) in the main manuscript are proved by ( 15) and ( 18) in Section B, respectively. Then, by iteratively generalizing the similar derivations to deep models and utilizing the fact that the GO gradients with variable-nablas in expectation obey the chain rule for models with continuous internal variables, Theorem 2 could be readily verified.

E PROOFS FOR COROLLARY 1:

When all q γ (i) y (i) |y (i-1) s are specified as Dirac delta functions, namely

q γ (i) y (i) |y (i-1) = δ(y (i) -σ(γ (i) , y (i-1) )),

where σ(γ (i) , y (i-1) ) denotes the activated values after activation functions, the objective becomes

EQUATION

where

γ = {γ (1) , • • • , γ (N ) }.

Back-Propagation. For the objective in ( 28), the Back-Propagation is expressed as

EQUATION

where

∇ y (i) f (•) = [∇ y (i) y (i+1) ][∇ y (i+1) f (•)].

Deep GO Gradient. We consider the continuous special case, where

D y (L) f (y (L) ) = ∇ y (L) f (y (L) ).

With q γ (i+1) y (i+1) |y (i) s being Dirac delta functions, namely,

q γ (i+1) y (i+1) k |y (i) = ∞, [σ(γ (i+1) , y (i) )] k = y (i+1) k 0, [σ(γ (i+1) , y (i) )] k = y (i+1) k we have Q γ (i+1) y (i+1) k |y (i) = 1, [σ(γ (i+1) , y (i) )] k ≤ y (i+1) k 0, [σ(γ (i+1) , y (i) )] k > y (i+1) k

Taking derivative wrt y

(i) v , we got ∇ y (i) v Q γ (i+1) y (i+1) k |y (i) = ∇ [σ(γ (i+1) ,y (i) )] k Q γ (i+1) y (i+1) k |y (i) • ∇ y (i) v [σ(γ (i+1) , y (i) )] k = -δ [σ(γ (i+1) , y (i) )] k -y (i+1) k • ∇ y (i) v [σ(γ (i+1) , y (i) )] k

Accordingly, we have

g q γ (i+1) (y (i+1) k |y (i) ) y (i) v = -1 q γ (i+1) (y (i+1) k |y (i) ) ∇ y (i) v Q γ (i+1) (y (i+1) k |y (i) ) = ∇ y (i) v [σ(γ (i+1) , y (i) )] k = ∇ y (i) v y (i+1) k

By substituting the above equation into (12) in Theorem 2, and then comparing it with(29), one can easily verify Corollary 1.

F PROOF FOR THEOREM 3:

Based on the proofs for Theorem 1 and Theorem 2, it is clear that, if one constrains all internal variables to be continuous, the GO gradients in expectation obey the chain rule. Therefore, one can straightforwardly utilizing the chain rule to verify Theorem 3. Actually, Theorem 3 may be seen as the chain rule generalized with random variables, among which the internal ones are only allowed to be continuous.

G DERIVATIONS FOR HIERARCHICAL VARIATIONAL INFERENCE:

In Hierarchical Variational Inference, the objective is to maximize the evidence lower Bound (ELBO) ELBO(θ, φ; x) = E q φ (z|x) [log p θ (x, z) -log q φ (z|x)].

For the common case with z = {z (1) , • • • , z (L) }, it is obvious that Theorem 3 of the main manuscript can be applied when optimizing φ.

Practically, there are situations where one might further put a latent variable λ in reference q φ (z|x), namely q φ (z|x) = q φ z (z|λ)q φ λ (λ)dλ with φ = {φ z , φ λ }. Following Ranganath et al. (2016b) , we briefly discuss this situation here.

We first show that there is another unnecessary variance-injecting "0" term.

EQUATION

where the second "0" term is straightly verified as

q φ (z|x)∇ φ log q φ (z|x)dz = ∇ φ q φ (z|x)dz = ∇ φ q φ (z|x)dz = ∇ φ 1 = 0.

Eliminating the "0" term from (31), one still has another problem, that is, log q φ (z|x) is usually non-trivial when q φ (z|x) is marginal. For this problem, we follow Ranganath et al. (2016b) to use another lower bound ELBO2 of the ELBO in (30).

-log q φ (z|x) = q φ (λ|z, x)[log q φ (z|x)]dλ = q φ (λ|z, x) -log q φ (z, λ|x) q φ (λ|z, x) dλ = E q φ (λ|z,x) -log q φ (z, λ|x) r ω (λ|z, x) + log q φ (λ|z, x) r ω (λ|z, x)

≥ E q φ (λ|z,x) -log q φ (z, λ|x) r ω (λ|z, x)

where r ω (λ|z, x), evaluable, is an additional variational distribution to approximate the variational posterior q φ (λ|z, x). Accordingly, we get the ELBO2 for Hierarchical Variational Inference as ELBO2(θ, φ; x) = E q φ (z,λ|x) log p θ (x, z) -log q φ (z, λ|x) + log r ω (λ|z, x) .

Note similar to (31), the unnecessary "0" term related to log q φ (z, λ|x) should also be removed. Accordingly, we have ∇ φ ELBO2(θ, φ; x) = ∇ φ q φ (z, λ|x) log p θ (x, z)-log q φ (z, λ|x)+log r ω (λ|z, x) dλdz.

Obviously, Theorem 3 is readily applicable to provide GO gradients.

H DETAILS OF GAMMA AND NB ONE-DIMENSIONAL SIMPLE EXAMPLES:

We first consider illustrative one-dimensional "toy" problems, to examine the GO gradient in Theorem 1 for both continuous and discrete random variables.

The optimization objective is expressed as

max φ ELBO(φ) = E q φ (z) [log p(z|x) -log q φ (z)] + log p(x),

where for continuous z we assume p(z|x) = Gam(z; α 0 , β 0 ) for set (α 0 , β 0 ), with q φ (z) = Gam(z; α, β) and φ = {α, β}; for discrete z we assume p(z|x) = NB(z; r 0 , p 0 ) for set (r 0 , p 0 ), with q φ (z) = NB(z; r, p) and φ = {r, p}. Stochastic gradient ascent with one-sample-estimated gradients is used to optimize the objective, which is equivalent to minimizing KL(q φ (z) p(z|x)).

Figure 8 shows the experimental results. For the nonnegative continuous z associated with the gamma distribution, we compare our GO gradient with GRep (Ruiz et al., 2016) , RSVI (Naesseth et al., 2016) , and their modified version using the "sticking" idea (Roeder et al., 2017) , denoted as GRep-Stick and RSVI-Stick respectively. For RSVI and RSVI-Stick, the shape augmentation parameter is set as 5 by default. The only difference between GRep and GRep-Stick (also RSVI and RSVI-Stick) is the latter does NOT analytically express the entropy E q φ (z) [-log q φ (z)]. One should apply sticking because (i) Figures 8(a )-8(c) clearly show its utility in reducing variance; and (ii) without it, GRep and RSVI exhibit high variance that unstabilizes the optimization for small gamma shape parameters, as shown in Figures 8(d )-8(f). We adopt the sticking approach hereafter for all the compared methods. Since the gamma rate parameter β is reparameterizable, its gradient calculation is the same for all sticking methods, including GO, GRep-Stick, and RSVI-Stick. Therefore, similar variances are observed in Figures 8(b ) and 8(e). Among methods with sticking, GO exhibits the lowest variance in general, as shown in Figures 8(a ) and 8(d). Note it is high variance that causes optimization issues. As a result, GO empirically provides more stable learning curves, as shown in Figures 8(c ) and 8(f). For the discrete case corresponding to the NB distribution, GO is compared to REINFORCE (Williams, 1992) . To estimate gradient, REINFORCE uses 1 sample of z and 1 evaluation of the expected function; whereas GO uses 1 sample and 2 evaluations. To address the concern about comparison with the same number of evaluations of the expected function, another curve of REINFORCE using 2 samples (thus 2 evaluations of the expected function) is also added, termed REINFORCE2. It is apparent from Figures 8(g)-8(l) that, thanks to analytically removing the "0" terms, the GO gradient has much lower variance and thus faster convergence, even in this simple one-dimensional case.

I DETAILS OF THE DISCRETE VAE EXPERIMENT:

Complementing the discrete VAE experiment of the main manuscript, we present below its experimental settings, implementary details, and additional results.

Since the presented statistical back-propagation in Theorem 3 of the main manuscript cannot handle discrete internal variables, we focus on the single-latent-layer settings (1 layer of 200 Bernoulli random variables) for fairness, i.e., p θ (x, z) :

EQUATION

Referring to the experimental settings in Grathwohl et al. (2017) , we consider

The used datasets and other experimental settings, including the hyperparameter search strategy, are the same as those in Grathwohl et al. (2017) .

For such single-latent-layer settings, it is obvious that Theorem 3 (also Theorem 1) can be straightforwardly applied. However, since Bernoulli distribution has finite support, as mentioned in the main manuscript, we should analytically express some expectations for lower variance, as detailed below. Notations of ( 8) and ( 9) of the main manuscript are used for clarity and also for generalization.

In fact, we should take a step back and start from (8) of the main manuscript, which is equivalent to analytically express an expectation in (9), namely

EQUATION

where

Q γ (y v ) = 1 -P v (γ) y v = 0 1 y v = 1 with P v (γ)

being the Bernoulli probability of Bernoulli random variable y v . Accordingly, we have

EQUATION

) For better understanding only, with abused notations

∇ γ P = [• • • , ∇ γ P v (γ), • • • ] T , ∇ P y = I, and ∇ y f (y) = [• • • , {f (y -v , y v = 1) -f (y -v , y v = 0)}, • • • ] T ,

one should observe a chain rule within the above equation.

To assist better understanding of how to practically cooperate the presented GO gradients with deep learning frameworks like TensorFlow or PyTorch, we take (34) as an example, and present for it the following simple algorithm.

Algorithm 1 An algorithm for (34) as an example to demonstrate how to practically cooperate GO gradients with deep learning frameworks like TensorFlow or PyTorch. One sample is assumed for clarity. Practically, an easy-to-use trick for changing gradients of any function h(x) is to define ĥ(x) = x T StopGradient[g] + StopGradient[h(x) -x T g] with g the desired gradients. # Forward-Propagation 1. γ → P (γ): Calculate Bernoulli probabilities P (γ ) 2. P (γ) → y: Sample y ∼ Bern(P (γ))

Change gradient: ∇ P (γ) y = I 3. y → f (y): Calculate the loss f (y)

Change gradient:

[∇ y f (y)] v = f (y -v , y v = 1) -f (y -v , y v = 0) # Back-Propagation 1.

Rely on the mature auto-differential software for back-propagating gradients For efficient of D y f (y), one should exploit the prior knowledge of function f (y). For example, f (y)s are often neural-network-parameterized. Under that settings, one should be able to exploit tensor operation to design efficient implementation of D y f (y). Again we take (34) as an example, and assume f (y) has the special structure

EQUATION

) where σ(•) is an element-wise nonlinear activation function, and r(•) is a function that takes in a vector and outputs a scalar. One can easily modify the above f (y) for the considered discrete VAE experiment.

Since y v s are now Bernoulli random variables with support {0, 1}, we have that

EQUATION

where Then to efficiently calculate the f (y -v , y v + a v )s, we use the following batch processing procedure to benefit from parallel computing.

a v = 1 y v = 0 -1 y v = 1 is the v-th

•

Step 1: Define Ξ y h as the matrix whose element [Ξ y h] vj represents the "new" h * j when input {y -v , y v + a v } in (35). Then, we have

[Ξ y h] vj = σ(y T W :j + b j + a v W vj )

where W :j is the j-th column of the matrix W. Note the vth row of Ξ y h, i.e., [Ξ y h] v: , happens to be the "new" h * when input {y -v , y v + a v }.

• Step 2: Similarly, we define Ξ y f as the vector whose element

[Ξ y f ] v = f (y -v , y v + a v ).

Utilizing Ξ y h obtained in Step 1, we have

Ξ y f = r([Ξ y h]ϑ + c T ),

where r(•) is applied to each row of the matrix ([Ξ y h]ϑ + c T ).

Note the above batch processing procedure can be easily extended to deeper neural networks. Accordingly, we have D y f (y) = a [Ξ y f -f (y)], where represents the matrix element-wise product. Now we can rely on Algorithm 1 to solve the problem whose objective has its gradient expressed as (34), for example the inference of the single-latent-layer discrete VAE in (32).

All training curves versus iteration/running-time are given in Figure 9 , where it is apparent GO provides better performance, a faster convergence rate, and a better running efficiency in all situations. The average running time per 100 iterations for the compared methods are given in Table 4 , where GO is 2 -4 times faster in finishing the same number of training iterations. We also quantify the running efficiency of GO by considering its running time to achieve the best training ELBO (within 1, 000, 000 training iterations) of RERAR/RELAX, referring to the black lines shown in the second-column subfigures of Figure 9 . It is clear that GO is approximately 5 -10 times more efficient than REBAR/RELAX in the considered experiments.

As shown in the second and fourth rows of Figures 9, for the experiments with nonlinear models all methods suffer from over-fitting, which originates from the redundant complexity of the adopted neural networks and appeals for model regularizations. We detailedly clarify these experimental results as follows.

• All the compared methods are given the same and only objective, namely to maximize the training ELBO on the same training dataset with the same model; GO clearly shows its power in achieving a better objective.

• The "level" of over-fitting is ultimately determined by the used dataset, model, and objective; it is independent of the adopted optimization method. Different optimization methods just reveal different optimizing trajectories, which show different sequences of training objectives and over-fitting levels (validation objectives).

• Since all methods are given the same dataset, model, and objective, they have the same over-fitting level. Because GO has a lower variance, and thus more powerful optimization capacity, it gets to the similar situations much faster than REBAR/RELAX. Note this does not mean GO suffers more from over-fitting. In fact, GO provides better validation ELBOs in all situations, as shown in Figure 9 and also Table 1 of the main manuscript. In practice, GO can benefit from the early-stopping trick to get a better generalization ability.

J DETAILS OF THE MULTINOMIAL GAN:

Complementing the multinomial GAN experiment in the main manuscript, we present more details as follows. For a quantitative assessment of the computational complexity, our PyTorch code takes about 30 minutes to get the most challenging 4-bit task in Figs. 4 and 11, with a Titan Xp GPU.

Firstly, recall that the generator p θ (x) of the developed multinomial GAN (denoted as MNGAN-GO) is expressed as ∼ N (0, I), x ∼ Mult(1, NN P ( )), where NN P ( ) denotes use of a neural network to project to distribution parameters P. For brevity, we integration the generator's parameters θ into the NN notation and do not explicitly express them. Multinomial leaf variables x is used to describe discrete observations with a finite alphabet. To train MNGAN-GO, the vanilla GAN loss (Goodfellow et al., 2014) is used. A deconvolutional neural network as in Radford et al. (2015) is used to map to P in the generator. The discriminator is constructed as a multilayer perceptron. Detailed model architectures are given in Table 5 . Figure 10 illustrates the pipeline of MNGAN-GO. Note MNGAN-GO has a smaller number of parameters, compared to BGAN (Hjelm et al., 2018) .

For clarity, we briefly discuss the employed data preprocessing. Taking MNIST for an example, the original data are 8-bit grayscale images, with pixel intensities ranging from 0 to 255. For the n-bit experiment, we obtain the real data, such as the 2-bit one in Figure 10 , by rescaling and quantizing the pixel intensities to the range [0, 2 n -1], having 2 n different states (values).

For the 1-bit special case, the multinomial distribution reduces to the Bernoulli distribution. Of course, one could intuitively employ the redundant multinomial distribution, which is denoted as 1-bit (2-state) in Table 2 . An alternative and popular approach is to adopt the Bernoulli distribution to remove the redundancy by only modeling its probability parameters; we denote this case as 1-bit (1-state, Bernoulli) in Table 2 .

Figure 11 shows the generated samples from the compared models on different quantized MNIST. It is obvious that MNGAN-GO provides images with better quality and wider diversity in general. With the learned variational inference nets, one could efficiently project the observed x to its latent variables z, c during testing. For applications requiring realtime processing, this is a clear advantage. For demonstration, Figure 12 shows test data samples x and their reconstruction x = Φ (1) ẑ(1) , ẑ(1) ∼ Gam NN (1) αz (x), NN

(1)

β z (x) ,

where one Monte Carlo sample is used to calculate ẑ(1) . 

:

A Bernoulli random variable y ∼ Bernoulli(P ) is identical to y = 1 <P with ∼ U (0, 1). But y is called non-reparameterizable because the transformation y = 1 <P is not differentiable.

In the NB experiments, REINFORCE uses 1 sample and 1 evaluation of the expected function; REIN-FORCE2 uses 2 sample and 2 evaluations; and GO uses 1 sample and 2 evaluations.

github.com/duvenaud/relax

github.com/duvenaud/relax

Figures and Tables:

Figure1is presented for better understanding. With variable-nablas, one can readily verify that the deep GO gradient in Theorem 2 in expectation obeys the chain rule.

Figure 2: Gamma (a-c) and NB (d) toy experimental results. (a) The gradient variance of gamma shape α versus iterations, with posterior parameters α0 = 1, β0 = 0.5. (b)-(c) The gradient variance of α and ELBOs versus iterations respectively, when α0 = 0.01, β0 = 0.5. (d) The gradient variance of NB r versus iterations with r0 = 10, p0 = 0.2. In each iteration, gradient variances are estimated with 20 Monte Carlo samples (each sample corresponds to one gradient estimate), among which the last one is used to update parameters.

Figure 3: Training curves for the discrete VAE experiment with 1-layer linear model (see Appendix I) on the stochastically binarized MNIST dataset(Salakhutdinov & Murray, 2008). All methods are run with the same learning rate for 1, 000, 000 iterations. The black line represents the best training ELBO of REBAR and RELAX. ELBOs are calculated using all training/validation data.

Figure 4: Generated samples from BGAN (top) and MNGAN-GO (bottom) trained on 4-bit quantized MNIST.

Figure 3 (also Figure 9 of Appendix I) shows the training curves versus iteration and running time for the compared methods. Even without any variance-reduction techniques, GO provides better performance, faster convergence rate, and better running efficiency (about ten times faster in achieving the best training ELBO of RERAR/RELAX in this experiment). We believe GO's better performance originates from: (i) its inherent low-variance nature; (ii) GO has less parameters compared to REBAR and RELAX (no control variant is adopted for GO); (iii) efficient batch processing methods (see Appendix I) are adopted to benefit from parallel computing. Table1presents the best training/validation ELBOs under various experimental settings for the compared methods. GO provides the best performance in all situations. Additional experimental results are given in Appendix I.

Figure 5: ELBOs of HVI for a 128-64 DEF on MNIST. Except for different ways to calculate gradients, all other experimental settings are the same for compared methods, including the sticking idea and one-MC-sample gradient estimate.

), we first introduce Abel transformation. Abel transformation. Given two sequences {a n } and {b n }, with n ∈ {0, • • • , N }, we define B 0 = b 0 and B n = n k=0 b n for n ≥ 1. Accordingly, we have

Figure 10: Illustration of MNGAN-GO.

Figure 11: Generated images from BGAN (top) and MNGAN-GO (bottom). Columns correspond to 1bit(Bernoulli), 1-bit, 2-bit, 3-bit, 4-bit tasks, respectively.

Figure 12: (a) Test data samples and (b) their reconstruction via the learned 128-64-32 DLDA.

Best obtained ELBOs for discrete variational autoencoders. Results of REBAR and RELAX are obtained by running the released code 4 fromGrathwohl et al. (2017). All methods are run with the same learning rate for 1, 000, 000 iterations.

Table content: <table><tr><td colspan="2">Dataset Model</td><td>Training</td><td/><td>Validation</td></tr><tr><td/><td/><td>REBAR RELAX</td><td>GO</td><td>REBAR RELAX</td><td>GO</td></tr><tr><td>MNIST</td><td colspan="4">Linear 1 layer -112.16 -112.89 -110.21 -114.85 -115.36 -114.27 Nonlinear -96.99 -95.99 -82.26 -112.96 -112.42 -111.48</td></tr><tr><td>Omniglot</td><td colspan="4">Linear 1 layer -122.19 -122.17 -119.03 -124.81 -124.95 -123.84 Nonlinear -79.51 -80.67 -54.96 -129.00 -128.99 -126.59</td></tr></table>

Inception

Table content: <table><tr><td/><td/><td/><td colspan="3">scores on quan-</td></tr><tr><td>tized</td><td colspan="2">MNIST.</td><td>BGAN's</td><td>results</td><td>are</td></tr><tr><td>run</td><td>with</td><td>the</td><td colspan="2">author-released</td><td>code</td></tr><tr><td colspan="6">https://github.com/rdevon/BGAN.</td></tr><tr><td colspan="2">bits (states)</td><td/><td>BGAN</td><td colspan="2">MNGAN-GO</td></tr><tr><td/><td>1 (1)</td><td/><td/><td/></tr></table>

Variable-nabla examples. Note y is a scalar random variable.

Table content: <table/>

element of vector a.

Table content: <table><tr><td/><td/><td colspan="3">-110</td><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>-110</td></tr><tr><td/><td/><td colspan="3">-115</td><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>-115</td></tr><tr><td colspan="2">ELBO</td><td colspan="3">-120</td><td/><td/><td/><td/><td/><td colspan="2">REBAR train REBAR valid</td><td colspan="2">ELBO</td><td>-125 -120</td><td>REBAR train REBAR valid</td></tr><tr><td/><td/><td colspan="3">-125</td><td/><td/><td/><td/><td/><td>RELAX train RELAX valid GO train</td><td/><td/><td/><td>-130</td><td>RELAX train RELAX valid GO train</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>GO valid</td><td/><td/><td/><td>GO valid</td></tr><tr><td/><td/><td colspan="3">-130</td><td colspan="2">0</td><td>2</td><td>4</td><td>6</td><td>8</td><td>10</td><td/><td/><td>-135</td><td>0</td><td>2000</td><td>4000</td><td>6000</td><td>8000</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td colspan="2">Iteration</td><td colspan="2">10 5</td><td/><td/><td>Time (seconds)</td></tr><tr><td/><td/><td/><td/><td/><td/><td colspan="5">(a) MNIST Linear Iteration</td><td/><td/><td/><td>(b) MNIST Linear Running-Time</td></tr><tr><td/><td/><td>-80</td><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>-80</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>-90</td></tr><tr><td/><td colspan="2">-100</td><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td colspan="2">-100</td></tr><tr><td>ELBO</td><td colspan="2">-120</td><td/><td/><td/><td/><td/><td/><td/><td>REBAR train REBAR valid</td><td/><td>ELBO</td><td colspan="2">-120 -110</td><td>REBAR train REBAR valid</td></tr><tr><td/><td colspan="2">-140</td><td/><td/><td/><td/><td/><td/><td/><td>RELAX train RELAX valid</td><td/><td/><td colspan="2">-130</td><td>RELAX train RELAX valid</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>GO train</td><td/><td/><td colspan="2">-140</td><td>GO train</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>GO valid</td><td/><td/><td/><td>GO valid</td></tr><tr><td/><td colspan="2">-160</td><td colspan="3">0</td><td/><td>2</td><td>4</td><td>6</td><td>8</td><td>10</td><td/><td colspan="2">-150</td><td>0</td><td>2000</td><td>4000</td><td>6000</td><td>8000</td><td>10000</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td colspan="2">Iteration</td><td colspan="2">10 5</td><td/><td/><td>Time (seconds)</td></tr><tr><td/><td/><td colspan="9">(c) MNIST Nonlinear Iteration</td><td/><td colspan="3">(d) MNIST Nonlinear Running-Time</td></tr><tr><td/><td/><td colspan="2">-120</td><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>-120</td></tr><tr><td/><td/><td colspan="2">-125</td><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>-130</td></tr><tr><td colspan="2">ELBO</td><td colspan="2">-135 -130 -140</td><td/><td/><td/><td/><td/><td/><td colspan="2">REBAR train REBAR valid RELAX train RELAX valid</td><td colspan="2">ELBO</td><td>-140 -150</td><td>REBAR train REBAR valid RELAX train RELAX valid</td></tr><tr><td/><td/><td colspan="2">-145</td><td/><td/><td/><td/><td/><td/><td>GO train</td><td/><td/><td/><td>GO train</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>GO valid</td><td/><td/><td/><td>GO valid</td></tr><tr><td/><td/><td colspan="2">-150</td><td colspan="3">0</td><td>2</td><td>4</td><td>6</td><td>8</td><td>10</td><td/><td/><td>-160</td><td>0</td><td>2000</td><td>4000</td><td>6000</td><td>8000</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td colspan="2">Iteration</td><td colspan="2">10 5</td><td/><td/><td>Time (seconds)</td></tr><tr><td/><td/><td/><td colspan="8">(e) Omniglot Linear Iteration</td><td/><td/><td colspan="2">(f) Omniglot Linear Running-Time</td></tr><tr><td/><td/><td colspan="4">-50</td><td/><td/><td/><td/><td/><td/><td/><td/><td>-50</td></tr><tr><td/><td/><td colspan="4">-100</td><td/><td/><td/><td/><td/><td/><td/><td/><td>-100</td></tr><tr><td/><td/><td colspan="4">-150</td><td/><td/><td/><td/><td/><td/><td/><td/></tr><tr><td colspan="2">ELBO</td><td colspan="4">-250 -200</td><td>-135 -130 -125</td><td/><td/><td/><td colspan="2">RELAX train REBAR train REBAR valid</td><td colspan="2">ELBO</td><td>-200 -150</td><td>-135 -130 -125</td><td>RELAX train REBAR valid REBAR train</td></tr><tr><td/><td/><td colspan="4">-300</td><td>-140</td><td>1</td><td>2</td><td>3 10 5</td><td colspan="2">RELAX valid GO train GO valid</td><td/><td/><td>-250</td><td>-145 -140</td><td>500</td><td>1500</td><td>2500</td><td>GO valid GO train RELAX valid</td></tr><tr><td/><td/><td colspan="4">-350</td><td>0</td><td>2</td><td>4</td><td>6</td><td>8</td><td>10</td><td/><td/><td>-300</td><td>0</td><td>2000</td><td>4000</td><td>6000</td><td>8000</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td colspan="2">Iteration</td><td colspan="2">10 5</td><td/><td/><td>Time (seconds)</td></tr><tr><td/><td/><td colspan="9">(g) Omniglot Nonlinear Iteration</td><td/><td colspan="3">(h) Omniglot Nonlinear Running-Time</td></tr></table>

Average running time per 100 iterations for discrete variational autoencoders. Results of REBAR and RELAX are obtained by running the released code 5 fromGrathwohl et al. (2017). The same computer with one Titan Xp GPU is used.

Table content: <table><tr><td colspan="2">Dataset Model</td><td colspan="2">REBAR RELAX</td><td>GO</td></tr><tr><td>MNIST</td><td>Linear 1 layer Nonlinear</td><td>0.94s 0.99s</td><td>0.95s 1.02s</td><td>0.25s 0.54s</td></tr><tr><td>Omniglot</td><td>Linear 1 layer Nonlinear</td><td>0.89s 0.97s</td><td>0.92s 0.98s</td><td>0.21s 0.45s</td></tr></table>

Peer review (iclr2019):
 * Summary The paper proposes an improved method for computing derivatives of the expectation. Such problems arises with many probabilistic models with noises or latent variables. The paper proposes a new gradient estimator of low variance applicable in certain scenarios, in particular it allows training of generative models in which observations and/or latent variables are discrete. The submission clearly improves the state-of-the-art, experimentally demonstrates the method on several problems comparing with the alternative techniques. In what concerns the optimization, the method achieves a better objective value much faster, confirming that it is a lower variance gradient estimator. The clarity of the presentation (in particular the description of when the method is applicable) and the technical correctness of the paper are somewhat lacking. In terms of applicability, it seems that many cases where discrete latent variables would be really interesting are not covered (e.g. sigmoid belief networks); the paper demonstrates experiments with discrete images (binary or 4-bit) not particularly motivated in my opinion. It also contains lots of additional technical details and experiments in the appendix, which I unfortunately did not review. * Clarity In the abstract the paper promises more than it delivers. Many problems can be cast as optimizing an expectation-based objective. The result does not at all apply to all of them. The reparameterization trick does not apply to all continuous random variables, only to such that the reparameterization satisfies certain smoothness conditions. Discrete variables are supported by the method only in the case that the distribution factors over all discrete variables conditionally on any additional “continuous variables” (to which the reparameterization trick is applicable). This very much limits the utility of the method. In particular it is not applicable to learning e.g. sigmoid belief networks [Neal, 92] (with conditional Bernoulli units) and many other problems. “reparametrizable distributions” A Bernoulli(p) random variable is discrete, yet it is reparametrizable as [Z>p] with Z following standard logistic distribution, whose density and cdf is smooth. Because of the above many discussions about discrete vs. continuous variables are missleading. Section 2. The notation of the true distribution as “q” the model as p and the approximate posterior of the model as “q” again is inconsistent. I find the background on ELBO and GANs unnecessary occluding the clarity at this point. For the purpose of introduction, it might be better to give examples of expectation objectives such as: - dropout: q is the distribution of NN outputs given the input image and integrating out latent dropout noises, gamma are parameters of this NN. - VAE, GAN: q is the generative model defined as a mapping of a standard multivariate normal distribution by a NN. - sigmoid belief networks: q is a Bayesian network where each conditional distribution is a logistic regression model. Then to state to which of these cases the results of the paper are applicable, allow for an improvement of the variance and at what additional computational cost (considering the cost of evaluating the discrete derivatives). Section 3. Contrary to the discussion, there are examples of non-negative distributions to which the reparameterization trick can be applied, including log-Normal and Gamma distributions. Method: In the case when Rep trick is applicable, is it identical to GO? The difference seems to be only in that the mapping tau may be different from Q^-1. However, this only affects the method of drawing the samples from a fixed known distribution and should have no more effect on the results than say a choice of a pseudo-random number generator. Yet, in Fig.1 some difference is observed between the methods, why is that so? Sec 7.1 “We adopt the sticking approach hereafter”. Does it mean it is applied with all experiments with GO? * Related Work The state of the art allows combining differentiable and non-differentiable pieces of computation: [Schulman, J., Heess, N., Weber, T., Abbeel, P.: Gradient estimation using stochastic computation graphs.] I believe it should be discussed in related work. Limitations / where the proposed method brings an improvement should be highlighted. * Technical Correctness Equations (5) and (6) require a theorem of differentiating under integral (expectation), such as Leibnitz rule, which in case of (6) requires q_gamma(y)f(y) to be continuous in y and q_gamma(y) continuously differentiable in gamma. Equation (7) (integration by parts) holds only with some additional requires on f. Theorem 1 does not take account for the above conditions.