Generalized Gumbel-Softmax Gradient Estimator for Generic Discrete Random Variables
Authors: Weonyoung Joo
Dongjun Kim
Seungjae Shin
Il-chul Moon
Venue:  (2021)
Keywords: Deep learning
Deep generative model
Unsupervised learning
Gradient estimator
Reparameterization trick
Discrete distribution
Gumbel-Softmax

Abstract: Estimating the gradients of stochastic nodes, which enables the gradient descent optimization on neural network parameters, is one of the crucial research questions in the deep generative modeling community. When it comes to discrete distributions, Gumbel-Softmax trick reparameterizes Bernoulli and categorical random variables by continuous relaxation. However, gradient estimators of discrete distributions other than the Bernoulli and the categorical have not been explored, and the the Gumbel-Softmax trick is not directly applicable to other discrete distributions. This paper proposes a general version of the Gumbel-Softmax estimator with a theoretical basis, and the proposed estimator is able to reparameterize generic discrete distributions, broader than the Bernoulli and the categorical. In detail, we utilize the truncation of discrete random variables and the Gumbel-Softmax trick with a linear transformation for the relaxed reparameterization. The proposed approach enables the relaxed discrete random variable to be reparameterized through a large-scale stochastic computational graph. Our experiments consist of (1) synthetic data analyses and applications on VAE, which show the efficacy of our methods; and (2) topic models, which demonstrate the value of the proposed estimation in practice.

Under review as a conference paper at ICLR 2021
GENERALIZED GUMBEL-SOFTMAX GRADIENT ESTI- MATOR FOR GENERIC DISCRETE RANDOM VARIABLES
Anonymous authors Paper under double-blind review
ABSTRACT
Estimating the gradients of stochastic nodes, which enables the gradient descent op- timization on neural network parameters, is one of the crucial research questions the deep generative modeling community. When it comes to discrete distributions, Gumbel-Softmax trick reparameterizes Bernoulli and categorical random variables by continuous relaxation. However, gradient estimators of discrete distributions other than the Bernoulli and the categorical have not been explored, and the the Gumbel-Softmax trick is not directly applicable to other discrete distributions. This paper proposes a general version of the Gumbel-Softmax estimator with a theoret- ical basis, and the proposed estimator is able reparameterize generic discrete distributions, broader than the Bernoulli and the categorical. In detail, we utilize the truncation of discrete random variables and the -lumee--Softmax trick with linear transformation for the relaxed reparameterization. The proposed approach enables the relaxed discrete random variable be reparameterized through a large-scale stochastic computational graph. Our experiments consist of (1) synthetic data analyses and applications on VAE, which show the efficacy of our methods; and (2) topic models, which demonstrate the value of the proposed estimation in practice.
INTRODUCTION
Stochastic computational graphs, including deep generative models such as variational autoencoders, are widely used for representation learning. Optimizing the network parameters through gradient methods requires an estimation of the gradient values, but the stochasticity requires the computation of expectation, which differentiates this problem from the deterministic gradient of ordinary neural networks. There are two common ways of obtaining the gradients: score function- bbsed methods and reparameterization methods. The score function-base based estimators tend to result in unbiased gradients with high variances, while the reparameterization estimators seem to result in biased gradients with low variances (Xu et al., 2019). Hence, the core technique of the score-function based estimators becomes reducing the variances of gradients to achieve stable and fast optimizations. Meanwhile, utilizing the reparameterization estimators requires the differentiable non-centered parameterizution (Kingma & Welling 2014a) of random variables.
we focus on the reparameterization estimators, one of the most popular examples is the reparame- terization the Gaussian variational autoencoder (VAE) (Kingma & Welling (014b)) which has an exact reparameterization form. Other VAEs with explicit priors suggest the reparameterization tricks with approximations (Nalisnick & Smyth 2017 Joo al. 2020). For continuous random variables, it is feasible to estimate gradients with automatic differentiation by utilizing a transport equation (Jankowiak & Obermeyer, 2018) an implicit reparameterizztion (Figurnov et al.) 2008). However, these methods are not applicable to discrete random variables, due to the non-differentiabiiity
Recently, some discrete random variables, such as Bernoulli or categorical random variables, have been well-explored in terms of the reparameterization method by overcoming such difficulty through continuous relaxation (Jang et al.) 2017 Maddison et al., 2017). However, other discrete distributions have not been explored from the learning perspective in the deep generative modeling community, for example, Poisson, binomial, multinomial, geometric, negative binomial distributions, etc. Prior works graphical models, such as Ranganath et al. (2015) 2016). utilized Poisson latent variables for the latent counting. Another line of work (Wu et al.) (2020) utilized the Gaussian approximation on the Poisson latent variable to count the number of words, which can be poor approximation

---

Under review as a conference paper at ICLR 2021
Fta
PMF Valuen
Rg Ry (Gg)
Relascal Dise Sample
Reparametertadd
Fg myfro)
(ogg - d
Rinax
T(e) E O
FEA
PC
Trincation Levell Gambel Random Samplax
Temperature
Outecam Verter
(a) Reparameterization trick.
(b) Visualization of GENGS.
Figure 1: Reparameterizztioon in stochastic computational graphs.
if the rate parameter is small. In this sense, study on the stochastic gradient estimator for discrete distributions required in the deep generative modeling community, which broadens the choice of prior assumptions and the utilization of various distributions.
This paper proposes a generalized version of the Gumbel-Softmai reparameterization trick, which can be utilized to generic discrete random variables through continuous relaxation, namely Generalized Gumbel-SSofumn (GENGS). The key ideas of GENGS are (1) conversion of the sampling process to one-hot categorical selection process; (2) reversion of the selected category in the one-hot form to the original sample value; and (3) relaxation of the categorical selection process into the continuous form. To implement these steps, GENGS first truncates discrete random variables to approximate the distribution with the finite number of possible outcomes. Afterward, GENGS utilizes the Gumbel-Softma trick together with special form of a linear transformation. Our main theorem supports that the proposed GENGS is applicable to general discrete random variables, other than the Bernoulli and the categorical. The GENGS experiments show the efficacy with synthetic examples and VAFs, as well as the usability in topic model application.
2 PRELIMINARY: REPARAMETERIZATUU TRICK & GUMBEL-SOFTMAX
2.1 BACKPROPAGATION THROUGH STOCHASTIC NODES WITH REPARAMETERIZATTIN TRICK
Suppose we have stochastic node, or latent variable, : p(= (), where the distribution depends on parameter 0. The goal is optimizing the loss function L( 17) = E (f, (f)], where fa is con- tinuous and differentiable function with respect to 17. i.e. neural networks. To optimize the loss func- tion in terms of 0 through the gradient methods, we need to find VoL(O,)) = Vo [z)), which can not be directly computed with its original form.
To compute V vo 7). the reparameterization trick alternatively introduces an auxiliary variable e ~ p(e), which takes over all randomness of the latent variable 2, so the sampled value a can be re-written as = g(0, e). with deterministic and differentiable function g in terms of 9. Figure ((a) illustrates the reparameterization trick: the shaded nodes indicate random nodes, and the dotted lines denote sampling processes. Here, the gradient of the loss function with respect to A is derived as
VoC = V (0
Fe-p(() VoS )) Eer-pte) (V (0, Vog(0
(1)
where the last term of Equation is is now achievable. A condition on enabling the reparameterization trick is the assumption of the continuity of the random variable z, the distribution of is limited to a class of continuous distributions. To utilize the differentiable reparameterization trick on discrete random variables, continuous relaxation can be applied: for example, relaxation from the categorical distribution to the Suubel-SSofumn distribution, described in the next subsection.
2.2 REPARAMETERIIAATI TRICK ON CATEGORICAL RANDOM VARIABLE
Gumbel-Max trick (Gumbel| 1948 is procedure for sampling a one-hot categorical value from the Gumbel distribution, instead of direct sampling from categorical distribution. This implies that the categorical random variable X ~ Categorical( (T), where lies on the (n 1)-dimensional simplex n-1, be reparameterized by the Gumbel-Max trick: (1) sample uy Uniform(0, 1)
2

---

Under review as a conference paper at ICLR 2021
to generate a gumbel sample 9j - -log(- log u,) for each 1,... ... n. and (2) compute k argmax)=1 (log / + gz1, where и is a categorical parameter. This procedure generates a one-hot sample x such that 2jj = 0 for j * and T3 = with P(X; = 1) = Tk. We denote (z) to be the distribution whose samples are generated by the Gumbel-Max trick.
A Gumbel-Softmas trick (Jang et al. 2017 Maddison et al. 2017) is an alternative to the Gumbel- Max trick that continuously relaxes categorical random variable. The Gumbel-Softmax utilizes the sof with a temperature > 0, instead of the argmax in the sampling process, which enables (1) relaxing the discreteness of the categorical random variable to the one-hot-like form xx(T) = soft log lox=+)) in the continuous domain; and (2) approximating the Gumbel-Max by laking T small enough. Lately, the Gunbel-Softmax estimator has been widely used to reparameterize categorical random variables, such as RelaxedOneHotee egor ical in TensorFlow (Abadi) trick.
et al. 2016
We denote GS 7) to be the distribution generated by the
3 PROCESS OF GENGS REPARAMETERIZATION
This section discusses the process of GENGS help understand the concept with minimal theoretical details, and Section
| provides the theoretical background of GENGS. The three steps of GENGS are the following: (1) approximate discrete distribution by truncating the distribution; (2) reparameterize the truncated distribution with the Gumbel-Mi trick and the linear transformation T, which will be introduced below; and (3) relax the discreteness by replacing the Gumbel-Max trick Step 2 with the Gumbel-Softuax trick. Figure [(b) illustrates the full steps of the GENGS trick.
Step 1. Truncate the discrete distribution to finitize the number of possible outcomes. Suppose X ~ Poisson( 100), which has mode near at x 100, and near-zero probabilities at x < 50 and x > 150. The key idea of the first step ignoring the outcomes of near-zero probabilities at certain levels (ex. < 50 and x > 150) and only focusing on the probable samples of meaningful probabilities (ex. 50 S x I 150), i.c., truncating the distribution, which finitizes the support of the distribution. Now, suppose we have a discrete random variable X D(A), and its truncated random variable Z ~ TD(A, R). where R denotes the truncation range that needs to be pre-defined.
Proposition 3 in Section 4 provides theoretical reason that Z approximates X. Since we finitized
the support by the truncation, we may assume Z has a support C = (co, ... ,Cn-1] } of n possible
outcomes and its corresponding constant outcome vector c = (Co, ... Ca-1). 1) Note that the ordering of C is not significant, and Appendix E provides examples of the setting on c.
Step Divide sampling process of Z into two-fold: select a one-hot category of Z. and revert the selected one-hot category into the original value. For example, if the sampled value of Z is 02 E C we will first focus on the one-hot category class vector one one_hot (c2) = (0, 0, 0. .0,, rather than the sampled value C2. Such a one-hot categorical selection process is possible by utilizing the categorical selection w ~ Categorical(z) or its reparameterized version, the Gumbel-Max trick GM(r).. Here, the categorical parameter π (To. -1) can be directly calculated by the explicit probability mass funciton (PMF) of the distribution, i.e., T = P(Z cx). However, the PMF of the truncated distribution requires modification from the PMF of the original distribution, which is determined by how we define / from X. See Definition | 2 and Appendix A for detailed configuration of и. Suppose parameter я. Afterward, we revert the one-hot selected categorical vector w = ( (Wy, wn-1) ) into the original sample value with linear transformation T(w) Σ Wg.Ck = Σ.
now have one-hot categorical sample
from the categorical
Proposition
4 shows the validity of the alternative sampling process in Section 4
Step 3. Relax the one-hot categorical selection into the continuous form by utilizing the Gumbel- Softmax trick. Up to now, the sole shortage of the reparameterivation trick is the differentiability due the one-hot categorical selecting Gumbel-Max process. Then, as in Section the process can be continuously relaxed with the Gumbel-Softmax trick GS(T, 7) for some temperature T. Theorem 5 in Section 4 shows that the alternative sampling process still holds under the continuous relaxation.
4 THEORETICAL BACKGROUND OF GENGS
To support our reparameterization methodology, this section provides the main theorem on the repa- rameterizations. The first proposition approximates original discrete distribution with its truncated
3

---

Under review as a conference paper at ICLR 2021
Gen
Temgerature
Truncation Level
T(GM(s)) TD(A.)
D(2)
Panson?7
(a) Effects of truncation level and temperature
(b) Concept GENGS approximation
Figure 2: (a) Approximation of GENGS in terms of choices of the truncation level n and the temperature in Poisson(7). As sub-figures go from left to right, the truncation level grows. Hence, the popped-out sticks, implying remaining probability of the right side, disappears if the truncation level is large enough. As sub-figures go from top to bottom, the temperature decreases, and the PMF of truncated distributions and the original distributions becomes similar. Appendix Deieeddixxpppooo proovidess the fine-graite fine-grained PMF of GENGS. (b) On the x-axis, TD(A, n) + D(A) as truncation level n f oc, according to Proposition
3 Then, (A,A, n.) can be reparameterized by the Gumbel-Max trick with 0, GENGS(T, -
linear transformation T in Proposition On the y-axis, as temperature TD(A, n2,, where is computed PMF value of (D,,, n),, according to Theorem
version. Next, the second proposition enables the truncated distribution to be reparameterized by the Gumbel-Max trick. Finally, the main theorem shows that the Gumbel-Softmax trick converges the Gumbel-Max trick under assumption of the linear transformation. Through these steps, we note that our proposed reparameterization trick generalized and grounded, theoretically.
4.1 FINITIZING THE SUPPORT BY TRUNCATING DISCRETE RANDOM VARIABLES
Definition I specifies truncated discrete random variable for truncating the right-hand side. Note that Definition D can be easily extended to truncate the left-hand side or both sides of distributions. Definition 2 is truncating both side version, and Appendix didicussess its variation. However, focus on the non-negative distribution in the remainder of this subsection, since most of the popularly used discrete random variables have the support of NNo-
Definition 1. (A special case of srrighthhadd-ie
truncation for non- negaetettvv discrete random variables) A fruncated discrete random variable Zn, of non-negative discrete random variable X D(X) is discrete random variable such that Zn = X if X <n 1, and Zn = n if X > n. The random variable Zn is said to follow truncated discrete distribution TD(X, R) with a parameter and truncation range R = [0, n). Alternatively, we write as truncation level I = n if the left truncation is zero in the non negaaivee case.
Definition 2. (Both-side truncation for general discrete distributions) A truncated discrete random variable Zm,.. of discrete random variable X D(A) is discrete random variable such that (1) Zm, X if m < X < n; (2) Zm,, n 1 if X 2 n; and (3) Zm.a = m + 1 if X S n. The random variable Zm,x is said follow truncated discrete distribution TD(A, R (m, n)) with a parameter À and a truncation range R = (m, n²).
As we discussed in Section truncating the distribution intends finitize the number of possible outcomes to utilize the categorical selection. From the definition, the samples of finitized support can be simply considered c in this special non-negative case. Furthermore, due the definition, the modification on the PMF only exists in the last category Cn-1 = n 1, and the modified PMF can be computed by injecting the near-zero remaining probability summation the last category, right before the truncation level. In other words, π = P(Z, cx) = P(Z,. k) = P(X k) for k = 0....
2 and nT 1- En Tkg hence, the sum-to-one property remains satisfied. Here, the idea which leads to Proposition 3 that if we take the truncation level far enough from zero,
cover most of the possible outcomes that be sampled from the original distribution. Note that the truncation step can be omitted if the original distribution already has a finite support, but one can utilize the truncation ignore the unlikely samples.
Proposition 3. With Definition ! Z,, converges to X almost surely as n - oo. Also, with Definition 7mm converges to X almost surely as m and ->
4

---

Under review as a conference paper at ICLR 2021
The almost sure convergence property of PropositionSssug supports the theoretical basis of approximating a discrete random variable D(A) with the truncated random variable TD(,, n), and Appendix [A] shows the detailed proof. Through the truncation, the discrete distribution is approximated with finitized support, and the Gumbel tricks are ready to he utilized.
4.2 REPARAMETERIZATION BY GENERALIZED GUMBEL-SOFTMAE TRICK
Next, we select one-hot categorical sample from the finitized categories and revert the one-hot selection to the original sample value. The widely utilized discrete distributions have the explicit forms of PMF, so we can directly compute the PMF values for the truncated support with a pre-defined truncation range. Once the distribution and the truncation range are fixed as ((A) and R. respectively. we have the corresponding constant outcome vector c (Co. ... .n- 1) and the computed PMF value vector æ (To,- -x-1) 1) of truncated distribution TD(A, R), where Tx = TD(cx: X, R). Afterward, we define a transformation T(w) = E. W&CC = >n w O c. Additionally, we denote the distributions, generated by applying T on the samples of GM and GS, as T(GM) and 7 T(GS), respectively. Then, we can generate a sampled value of TD(A, (A, R) with a linear transformation and a Gumbel-Max sample, stated in Proposition 4 proved in Appendix(B
Proposition 4. For any truncated discrete random variable Z ~ (D(A, R) of discrete distribution D(A) and transformation T, Z can be reparameterized by T(GM(=)) we TT& P(Z Ck)-
Due to the reparameterization, the randomness of TDA,, R) with respect to the parameter À moves into the uniform random variable in the Gumbel-Max, since T is continuous and deterministic function. Then, TD R) can be approximated by replacing the Gumbel-Max with the Gumbel-Softmen in T, as stated in Theorem[5] proved in Appendix C
We define ((GS(T,, T)) as GENGS(TT, (T, 7).
Theorem 5. For a transformation T and a given categorical parameter π E A the convergence property of Gumbel-SSftmma Softtmxx to Gumbel-Max still holds under the linear transformation T, i.e., GS(T, T) + GM(T) as T + 0 implies GENGS(T, - (TGM(=)) as ->
Theorem 5 implies that we can relax and approximate the truncated discrete random variable TD(A, R) by the Gumbel-Softtax and the linear transformation. The assumption of the theo- - 0 has not been mathematically proven in the original literature (Jang et al. 2017 Maddison et al. 2017). Instead, the authors have empirically shown that GS(T,T) eventually becomes GM(z) π as 5 + O. Figure 2 shows how GENGS gets closer to the original distribution by adjusting the truncation range and the temperature.
rem that GS TT, - GM(m)
We can observe that the approximation becomes closer to the original distribution as we widen the truncation range R. However, the increment of R is technically limited due to the finite neural network output for the inference. Note that the choice of truncation range crucial in terms of covering many probable samples. Therefore, we set the truncation range to cover all but less than 18-5% of probability with respect to the prior distribution, or arbitrary large. 7
Truncation Range.
Temperature. The decrement T from so softmax log results in the closer distribution to the original distribution. However, the initially small leads to high bias and variance of gradients, which becomes problematic at the learning stage on л. Therefore, the annealing of r from a large value to a small value is necessary provide a learning chance of n.
Note that there is no condition on the distribution to be reparameterized with GENGS in the theoretical analysis. Hence, once discrete distribution has explicit PMF, GENGS can be easily utilized to the reparameterization. Appendix E suggests examples of GENGS utilization, including one that shows the regular Gumbel-Softman special case of GENGS. Appendix [P provides the algorithm of GENGS, and Appendix G gives discrete distributions in TensorFlow, which can utilize GENGS.
5 EXTENTION OF GENGS: IMPLICIT INFERENCE & DISCRETIZATION
Implicit Inference. Unlike continuous random variables, discrete random variables have countable outcomes. Instead of inferring the distribution parameter X and then computing the PMF values through the fixed PMF and X, we can directly infer the PMF values æ of possible outcomes with softmax, which becomes the input of the Gumbel tricks, by loosening the assumption approximate posterior PMF shape. This implicit inference on the PMF values becomes possible due to truncating distribution by finitizing the possible outcomes. However, this inference approach needs
the
5

---

Under review as a conference paper at ICLR 2021
regularizer, such the KL divergence term in the objective function of VAEs, which ensures the distribution shape to be similar to a prior distribution with a pre-defined parameter. We found that loosening the approximate posterior assumption leads to significant performance gain in our VAE experiments. See Appendix
F for the algorithm of the implicit inference.
Discretization of Continuously Relaxed Sample. GENGS outputs continuously reparameterized sample value since we are relaxing the discrete random variable into a continuous form. Utilizing the Straight-Through (ST) Gumbel-Softmax estimator (Bengio et al. 2013 Jang et al.) (017), instead of the naive Gumbel-Softnax, we can obtain the discrete sample as well. Since ST Gumbel-Softne discretizes the relaxed Gumbel-Softual output with arg.ax, ST Gumbel-Softmax uses the gradients obtained from the relaxed ones, which could result in significant performance degradation.
6 RELATED WORK
GENGS is basically single-sample gradient estimator like other reparameterization gradient estima- tors. Though GENGS could use multiple samples to obtain the stable gradients, we compare GENGS with the other estimators using single sample to test the fundamental performance of gradient estimators. RF denotes the basic REINFORCE (Williams) 1992). NVIL (Mnih & Gregor] 2014)
utilizes neural network to introduce the optimal control variate. MUPROP (Gu et al.]]]]11010 (2016) utilizes the first-order Taylor expansion on the loss term as a control variate. VIMCO(K) (Mnih & Rezende 2016) designed as k-sample gradient estimator. REBAR (Tucker al. 2017 and RELAX (Grath wohl et al. 2017) utilize reparameterization trick for constructing the control variate. Deterministic RaoBlack (DETRB) (L.iu et al., 2019) uses the weighted value of the fixed gradients from m-selected categories and the estimated gradients from the remaining categories with respect to their odds to reduce the variance. The idea Stochastic RaoBlack (STORB) (Kool et al. 2020) is essentially same as that of DETRB, but STORB randomly chooses the categories at each step, instead of using fixed categories. Kool et al. (2020) also suggested an unordered set gradient estimator (UNORD), which also uses the multiple sampled gradients, utilizing the sampling without replacements. For DETRB, STOR B, and UNORD, we use m = category that utilizes the fixed gradient for the fair comparison. Note that if there are K >> 1 dimensions be inferred, the models require computing mk gradient combinations, which has higher complexity than GENGS, The symbol denotes a variation that utilizes built-in control variate introduced in the work of Kool et al. (2020).
7 EXPERIMENT
7.1 SYNTHETIC EXAMPLE
Experimental Setting. In this experiment, we expand the toy experiments from Tucker al, (2017);
Grathwohl et al. (2017) to various discrete distributions. We first fix constant i., th. and then optimize the loss function Ex-p(=]\) Et =4 with respect to X. Here, we set p(=|A) as Poisson (20). Binomial (20, .3). Multinomial(3, [.7,.2..1] and NegativeBinomial3,, .4) in this experiment. We also adapt the Rao- Blackwellization (RB) idea in GENGS by utilizing m 1 in calculating the selected gradient, so this adaptation results in GENGS RB that estimates the remaining gradients by GENGS. See Appendix for the detailed experimental settings.
Experimental Result. Figure Tiguueeccome compares the log-loss and the log-variance of estimated gradients from various estimators. the figure, the log-loss needs to be minimized to correctly estimate the backpropagated gradient value in the learning process. Additionally, the log-variance requires being minimized to maintain the consistency of the gradients, so the gradient descent can be efficient. GENGS shows the best log-loss and the best log- varvanri if GENGS keeps the continuous relaxation of the modeled discrete random variable. For the Poisson case, the exact gradient has a closed-form solution, as in Appendix J and GENGS shows the lowest bias among all gradient estimators. See Appendix | for the curves with confidence intervals and the curves without smoothing.
7.2 VAE: SYNTHETIC EXPERIMENT ON DEEP GENERATIVE MODELS
Experimental Setting, We follow the VAE experiments Figurnov et al. (2018) with discrete priors, which diversifies the choice of prior assumption, as the latent factor count in the discrete case. This

---

Under review as a conference paper at ICLR 2021
Broeca 0.)
H:l 1.1.5.
SeppCivdl
2.0
10000
1000
Mall 2 ab
Segitrod foceat
Highh GonGS High Temp
Temp
Highh Temp
DecRa*
StoR3
Tomp
Tomp
Temp
Figure 3: Synthetic example performance curves in log scale: (Top Row) losses, variances, and biases of gradients for Poisson; (Middle Row) losses for Binomial, Multinomial, and NegativeBinomial: (Bottom Row) variances of gradients for Binomial, Multinomial, and NegativeBinomial. We utilize the cumulative average for smoothing the curves, and the curves with confidence intervals and the curves without smoothing are in Appendix g
Table I: Test negative ELBO on MNIST and OMNIGLOT datasets. The lower is better for the negative ELBO. We provide a full table including baselines with insignificant results and variations of GENGS in Appendix K Symbol " indicates convergence.
MNIST
RF
NVH.
MuPror
VIMCO()
REBAR
RELAX
STORB
GENGS (Ex.)
GENGS (1m.)
Pais(2)
+2.44
24.72
125.43
122.55
123.44 123.44g224
122.71.1.2
124 22..02
103.18
96.0442144
Pais(3)
123. 123.12
130 24
125
121.15
120 220.0222m
119.8422.8
124.4125.0
105.15
96.01 21
Geamt 25)
127.9021
135. 30
137
137.90
127 121.21
196. 12
136. g0
131.09
98.49_ost
92.52g1.
Geumi 5)
129.20
138. 10.30
138.40_am
129.91.a.0
138.37
139.41 139.41
139.67
100.92,3.84
98.81g1.0
N5gBinc3..5)
116.07a..77
28.7.00
131.96...4
112.60.4.00
114. 4.33
98 28.58
94.52.3.5
NegBing5..3)
130.03
133. 44.4.27
144.05gs.10
124.48-272
128 128.02e2.00
100.88
95.37
OMNIGLOT
RF"
NVIL
MuPRoe
VIMCO(S)
REBAR
RELAX
STORB"
GENGS (Ex.)
GENGS (1m.)
Pois(2)
139. 139.47
148 01.4.0
142 22.35
138, 22.445
138. 12
197 2..
139 61.57
127 20.44
118.
Puis(3)
1d0.54e2.3
13
153.85
139.0 139.97
42
140 140.05e3.0
151.53=1.
119.15
GoumL.2)
102. 68
153 6922
152. 152.17
122.91
100. 78
148 s1sss
143. 10
115.28=20
107 79
Geomr. 5)
153 20
149.76g2210
142.06 142.06:28.0
149.63as
161.99 161.97
142. 26
115.14ga.0
108. 48
NogRinc3.55 5)
141 141.442220
144 44.2.75
147.75 147.78.4.0
141.89 141.8843334
129 24.54
118.57=2m
117 117.0222.0
NogBinc5..3)
144.44gsss
150.40 150.40
15281
150.40gcm
151. 151.90gxx
119.57
117.54g2
experiment utilizes the Poisson, the geometric, and the negative binomial distributions. The evidence lower bound (ELBO) L(x)= Eggenee Po (x|z)] (z)xxx))))[lp (z)). which consists of the reconstruction part and the KL divergence part, is minimized during the training period. Optimizing the ELBO of VAEs requires computing the KL divergence between the approximate posterior and the prior distributions. In GENGS, by truncating the original distribution, the KI divergence becomes the derivation with categorical distributions. See Appendix H for the detailed statement and proof.
Note that the purpose of VAE experiments is not to compare the performance across various prior distributions. The VAE is considered as a more challenging task than the synthetic example in the last subsection, since (1) this task requires computing the gradients of the encoder network parameters through the latent distribution parameter X; and (2) each stochastic gradient of the latent dimension affects every encoder parameter since we utilizing the fully-connected layers. Hence, a single poorly estimated gradient with respect to the latent distribution parameter À could negatively affect the learning of encoder parameters. See Appendix K for the detailed experimental settings.
7

---

Under review as a conference paper at ICLR 2021
WUST
0
2 3
5
6
7
8
9
M3
NG
0F
0
2
2
9
5
6
7
0
9
1
0
2
0
5
6
7
8
9
P
SI
FF
THE
10003
0
2
3
5
6
7
8
1
UF
Figure 4: Reconstructed images by VAEs with various gradient estimators. clearest images among other gradient estimators with better reconstruction.
GENGS shows the
Experimental Result. Table 1 shows the negative ELBO results of the VAE experiments. We found that some baselines failed to reach the optimal point, so excluded those estimators in such suboptimal cases. The variants of GENGS showed the lowest negative ELBO in general, and loosening the PMF condition idea (i.e., the implicit inference) reached the better optimal points. Figure 4 the reconstructed images by VAEs with various gradient estimators on MNIST and OMNIGLOT. GENGS draws the clearest images and better reconstructions, which aligns with the quantitative result in Table I
See Appendix K for the full table and additional discussion.
7.3 TOPIC MODEL APPLICATION
Experimental Setting. This experiment shows another application of GENGS in the topic mod- cling. The Poisson distribution is one of the most important distribution for counting the number of outcomes among all discrete distri- butions. The authors of Deep Exponential Fami- lies (DEFs) |Ranganath et al. 2015) utilized the exponential family, including the Poisson distri- bution, on the stacked latent layers. Therefore, we focus on the Poisson DEF, which assumes the Poisson latent layers to capture the counting numbers of latent super-topies and sub-topies; and we convert the Poisson DEF into a neural varia- tional form, which resembles NVDM (Miao et 2016). Figure 3 shows the neural network and its corresponding probabilistic modeling structure, We utilize GENGS on the Poisson DEF to sample the values in the latent variable, namely the neural variational Poisson DEF (NVPDEF), See Appendix for further description of DEFs, NVPDEF, and detailed experimental settings.
P
Figure 5: (Left) graphical notation of NVPDEF with generative process (0) and inference network (Φ). The multi-stacked latent layers have A as a prior distribution parameter. (Right) A neural network diagram of NVPDEF: diamond nodes in- dicate the auxiliary random variable for the repa- rameterization trick.
We enumerate the baselines and the variants of NVPDEFs Table and confirmed that NVPDEF shows the lowest perplexity overall with 20Newsgroups and RCVI-V2. Since NVPDEF and the original DEFs have dif- ferent training and testing regimes, we com- pare NVPDEF to representative neural variational topic (document) models, which are listed in Table Additionally, Appendix
Experimental Result.
Table 2: Test perplexity on 20News and RCV dataset.
Model
20News (Dim.)
RCV (Dim.)
LDA (Bis et
2003
1082a12. (50)
1187 (200)
NVDE
8034
574 (200)
GSM Maare A0
854gn. (50)
80l (200)
NVLDA Smiasiia & Mation 2017
1150 (50)
1574 (200)
PannI (Snvislav & Smron 2017
1145 (50)
1425 (200)
NVPDEF MULTI-STACKED XVPDEF
759 (50)
502a11.
783 (20-50)
570 (50-200)
L sshows the qualitative results from topic models.
8 CONCLUSION
This paper suggests a new gradient estimator for general discrete random variables, a generalized ver- sion of the Softtax estimator. To strengthen the practical usage of reparameterization tricks with the Gumbel-Softmax function, we provide theoretical background of our reparameterization trick. Our finding claims that discrete random variable can always be reparameterized through the proposed GENGS algorithm. The limitation of GENGS is the setting of the truncation level and the Gumbel-Softmall temperature, which becomes the trade-off between the gradient estimation accuracy and the time budget. Subsequently, we show the synthetic analysis and the VAE experiment, well as topic model application of GENGS, With the generalization, we expect that GENGS can diversify the options of distributions in the deep generative model community.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
M. Abadi, P. Barham, J. Chen, 7. Chen, A. Davis, J. Dean, M. Devin, S. Ghemawat, Ge. Irving, M. Isard, and et al. Tensorflow: system for large-scale machine learning. USENIX Symposium on Operating Systems Design and Implementation, 2016.
Y. Bengio, N. Leonard, and A. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint orXiv:1308.34333 2013.
D. M. Blei, A. Y. Ng, and M. Jordan. Latent dirichlet allocation. Journal of Machine Learning Research, 2003.
M. Figurnov, S. Mohamed, and A. Mnih. Implicit reparameterization gradients. Advances in Neural Information Processing Systems, 2018.
W. Grathwohl, D. Choi, Y. Wu, G. Roeder, and D. Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. International Conference on Learning Representations, 2017.
S. Gu, S. Levine, Sutskever, and A. Mnih. Muprop: Unbiased backpropagation for stochastic neural networks. International Conference on Learning Representations, 2016.
E. J. Gumbel. Statistical theory of extreme values and some practical applications: a series of lectures (vol. 33). US Government Printing Office, 1948.
E. Jang, S. Gu, and B. Poole. Categorical reparameterization with gumhel-softmax. International Conference on Learning Representations, 2017.
M. Jankowiak and F. Obermeyer. Pathwise derivatives beyond the reparameterization trick. Interna- fional Conference on Machine Learning. 2018.
W. Joo, W. Lee, Park, and I. C. Moon. Dirichlet variational autoencoder. Pattern Recognition, 107, 107514, 2020.
D. P. Kingma and M. Welling. Efficient gradient-based inference through transformations between bayes nets and neural nets. International Conference on Machine Learning, 2014a.
D. P. Kingma and M. Welling. Auto-encoding variational bayes. International Conference on Learning Representations, 2014b.
W. Kool, H. van Hoof, and M. Welling. Estimating gradients for discrete random variables by sampling without replacement. International Conference Learning Representations, 2020.
R. Liu, J. Regier, N. Tripuraneni, M. I. Jordan, and J. McAuliffe. Rao-blackwellized stochastic gradients for discrete distributions. International Conference on Machine Learning, 2019.
C. J. Maddison, D. Tarlow, and T. Minka. A sampling. Advances in Neural Information Processing Systems, 2014.
C. J. Maddison, A. Mnih, and Y. W. Teh. The concrete distribution: A continuous relaxation of discrete random variables. International Conference on Learning Representations, 2017.
Y. Miao, L. Yu, and P. Blunsom. Neural variational inference for text processing. International Conference Machine Learning, 2016.
Y. Miao, E. Grefenstette, and P. Blunsom. Discovering discrete latent topies with neural variational inference. International Conference on Machine Learning, 2017.
A. Mnih and K. Gregor. Neural variational inference and learning in belief networks. International Conference Machine Learning, 2014.
A. Mnih and J. Rezende. Variational inference for monte carlo objectives. International Conference on Machine Learning, 2016.
E. Nalisnick and P. Smyth. Stick- breaking variational autoencoders. International Conference on Learning Representations, 2017.

---

Under review as a conference paper at ICLR 2021
R. Ranganath, L. Tang, L. Charlin, and D. Blei. Deep exponential families. Artificial Intelligence and Statistics, 2015.
R. Ranganath, A. Perotte, N. Elhadad, and D. Blei. Deep survival analysis. Machine Learning for Healthcare Conference, PMLR, 2016.
A. Srivastava and C. Sutton. Autoencoding variational inference for topie models. International Conference on Learning Representations, 2017.
G. Tucker, A. Mnih, C. J. Maddison, J. Lawson, and J. Sohl-Dickstein Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. Advances in Neural Information Processing Systems, 2017.
R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning, 8(3-4), 229-256, 1992.
J. Wu, Y. Rao, Z. Zhang, H. Xie, Q. Li, F. L.. Wang, and Z. Chen. Neural mixed counting models for dispersed topic discovery. Annual Meeting of the Association for Computational Linguistics, 2020.
M. Xu, M. Quiroz, R. Kohn, and S. A. Sisson. Variance reduction properties of the reparameterization trick. International Conference Artificial Intelligence and Statistics, 2019.
10

---

Under review as a conference paper at ICLR 2021
A PROOF OF PROPOSITION 3 & TRUNCATING BOTH SIDES
A. PROOF OF PROPOSITION 3 FOR TRUNCATING HIGHT-HAND-S5
Before reminding Definition and Proposition 3 of the main paper, we first introduce the definition of almost sure convergence of sequence of random variables.
Definition. (Almost Sure Convergence) For sequence random variables ( Xn converges almost surely random variable X, if
P
lim X, = X) P((w E E lim (n(w) = x(w))))) 10000
n-+00
for sample space S2
Definition, (A special case of right-hand-side truncation for non-negative discrete random variables) fruncated discrele random variable Zn of non-negative discrete random variable X D(A) is discrete random variable such that Zn = if X <n 1, and Zn = 1 if X > n. The random variable Zn said to follow truncated discrete distribution TD(A, R) with a parameter and truncation range R = [0, n). Alternatiell,, we write as truncation level R if the left truncation is zero in the non- non-negat case.
With this definition,
discussed in the main text, the constant outcome vector can be defined
as c = (0, 1,2,.... .................. n - 1). Note that for Z, = 0, 1, n 2, the sample spaces of Z, and X equal, hence, P(Z.) can be computed by using the PMF of X = k2' in such cases, i.e., P(Zn = k) = P(X = k) for
P(X) in such cases. This implies that the modified PMF of Zn k
k = 1. .n-2. 2. Then, consequently, P(Zn = n- P(X 2 n n-1) = =1 ルーー due the sum-to-one property of probability. Hence, the corresponding categorical parameter of the constant outcome vector can be computed using the PMF of the original discrete random variable X: n. P(Zn
k= P(Z.. k)
k) = P(X- = k) for k = 0,1, n 2, and Wn-1 = P(Zn = n1 -
P(X 2 n 1) =1 Σ 0= Tk-
Proposition. With Definition
7
Zn converges to X almost surely as n + oc.
Proof. Note that {w (l x (())) (w (N)((w) < n... Then, we have the following:
P( lim 7n X) = P([w E Ω2 lim (m( X(w)))
-P([ E R lim (n(w)= = lim X(w) < lim = oo]))
(4)
=P([w E ([)(w)) <o))
(5)
= P(X < o0)
(6)
=1
(7)
since is a non-negative discrete random variable, and its cumulative distribution function Fx (a)
has the following properties: non-decreasing as x oo, and 0 < Fx < 1.
A.2 PROPOSITION FOR TRUNCATING BOTH SIDES
The below is truncating both side version of Definition 2 and Proposition 3 in the main paper. For simplicity, we assume the distribution has support Z.
Definition. (Both-side truncation for general discrete distributions) truncated discrete random variable Zm,. of a discrete random variable X ~ D(A) is discrete random variable such that X, if m X
Zmm.
n1 1, if 2 n
The random variable Zm, is said to follow a truncated
m 1, if & m
discrete distribution TD(A, R = (m, n)) with a parameter À and a truncation range = (m, n2).
With this definition, the constant outcome vector can be defined as c (m + 1,... 1). Note that for Zm, = m + 2, . 2, the sample spaces of Zm... and X are equal, hence,
Note that using the distribution with eplicitly known PMF for discrete distribution X.
11

---

Under review as a conference paper at ICLR 2021
P(Zm,n) P(X) in such cases. This implies that the modified PMF of Zm = k can be computed by using the PMF of X = k in such cases, ie.. P(Zm.n = k) = P(X = k) for k == m 2, 2. n 2. Then, with the definition, P(Zm.. = m + = P(X I m 1) and P(Zm,n = n 1) = P(X 2 n - 1). However, during the implementation, this definition can be problem since the configuration implies that there might be case that we have to compute
two infinite sums, which are P(Zm,n = m + 1) = P(X S m + +1) = Ensm+1 I P(X = k) and
P(Zm,n - = P(X 2 n n-1) 1) = Eien-1 P = k)² Hence, we also provide an alternative configuration of 2m,, n in the later of this section.
Proposition. With definition above for the truncating both sides, Zm,, converges to X almost surely as m 7- -00 and n + DC.
Proof. Note that fw E S2]2m.n(w) = X(w)] = {w E 9/m < X(w) < n2. Then, we have the following:
P( lim
Zm, = X)= = P(fw E Q2|
lim 12--00
Zm, (w) = X(w)))
(8)
P((w E O/ - xx= = lim < X(w) < lim -
n-00
m=-00
P([w Q < X(w) < x0)))
(10)
P(-20 < X xx)
(11)
= P(X < <0))- P(-00 x)
(12)
(13)
since X is a discrete random variable. and its cumulative distribution function Fx(x) has the following properties: non-decreasing as x + oo, non-increasing as x + -00, and 0 S Fx S 1. O
As we discussed above, during the implementation. computing the small probabilities of both left and right tail can cause a problem, either it can have high complexity even impossible.. Hence,
X, if < X < n
- 1, if 2 or X m when we can consider the alternative definition such as Zm, n
which has simpler PMF computation. In this case, we can simply add the remaining probability of unlikely samples at the right-most value, and moreover, the alternative configuration does not harm the proof of the proposition. Hence, with the alternative defitnition, the the corresponding categorical parameter π = (mm-+1.-
1n of the constant outcome vector c = (m + 1, ... n 1) can be computed using the PMF of the original discrete random variable X: T P(Z,, k)) P((X k)
for l m + 1, n 1 Ik=m+1 Tg.
2, and nn- = P(Zn = n 1) = P(X 2 n - 1) + P(X < m + 1) =
B PROOF OF PROPOSITION! 4
Proposition. For any truncated discrete random variable Z TD(X, R) of discrete distribution D(A) and a transformation T, Z can be reparameterized by T(G((=)) if we set Thg P(Z = cz).
Proof. Note that Z has two parameters: the distribution parameter À and the truncation range R. Assume that we have possible outcome set C = {Co,... , Cn-1] of n possible outcomes by truncating the distribution with truncation range R. Note that the transformation
defined as T(w)= = - W&Ck- = Σx w O c. By pre-defining the truncation range R as a hyper- syper-parameter, p the randomness of Z is fully dependent on the distribution parameter X. Now, we introduce the Gumbel random variable G = - log(- log U) where U ~ Uniform(O, 1) as an auxiliary random variable. Then, given a categorical parameter π E A n-1, any n-dimensional one-hot vector ej-
2Or there might be case that one of P(Zm.n = m + 1) P(X S m + 1) and P(Zm.n = 1)
P(X 2 1) requires summation of high complexity, even though it is finite summation.
For example, we truncate Poisson(1000 (000) with truncation range R = [900, 1100), we have to compute PMF values for 0 I x < 900 to sum-up the left-hand meaningless probability. and it causes high complexity. Foor example, the case when the support has infinite support for both left (-oo) and right (co) sides.
12

---

Under review as a conference paper at ICLR 2021
(0.... 00.1,0.... which has 1 in 5th entry and 0 in all other entries, can be reparameterized by Gumbel-Max trick.
Suppose have sample cm from Z, and note that we have known PMF values (mo. Tn-1) of by the definition of Z. Then, with the transformation and the constant outcome vector (co," , Cn-1), the following holds:
Cm Σ C4C+ = T(em)
(14)
Since the transformation is also deterministic function, by introcuding the Gumbel random variable as an auxiliary random variable, we can replace the randomness of Z from À (or i in the implicit inference case) with the uniform random variable composing the Gumbel random variable. Hence, the truncated discrete random variable Z can be reparameterized by the Gumbel-Max trick and the transformation T.
C PROOF OF THEOREM 5
Theorem. For transformation TT and a given categorical parameter π E A=-1, the convergence property of Sumbel-Softmax to Gumbel-Max still holds under the linear transformation T, i.e., GS( GM(TT) - 0 implies GENGS TT) (TGM(z)) as - 0.
Proof. Suppose we have given a categorical parameter π E A-11, Define fM be a Gumbel-Max trick function, and f be Gumbel-Softmat trick function with temperature > 0 that both functions take the categorical parameter and a Gumbel sample as inputs. Note that far returns a one-hot vector which has in the argmax entry after the Gumbel perturbation, and fá returns one-hor-like so sof faaa activation value with the temperature with the Gumbel perturbation.
Draw random sample perturbation. Assume that log mm log (- log um) > log π - log (- log 4z) for all j * m, i.c., the selected sample as category cm out of possible outcome set {co.
Uniform(0, 1)a which defines the Gumbel sample g for the Gumbel
Cn Therefore, for the Gumbel-Max trick, clear that T(/as (TT.g)) g)) E [em] Cx = Σt= em ④ c = Gm for c= (co. ,C--1) where E (0, 0, 1, 0....... 0) n--imensional one-hot vector, which has 1 in the mh entry and O in all other entries. Then, the statement that GS(T, r) GM(x) ( as T + 0
implies (z, g) + [M( g) i.e., the following:
log log( log u.
exp
[f5(F,9))
(15)
log log(
log m6))
k= exp
if = m
as T +
(16)
if + m
Then, f3(T,, g) Em for some relaxed one-hot vector em by introducing the softmax relaxation. As consequence,
log m. log
log u.
exp
[S3(E, x x
X
(17)
log log( log MK)
Ek=0 exp
Cm if = m
10 if + m
as T T
(18)
since the constant multiplication gives no harm to the approximation. Hence, by taking the summation, which also gives no harm to the approximation, ((f3(T.g)) = Em=dlémle + Cm = Σビ=0 em
C. (f.))))
D PMF SHAPE OF GENGS
Note that Figure 2 in the main paper is drawn by rounding-up continuous values into integers. Since PMF for discrete distributions (Poisson(7) in Fig) and PDF for continuous distributions (GENGS for Poisson(7) in Figure 2(b) cannot be directly compared within the same figure due to their scale, provide Figure of the fine-grained PMF by rounding-up in small decimals.
13

---

Under review as a conference paper at ICLR 2021
GenGS Poisson(/) Low Temperature, Truncation
GenGS Paisson(?), Tempersture, Truncation Level 15
Palsson())
Figure 6: Fine-gres aeneggrined PMF of GenGS for Poisson( (7).
E EXAMPLES OF GENGS
E. TRUNCATING RIGHT- SIDE FOR POISSON DISTRIBUTION
For the distributions that the left-side truncation needs to be zero, such as the Poisson with small rate parameter, the geometric, the negative binomial, etc., as we mentioned in the main text, we can simply set the constant outcome vector с = (0, 1, ... , - 1). Note that the ordering of e not crucial, for example, we can also set (2, 0, 3, 4,... n 1). Then, the corresponding PMF value needs to
be computed as = (P(Z 2), P(Z 1), P(Z = 0), P(Z = 3), P(Z = ... = 1)) where Z is the truncated discrete random variable.
E.2 TRUNCATING BOTH SIDES FOR POISSON DISTRIBUTION
For the case when the distribution requires left-side truncation, for example, Poisson(100)s,
set the constant outcome vector such as c = (50, 51, 51..... .... 149, 150). Again, the ordering of e is not erucial, hence, we can also set (148, 150, 149, 147, 53, 50, 52, 51), for example. Afterward, the PMF value π is naturally computed in the same order as the constant outcome vector c.
E.3 GUMBEL-SOFTMAX IS A SPECIAL CASE OF GENGS.
The Gumbel-Softmax estimator of categorical random variables is trivial case of GENGS, Assume the number of dimensions 3 and categorical parameter Then, the poissble outcomes of Categorical( (T) are c = [1, 0, 0)" c1 = [0, 1, 0) and 0o = [0, 0, 1]T. Afterward, draw a sample from GS( r) for some temperature a relaxed one-hot form, for example, (0.95, 0.04, 0.01). If we construct
(0.5, 0.3, 0.2) in this example.
O, and the value will become
[co, C1, c2]
(19)
[ 0
0
0 1
Tg
(20)
0
0
1
then T(w)= w 3=== En WaCk Hence, the Gumbel-Softmax trick can be written in the
form of GENGS with the identity matrix in the linear transformation. Note that we abuse the symbol of Hadamard product (3) in terms of the dimension, where the last term of T(w) >x
En W&Cr is actually the multiplication of scalar wt E [0, 1] and vector e₁ E R in this case.
E.4 GENGS CAN BE APPLIED TO MULTINOMIAL DISTRIBUTION.
If we go one step forward from the example above, we can reparameterize multinomial distribution
with GENGS trick, also. For example, assume the number of trial m and the probability vector
p P2. Pa] = [.7. 2, .2.11. Then, the possible outcomes are co = [3. 0. 0)
[0. 3, 0)"",
C2 [0,0,3 C3 [2, 0) C4 = [2, 0, 1]", cs [1, 2, 0]" Gs [1, 0.2 C7
C3
[0, 1, 2)" and cg = 1, where the corresponding probability is buging! P1 "1 T2 "2 for
Note that Poisson(100) has inprobable samples at 50 and x: 150.
14

---

Under review as a conference paper at ICLR 2021
outcome [n1, n2, ns]. Construct the linear transformation constant as the following:
[Co
C3
Cs
Cs
Co]
(21)
3
0
0
2
2
0
0
17
0
2
1
(22)
0
3
0
0
2
1
2
1
where the categorical parameter π
tma "1 P2
Mylngina! Pi
"2
P2
If sample Gumbel-
[ n2,33)
Softmax sample w from GS( 7) and compute T = Σx ④ the result will be in the relaxed
form of selected samples Cm- This example shows how the linear transformation constant с can be generalized to the matrix form. However, if we recall that the equation To0 +... Ri-1 n has ("*k=l) solutions of non-negative tuple (no. mk-11), where ("=k-1 I 0((2mn((knn1)) the relaxed categorical selection through the Gumbel-Softmax can become problematic due to the high complexity when n & has large value. Hence, in this situation, reducing the possible outcomes by disregarding unlikely samples by user guidance can be a remedy, but the treatment can not a fundamental solution, and handling such case can be open research question.
INFERENCE STEP, ALGORITHM & COMPLEXITY OF GENGS
VISUALIZATION OF GENGS REPARAMETRIZATION STEPS
Figure 7 and Figure represent the reparameterization stpes of explicit inference and implicit inference of GENGS, respectively. For both figures, the shaded nodes indicate the auxiliary Gumbel samples, composed of uniform samples, which enable the reparameterized variable be deterministic with respect to the parameter of the target distribution.
Forward Pass
PMF Values
No (2()))
Relaxed One-Hot Sample
Reparameterized Sample
Pa(-)
n1 P2(c1)
log +
A
softmax
T(w) Σ Www O
r
n- 1 - ExTt
No
CT-1
Truncation Level
Gumbel Random Samples
Temperature
Outcome Vector
Figure 7: Visualization of GENGS reparameterization of explicit inference version. Note that compute PMF values from the infered parameter X.
F.2 ALGORITHM ADDITIONAL COMPUTATIONAL COMPLEXITY OF GENGS
Algorithm D and Algorithm 2 provides the algorithm of explicit and implicit inference steps of GENGS, respectively. The explicit inference has further computation on PMF value calculation in $ 4 of Algorithm and the linear transformation computation in Line of compared to the original Gumbel-Softmax reparameterizer of the categorical distribution. Note that this additional computation complexity may not be represented as (n)) if we assume that there are n possible outcomes by truncating the distribution, since the computation on PMF values from the inferred distribution parameter depends on the PMF distribution. Meanwhile, the implicit inference only has extra linear transformation computation in Line 5f Algorithm 2 compared the original Gumbel- Softtma reparameterizer of the categorical distribution. by inferring the logit of the PMF values directly in Line of Algorithm Hence, in this case, the additional computational complexity
Line
15

---

Under review as a conference paper at ICLR 2021
Forward Pass
PMF Logit Values
PMF Values
Va
Relaxed One-Hot
Reparameterized Sample
V1
Softmax
1
log g
softmax
(w) Σ O
Vn-1
T
9n
91
On--1
C₁
Cp--1
Truncation Level
Gumbel Random Samples
Temperature
Outcome Vector
Figure 8:
Visualization of GENGS reparameterization of implicit inference version. Note that we do infer the distribution parameter and the PMF values are computed directly, instead.
of GENGS is (n), compared to the original Gumbel-Softmax. Note that this complexity does affect on the total complexity under the O-notation, since the inference step of the original Gumbel- Softmax also contains the inference on the logit of the probability which has (2(n) complexity.
Algorithm 1 Explicit GRNGS: Inference step the parameter À of distribution px
Input: PMF px(-) of distribution D(A), loss function f(-). Gumbel-Softmat trick GS, tem- perature > O, truncation range R, linear transformation T(+), constant outcome vector c.
2: Infer distribution parameter X.
do 4: Compute = (c())).
3: for ==
5: end for
6: Compute -- = k= Th.
7: Sample n-dimensiinnl one-hot hike w ~ GS(T, T).
8: Compute transformation = T(w) E, WyCj.
9: Compute loss (==).
Update À via stochastic gradient method.
Algorithm 2 Implicit GRNGS: Inference step on the PMF values v of distribution p>
1: Input: PMF px(-) of distribution D(A), loss function f(-), Gumbel-Softnax trick GS, tem-
perature >> 0, truncation range R, linear transformation T(+), constant outcome vector c.
2: Infer PMF logit-value i.
3: Compute = soft
4: Sample re-dimensional one-hot-liie hik- GS(T, 7).
5: Compute transformation 2 = T(w) = Σ, Wyeg.
6: Compute loss f(z).
7: Update e via stochastic gradient method.
G DISTRIBUTIONS TENSORFLLOW PROBABILITY
We provide the table of discrete distributions where GENGS be applied in Table We list up the discrete distributions which are available in TensorFlow Probability 0.8.0 in lexicographical order. Note that Bernoull 1i and Categori cal different from OneHot cal and ReLaxxedOneeorrCtteerrie.. The original Gumbel-Softmax are available only for
tyy/ppi pacss/pytoon/tfp//c
16

---

Under review as a conference paper at ICLR 2021
(I) Oneliotattegorical ,which is relaxed RelaxeeoneeHet and (2) Bernoulli, which is relaxed as RelaxedBernoul1, but can not be applied to Categorical However, by generalizing the Gumbel-Softmax, ie.. GENGS, other distributions listed below, including Categor cal. Also, Empirt ra which is a user-defined distribution in TensorFlow Probability, utilize GENGS which is not in the list. Note that for Bí nomi a if n. is large p has extreme value close to 0 or 1, one can truncate left, right, or both. Also, for Mult innniiialli aoomiiallaa as an extension of the Binomi al case, one can disregard unlikely samples.
Table 3: A list of distributions which can be reparameterized by GENGS with their distribution parameters in TensorFlow Probability 0.8.0.
PMF
Reparameterized Parameter
Infinite Support
Truncation Side
Distribution
P(X = k)
Rernou11i(p) e [0, (0,1}
p²11 -p)
None
Binomial(m.p) n NENN.. E 0. 1) (1, (1, m
(2)² (1 - p))
P
None
Categorical(p= pi.... PK)
e [0. p: 1, k [1, [1..... K).
II P/
None
e [0.1). k; 11
Dtr.chhletMullinnmmaall...
(a,,"",,x ,ak])
n N 04 0
mr 0zم rim+) a,) I 1or)
None
kj))(a))
E "n.
Geometric(p) e k0,1), I2
(1-p)"p
Right
Multinomialle an..iaannn...p...... (p)) N.N. E Ep, e (,,0,,-.,,,,,, kk=n
(e " ) [1
None
E N, IN
NegativeBinomisl/ (m.
( '1 7x
k
p
Right
OncHiotCateppriealp-ppp ........ .PK)
p. E 00, 1), 1 =
11 P3
p
None
KE <o k
Poisson(2) E Kyok Nzo
A exp
Right, Both
2ipf(n,s) $2 N f1,-..,m)
None
Σ/
H KL DIVERGENCE BETWEEN TWO TRUNCATAED DISTRIBUTIONS
Theorem. Assume two truncated distributions X ~ n) and Y ~ TD(), n) where Th P(X = k). ft. P(Y k). Then, the KL divergence between X and Y can be represented in the KL divergence between the categorical distributions where KL(Y(| [[ateggorriclll Categooocl))))))))
17

---

Under review as a conference paper at ICLR 2021
Proof.
P(Y k)
XL(Y||X) = Σ P(Y = =b)
(23)
P(X = k)
it t log
Σ
(24)
KL KL (Categorriaal [Categorrl (T))
(25)
I EXPERIMENT: GENERAL SETTING
For all experiments, use Intel Core i7-6700K CPU, 32GB RAM, and Titan X. For the dependency, we use TensorFlow version 1.15.0, TensorFlow Probability version 0.8.0, and PyTorch version 1.0.1. Also, we the experiments over 10 times for each experiment.
J SYNTHETIC EXAMPLE
J.1 EXPERIMENTAL SETTING
In this experiment, we first sample t1, ... ty i..dd from a discete distribution D(#) for a fixed > 0, and optimize the loss function Ex-((=))) Et [ with respect to À where p(=(A) is D(A). We use Poisson(20) Binomial(20,, .3), Multinomial(3, 1,7,,2,.11 .]],, and NegativeBinomial33 .4) in this experiment, and the distribution parameter which we want infer in each distribution is À in Poisson(A), (A), Binomial(20, A). Multinomial A), and NegativeBinomial A). For GENGS, we use truncation level (7, 36) and 12 for the Poisson and the negative binomial, respectively. Note that the binomial case does not require truncation of the distribution. We use k = 5 sampled targets for the Poisson and the binomial cases, and = 1 for the negative binomial case. this experiment, separately utilize the temperature т the low-temperature case. To compute the variance of gradients, we sampled 100 gradients for the Poisson and binomial, and 500 gradients for the negative binomial. For fair comparisons, m =1 fixed category for gradients for the RBs. Whereas it is able to use more than one fixed gradient in the synthetic example, if there is more than one latent dimension, K, it requires to compute m gradient combinations, which has high complexity. We also adapt the Rao- Hao-Blackwellization idea in GENGS, which is utilizing m fixed gradient and utilizing GENGS for the remainings, namely GENGS-RB, We exclude UNORD snice UNORD fails to converge to the optimal parameter because of approximation accuracy problem with single gradient sample.
1. for the high-temperature case, and = .25 for
Closed-form True Gradient Derivation for the Poisson Synthetic Example. Throughout the synthetic example, we compare the quality of gradient estimators by the convergence of losses, variances of estimated gradients, and biases between true gradient and estimated gradient. To compute the bias between the true gradient and estimated gradient, we need the closed-form solution of the true gradient. We find that the Poisson case has the closed-form true gradient, and the derivation is as follows.
Proposition. p(z)A) is Poisson distribution with a rate parameter X, the gradient of f. = E Ee>p(e/A)\((= 1)2] with respect X has closed-form solution, 85 2X 2/ 1.
18

---

Under review as a conference paper at ICLR 2021
40-
2.2
360
1404
1906
3006
2740
300
390
180
1008
2014
2518
3.00
5
1006
-50
3909
7504
Makrsen JB
4.-
1
10-
350
3500
3082
10:00
1000
60000
150
3000
1400
Branac PA, 3X
Vidloni (S1) 2 3b
600
2003
2001
10:00
20/00
305
2002
1205
2300
DetRE
Tomp
GonGS Low Tomp
GonGS Tomp
Low Temp
Figure 9: Synthetic example performance curves in log scale: (Top Row) Losses, variances, and biases of gradients for Poisson; (Middle Row) Losses for Binomial, Multinomial, and NegativeBinomial: (Bottom Row) Variances of gradients for Binomial, Multinomial, and NegativeBinomial. We utilize the cumulative average for smoothing the curves, and we also provide confidence intervals together.
Proof. Note that Poisson distribution with the rate parameter À has mean À and variance A. Hence, the Poisson distribution has the first moment μك = x, and the second moment H2 = X2 + A.
Oc
E. ~pp=||)]((= - [)2]
(26)
ax
= Σ C(<))A((z - t)2
(27)
ax
#20
Σ
p(2|A)(= t)²]
(28)
=20
-2
(2 1)2
(29)
z!
220
0 7
E
(2 1)2
(30)
x!
#20
0
2tz + +82) 22 - - N=-
=?
Σ
2
(31)
0x
z!
121
(Xe--))
= -P- + Σ
((2 + 1)2 - 22 + (1/2+11++ 12)
Σ (22 2tz + t2)
z!
z!
#20
121
(32)
= --f
Σ
(22 - 2(t - 1)2 (t - p(1))())))) -
(22 21 + )(=
(33)
220
19

---

Under review as a conference paper at ICLR 2021
=-fe- -x + (M2 - 2(t 1)11 + (t (t)) -N12 2tp1 + +t²(1 c e))))
(34)
2X- 2/ +
(35)
J.2 EXPERIMENTAL RESULT
We compare the log-loss and the log-variance of estimated gradients from various estimators in this experiment. We also compare the log-bias in the Poisson case. We additionally provide Figure to to report the confidence interval, and Figure 10 to show the convergence of loss which may not be seen Figure of the main paper.
34
1.5
2.4
-50
(S15
3.5
03-
-3%
2000
20/83
1049
1.580
1.0
800
2000
18:00
3000
40000
1P:00
DeIRB
StoR8
GenGS High Temp
GenGS ST Low Temp
GenGS High Temp
Gen Low Temp GanGS RB Low Temp
Rp*
OugRg
StoRg*
GandS RB ST Tomheemp
GandS RS ST Low Teme
GandS RB High HighTome
Figure 10: Synthetic example performance curves in log scale: (Top Row) Losses, vari- ances, and biases of gradients for Poisson; (Middle Row) Losses for Binomial, Multinomial, and NegativeBinomial; (Bottom Row) Variances of gradients for Binomial, Multinomial, and NegativeBinomial. We do not smoothen the curves to show the convergence of losses in this figure.
VAE: SYNTHETIC EXPERIMENT ON DEEP GENERATIVE MODELS K
K.1 EXPERIMENTAL SETTING
We utilize the (truncated) Poisson, the (truncated) geometric, and the (truncated) negative binomial distributions in this experiment. Both MNIST and OMNIGLOT" are hand- written gray-scale datasets of size 28 x 28. We split MNIST dataset into (train validation : test] = {45, 000: : 5, 000 : 10, 000). and OMNIGLOT dataset into {22, 095 2, 250 : 0700.
We construct two fully-connected hidden layers of dimension 500 for the encoder and the decoder, and we the latent dimension K 50 for both MNIST and OMNIGLOT datasets. activation function, learning rate 5e se-4 4, training epoch 500, and batch size 100, 45 for MNIST and
tanh
Thttps:/ http//gtttttt thub. couub.com/ybubbae http:///eeeeeaaaaeeaaaeteaee
OMN GLOTT
20

---

Under review as a conference paper at ICLR 2021
OMNIGLOT, repectively. For GENGS, we exponential temperature annealing" from 1. to .5. and truncation levels (1) 12 for Poisson( 22;; (2) 15 for Poisson(3); (3) 25 for Geometric( .25;; (4) 15 for
Geometric(.5); (5) 30 for NegativeBinommal(3, 5); and (6) 30 for NegativeBinomia((5,
K.2 EXPERIMENTAL RESULT
Table shows the negative ELBO results on the VAE experiments with a full range of gradient estimators. The variants of GENGS show the lowest negative ELBO in general. We empirically found that the extreme probability imbalance, due to the explicit PMF restriction, induces unstable learning which leads the performance degradation of RELAX REBAR. We found that utilizing Straight-Through for the discretization degrades the performance. Mcanwhile, the implicit inference methodology further leads to better optimal points, enabled by loosening the PMF condition. The empirical reason why the implicit version is better than the explicit version is that the inferred PMF shape is thinner the implicit case.. Hence, the implicit distribution has lower variance than the explicit one, and consequently samples consistent values that lead to better trained neural network parameters, although it looses the original PMF shape.
Table 4: Test negative ELBO MNIST and OMNIGLOT datasets. The lower is better for the negative ELBO. Symbol indicates no convergence.
MNIST
RF
RF
NVII
MUPROP
VIMCO(5)
REBAR
RELAX
Pois(2)
153.36 263.36
122.81=2.41
129.34,4.72
125. 125.43+2
122. 122.55
123 123.44
122.71
Pois(3)
153. 2.
123. m2
130.24,5.2
125.92
121.152257
120.62 2.:31
10.8442.m
Geomí 25)
168. 36
127.90_1.07
135.90
137.90
127. 21 22.55
135. 1242.74
136. 136.8043.00
Goomi
165.00
129.20=2.05
288.07
136-40+1.78
129.91 129.91 =2.00
138.37=200
139.41
NegBin(3..5)
116.67
119.28,7.50
131 131.9646
112.60
NegBin(5..33
130.03
133.4424.27
144.05
124.48
MNIST
STORB'
UNORd
GENGS ST (Fx.)
GENGS (Fx.)
GENGS ST (1m.)
GENGS (Im.)
Pois(2)
124.22 14.00
153.9115.:
130.481181
103.18
106.95 12000
96.04111.4
Pois(3)
124.41.5.06
153.65
129. 129.86
105.15
106. 106.66ia2
96.01
Gcomí.25)
14.5
120.1611.r
98. 43 Lo.s1
108.1412.7
92. 11.02
Geoml.5) (5)
139.67 299.67_2.42
119.10
100.92
106. 106.44....7
931.81
NegBin(3..5)
114.36-4.12
117.23+1
98.58
232.44
94.52+1.52
(Sg)
128 128.022.2.60
122.614231
100.88
102.70g3.64
95.37
OMNIGLOT RF
RF
NVIL
MUPROP
VIMCO(S)
REBAR
RELAX
Pois(2)
165.77
139.47.3329
148.01
142.95
138.78.3:42
138. 12
137.56
Pois(3)
164.8644.19
140.5442.36
148.13_a.0s
143.85
139.37
137 337.92233.0
137.42 137.424s
Geom(.25)
171.7944.50
142.68.22.0
153.69-2.82
152.17
142. 2.200443.00
146.7845.62
148.9 24.00
Goomí.5)
170.96.5.10
142.70 20
153.20
149.76g2..
142.05gs..0
149.63
151. 1:.00
(egBin(3..5)
24
144.44-2.78
147.78+4.40
141. 20.4
(ogBin(5..3)
14
159.40ma.1s
152.81
130.49g4.00
OMNIGLOT
STORB"
UNORD
GENGS ST (Ex.)
GENGS (Ex.)
GENGS (1m.)
GENGS (Im.)
2)is(2)
139.61a5.87
166.64_5.1a
148.60
127. 227.89.144
134.76ga.s
118.174a.99
Pois(3)
140.05=2.68
166.78_6.7
147.79
131.53
135. 14
119.15g1.3
Geom(.25)
143.10g3.1
146.54g1.1
115.28ga.0
135.61
107.794a.s4
Geom(.5)
142.56 142.56.2297
141.23
115.1442.43
136.02+3.6
108.48
3egSin(3..5)
129. 129.481134
142.33
118.57
135..12 135.02243.0
117.02
NegBin(S..3)
151.30=3.08
145.27+2.18
119.57
134.54
117.5442.00
L TOPIC MODEL APPLICATION
L. EXPERIMENTAL SETTING
Deep Exponential Families (DEFs) (Ranganath et al. 2015) are probabilistic graphical model which utilize the stacks of exponentail family distributions. If we assume the Poisson distribution, which is included in the exponential family, each 4th Poisson latent variable counts the number of sub-topics occurrence.
The relationship between the super-topic and the sub-topic is modeled with the linked weights, which has positive values. Hence, with Poisson DEF, model hierarchical relations between
For GENGS ST, the temperature annealing is unnecessary the ST Gumbeel-Sftman estimator does Jang 2017
21

---

Under review as a conference paper at ICLR 2021
super-topics and sub-topics including the vocabularies. Here, utilize the idea of Miao al. (0016) 2017); (rivasstava & Sutton (2017), the neural variational architecture, to extract the latent document representation as (relaxed) counts of supermost-topic, and consequently capture sub-topic counts. To ensure the positive linked weights between super-topics and sub-topics, we utilize absolute value function.
The generative process of NVPDEF is
21 Poisson Ao). 22 ~ Poisson(A1), ,.. 2K ~ Poisson(AK-1)
(36)
caltiiomiallooiistieeeererr
(x)
(37)
where we adopt multinomial logistic regression from NVDM (Miao et al. 2016), and the inference process of NVPDEF is
o MLP(x), d1 = Wid,, Ax WK WK-1
(38)
so that the approximate Poisson posterior q( (z-1 1) has Å distribution parameter. Here, each 2k ~ Poisson(Ax-1) represents the count distribution of topics from the super-topic. Each component of Wk. Wk,i, is positive, and Wk,i.j captures the positive weight of relationship between super-topic i of the hth layer and sub-topic j of the (k 1)th layer. The objective function, ELBO, of NVPDEF is
L Bg(xx,- .n [lo
21)
(L(g(=)/==)))l(xx)))
(39)
where 2g = for simplifying the equation.
20Newsgroups and RCVI-V2 datasets are used in this experiment. 20Newsgroups dataset has (train test] {11, 258 7, 487) split with the vocabulary size of 2, 000, and RCV1-V2 has (train test]= = (794, 414 : 10, 000} split with the vocabulary size of 10, 000. For the data pre-processing. stopwords are removed and the most frequent vocabularies are chosen. Especially for 20Newsgroups, we use the vocabulary from Srivastava & Sutton (017).
For the single-stacked version of NVPDEF, temerature r = .5. For the multi-stacked version of NVPDEF, i.e., MULTI-STACKED NVPDEF, we utilized 10-sample on the latent layers for the stable optimization of consecutive sampling. Also, to have better chances of learning, we utilize linear temperature annealing from T = 3. to T .5 during the training period. For all neural network models, we utilize two 500-dimensional hidden layers for the encoders. We use 50, 20-50 stacked layers for 20Newsgroups, 200, 50-200 stacked layers for RCV1-V2 dataset. We set A = .75 with truncation level 15 for the single-stacked case, and A = 1.1, >2 1. with truncation level 15 for the multi-stacked case. We train NVPDEF for 100 epochs with batch size 256 and learning rate 3. We also iteratively update encoder parameters and each linked weight parameter of latent variable. As a performance measure, utilize perplexity perp exp(- 1 En log where Na is the number of words in document d, and D is the total Na
do not anneal the temperature, instead, we set
number of documents.
L.2 EXPERIMENTAL RESULT
We provide super-topic, sub-topic, and word relationship obtained from two-layer-stacked NVPDEF 20Newsgroups dataset by listing up the top-weighted sub-topics and words Table
M OPEN RESEARCH QUESTION: NON-PARAMETRIC REPARAMETERIZATION TRICK
In GENGS, to reparameterize discrete distributions, convert sumpling process to categorical selection process by finitizing the support of the distribution with truncation. Here, truncating distribution converts categorical selection on countably finite number of categories to categorical selection on finite number of categories, i.c., turn non parmmarmttrric problem into parametric problem.
22

---

Under review as a conference paper at ICLR 2021
Table 5: Super-topic, sub-topic, and word relationship obtained from two-layer-stacked NVPDEF in 20Newsgroups dataset.
Super Topie
Super Topic 2
Super Topic
Super Topic 4
Super Topic 5
Super Topic 6
Topic 1
Topic 2
Topic 3
Topic 4
Topic
Topic 6
Topic 7
Topic 8
Topic
Topic 10
Topic 11
Topic 12
lebanese
hitler
knife
tire
water
probe
brand
honda
flyers
pitcher
boly
bible
lcbanon
jewish
gun
helmet
air
spacceraft
outlet
dealer
braves
montreal
resurrection
biblical
palestinian
nazi
police
rider
hear plane
sale
offer
hitter
score
passage
faith
arabs
religion
weapon
bike cold shuttle shipping
sell condition purchase market
philadelphia
season
jesus
prayer
isracli
territory
firearm
gear
oil
insurance
dctroit game
worship
doctrine
islamic
sentence
handgun
motorcycle
gas launch price
rangers
player
christ
regulation
moral
officer
wheel
noise
fuel
supply
minnesota
tie
sin
god
The categorical selection on finite number of categories by disregarding the samples of extremely small probability might cause a problem if we need to utilize full range of possible outcomes. For example, in the multinomial case in Appendix [E,4] as n and k grows, proposed GENGS should ignore numerous probable samples due to the high complexity on the number of possible outcomes.
A sampling (Maddison et al. 2014) a non-parametric version of Gumbel-Max trick which we also utilize GENGS, since A sampling searches maximum Gumbel sample among countably infinite Gumbel samples by A algorithm. Utilizing A sampling concept in reparameterizing the distribution with countably infinite support could lead better reparameterization terms of reparameterizing the exact distribution instead of approximate distribution. However, we utilize the truncated distribution in proposed GENGS to convert countably infinite categorical selection into finite categorical selection for the following reasons.
First, while adapting non-parament methodology into a neural network, which has a fixed number of parameters, people usually give limit certain point by utilizing the truncation level. An example of such case is a Stick bick-breaking VAE (Nalisnick & Smyth 2017) which utilizes Dirichlet process in the latent variable of VAE, and the authors fimitized the number of sticks by the human guidance. Second, while there is no previous work on reparameterization trick for fully-non-paramenn connmparametric categor- ical selection, if we finitize the number of categories with the truncated distributions suggested in the paper, we can utilize Gumbel-Softmat reparameterizer which already verified in deep generative model community and widely used by the implementation in the deep learning framework such as TensorFlow, i.e., Re Hel.aaddeeooottteeeee parametric model and the non-pat non-parametric model, the choice depends on the situation that face up to. For example, we can compare Gaussian mixture model (GMM) and Dirichlet process Gaussian mixture model (DPGMM). we have a clue on the number of clusters, we could directly apply GMM instead of DPGMM. However, if we know nothing about the data, utilizing DPGMM can be good choice. Also, we can not directly compare GMM and DPGMM along the same line, since the experimental result differs from data to data.
Finally, if we have to choose one between the
In summary, as other non-parametric models do, we turn the non-parametrie problem into the parametric problem, especially by utilizing the truncated distribution, and this kind of treatment a natural way of solving such difficulty. However, we believe that investigating non normmaraetrtic reparameterizer, partieularly utilizing * sampling which is theoretically solid, is crucial and open research question in the deep generative model community.
23


Reviews:
[{"review_id": "wFozcC8PTnN", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**summary** \nthe paper proposes a method of learning discrete approximate posterior distributions with potentially unbounded support. The idea is to truncate it to a finite set of states and use Gumbel-Softmax relaxation for samples from the truncated distribution. \nThe approach is illustrated on VAE example as well as topic modelling.\n\n**pros**\nThe idea of using Gumbel-Softmax to relax samples from another discrete distribution isn't new, however the application of this idea to distribution discussed in this paper does look novel. Gumbel-Softmax estimator is known to be computationally efficient and performing well in many situations, so it's perhaps not surprising that it performs well in the current context. Experimental results in topic modelling look very impressive.\n\n**cons**\n* I found it hard to follow the experimental section. In Section 7.2 the authors reference Figurnov et al (2018) for the setup however it seems that the reference only discusses continuous distributions. It is not clear what is the form of the prior and approximating posterior used in VAE experiments. This makes it hard to evaluate the paper.\n* I found the proofs that proposed approximation converges to true samples to be of little value: they do not give quantitative bounds on what the truncation region should be for given sample quality.\n* I think that comparison with StoRB and UnOrd estimators using only m=1 is not enough, it is possible that increasing m, although adding more computation, will significantly improve overall performance of the final LL.\n\ncomments:\n* My suggestion for improving the paper would be to provide much more experimental details, regarding the setup, datasets and baselines in the main text of the paper. At the same time the theoretical part can be reduced given that the idea is very straightforward.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "k2c67mCeww", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: The paper presents a generalization of the Gumbel-Softmax gradient estimator. The original Gumbel-Softmax is usually applied to Bernoulli and categorical random variables. The method proposed in the paper attempts to extend it applicability to other discrete distributions, such as Poisson, multinomial, geometric, among others. The main ideas of the approach are: (1) Random variables that may take countably infinite values are truncated, (2) The sampling process of the random variable is converted to a one-hot scheme (where Gumbel-Softmax relaxation is applied), (3) ``One-hot'' samples are reverted to the original sample space.\n\nClarity: While there's some room for improvement, I think the paper is clear.\n\nPros:\n- The paper tackles a very relevant problem; namely, estimating gradients with respect to the parameters of general discrete random variables. These distributions are used in many scenarios, and some of the existing methods are not always applicable. This paper aims to extend some of the current methods (specifically, Gumbel-Softmax) to general discrete distributions.\n\nCons:\nI have concerns about the method's significance.\n- The main claim is that the proposed method generalizes the Gumbel-Softmax estimator to general discrete distributions. However, I don't see how the method leads to an efficient algorithm for distributions with some combinatorial component. For instance, take the multinomial distribution with parameters (n, k). The number of possible outcomes is finite, but grows fast with n and k. Without truncation, the method can be applied by assigning each possible outcome to a one-hot vector, which becomes very large for moderate values of n and k. An alternative would involve removing potential outcomes with very low probability (i.e., truncating, as suggested in the paper). However the paper does not explain how to apply the truncation in cases like this. The only experiment in the paper with multinomial distributions involves small values of n and k (3 for both), and thus this issue does not arise.\n- The multinomial distribution is just one example where the issue mentioned above occurs. There are several other distributions for which this happens (e.g. a discrete distribution over tree structures), and it seems that no solution is proposed. If I understood the approach correctly, and this is the case, I think that a discussion of this should be included in the main paper. Is there a way to deal with this using the proposed approach?\n- For other non-combinatorial cases (e.g. Poisson, geometric), where the proposed approach may be applied efficiently, the method reduces to truncating the support of the random variable, which may be seen as a simple heuristic.\n\nIn summary, as presented, I think the paper does not provide an efficient approach to deal with general discrete distributions, and reduces to a simple heuristic for some of the simpler cases mentioned above. I think extending the approach to other more challenging distributions (multinomial, etc) could be an interesting direction to explore.\n\nRecommendation: Reject.", "extra_scores": {}, "extra_text": {}}, {"review_id": "W5RJ3FsxO_t", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n* This paper addresses the issue of applying continuous relaxations to simple countably infinite discrete distributions such as the Poisson or Geometric distributions.\n* The proposed approach truncates discrete distributions and applies a continuous relaxation to the truncated subset, i.e. the Gumbel-Softmax trick.\n* The experiments demonstrate the method outperforms other estimators that do not apply continuous relaxations.\n\nStrengths:\n* The method is straightforward and is a good first step towards applying continuous relaxations to high-dimensional discrete distributions.\n* The approach demonstrates improvements in the test likelihood in the presented experiments. \n\nWeaknesses:\n* A contiguous range must be chosen. This has clear drawbacks if there is no obvious structure as in the simple distributions considered in this paper (Pois, Geom, NegBin).\n* Too much space was dedicated to explaining the method twice, resulting in a lack of space for experiments. The MNIST/Omniglot experiments clearly show that the proposed estimator performs well by fixing the model and varying over different gradient estimators. What is the purpose of the topic modeling experiments? Are the baselines comparable to the proposed model?\n\nDecision: Reject\n* The method is simple and extends the Gumbel-Softmax trick to countably infinite discrete distributions.\n* The empirical results appear to be quite good compared to the reported baselines, although I am not familiar with the datasets used or whether the baselines are fair comparisons.\n* However, the writing needs a lot of work. Most of the experimental details have been relegated to the appendix because of space used to explain the method twice.\n\nSuggestions:\n* The truncation seems unnecessary. Is there a tweak to the method similar to Liu et al [1], such as applying the Gumbel-Softmax trick to the largest atoms (such that their total mass > 1 - 0.0005)? Or exploring learned partitions? Additionally, Liu et al [1] correct for bias via sampling. Is there a similar extension here? For example the (n+1)-th value of c can be sampled randomly from the complement of the support subset C.\n* Sections 3 and 4 are redundant. Get rid of section 3 and, if necessary, simplify the formal description of the method so that it is understandable without section 3. Find the minimal way to convey the method while remaining comprehensive. This should give more room to fully and clearly explain the experiments.\n* Given the simplicity of the method (which is a good thing), is all of the notation necessary?\n* In addition to the current results on topic modeling, I would like to see the fixed NVPDEF model trained with some of the gradient estimators from the MNIST/OMNIGLOT experiments.\n\n[1] R. Liu, J. Regier, N. Tripuraneni, M. I. Jordan, and J. McAuliffe. Rao-blackwellized stochastic gradients for discrete distributions. International Conference on Machine Learning, 2019.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "4cP2YiYgay", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a new methodology for inference of discrete variables in computational graphs.\n\nOverall the paper represents a substantial amount of work, and it is clearly written. Nonetheless, I believe it is not significant enough so I can recommend acceptance.\n\n1)Theoretical analysis, while pertinent, is mostly elementary and unsurprising. I would have expected something closer to rates of convergence but the kind of analysis presented (almost sure convergence when the truncation size goes to infinity) is essentially  obvious for anyone who was taken an elementary course on probability/measure theory. I would expect theoretical results to enlighten some aspects of their methods, but that is absent.\n\n2)I am in general unconvinced by the method. I don't think there is anything specially novel besides truncation. But just proposing truncation is not enough merit, I believe, as truncation is a natural choice. I would have expected an analysis of how truncation affects the results.\n\n3)I am not convinced by experimental results. Authors don't include Gumbel-Softmax as a baseline and I wonder whether any goodness in reported results (Besides the topic model) is a consequence of the action of Gumbel-Softmax instead of their method.\n\n4)The topic modeling experiment is interesting, since it involves a poisson distribution, with infinitely many possible values. However,  there, the relevant comparison is missing. This alternative could be, for example, the detr method (Liu et al, 2019). Notice that I disagree with the authors claim that this method is doing something completely unique. The Liu et al paper can be used with similar purposes.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers agree the paper does not meet the acceptance bar, and an authors' rebuttal is not available. Therefore, I recommend rejection.
Decision: Reject