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 -Mrree (ARM) (Ynn and Zhou 2019) estimator for binary variables and 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()) logooo( + 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) - ((3) (Vo, log (z())) - Va, Ax 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)))((²z))
o( o(-ar,)) ( (-amiu(zu,saa) 1 1- (=2a))
II
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-brea 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)]] ELs 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 implies 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.,,)....,.
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, 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 Attriuttonn 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
-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
SO0K
250K.
300K
3006 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-breaking 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. (0211).
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-RRENNFOREEE.M 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-2555
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 1505.0055...0
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---aaaap-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": "mGewEAJqu3F", "reviewer": "Reviewer_AbN2", "summary": "This paper considers the problem of designing a gradient estimator for categorical latent variables that is unbiased and does not depend on continuous relaxation. It proposes three different types of such estimators (called DisARM-IW, DisARM-{SB,Tree}, and {ARS,ARSM}+). DisARM-IW is based on a new derivation, which uses a coupling, of a previous estimator called DisARM; DisARM-{SB,Tree} are obtained by reparameterizing with binary variables and then applying DisARM; and {ARS,ARSM}+ are obtained by Rao-Blackwellizing previous estimators called {ARS,ARSM}. DisARM-{IW,SB,Tree} are shown to outperform a state-of-the-art estimator called RLOO on several benchmarks.\n", "questions": "", "limitations": "", "rating": 6, "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 Review\n\nThe paper considers an important problem and proposes new methods, some of which outperform the state-of-the-arts. One of the proposed methods (DisARM-IW) looks novel to me, but the other methods look less novel. The paper is technically sound. The paper needs more improvement in writing.\n\n## Strengths\n\n(1) This paper considers an important problem and proposes several solutions. Among the proposed solutions, DisARM-IW and DisARM-{SB,Tree} are shown to outperform a state-of-the-art estimator.\n\n(2) DisARM-IW is novel to me in that its derivation looks new and it can be seen as a generalization of a previous estimator called DisARM. \n\n## Weaknesses\n\n(1) The estimators other than DisARM-IW look incremental and much less novel to me. DisARM-{SB,Tree} are obtained by applying reparameterization based on the stick-breaking or tree construction and then simply passing the reparameterized form to the well-known estimator DisARM. {ARS,ARSM}+ are obtained by simply following an approach similar to DisARM. I think little novelty in DisARM-{SB,Tree} is less problematic as they outperform RLOO.\n\n(2) The writing needs more improvement. First, even though an antithetic Bernoulli coupling is used importantly in several parts of this paper, it is not explained at all in detail. Second, it was difficult to grasp the overall picture of, and connections between, proposed estimators. I was misled that DisARM-IW and DisARM-{SB,Tree} are similar types of estimators, which is not the case: the former does not use DisARM at all, while the latter uses DisARM importantly. In addition I felt that {ARS,ARSM}+ and Sec 3.2 are completely separate from the other estimators DisARM-* and Sec 3-3.1. This makes the paper less coherent and makes it difficult to understand the main point of the paper.\n\n## Questions and comments\n\n- Line 127-128. p(b_k,i) > 0.5 should be p(b_k,i=1) >= 0.5 (note >= instead of >)? If this is the case, isn’t it a problem if q(z_k=i) / \\sum_{j=i}^C q(z_k=j) = 0.5 (Line 131), as this implies p(z_k>i, \\tilde{z_k}>i) = 0 and violates the support condition?\n- Line 106. Include the proof or calculation (to text or appendix) showing that Eq (3) recovers DisARM for a particular coupling, as this is one of the major contributions of this paper.\n\n- Line 78-79. \\sigma is undefined.\n- Line 86-87. j is undefined.\n- Line 133. \\sum_{j=i+1}^C … should be \\sum{j=i}^C?\n\n-----\n**Updates after the author response.** Thank the authors for answering my questions. The response gives a more clarified view of the paper. I am still positive on the paper and will keep the same score. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations and societal impact are discussed in the paper.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "jI-_o-7ssHV", "reviewer": "Reviewer_ozdD", "summary": "The authors consider optimization problems that can be formulated as maximizing parameterized functions that can be expressed as an expectation of a function f with respect to some factorial catergorical distribution.  These kinds of optimization problems arise naturally in standard mean-field ELBO optimization problems.  The proposed approaches aim to go beyond existing work by introducing a stick-breaking based coupling technique to reduce the variance of the sampling based estimator.  This is done by first presenting an importance-sampling characterization of DisARM, from which the newly proposed techniques are derived.", "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": "Originality:  While the work does build on top of existing work, the perspective is new.\n\nQuality:  1) The claim that, \"In the case of nonlinear models (Figure 1), all three proposed estimators perform similarly, with lower\n240 gradient variance and better performance than the baseline estimator (RLOO) across all datasets,\" does not seem to match the figures, e.g., not on Omniglot.\n\n2) The data sets selected for comparison are all quite similar.  It is unclear to me that similar behavioral trends would hold for a wider variety of data sets.  Can the authors explain if they looked at other data sets and why these data sets were chosen?\n\n3) The Rao-blackwellized portion of the experiment section feels like a bit of an afterthought and most of it is relegated to the appendix.  I think I'd prefer a more detailed study of disarm variants.\n\nClarity:  1) While this paper is a bit out of area for me, it is incredibly difficult to read as many assertions are made without concreate explanations or citations of why something is true, e.g., \"Monte Carlo estimates of the first term can have large variance.\"  This should have a citation.  There are many more examples of this, too many to reasonably include here which greatly limited my ability to evaluate and certify the results.\n\n2) The whole discussion of related work seems like a bit of a mess to me as the paper bounces back and forth from discussing related work, a background section on related work, a section that builds on related work, and then a separate related work section. I think these sections can be combined and streamlined so as to minimize repetition of information and allow for additional explanations where needed.\n\nSignificance:  The performance improvements of the new approach range from mild to moderate, but the perspective may be of interest more broadly.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "In my opinion, no detailed discussion of limitations is provided.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 hours"}, "extra_text": {}}, {"review_id": "XvrZNNKD9O", "reviewer": "Reviewer_u8r7", "summary": "The paper proposes three versions of DisARM/U2G estimators.\nThe authors extend the binary version to the categorical random variables, and the extentions are in the importance sampling sense (IW) and stick-breaking coupling perspecitve (SB and Tree).\nThey also verify their method through VAE experiments.", "questions": "", "limitations": "", "rating": 4, "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": "Originality:\nExtending the idea of binary random variable gradient estimator to the categorical case is natural.\nAlso, constructing categorical variable through stick-breaking or tree construction is also a natural idea.\nHowever, combining those two ideas in the research on the gradient estimotors for the categorical variable seems novel.\n\nQuality:\nThe paper seems technically sound.\nThe appendix substantiate the unbiasness of the proposed gradient estimators.\n\nClarity:\nI can't say that the paper is well-written (it seems that it can be written in some better way), the paper is written in relatively clear manner.\n\nSignificance:\nMy concern on this paper is in two parts.\nThe first one is that the experiment is too limited.\nThe baselines are limited and the experimental settings, where the reparameterization estimators are excluded and a lot of score function estimators are missing. (including UnOrd [1])\nAlso, the experimental task is limited to VAE case.\nAs the research on the gradients over categorical random variables grows, more extreme experiments are gradually required to differentiate the effect of the proposed methods.\nI suggest authors demonstrate the idea in the broader experimental settings.\nThe second concern is that the method (or the implementation) is more complex than the sole baseline, RLOO, but the performance gap is too marginal.\nI guess the authors should provide an example that the proposed complex method gives a clear benefit compared to the baseline.\nTherefore, I lean my score to reject.\n\nQuestions:\n- How does the running time of each model as the number of latent categories changes? (compared to RLOO) Especially, the stick-breaking construction might take a long time.\n- From the learning curves in Fig 1 and 2, it seems that the models are not reached the optimal point, and it seems they still have a learning chance. Am I missing something?\n- From the right-most figures of Fig 1, the variance of gradients reduces up to a certain point and keeps increases later. What causes such curve shapes?\n- How does the performance change (including RLOO) if one sets the extreme prior in VAE?\n\n[1] Estimating gradients for discrete random variables by sampling without replacement., Kool et al., ICLR 2020.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors argue that they pointed out the limitation of their work in the conclusion (but I can not find it), and there is no negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "UDH59Yz7rWE", "reviewer": "Reviewer_KiRx", "summary": "This problem addresses the problem of differentiating expectations of discrete random variables. The authors proceed from a derivation of DisARM that applies importance sampling to the RLOO gradient, extending this estimator beyond 2 categories. A core component of this new multi-sample gradient estimator is a coupling which leads to (negatively) correlated samples of the categorical. The main proposed coupling uses stick-breaking. The authors propose two estimators that do not use importance sampling, but extend DisARM to the categorical case by reparametrising a categorical using binary random variables. This approach is amenable to standard DisARM. The authors consider two such binarizations: first using the stick-breaking construction again, second using a decision tree. Finally, the authors propose ARS+ and ARSM+ plus, which are Rao-Blackwellized versions of ARS and ARSM. Experimental evaluation focuses on categorical latent variable VAEs.", "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": "Novelty/related work\n-----------------------------\nThe work constitutes a natural extension of existing work, focusing on moving from binary to categorical random variables. Regarding related work, a piece that seems to be missing is ARMS [1]: a recent (ICML 2021) paper that generalizes RLOO and DisARM, but remains focused on the binary case. To me, the key mathematical contributions seem to be\n1. Extending antithetic coupling based methods to categoricals (DisARM-IW)\n2. Proposing an explicit coupling using stick breaking (DisARM-IW)\n3. Alternatively, two binarizations of categoricals that can then use vanilla DisARM (-SB and -Tree)\n4. Rao-Blackwellization of ARS and ARSM to give ARS+ and ARSM+\n\nWith all these different variations floating around, I do think the authors should do a little bit more work to clarify the relations between these different methods. For instance, it wasn’t 100% clear to me whether there is any relationship between ARS+ and ARSM+ and the DisARM-* methods. My intuition was no, because the DisARM-* methods don’t use any continuous augmentation that could be integrated out. I think clarifying this would be very helpful, perhaps with some kind of table.\n\nQuality\n---------\nThe paper is of good quality. Although I haven’t read into the appendix in great detail, I have good confidence in the underlying mathematics of the approach. One thing that could have further grounded this work is stronger theoretical analysis. When considering the properties of a good coupling, the authors say (112-114) “moving mass away from this configuration will reduce variance. Furthermore, for the estimator to be valid, the coupling must put non-zero mass on all $z_k \\ne \\tilde{z}_k$ configurations.” But this is not followed up by more careful analysis. For instance, the variance of the new estimator is not calculated / bounded. Furthermore in 279-280 “DisARM-IW is simpler to implement, more natural to understand, and easier to generalize, we recommend it in practice”. This kind of comparison between different proposed methods would again have been stronger with some kind of analysis or calculation of the actual variance, to guide a selection of the right estimator. \n\nOn the empirical side, I would say the authors have been fairly detailed here. They have looked at three sensible datasets, and compared against suitable baselines. It was a tiny bit confusing to separate out the comparison as DisARM-* versus RLOO, and secondarily RLOO versus ARS(+) and ARSM(+). But I think the logic makes sense here: the estimators in the first group only require 2 function evaluations. One direction that the experiments could have been pushed is to consider other models beyond the VAE. Could these estimators be applied to belief nets, or would be number of discrete latents be problematic?\n\nOn strengths and weaknesses, one thing that could be discussed in more detail is the relative computational cost of all of the estimators being discussed in this paper (as there are many). For the DisARM-Tree estimator, is there any cost to calculating the tree probabilities (if C was extremely large, perhaps this becomes problematic)?\n\nClarity\n---------\nI found the paper to be well-written. However, there are some places where I think the authors could be clearer. As mentioned previously, a more precise and self-contained comparison between the different estimators proposed in this paper (plus some baselines) would help. I understand that this information is present in the paper, but it can be hard to locate. For instance, a table with concrete information such as complexity, continuous augmentation (y/n) and so on could be a real asset. One thing that should be stressed is the distinction between DisARM-IW with a stick breaking coupling and DisARM-SB, there is definitely a possibility for confusion here. \n\nSignificance\n----------------\nThis paper deals with an important area of ML research. It makes some sensible and logical next steps in the development of low variance estimators for discrete random variables. The empirical gains are not colossal, but they do provide new state of the art performance. I would say this paper is definitely a valuable contribution, without being a huge leap forward.\n\n[1] Dimitriev, Alek, and Mingyuan Zhou. \"ARMS: Antithetic-REINFORCE-Multi-Sample Gradient for Binary Variables.\" arXiv preprint arXiv:2105.14141 (2021).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I think the authors have addressed wider societal impact adequately. On limitations, see my main review.", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "00mo_7lLZmr", "reviewer": "Reviewer_Msvo", "summary": "The paper reformulates the DisARM/U2G (Dong et al., 2020; Yin et al., 2020) gradient estimators using importance sampling. While these estimators are originally applicable to binary latent variables, the new formulation allows for the extension to the categorical latent variables case by making use of couplings (common random numbers). The paper also shows that another family of estimators ARS & ARSM (Yin et al., 2019) is not performant compared to RLOO (leave-one-out REINFORCE) (Kool, et. al, 2009) even when its variance is reduced further with Rao-Blackwellization.\n", "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": "**Originality:** While the paper builds on previous work, the reformulation of DisARM/U2G (Dong et al., 2020; Yin et al., 2020) with importance sampling and the use of couplings are novel. The experimental evaluation is standard for this type of problem and uses the usual datasets and models. The paper does a decent job at positioning its contributions within the literature and differentiates itself from its direct comparisons; however, it still misses some relevant references, e.g. Lee et al. 2018, Peters and Welling 2018, Shayer et al. 2018, Cong et al. 2019, Richter, et. al. 2020.\n\n**Quality:** I think this work is technically solid. The mathematical derivations seem to be correct to me (although I did not verify all of them in detail) and the motivation was clear. Although, I think that the paper is lacking some more explanation and intuition on why this type of estimator reduces the variance (I had to go to Owen 2009 to get some intuition). \n\nOne thing I was not able to verify is how computationally expensive this estimator is? In the experiments section, the authors state that: “These estimators require only 2 expensive function evaluations regardless of the number of categories (C) ”. How expensive is this compared to the 2-sample RLOO estimator?  Does the SB and Tree coupling constructions have any additional overhead? I think these questions should be addressed in the experimental section. Currently, the experiments only show the ELBO with respect to the training steps. I think a similar comparison needs to be performed but with respect to wall-clock time to show the effectiveness of this estimator in practice. \n\nFurthermore, it is also good to include a comparison to other popular gradient estimators for categorical variables such as REBAR (Tucker, et. al. 2017) and RELAX (Grathwohl, et. al. 2018). While these estimators are expected to perform well because of their adaptivity, they are quite expensive to compute. I think that benchmarking w.r.t. wall-clock time might reveal that DisARM-IW is competitive in comparison.\n\nFinally, I found the benchmarking of ARS+ & ARSM+ separately to be strange, since the main thesis of this paper is that these do not close the gap to RLOO motivating the need for DisARM-IW. I think the comparison with DisARM needs to be direct in this instance. Please see further comments below on this point.\n\n**Clarity:** I think the paper is well-written and the technical details are mostly clear. However, I found the discussion of ARS & ARSM a little bit arbitrary. These estimators are only considered in the experiments and in section 3.2 on their Rao-Blackwellization. \n\nFurther, these are not compared to the contributed estimator (DisARM-IW/SB/Tree). I understand that their inclusion is to demonstrate they are not as effective in reducing the variance as the proposed estimator or RLOO, but couldn’t this comparison be made directly? At present, this paper reads as two distinct contributions (A major one: DisARM-IW and a minor one: Rao-Blackwellization of ARS & ARSM) in a single manuscript. I think the manuscript will be much clearer if the discussion of ARS & ARSM is integrated better in the exposition (including the experiments). Alternatively, this discussion could be demoted to the appendix or submitted as a workshop paper in the future.\n\nSince this paper is relatively heavy on notation, I would suggest the following three points to improve its clarity:\n* Needed: An algorithmic description (pseudocode) of DisARM-IW and the computation of the different couplings (some of this can be added to the supplement).\n* Nice to have: following the standard notation of representing vectors in boldface.\n* Nice to have: definitions of undefined symbols, e.g. $1_C$.\n\n**Significance:** I think this paper provides a significant contribution to a few ML communities (mainly probabilistic modelling and RL). The introduction of the idea of coupling for variance reduction is an important contribution to these communities and can create many avenues for more research into variance reduction for Variational Inference and Policy Gradients. However, the impact of this paper is hampered by the limited experimental evaluation. That being said, I am willing to reconsider my evaluation based on the authors’ rebuttal/clarifications and the discussion with the other reviewers.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "There was not much discussion of the technical limitations of this work in the manuscript. It would be great if the authors can comment on those, especially if there is any extra computational overhead.\n\nThe potential societal impact of this work was sufficiently considered.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The submission is concerned with low-variance gradient estimators for discrete random variables, extending recent work on binary to categorical. The reviewers agreed the work here is novel and interesting, but shared concerns around the lack of baselines and clarity.  The authors promised to improve clarity, and add more experimental results [but did not post these in the rebuttal]. An additional expert review was sought and the AC agrees with the expert reviewer that the submission’s overall contribution is enough to justify acceptance. The authors are strongly encouraged to improve the writing and add more experiments, as promised to the reviewers.
Decision: Accept (Poster)