Coupled Gradient Estimators for Discrete Latent Variables
Authors: Zhe Dong
Andriy Mnih
George Tucker
Venue: NeurIPS 2021 Poster (2021)
Keywords: discrete latent variable
antithetic
coupling
variance reduction
vae

Abstract: Training models with discrete latent variables is challenging due to the high variance of unbiased gradient estimators. While low-variance reparameterization gradients of a continuous relaxation can provide an effective solution, a continuous relaxation is not always available or tractable. Dong et al. (2020) and Yin et al. (2020) introduced a performant estimator that does not rely on continuous relaxations; however, it is limited to binary random variables. We introduce a novel derivation of their estimator based on importance sampling and statistical couplings, which we extend to the categorical setting. Motivated by the construction of a stick-breaking coupling, we introduce gradient estimators based on reparameterizing categorical variables as sequences of binary variables and Rao-Blackwellization. In systematic experiments, we show that our proposed categorical gradient estimators provide state-of-the-art performance, whereas even with additional Rao-Blackwellization previous estimators (Yin et al., 2019) underperform a simpler REINFORCE with a leave-one-out-baseline estimator (Kool et al., 2019).

Coupled Gradient Estimators for Discrete Latent Variables
Zhe Dong Google Research, Brain Team zhedong@google. com
Andriy Mnih DeepMind andriy@deepmind.comm
George Tucker Google Research, Brain Team gjt@google.co com
Abstract
Training models with discrete latent variables is challenging due the high vari- ance of unbiased gradient estimators. While low-variance reparameterization gradients of continuous relaxation can provide effective solution, contin- uous relaxation is not always available or tractable. [Dng et al. (2020) and Yin et al. (2020) introduced performant estimator that does not rely on continuous relaxations; however, it is limited to binary random variables. We introduce a novel derivation of their estimator based on importance sampling and statistical couplings, which we extend to the categorical setting. Motivated by the con- struction of stick-breaking coupling. we introduce gradient estimators based reparameterizing categorical variables as sequences of binary variables and Bao-BBaacwelllization. In systematic experiments, we show that our proposed categorical gradient estimators provide hale-of-the-art performance, whereas even with additional Rao-Black Blackwellization, previous estimators (Yin et al. 2019) un- derperform simpler REINFORCE with leave-one-out out-baseline estimator (Kool ct al. 2019).
1 Introduction
Optimizing an expectation of a cost function of discrete variables with respect to the parameters of their distribution is frequently encountered problem in machine learning. This problem is challenging because the gradient of the objective, like the objective itself, is an expectation over an exponentially large space joint configurations of the variables. As the number of variables increases, these expectations quickly become intractable and thus are typically approximated using Monte Carlo sampling, trading reduction in computation time for variance in the estimates. When such stochastic gradient estimates are used for learning, their variance determines the largest learning rate that can be used without making training unstable. Thus finding estimators with lower variance leads directly to faster training by allowing higher learning rates. For example, the use of the reparameterization trick to yield low-variance gradient estimates has been essential the success of variational autoencoders (Kingma and Welling 2014 Rezende et al. 2014). However, this estimator, also known as the pathwise derivative estimator Glasserman) 2013), can only be applied continuous random variables.
For discrete random variables, there are two common strategies for stochastic gradient estimation. The first one involves replacing discrete variables with continuous ones that approximate them as closely as possible [Maddison al. 2017] Jang et al. 2017) and training the resulting relaxed system with the reparameterization trick. However, as after training, the system is evaluated with discrete variables, this approach not guaranteed
Jang
perform well and requires careful choice of the continuous relaxation. Moreover, evaluating the cost function at the relaxed values instead of the discrete ones is not always desirable or even possible. The second strategy, involves using the REINFORCE estimator
Code and additional information: htttss://iitette http://siiss.ogggleeeeee www.go.ooggooooooo wtt.ee...o.iviieeeeeeeee
estimator
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

(Williams 1992). also known as the score-function (Rubinstein and Shapiro 1990) or likelihood-raiie (Glynn| 1990) estimator, which, having fewer requirements than the reparameterization trick, also works with discrete random variables. As the simplest versions of this estimator tend to exhibit high variance, they are typically combined with variance reduction techniques. Some of the most effective such estimators (Tucker et al. 2017;; Grrthwohl et al. 2018), incorporate the gradient information provided by the continuous relaxation, while keeping the estimator unbiased w.r.t. the original discrete system.
The recently introduced Augment- -EW--RRNNNOOREEREE -Mrrge (ARM) (Ynn and Zhou 2019) estimator for binary variables and Augment- Augment-REINFORCEE REINFORCE- (ARS) and Augment-REINFORCE REINFORCE SaaEE--Sapp-- Swap-Merge (ARSM) estimators (Yin et al. 2019) for categorical variables provide a promising alternative to relaxation-basee estimators. However, when compared to simpler baseline approach (REINFORCE with a setee-one-out-baseee
(RI.O;; Kool et al. 2019)), ARM underperforms in the binary setting (Dong et al., 2020), and we similarly demonstrate that ARS and ARSM underperform in the categorical setting. Dong al. (2020) Yin et al. (2020) independently developed an estimator that uses Rao-Blac Noo-BBackkellizationn to improve ARM and outperforms RLOO, providing state-of-the-art performance in the binary setting.
In this paper, we explore how to devise a performant estimator in the categorical setting. A natural first approach to apply the ideas from (Dong al., 2020 Yin et al., 2020) ARS and ARSM. However, empirically, we find that this is insufficient to close the gap between ARS/ARSM and RLOO. Due to space constraints, we defer the details to Appendix A.6 Instead, we develop a novel derivation of DisARM/U2C (Dong et al. 2020) Yin et al. (020)) using importance sampling which is simpler and more direct while providing natural extension to the categorical case. This estimator requires constructing coupling on categorical variables, which we do using stick-breaking process and antithetic Bernoulli variables. Motivated by this construction, we also consider estimators based on reparameterizing the problem with sequence of binary variables. We systematically evaluate these estimators and their underlying design choices and find that they outperform RLOO across a range of problems without requiring more computation.
2 Background
We consider the problem of optimizing
(=) [fo(z]],
(I)
with respect to the parameters 0 of a factorial categorical distribution (o(2) = He Cat( o0.e) where indexes dimension and ogg,k the vector of logits of the categorical distribution with C' choices | This situation covers many problems with discrete latent variables, for example, in variational inference fz( could be the instantaneous ELBO (Jordan et al. 1999) and (o(2) the variational posterior.
The gradient with respect to 0 is
Vo [fo(=)] Ey)(=) (fo()) logoooo + Vole(=)).
(2)
It typically suffices to estimate the second term with single Monte Carlo sample, so for notational clarity, we omit the dependence of on 6 in the following sections. Monte Carlo estimates of the first term can have large variance. Low-variance, unbiased estimators of the first term will be our focus.
2.1 DiSARM/U2G
Dong al. 2020) and Yin et al. (2020) derived Bao-Blackwelliized estimator for binary variables based on the coupled estimator of Yin and Zhou (211)). In particular, if qo (b) ~ II, Bernoulliç (0.,)) with 0 parameterizing the logits oo, of the Bernoulli distribution, and and b are antithetic sample from go (b) (independent across dimension k), then the estimator
2 (f(b) (1-1)01 1e a (((ag,)))
To simplify notation, we omit the subscripted 0 on a when it is clear from context.
2 Antithetic Berouuli samples can be defined by the following process: sample Uniform(O, 1), then set b p and i (1 u) < p. where is the probability parameter of the Bernoulli variable.

---

is an unbiased estimator of the gradient V V (1) [f()]]. This estimator has been shown to outperform RLOO, but is limited to the binary variable case.
Methods
First, provide novel derivation of DisARM (Dong al. 2020) / U2G Yin al. 2020) from an pepoortane-saappling perspective, which naturally extends to the categorical case. Starting with the (2 (2-sample) REINFORCE LOO estimator (Kool al. 2019) and applying importance sampling, have
(1,00(2,5)
Vor E 20(2) [f(z)] = B (=))
(f(z) - f(ž)) (Vo, log (z())) Vo log go (2x))
go(2)g0
(f(z) - f(z)) (Ve, log (o(x)) - Vo, log qo (2x))
2 pa (z.5)
To reduce the variance of the estimator, we can use the joint distribution (= 8) to emphasize terms that have high magnitude. However, controlling the weights can be challenging for high dimensional
2. We can sidestep this issue by taking advantage of the structure of the integrand and requiring that coppplinecouupiigg
Pe be a properly supported coupling that independent across dimensions (i.e., a joint distribution
p(2, 2) II p(z). 2x) such that the marginals are maintained p(=x.) p(2)) = 9(2x)). Then with
(z, E) ~ p, the following will an unbiased estimator of the gradient (see Appendix A.3 for derivation)
1 qo (E)) (f(2) - (()) (Vo, log (z())) - Va, Va log qa (āx)).
SDSARM-IW 2 (ZN: =b)
(3)
Critically, because the coupling maintains the marginal distributions, we only need to importance weight a single dimension at time, which ensures the weights reasonable. The estimator is unbiased, so the coupling can be designed to reduce variance. In the binary case, with an antithetic Bernoulli coupling, it is straightforward to that this precisely recovers DisARM/U2G. We note that Dimitriev and Zhou 2021) independently and concurrently discovered a similar result. Conveniently, this estimator is also valid in the categorical case, however, the choice of coupling affects the performance of the estimator.
Stick-breaking Coupling
The ideal variance-reducing coupling would take into account the magnitude of (f(z) f())) (Vo, log qo (=x)) - Vo, log (xx)) which we do not know a priori. However, we know that when 2g 2g. the last multiplicative term of Eq. B] vanishes, moving away from this configuration will reduce variance. Furthermore, for the estimator to be valid, the coupling must put non-zero mass all 2 # z configurations [
In the binary case, there is a natural construction of an antithetic coupling which minimizes the probability of z ãµ. We can extend this construction using formulations of categorical variables as a function of a sequence of binary decisions. The stick-breaking construction for categorical variables (Khan et al. 2012) provides such approach. Suppose we have categorical distribution ((zx) and we construct a sequence of independent binary variables bb. bb.c with by.; ~ Bernoulli(g(z,, i)/ E5= ((= =))) and define then 2 ((x)). Given each of .., bb.C. we can independently sample antithetic Bernoulli variables by b........ b and let = SB(6b.1.... 1..b......... b,. This process defines joint distribution on (ZR, ),), which we call the stick bick-breaking coupling. By construction, the joint distribution preserves the marginal distributions. Note that the stick-breat beeckbre coupling depends on the an arbitrary ordering of the categories.
SB(bb,,,....... ,b,c) == min i s.t. bR = 1;
As is standard for importance sampling, ensure unbiasedness, require that either g(2, 8) 0 the integrand is whenever p(2, ;) = 0.
'Technically, the coupling can zero mass configurations as long the expectation is zero across those configurations, however, this hard to ensure for general f.

---

Without reordering the categories, the coupling arising from the default ordering may not put non-
zero mass on all 2 * Σn configurations. When p(bx.a) : 2 0.5, then p(bg. = 0, ba. 0) = 0. which implies that p(2 > i, E > i) 0 by definition. So if p(be,s) > 0.5 for i < C, this violates the condition that the coupling put mass on all * Es configurations. We can avoid this situation by relabeling the categories in the ascending order of probability (as this guarantees that g(=x = / Σj= g( = j) & 0.5). With this ordering, computing the required importance weights with the stick- breaking coupling is straightforward when 2& & &k
min(*k,8k)) -1
(e(2))90 (0x)))((2z))
o( o(-ar,))
II
( (-amiu(zu,saa) 1 1- (=2a))
po [2k, kx.)
where Cue, = logit g(=6=1)
((x= (see Appendix A.4 for details). Note that due to the ascending
ordering. the ((a,,) term in the product is always 3. ensuring that the weights are bounded.
Finally, we note that mixture of couplings is also a coupling, so we can take any coupling and mix it with the independent coupling ensure the support condition is met. Because the estimator
is unbiased for any choice of nonzero mixing coefficient, the mixing coefficients can be optimized following the control variate tuning process described in (Ruiz et al., 2016; Tucker et al. [017]] Grathwohl et al. 2018). This also ensures that the performance of the estimator will be at least as strong RLOO. Exploring this idea could be an interesting future direction.
3.1 Binary Reparameterization
Motivated by the stick biek-breaking construction, we can eschew importance sampling and directly reparameterize the problem in terms of the binary variables and apply DisARM/U2G. In particular, if we have binary variables 6u,1 (erno .b.c ~ Bernoull((((ww.c) such that Zs == h(bb, I, J.... bx.c) ~ (zx) where a deterministic function that does not depend on 0, then can rewrite the gradient
V V go (4) (f (M(b)) = =V as. Be]s(h(D)]] EL log (b)) (bbl( or.. log ge (bx.c)].
which is expression that DiSARM/U2G can be applied to.
We can additionally leverage the structure of the stick-breakin construction to further improve the estimator. Intuitively, by the definition of in terms of SB(...). ), if we know that by.c 1, then the values of e for c' > c do not change the value of zg. which implics the gradient terms vanishing. In particular, given antithetically coupled Bernoulli samples, the following is an unbiased estimator (see Appendix A.4 for details):
(f(2) - f()) ((-1)mme Io tbe. a (((ax.e)))
< min [x)
9DLARMM-BBECC
1 (J(E) - [()) Von,e log go (br.c)
4 < S Es
(4)
(f(z) - -2)) Var.. log go (bk,c)
21 < 2k:
> max (2k kkk
In contrast to the previous section, any choice of ordering for the stick-breaking construction results in an unbiased estimator, and we experimentally evaluated several choices.
Alternatively, we can use a sequence of independent Bernoulli variables arranged as a binary tree to construct categorical variables, with each leaf node representing category (Ghahramani et al. 2010). For simplicity, we only consider balanced binary trees (i.c., C' a power of 2). The binary variables encode the internal routing decisions through the tree, defining a deterministic mapping from the binary variables to the categorical variable 2 == T(b 1 It is straightforward to derive the Bernoulli probabilities from (2( so that go (2).).
defer the details to Appendix A.5
As above, we can additionally leverage the structure of the tree construction to further improve the estimator. Given binary variables b), ....... bc- 1. let [bb...., bc- be the set of variables used in routing decisions (ITbb,.... bb--) = log, C). By the definition of 24 terms of T(...), we know that z is unaffected by binary decisions that occur outside (bb.), which implies the gradient terms

---

vanishing. With pair of antithetically sampled binary sequences, the following is an unbiased estimator (see Appendix A..5 for details):
(f(2) f(2)) be, (aaesbr..ol o(((ax,e))
E I(br.) ) n n I(br.))
32 (f(2) f(2)) Vaz,e log 90 (bk,c)
I I(bb,) I(bk,))
gDisARM-Trest T
(5)
(f(z) - f(8)) Var,o log go
I (bk,.) I(bk..)
E T(bR,)) ) U U(bE.))
Again, the choice of ordering affects the coupling, however, we leave investigating the optimal choice to future work.
Notably both estimators include terms that are reminiscent of the LOO estimator, however, they use dependent 2, and are still unbiased due to careful construction.
3.2 Multi-sample extension
We can extend the proposed estimators to the case when we have n-coupled pairs (= .,,........ ""). Dimitriev and Zhou (202T) note that the 2n-sample RLOO estimator can be written as average over all 2-sample RLOO estimators. In other words, given 2n independent samples z² 2?.
( 2n)
2n(2n 1) Σ [RL
If instead of 2n independent samples, we used n-coupled pairs in the estimator, most 2-sample RLOO estimators will be with independent samples, however, n estimators will use coupled pairs which introduces bias. We form the unbiased n-coupled pair estimator by adding a correction term that replaces those terms with DisARM estimator
GKLOO(2* (t.,2)....,
20 + 2n(2n 1)
(2', (
where we can use any of the previous DisARM-* estimators.
4
Related Work
Almost all unbiased gradient estimators for discrete random variables (RVs), including our proposed estimators, belong to the REINFORCE/ score function / likelihood-ratio estimator family (Williams) 1992 Rubinstein and Shapiro 1990, Glynn 1990). The primary difference between such estimators is the variance reduction techniques they incorporate, trading additional computation for lower variance of the estimate. Averaging over multiple independent samples is the simplest but relatively inefficient variance reduction technique, as the variance of the estimate inversely proportional to the number of samples and thus decreases slowly. REINFORCE with the lcave-out-out baseline (RLOO; Kool al., 2019) provides much more effective way of utilizing multiple independent samples, by using the average over all the other samples as the baseline for each sample. Kool et al. (2020) extend this approach by sampling without replacement, however, they find that the high-dimensional spaces we consider, we almost never see such duplicates, resulting in little to no improvement in sampling without replacement compared to sampling with replacement (Figures 2(b) and 4(b) (Kool et al.
2020))
The estimators we develop use technique called coupling (wwn 2013) to make better use of multiple samples by making them dependent. Recently, this kind of approach has been applied to binary RVs, resulting in the Augment-REin Sewteeee-RRNFORCE--MMeee
(ARM: Yin and Zhou 2019) and DisARM DDog et al. 2020) / U2G (Yin et al., 2020) estimators. ARM is constructed by reparameterizing Bernoulli RVs in terms of Logistic variables and applying antithetic sampling, which is the simplest kind of coupling, to the resulting REINFORCE estimator. DisARM U2G is obtained by Rao- Blackwellizing ARM estimator to eliminate its dependence on the values of the underlying Logistic variables, reducing its variance. We show how to employ DisARM with categorical RVs by representing them as sequences of binary decisions and applying DisARM to the resulting system. By considering the sequences obtained through
and tree-structured decisions, obtain
5

---

Table 1: Mean variational lower bounds and the standard error of the mean computed based on runs of 5 x 1te-steps training from different random initializations. The best performing method (up to the standard error) for each task is in bold. Top: comparing the performance of the proposed estimators and the RL.OO baseline on three datasets. Bottom: comparing the effect of arranging the categories the ascending or descending order to the default ordering for DisARM-SB.
Estimator Comparison with Train ELBO
DiSARM-IW
DiseRM-Tree
DisARM-SB
RLOO
RELAX
DynamicMNIST RishionMNIST
-103.41 + 0.23
-033. 103.1022 + 0.25
- 103..3 + 0.17
-104.03 + 0.23
-102.22 #: 0.34
-240.08 + 0.17
239 0.29
-239.60 66 + 0.20
-240.89 + 0.49
-338.81 0.34
Onniglot
-122 25 + 0.35
123.08 0.43
-121.99 + 0.17
-122.70 + 0.80
-120.40
0.16
Ordering Comparison with Train ELBO
Ascending
Descending
Default
DynanicMNIST MNShi
-103.35 1 0.17
-103.56 1 0.28
-103.48 46 0.17
-239. 6899.6612 1 0.20
-240.04 = 0.24
-239.00 + 0.20
Omniglot
-121.99 + 0.17
-122.79 0.12
-122.75 0.15
the DiSARM-SB and DisARM-Tree estimators respectively. We obtain DiSARM-IW by generalizing the two-sample version of RL.O to allow the two samples to be coupled and using an importance- weighting correction to keep the estimator unbiased. DisARM-IW reduces to DisARM in the case of Bemouuli RVs and antithetic coupling, and thus can be seen as generalization. Dimitriey and Zhou (202T) has concurrently and independently developed an estimator identical DisARM-IW, though they considered only the case of Bernoulli random variables.
Augment- Reegnenn-REINOORCE.C. Swap estimator (ARS; Yín et al.) 2019) is generalization of ARM to categorical RVs; it reparameterizes the categorical variables in terms of Dirichlet RVs and applies REINFORCE to them, using a coupling based on a swapping operation. Bngment-REINFORE Swap-Merge (ARSM; Yin et al., (019) goes one step further by averaging over all possible choices of the pivot dimension that ARS chooses arbitrarily. Our Boo-Blaakwellized ARS+ and ARSM+ estimators are obtained by analytically integrating out the randomness introduced by the augmenting Dirichlet dimensions from the ARS and ARSM estimators respectively, following the strategy that was used to derive DisARM from ARM (see Appendix for details).
Gao-Blackwelliztine has been popular variance reduction technique in the literature gradient estimation with discrete random variables, having been used in BBVI (Ranganath et al.) 2014). NVIL (Mnih and Gregor, 2) LEG (Titsias and Lázaro- Lázaro-Gredilla 2015), well as more recently by Liu jet al. (2019) and Paulus et al. (202)). Stick-breaking has previously been used in the context of VAES to allow inferring the latent space dimensionality (Nalisnick and Smyth 2016).
In this paper we use variational training of latent variable models as a benchmark for unbiased gradient estimators for general expectations rather than a goal in itself. As result, we do not compare to specialized techniques for training models with discrete latent variables, such as Reweighted Wake Sleep (Bornschein and Bengio| 2015) and Joint Stochastic Approximation (Ou and Song 2020). which do not provide such unbiased estimators and optimize different objectives.
5 Experiments
Our goal was to derive unbiased, low-variance gradient estimators for the categorical setting, which do not rely on continuous relaxation. We primarily compared against the REINFORCE estimator with a leave-one-out baseline (RLOO, Kool et al. 2019). In the high dimensional setting that we evaluated, RLOO has been shown to be a strong baseline algorithm, performing at least as well as more complex methods, such as RELAX, REBAR, STORB, and UnOrd, as shown in (Richter et al.) 2020 Kool et al. 2020). As in prior work Yin et al. 2019), we benchmark the estimators by training variational auto- encoders (Kingma and Welling 14 Rezende al.) 2014) (VAE) with categorical latent variables on three datasets: binarized MNIST LeCun et al. 2010), FashionMNIST (iao
*Creative Commons Attriuutonn Alike 3.0 license
MIT license

---

Table Evaluating the estimators and the effect of ordering on FashionMNIST with different model sizes. Mean variational lower bounds and the standard error of the mean computed based on 5 runs of 5 -te-steps training from different random initializations. We consider three combinations of the number of categories (C) and the number of latent variables (V). The best performing method (up to the standard error) for each task is in bold. Top: evaluating the performance of the proposed estimators and the RLOO baseline. Bottom: comparing the effect of arranging the categories in the ascending or descending order to the default ordering for DisARM-SB.
Estimator Comparison with Train ELBO
DisARM-IW
DisARM-Tree
DiSARM-SB
RLOO
C64
-244.75 + 0.53
-244.29 0.30
244.81
0.64
-245.27
0.30
C64 / V32
-240.08 + 0.17
-239.83 + 0.29
229..6 0339.66-2 + 0.20
-240.89 + 0.49
C16 / V32
-239.94 + 0.14
-240.52 + 0.22
-240.43 + 0.13
-240.99 + 0.25
Ordering Comparison with Train ELBO
Ascending
Descending
Default
C64 / V5
-244.81 1 0.64
2244.95 1 0.36
244.61 1 0.38
C64 V32
-239.66 0.20
-240.04 1 0.24
-239.99 1 0.20
C16 V32
-240.43 = 0.13
-240.39 + 0.21
-240.50 + 0.13
al., 2017), and Omniglof Lake al. 2015). As we sought evaluate optimization performance, we use dynamic binarization to avoid overfitting and largely find that training performance mirrors test performance. We use the standard split into train, validation, and test sets.
To allow direct comparison, we use the same model structure as in ( Yin et al., 2019). Briefly, the model has a single layer of categorical latent variables which are then represented one-hot vectors and decoded to Bernoulli logits for the observation through either a linear transformation or an MLP with two hidden layers of 256 and 512 of LeakyReLU units (Xu et al. 2013). The encoder mirrors the structure, with two hidden layers of 512 and 256 LeakyReLU units and outputs producing the softmax logits. For most experiments, we used 32 latent variables with 64 categories unless specified otherwise. See Appendix A.2 for more details.
5.1 Evaluating DiSARM-IW, DiSARM-SB
& DisARM-Tree
First, we evaluated the importance-weiihhed estimator (DisARM-IW),, and DisARM estimators with the stick breaking construction (DisARM-SBB) and the structured construction (DisARM-Tree)) These estimators require only 2 expensive function evaluations regardless of the number of categories (C), unlike ARS/ARSM which require up O(C)/O(C2) evaluations, so we use the 2-independent sample RLOO estimator as a baseline.
In the of nonlinear models (Figure 1. DiSARM-IW and DisARM-SB perform similarly, with lower gradient variance and better performance than the baseline estimator (RLOO) across all datasets. While for DisARM-IW, ensure unbiasedness, for DiSARM-SB any ordering results in an unbiased estimator. We experimented with ordering the categories in the ascending and descending order by probability as well as using the default ordering (Figure ( We found that the estimator with the ascending order consistently outperforms the other ordering schemes, though the gain is dataset- -deeedeedd However, when we varied the number of categories and variables (Figure (). we found that no ordering dominates. Given that ordering introduces an additional complexity, we recommend no ordering in practice.
have use the ascending order for the stick-breaking construction to
In the case of linear models, we found that all estimators performed similarly (Appendix A. Dong et al. (020) found that for multi-layer linear models in the binary case, DisARM showed increasing improvement over RLOO for models with deeper hierarchies. As nonlinear models are more common and expressive, we leave exploring whether this trend holds in the categorical case future work.
We further verified that the observed improvements of the proposed estimators w.n.t. the RLOO baseline are consistent for models with different sizes of the latent space. As shown in Appendix
MIT license
7

---

-10
SA10 3000
Coupal
Vie Giw.deeeee
PynamicMNNSS
Ton 1180
-104
-100
-102
-100
200K
200K
400K
400K
300K
230x
200K
200K 2305,
2028
320X
4006
300K
Traliing
-2.4
-230
FistiounNNIST
240
0400
0-400
Tes: 300 Sanels sound
-240
Venance of Gred weth
10.44
-241
-244
-284
2400K
2008
103K
100K
500K
-2400X
200K
100K
203K
450K
500K
-120
Tixt 100 Samplo Sound
-120
Cradhente
10-
Tran EL8O
-122
122
Omnigghot
-126
-128
103K
250K
3000
153K 400% 43.0K 5O0K
200k
250K
303K 3506 40DK
450K 500K
3806
203K
300K 4006 300K
Traning St.p
Stop
Figure Training non-linear categorical VAE with latent variables with 64 categories on dynamically binarized MNIST. FashionMNIST, Omniglot datasets by maximizing the ELBO. We plot the train ELBO (left column), test 100-sample bound (middle column), and the variance of gradient estimator (right column). For fair comparison, the variance of all the gradient estimators was computed along the training trajectory of the RLOO estimator. We plot the mean and one standard error based on 5 runs from different random initializations.
Drcll
D4all Jveesss
-24
Tran E180
241
-106
Train
-128
244
200k 253K.
300K 200K 400K Traiiing Step
459K
SOOK
200k
200K
B00K
400K
5O0K
250K.
300K
3005 400K 400k SOOK Training Step
Traising Step
DynAMicMNIST
MNShionMNIST
Omniglot
Figure 2: The effect of logit ordering on the performance of DiSARM-SB. We sort the encoder logits the ascending or descending order and compare against the default ordering.
Figure 6 we find that the best choice of the proposed estimators depends on the model size, however all three estimators outperform RLOO in all cases.
5.2 Evaluating Multi-sample Estimators
We evaluate the multi-sample extension of the DisARM-* estimators three benchmark datasets, by training a non-linear categorical VAE. The VAE has a stochastic hidden layer with 128 categorical latent variables, each with 16 categories. We run experiments with 5 antithetic pairs (10 samples) and 10 antithetic pairs (20 samples), and found that the proposed estimators outperform RLOO with comparable number of independent samples. Even with the increasing number of samples, the performance improvement of the proposed estimators w.r.t. RLOO still holds, as shown in Figure 3 Appendix Figure 7 and Appendix Table
Conclusion
We introduce a novel derivation of DisARM/U2G estimator based on importance sampling and statistical couplings, and naturally extend it to the categorical setting, calling the resulting estimator

---

-210
-154
-132
Tran 11.30
Pains
24
-235
-236
1409
15ox
105K
Traliing Stop
Trahiing Stop
Training Step
-2-
10 Pairs
110
Rico
-04
-231
112
-235
-150
2004
200K
300K
400K
500K
-240
200K
300K
401K
5
100K
200K
5028
400K
500K
DynAmicMNIST
FashionMNIST
Omniglot
Figure 3: Training a non-linear categorical VAE with 128 latent variables and 16 categories, using multi-sample estimators, on dynamically binarized MNIST, FashionMNIST, Omniglot datasets by maximizing the ELBO. The examples in the top row are using 5 pairs of samples, while the ones in the bottom using 10 pairs.
DisARM-IW. Motivated by the construction of tiick-breaiing coupling, we introduce two estimators, DiSARM-SB and DisARMM-Tee: by reparameterizing the problem with sequence of binary variables and performing Rao-Bi Blackwellization.
With systematic experiments, we demonstrate that the proposed estimators provide state-of-the-at performance. We find that the proposed estimators usually perform similarly and all outperform RLOO, with the winner depending on the dataset and the model. As the importance sampling estimator D.SA simpler to implement, more natural understand. and easier to generalize, we recommend it in practice. We expect that this estimator can be further improved through better couplings, which is something we intend to explore in the future.
limitation of the introduced categorical couplings is that they impose structure on the categorical space (i.e., ordering tree structure), which is not fully satisfactory because in most settings there is no such natural structure for categorical spaces. Developing coupling-based estimators that do not rely on such structure would be interesting future work which might lead further improvements. While we see that our coupling- -aaed estimators generally outperform or perform at least as well as RLOO in our experiments, using coupled samples instead of independent samples is not guaranteed to lead to better performance. Learning the couplings would provide a way of ensuring an improvement over RLOO. Finally, the estimators we propose this paper, like all nalli- estimators, can be used for RL only if the enviromment is simulated or we have model of it, as they require being able to perform multiple rollouts from the same state.
Discrete latent variables have particular applicability to interpretable models and sparse/conditional computation. Improving the foundational tools train such models will make them more widely available. While interpretable systems are typically viewed as a positive, they only give a partial view of complex system, and they be misused give false sense of understanding. While sparse/conditional models can reduce the environmental cost of learned models (Patterson et al. 2021), the unintended implications of large models must also be considered (Bender et al. 2021).
Acknowledgments and Disclosure of Funding
We thank Chris J. Maddison and Ben Poole for helpful comments. The authors did not receive any third-party funding or third-party support for this work, and have no financial relationships with entities that could potentially be perceived to influence what they wrote in the submitted work.
References
Bender, F. M., Gebru, T., McMillan-Majon A., and Shmitchell, (2021). On the dangers of stochastic parrots: Can language models be too big? In Proceedings of the 2021 ACM Conference

---

on Fairness, Accountabiliry, and Transparency, pages 610-623.
Bomschein, J. and Bengio, Y. (2015). Reweighted wake-sleep. ICLR.
Dimitriev, and Zhou, M. (2021). ARMS: Anithheti-RRENNFOREEEMM sample gradient for binary variables. In Proceedings of the 38th International Conference on Machine Learning.
Dong, Zz,, Mnih, A., and Tucker, G. (2020). DisARM: An antithetic gradient estimator for binary latent variables. arXiv preprint arXiv:2006. 10680.
Ghahramani, Z.,, Jordan, M., and Adams, R. P. (2010). Tree-structured stick breaking for hierarchical data. In Advances in Neural Information Processing Systems, volume 23.
Glasserman, P. (2013). Monte Carlo Methods in Financial Engineering. Springer Science & Business Media.
Glynn, P. W. (1990). Likelihood ratio gradient estimation for stochastic systems. Communications of the ACM, 33(10):75-84.
Grathwohl, W., Choi, D., Wu, Y., Roeder, G., and Duvenaud, D. (2018). Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In International Conference on Learning Representations.
Jang, E., Gu, S., and Poole, B. (2017). Categorical reparameterization with Gumbel-softmax. In International Conference Learning Representations.
Jordan, M. I., Ghahramani, Z., Jaakkola, T. S., and Saul, L. K. (1999). introduction to variational methods for graphical models. Machine learning, 37(2):183 183-233...
Khan, M., Mohamed, S., Marlin, B., and Murphy, K. (2012). A stick-breaking likelihood for categorical data analysis with latent gaussian models. In Artificial Intelligence and Statistics, pages 610-618. PMLR.
Kingma, P. and Welling, M. (2014). Auto-encoding variational Bayes. International Conference on Learning Representations.
Kool, W., van Hoof, H., and Welling, M. (2019). Buy 4 REINFORCE, samples, get baseline for free! In Deep RL. Meets Structured Prediction ICLR Workshop.
Kool, W., van Hoof, H., and Welling, M. (2020). Estimating gradients for discrete random variables by sampling without replacement. arXiv preprint arXiv:2002.06043.
Lake, B. M., Salakhutdinov, R., and Tenenbaum, J. B. (2015). Human-level concept learning through probabilistic program induction. Science, 350(6266))13333 1332-13333
LeCun, Y., Cortes, C., and Burges, C. (2010). MNIST handwritten digit database. ATT Labs [Online].
Available: http://yann.lecunt http:/aww..cu...oo/////////////e 2.
Liu, R., Regier, J., Tripuraneni, N., Jordan, M., and Mcauliffe, J. (2019). Rao-Blackwellized
stochastic gradients for discrete distributions. In International Conference on Machine Learning. pages 4023-4031. 4031. PMLR.
Maddison, C. J., Mnih, A., and Teh, Y. W. (2017). The Concrete Distribution: A Continuous Relax- ation of Discrete Random Variables. In International Conference on Learning Representations.
Mnih, A. and Gregor, (2014). Neural variational inference and learning in belief networks. In Proceedings of The 3 International Conference on Machine Learning, pages 1791-1799.
Nalisnick, F. and Smyth, P. (2016). Stick-breaking variational autoencoders. arXiv preprint arXiv: 1605.06197.
Ou, Z. and Song, Y. (2020). Joint stochastic approximation and its application to learning dis- crete latent variable models. In Proceedings of the 36th Conference on Uncertainty Artificial Intelligence (UAI).
10

---

Owen, A. B. (2013). Monte Carlo theory, methods and examples.
Patterson, D., Gonzalez, J., Le. Q. Liang, C.,, Munguia, L.-M., Rothchild, D., So. D., Texier, M., and Dean, (2021). Carbon emissions and large neural network training. arXiv preprint arXiv:2104. 10350.
Paulus, M. B., Maddison, C. J., and Krause, A. (2020). Rao-Blackwellizing the straight-through Gumbel-softmax gradient estimator. arXiv preprint arXi:20010.04838.
Ranganath, R., Gerrish, S., and Blei, D. M. (2014). Black box variational inference. In AISTATS, pages 814-822.
Rezende, D. J., Mohamed, S., and Wierstra, (2014). Stochastic backpropagation and approximate inference in deep generative models. In Proceedings of The 31st International Conference on Machine Learning, pages 1278-1286.
Richter, L., Boustati, A., Nüsken, N., Ruiz, J., and Akyildiz, Ó. D. (2020). VarGrad: low-variance gradient estimator for variational inference. arXiv preprint arXiv:2010.10436
Rubinstein, R. Y. and Shapiro, A. (1990). Optimization of static simulation models by the score function method. Mathematics and Computers in Simulation, 32(4):373-39)
Ruiz, F. J., Titsias, M. K., and Blei, D. M. (2016). Overdispersed black-box variational inference. In Proceedings of the Thirty-Second Conference on Uncertainty in Artificial Intelligence, pages 647-656. AUAI Press.
Titsias, M. K. and Lázaro-Gredilla, M. (2015). Local expectation gradients for black box variational inference. In Advances Neural Information Processing Systems, pages 2638-2646.
Tucker, G., Mnih, A., Maddison, C. J., Lawson, J., and Sohl-Dickstein, (2017). REBAR: Low- variance, unbiased gradient estimates for discrete latent variable models. In Advances in Neural Information Processing Systems 30.
Williams, R. (1992). Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Machine learning, 8(3-4):229-255)
Xiao, H., Rasul, K., and Vollgraf, R. (2017). Fashion-MNIST: novel image dataset for benchmarking machine learning algorithms. arXiv preprint 178.077708
Xu, B., Wang, N., Chen, T., and Li, M. (2015). Empirical evaluation of rectified activations in convolutional network. arXiv preprint 150.00335.
Yin. M., Ho, N., Yan, B., Qian, X., and Zhou, M. (2020). Probabilistic Best Subset Selection by Muddient-Based Optimization. arXiv e-prints.
Yin, M., Yue, Y., and Zhou, M. (2019). ARSM: Augment- Seteeenn-RENNoORC---aaa-p-eeeeee estimator for gradient backpropagation through categorical variables. In Proceedings of the 36th International Conference Machine Learning.
Yin, M. and Zhou, M. (2019). ARM: Augment- REINFORCE merge gradient for stochastic binary networks. In International Conference on Learning Representations.


Reviews:
[{"review_id": "ymJ3jPbpGW", "reviewer": "Reviewer_Zj7z", "summary": "The paper concerns the estimation of gradients of expectations when sampling categorical random variables.\nThe work is closely related to recent works ARM and DisARM for binary variables, and ARS, ARSM for categorical random variables.\nARM was an estimator that samples a pair of binary random variables and couples them via antithetic sampling. This was achieved by sampling a continuous variable from a logistic distribution in an antithetic manner, and then transforming this continuous variables via a step function to obtain binary variables. The final form of the estimator in ARM still depended on the sampled continuous variable. DisARM noticed that this is unnatural, as in principle only the final binary variable should matter, and they proposed a Rao-Blackwellized estimator that integrates out the continuous random variable and obtains a strictly better gradient estimator compared to ARM. Both ARM and DisARM were meant for binary variables. \n\nARS and ARSM instead worked for categorical variables. The current work first proposes to Rao-Blackwellize ARS and ARSM leading to ARS+ and ARSM+. Then they propose 3 ways to extend DisARM to the categorical case: 1) they perform a new derivation of DisARM based on importance sampling, then perform an antithetic coupling of the categorical variables by constructing them via a stick-breaking process from binary variables, and computing the importance weight; 2) instead of computing probabilities in the categorical space, the problem is reparameterized into a vector of binary variables motivated from the stick-breaking construction, and the DisARM (recall that the original one could be used with binary variables) estimator is applied in this binary space together with a stick-breaking coupling (they also use knowledge of the stick-breaking coupling to further improve the estimator); 3) construct the categorical variable via a binary tree, and apply the DisARM estimator on these binary variables together with similar improvements as in the stick-breaking construction case.\n\nMoreover, they compared with REINFORCE with a leave-one-out baseline (RLOO). Previously, it was shown that ARM actually does not improve over the standard REINFORCE with a leave-one-out baseline, while DisARM did improve. In the current work, the authors showed that ARS and ARSM also do not improve over RLOO when the same number of function evaluations are used (5 evaluations for ARS and 11 for ARSM on a 5 category task), while the new categorical DisARM-IW, DisARM-SB and DisARM-tree methods did improve compared to a 2-sample RLOO estimator.\n\nThey performed experiments of training categorical variational autoencoders on binarized MNIST, FashionMNIST and Omniglot. The experiments showed a slight improvement in terms of learning speed and gradient accuracy (maybe around 25% -- 50% reduction in gradient variance compared to the 2-sample RLOO). They also tested the orderings of the binary variables in the stick-breaking estimator, and found that the dependence on ordering is inconsistent.", "questions": "", "limitations": "", "rating": 7, "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": "In general the quality of the work was good. The experimental analysis seemed fair, though I would have liked to see how the new estimators scale with the estimator sample size compared to RLOO. The paper presented a sensible technique for extending the coupling technique to categorical random variables. However, the improvement in performance seems small (25 - 50% reduction in gradient variance in practice is not a huge deal), so I am skeptical that that the method will make a major difference in practice. The clarity was mostly good, but could be improved by adding some more definitions, e.g. the antithetic binary variables are not properly defined, and some derivations in the appendix could be extended with the intermediate steps. Moreover, more related literature could have been added.\n\nCurrently, I gave it a 6 and may improve to 7 if some points in my review are addressed. \n\n**Experiments**\n\nI think the comparisons performed in the paper are fair, and it was nice to see analysis from multiple perspectives, e.g. the gradient variance, ordering of the categories, etc. However, for the ARS and ARSM estimator comparisons, RLOO used more samples than were used for DisARM, which makes me question whether RLOO would also not outperform DisARM when the sample size is increased. For example, currently DisARM uses one pair of samples (i.e. 2 samples), when we double this, we expect the variance to decrease by 2 times; however, the same is not true for RLOO---when we double the sample size for RLOO, the LOO baseline will now be using 3 samples, instead of just 1, so I would expect the variance of the gradient estimator to decrease by more than 2 times, due to the better baseline. Therefore, if the sample size is increased, it is plausible that RLOO will outperform the new DisARM estimators, even when matching the sample sizes between RLOO and DisARM. It would have been nice to see an experimental analysis of this. Of course, in this case one could try to construct a coupling between more than 2 samples to try to improve over the standard RLOO, and this could be left for future work.\n\n**Other literature**\n\nThe RLOO estimator was used earlier than Kool et al. (2019) by Mnih, A. and Rezende, D. (2016). \"Variational inference for monte carlo objectives\", and I find it inappropriate to only cite Kool. \n\nFor Rao-Blackwellization, the paper by Titsias et al.  (2015). \"Local expectation gradients for black box variational inference.\" was one of the first to use this method with binary random variables.\n\nRegarding importance sampling there are also several works, e.g.\nRuiz, F. J., Titsias, M. K., & Blei, D. M. (2016). Overdispersed black-box variational inference.\nParmas, P., & Sugiyama, M. (2021, March). A unified view of likelihood ratio and reparameterization gradients.\nJie, T., & Abbeel, P. (2010). On a connection between importance sampling and the likelihood ratio policy gradient.\n\nIn addition to the mentioned approaches to deal with discrete variables, biased estimators are also an approach that could be mentioned.\nSome methods:\nStraight-through estimators:\nYoshua Bengio, Nicholas L ́eonard, and Aaron Courville. Estimating or propagating gradients through stochas-tic neurons for conditional computation.\nPath-sample analytic estimators:\nShekhovtsov, A., Yanush, V., & Flach, B. (2020). Path Sample-Analytic Gradient Estimators for Stochastic Binary Networks.\n\n**Clarity**\n\nI think it would be useful to add:\n- A definition of antithetic binary random variables, and also explain how these are sampled and transformed from a logistic distribution.\n- What is SB? I was able to figure out that it means stick-breaking, but you\nshould explicitly define the abbreviation when you first use\n\"stick-breaking\".\n-Explicitly specify that z_k is an integer in [1,...,C]. Otherwise\none might assume a one-hot encoding.\n-In the equation between lines 78--79, please explain the notation in the\ntext.\n-And other points...\n\n----------------------------------------------------\nUpdate:\nThe authors added experiments looking at the scaling with number of samples. The performance improvement reduced, but the new estimator still outperformed RLOO. I think the paper is a good addition to the literature on coupled gradient estimators for discrete random variables, so I increased my score to 7.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6h"}, "extra_text": {}}, {"review_id": "k6C1iK0ZxSz", "reviewer": "Reviewer_kQGn", "summary": "**Summary**\n- The paper introduces a gradient estimator for discrete latent variable models. \n- The estimator is a natural extension of estimators that have been recently proposed for binary variables to categorical variables. \n\n**Edit after Author Response**\nThank you for taking the time to address all my concerns convincingly. I also read through the discussions with other reviewers. In light of your response, and the discussions I raise my score to 7. Even though, the empirical improvements are not very large, I think the insights, methods and experiments are a valuable contribution to the community and justify acceptance. This is a good paper. ", "questions": "", "limitations": "", "rating": 7, "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": "**Strengths**\n- Originality. The new derivation of DisARM/ U2G is insightful. \n- Quality. The submission is technically sound. Claims are well-supported, derivations are followable. \n- Clarity. The paper is well-written and has a clear, sensible structure. I enjoyed reading the paper. \n- Significance. Gradient estimators for categorical distributions are important. \n\n**Weaknesses**\n- Quality. The experimental section could be more extensive. \n- Significance. Experimentally, the improvements of the proposed method appear marginal to me. \n\n**Decision**\n- I like this submission and think it is above the acceptance threshold. \n- I wished the experimental improvements were larger and the experimental section was richer. \n- I also make some suggestions below which I think would be interesting to explore as part of the paper. \n\n**Suggestions and/ or questions**\n\n- Does the estimators favorable properties over RLOO arise solely from avoiding “duplicate” samples in expectation or is there more behind it? In the former case, it would be interesting to highlight, how often duplicate samples with RLOO arise in the experiments to highlight that it is a real issue. \n\n- Are there some settings where the gains over RLOO are expected to be larger? This would strengthen the experimental case for this method, to me the improvements appear marginal. \n\n- The model was only evaluated on categorical VAEs. If possible, I would like to see the estimators being evaluated in a RL setting, for example discrete-action policy optimization as in Yin et al. (2019), or on the structured prediction prediction task considered in various papers, e.g. Maddison et al. (2016).\n\n- The paper could also compare to other baselines, for example the unordered set estimator from Kool et al. (2020). In general, the case for multi-sample estimators could be made more strongly by comparing to single-sample relaxed estimators and keeping the number of function evaluations constant. However, I do think that RLOO is a strong baseline to compare against. \n\n- Would it be possible to do hypothesis testing (t-test) to highlight the statistical significance of the improvements of the proposed estimators over the baseline? Statistical significance of the improvements of the method is a concern I have, and this will hopefully highlight that the improvements are significant.\n\n- The authors use Rao-Blackwellization as in Dong et al. (2020) to improve their gradient estimator. Rao-Blackwellization has also been used for improving other gradient estimators for discrete distributions, for example in Liu et al. (2018), Kool et al. (2020) or Paulus et al. (2021), and I think it's worth highlighting that in related work.\n\n**References**\n- Liu et al. (2018). Rao-Blackwellized stochastic gradients for disccrete distributions. \n- Kool et al. (2020) Estimating gradients for discrete random variables by sampling without replacement. \n- Paulus et al. (2021). Rao-Blackwellizng the Straight-Through Gumbel-Softmax estimator. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "fegOLICP4Y3", "reviewer": "Reviewer_WpYg", "summary": "Training models with discrete latent variables is an important problem in machine learning. Among previous methods, REINFORCE with a leave-one-out-baseline (RLOO) outperforms Augment-REINFORCE-Merge (ARM) in the binary setting; DisARM/U2G (independently in Dong et al. (2020) and Yin et al. (2020)) outperforms RLOO, but is limited to binary random variables. This work develops a novel derivation of DisARM/U2G based on importance sampling and statistical couplings, which is simpler and more direct while providing a natural extension to the categorical case. Gradient estimators based on reparameterizing categorical variables as sequences of binary variables and Rao-Blackwellization are obtained. Experiments show the effectiveness of the proposed method.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "Strengths and weaknesses:\n\nThis paper extends DisARM/U2G, which is limited to the binary setting, to the categorical setting. It appears to be a new work.\n\nThe paper is generally well-written. But there are a few places, which need further clarification, as detailed below.\n\nA main concern is about performance evaluation and missing of related work. It has been shown in [a] that the joint stochastic approximation (JSA) method significantly outperforms a number of methods including VIMCO, Gumbel-Softmax and ARM/ARSM in both binary and categorical settings for learning discrete latent variable models. Evaluation experiments in [a] are conducted by training VAEs with discrete latent variables, which is also taken in the current paper. The model structures evaluated in the current paper are also similar to [a], which follows from the work of ARSM (Yin et al., 2019). It seems that the proposed method in the current paper is more complicated than [a], in both methodology and implementation. From a quick reading, the results reported in Table 1 are much worse than the Table 2 in [a]. Further, since the proposed method uses the Stick-breaking construction, the cardinality of the discrete latent variable needs to be the integer power of 2. Comments and comparisons are better to be added in this paper.\n\n[a] Z. Ou, Y. Song. Joint Stochastic Approximation and Its Application to Learning Discrete Latent Variable Models. UAI, 2020.\n\nOther problems:\n\nBetter to make clear the dimension/size of z, b, etc.\n\nLine 51: \"it is straightforward if more challenging to apply the ideas of Dong et al. (2020) to ARS and ARSM\"\nnot clear\n\nThe equation under Line 78 contains typo?\n\nFor the equation under Line 95, \\tilde{z} is used without definition.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "3PMlrn7bl7", "reviewer": "Reviewer_tJ9n", "summary": "This paper presents the multiple methods of estimating the gradients of discrete random variables. Particularly, they expand DisARM by adding the variations of importance sampling, stick-breaking construction, and tree-based decision makings. The performance is compared to previous methods, such as RLOO, with a standard benchmark dataset in VAE settings. ", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "1)\nGiven DisARM/U2G, the extension with the importance sampling in Eq 3 seems to be straightforward, and this variant is not congruent with the two other variants with SB and Tree. I think that this could be regarded extension of DisARM/U2G rather than considering this as a separate original version.\n\n2)\nMost development has been made to construct the coupling between $z_k$ and $\\hat{z}_k$. Utilizing the stick-breaking prior in this regard has been considered in the previous literature, as well, though they are not in the suggested form and purpose [1]. Anyhow, this could be briefly mentioned.\n\n3)\nThe two variations looks essentially same mechanism of creating the mapping function of $h$. Their differences are subtle, and it is worth enough to create a diagram explanation on their abstrate support sets. Eq 4 and Eq 5 are all same except the conditioned regions. These support regions could be visualized in both variations to express the difference between the two variations. Or, is it that different? I wonder.\n\n4)\nExperiments are not satisfactory. There have been many papers in this gradient estimations. The large categorization would be score-based and reparameterization-based approaches. There has been models like NVIL, REBAR, RELAX, STORB, etc. (These are easily searchable through acronym in google) Currently, RLOO is the only baseline in Table 1 and 2. Moreover, the original version of DisARM/U2G can also be compared in Table 1 and 2. \nDatasets are also limited to Omniglot and MNIST, which could be understandable given the prior papers with the same practice, but it seems that this line of work requires further expansion in the dataset, as well.\n\n[1] Nalisnick, Eric, and Padhraic Smyth. \"Stick-breaking variational autoencoders.\" arXiv preprint arXiv:1605.06197 (2016).\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Have the authors adequately addressed the limitations and potential negative societal impact of their work? If not, please include constructive suggestions for improvement.\n\n\nThere is no discussion at all. I understand that authors may not discuss the societal impact, but they can surely include the limitations that they expect from their models.\n\n\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}]

Author rebuttal:

Meta-review: Four knowledgeable reviewers recommend accepting this submission. I agree. This submission makes a valuable contribution by proposing new inference techniques for models with discrete latent variables.
Decision: Accept (Poster)