CARMS: Categorical-Antithetic-REINFORCE Multi-Sample Gradient Estimator
Authors: Alek Dimitriev
Mingyuan Zhou
Venue: NeurIPS 2021 Poster (2021)
Keywords: Discrete latent variables
gradient estimation
antithetic
unbiased
low variance
probabilistic modeling
variational autoencoder
copula

Abstract: Accurately backpropagating the gradient through categorical variables is a challenging task that arises in various domains, such as training discrete latent variable models. To this end, we propose CARMS, an unbiased estimator for categorical random variables based on multiple mutually negatively correlated (jointly antithetic) samples. CARMS combines REINFORCE with copula based sampling to avoid duplicate samples and reduce its variance, while keeping the estimator unbiased using importance sampling. It generalizes both the ARMS antithetic estimator for binary variables, which is CARMS for two categories, as well as LOORF/VarGrad, the leave-one-out REINFORCE estimator, which is CARMS with independent samples.  We evaluate CARMS on several benchmark datasets on a generative modeling task, as well as a structured output prediction task, and find it to outperform competing methods including a strong self-control baseline. The code is publicly available.

CARMS: Catepriial-Aniihee-RREEENOOO Multi-Sample Gradient Estimator
Alek Dimitriev
Mingyuan Zhou
McCombs School of Business
McCombs School of Business The University of Texas at Austin Austin, TX 78712
The University of Texas at Austin Austin, TX 78712
alek. dwwitriee@@ecconteeeiiii utexas, edu
wwngguan..zhuu@mccommbb utexas edu
Abstract
Accurately backpropagating the gradient through categorical variables a chal- lenging task that arises in various domains, such training discrete latent variable models. To this end, we propose CARMS, an unbiased estimator for categorical random variables based multiple mutually negatively correlated (jointly anti- thetic) samples. CARMS combines REINFORCE with copula based sampling avoid duplicate samples and reduce its variance, while keeping the estimator unbiased using importance sampling. It generalizes both the ARMS antithetic estimator for binary variables, which is CARMS for two categories, as well as LOORE/VarGrad, the leave-one-out REINFORCE estimator, which is CARMS with independent samples. We evaluate CARMS on several benchmark datasets on generative modeling task, as well as a structured output prediction task, and find it to outperform competing methods including strong self-control baseline. The code is publicly available
1 Introduction
When optimizing an expectation based objective of the form (f)= we sometimes require the gradients with respect the parameters & of the distribution. This challenging for discrete variables, because the commonly used reparameterization gradient does not directly work, unless the discrete distribution is approximated by a continuous and reparameterizable one [Jang al. 2017 Maddison al.) (017]. A significant part of this field has thus been score function (REINFORCE) based estimators [Glynn 1990 Williams 1992 Fu 2006), which are general and do not require to be the differentiable. le In this paper, we focus the case when z is high dimensional categorical variable with logits Φ. An common example this form the evidence lower bound (ELBO) [Jordan et al. 1998). which arises in variational inference and is used for training variational autoencoders Kingma and Welling 2014 Rezende et al. 2014]. Because their latent space consists of large number of categorical variables, they require Monte Carlo gradients with respect to the parameters the stochastic distribution, but have excellent performance, as a categorical VAE has in practice achieved state-of-the offtthe-r art zero-shot image generation [Ramesh et al. 2021
Our main contribution is a novel unbiased and low-variance gradient estimator for categorical watprraa--nnniieeee-EEIIIRRREETT
variables. The to generate any number of antithetic (mutually negatively correlated) [Owen] 2013) categorical samples, and constructs an unbiased estimator by combining them into baseline for variance reduction. This approach is inspired by the ARMS estimator [Dimitriev and Zhou 2021], which also uses multiple antithetic samples, but only works for binary variables. For two categories, CARMS
(CARMS) estimator uses copula
http:://ggiaabb.ccc//aaaaaaaaaaa
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

reduces to it, while for independent samples, CARMS reduces to the leave-one -wt-oone-------OO--- (LOORF) estimator. Our approach achieves higher ELBO with VAE models than the state of the art, as well as higher log likelihood for conditional image completion.
Related work One widely used group of gradient estimators for categorical variables is based on trading off bias for lower variance. This includes the straight through (ST) estimator [Bengio et al.. 2013], the direct argmax [orberrom et al. 2019], as well as the concurrently developed equivalent Gumbel-Softmax (GS) [Jang et al. 2017 or Concrete [Maddison al., [017] that uses a continuous relaxation. These were further improved by combining them with REINFORCE to obtain unbiased estimators, which includes REBAR [Tucker et al. 2017]. as well as RELAX [Grathwhll et al. 2018). which uses a free-form neural network.
In practice, REINFORCE is almost always augmented by baselines, e.g.. in variational infer- ence [Mnih and Gregor 2014, Ranganath et al. 2014, Paisley et al. [012 Ruiz et al.; 2016
Kucukelbir al. 2017 I. MuProp [Gu et 2016 uses a first order mean-field Taylor approx- imation, but requires f to be differentiable. Other estimators apply Ahoo-Backkwellizttion [Liu et al.
2019 Kool et al., 2020). e.g., to one latent variable at a time (Titsias and Lázaro- Gredilla, 2015 to the ST estimator Paulus et al. 2021), or to the reparameterized Dirichlet vector in ARSM [Dong al. 202T]. Besides REINFORCE and reparameterization type gradients, there is also the measure valued gradient [Rosca et al.) 2019] and finite differences [Fu (006), but are less common in practice. comprehensive review can be found in Mohamed et al. [202
The more recent approaches for categorical variables are unbiased and REINFORCE based, building on the idea of using multiple samples construct a baseline for variance reduction. One such baseline is LOORF KKool et al.] [209a]], originally introduced in Salimans and Knowles [2014) and also known VarGrad [Richter al. 2020]. where its theoretical properties are further analyzed. has a similar form, but is specific to the importance weighted multi sample bound [Burda al. 2016). A different approach is ARSM (Yin et al. 2019), which reparameterizes the gradient with a Dirichlet distribution and uses swaps to obtain mul multiple correlated samples. However, it adds some amount of variance due to the continuous reparameterization, and can require up to C(C - 1)/2 function evaluations per step, which can be computationally expensive. More recently, the unordered set estimator (UNORD) [Kool et al. 2020) uses the Gumbel top-k trick to sample without replacement, and then constructs a baseline using a multiplicative term to preserve its unbiasedness.
VIMCO MMiih and Rezende,
2 Background
Let z = (z1...., zD) denote D independent categorical variables, where za is a one hot encoded categorical sample from Za ~ [ (Za) = ((a)), with (xx); =
Σn ext being the softmax function. Let also i be the basis vector with its 5h coordinate set to one, and otherwise zero, such that P(= = a (Φ)). or equivalently E[z] (Φ)). Unless otherwise stated, superscripts denote the dimension, and subscripts denote different samples, with vectors and matrices being bold lowercase and bold uppercase symbols, respectively. We are interested in optimizing the following objective with respect to the logits & (Φ1...., .... ΦD):
c:: =
D
(())= = E [f(]],, (z(2)= = II
qou (za).
Although the score function gradient contains two terms:
VoL(Φ) E (x)
Vos(z) (et (n
the first term is easily estimated, so we omit the subscript in f for notational clarity, and we focus on the latter term in this work. Since CARMS generalizes the multisample LOORF estimator beyond independent samples, we review LOORF here. We also review the ARMS estimator, which uses a copula to generate antithetic samples, but is restricted to binary variables.

---

2.1 LOORF
Lawc-one--uURRIINFRREEE (LOORF) (Salimans and Knowles 2014 Kool al. 2019a]. also known as VarGrad Riicher et al. 2020 is a general score function based gradient estimator that uses N i..dd samples to construct baseline for variance reduction, and competitive with state of the art estimators. Its theoretical properties have recently been analyzed in Richter et al. (2020), where the same estimator results from minimizing the variance of the log ratio between the posterior and approximating distribution in variational inference. The only requirements for LOORF are the ability to sample z ~ qs and evaluate f(z). Given N samples 21, NN
sid Cat(or it has the form:
N Σ
N
(f(za) - [(z)) V go)) = - Σ
(((x)) (2x a(H)),, (b))),, (1)
1 N-T
9LOORD
==1
n=1
where [(z) - N I It is commonly used due to its simplicity and strong performance.
2.2 ARMS
The Antithetic- REINFORCE-
(ARMS) [Dimitriev and Zhou 2021] estimator is a recent
work that uses any number N of mutually negatively correlated samples. However, although unbiased, is limited to binary variables, which motivated us to extend it to the categorical case. Since any antithetic copula in two dimensions reduces to the pair (u, 1 - u), it generalizes DisARM [Dong et al., 2020], independently discovered as unbiased uniform gradient (U2G) (Yin et al. 2020)] which 2 samples. ARMS achieves this generalization by using a copula [Trivedi and Zimmer
2007
which is any multivariate distribution whose marginals are uniform random variables:
u (ui,.... uN) CN, Vi: U4 ~ Unif(0, 1).
ARMS uses a Dirichlet or Gaussian copula with strong negative dependence between each dimension of u. However any copula can be used instead, with the only two requirements being the ability to generate samples easily, as well as being able to evaluate the bivariate CDF o (u,, u,) of the copula, so that the debiasing term can be calculated. Given a copula sample u, it uses inverse CDF sampling to convert it into N antithetic Bern(p) samples, which are simply b; = I <p; Vi. For set of N antithetic D-dimensional Bernoulli variables b bi.... by with probabilities (()), the -amppl unbiased estimator has the following simple form:
N
N Σ f(bm) m=1
bn (())
JARMS N
f(bm))
(2)
1 1-p
n=1
with p = (PI. PD) and Pa corr(bf bg), which simple to compute given the bivariate CDF of the copula.
3 CARMS
The CARMS estimator has similar form to LOORF, but requires a multiplicative term to remain unbiased. It has two requirements: an easy way to sample antithetic categorical variables, and being able to compute the bivariate probability mass function (PMF), which should identical for any pair. For clarity, we begin by assuming the ability to do this, and derive the univariate version for two samples, which we then extend to samples. Next, we generalize CARMS to any number of categorical variables. Lastly, in Section 3.2 show two different ways of sampling that satisfy both conditions: inverse CDF sampling with an easily computable analytical bivariate PMF, and the Gumbel max trick combined with an empirical estimate of the PMF
The two sample version of CARMS can be obtained by replacing the two i.i.d. samples in 2-LOORF with an arbitrarily correlated pair of categorical variables and an added debiasing term. For N = 2 samples 2, z' (Ø)) LOORF has the following simple form:
gL.ORF(Z2 x²) 2 (f(z) f(x')) z²).
(3)
3

---

which, importantly, is unbiased [Kool et al. 2019a Dimitriev and Zhou 2021]. This means that using a simple importance weight preserves its unblasedness, for an arbitrary bivariate categorical distribution (z,z") ~
P( P( = 5)
P(= P(==i)P((Z = 2 = 9LOORF (?.
E 9LOORF 22)
P(z = 2, = = = Σ
Σ P(z= (opLOORr((.)) E x, Catlo (o)) (OLOORr(Z, z') VoL(Ø).
We summarize the derivation of the two sample version of CARMS, which we denote as the Calegorical-Antihhtti--- Anttettc-iicciCCEEERRCRCCCCEEC Two-Saplee (CARTS) estimator the following theorem.
Theorem Let z') be a sample from an arbitrary bivariate Categorical distribution with marginal distributions x' ~ Cat(o(o)), and known bivariate PME, An unbiased estimator of V (&Bs(()) is:
)()); 1 ()))) ((z))) - - Rij = p(z=i,z'=0) SCARTS 2') 2
(4)
The intuition behind using negatively correlated variables is that we want to avoid the case when x², because the sample is then "wasted." We formalize this intuition below, and defer the proof to the Appendix.
Theorem 2 Let (z,z') Cat(a
the bivariate PMF satisfies;
Vi Vi j P(z = =, 2 P(z = P 5), with strict inequality for at leust one pair, then Var[gCARTS(2. = VarlgnooRFFZ, =')].
We now extend CARTS to N samples, using the following identity, the proof of which can be found in Dimitriev and Zhou 2021 It states that N N-sample LOORF is equivalent to averaging 2-sample
LOORF over all pairs:
N
N Σ f(zm) m=l
(LOORF(Z1, 2N) =
Σ n=1
f(Zz)
(===)()))
N
N
(s(z)) (((m)))(Zzm))
)zn 2m) = N(N-I)
GLOORF En, Zm)-
N(N-1)
nzm
With the above identity it easily follows that given antithetic categorical samples = [21, xN IN. applying CARTS to all pairs results in an unbiased estimator, due linearity of expectations:
Σ SCARTS( Zn. Zm, R)
E. [SCARMSs22 Z2 R)]
E N(N
BLOORF (Zn, 2m)
N(N
17
nfm
nfm
EE [9LOOR
[)] = V C(Ø).
We summarize CARMS in the next theorem and also rewrite it in simpler matrix form used in our implementation. The matrix form can be obtained after some algebra, which can be found in the Appendix.
Theorem 3 Let Z [21,..., 2N be sample from arbitrary N-variate Categorical dis- tribution with identical marginal and bivariate distributions, such that z; ~ Cat(a(()), and Rij = ( P(zn. i 2m 5). An unbiased estimator of V (f(((() is:
gCARMS ( N(N-1
Σ (1(En)
2m)² R RT
Lemma Let f(2) (f(21), .... s(2N )) denote the Hadamard product, 1 NNNN matrix of
ones, and I the identity matrix. Define O x- (1NNN IN) ZRZ"), and D diag(O11

---

to be diagonal matrix. The CARMS estimator can equivalently be written in the following form:
9CARMS Z - (2 (2-1xo(0)))
(5)
Furthermore, for independent samples, ZRZ = =1NxN and LOORF has the form:
9LOORF 2) 1(2) (INX)
-INNN-N-> - (INAN 1x) (z (xT())))
(6)
3.1 Multivariate CARMS
In the univariate case, Monte Carlo estimators are not necessary, as the expectation has C' terms and can simply be analytically summed, but for many categorical variables, stochastic gradients required. For dimensions, we can sample Z ~ (at(or())) independently and combine them into D N x C tensor Z. with the corresponding D x C C importance ratio tensor R. Below, we use superscripts and subscripts to index the first and second dimension of the tensors, with dd denoting excluding the 4 row of Z Focusing the d dimension, we have:
V ab 15 [f(Z)]
= Bz
V SulEzz [f(z d .z")]]
=
Ez
Ezz
N(N-I)
Σ 2 (1(Z," 2 -1( _nni (2) 2nd 2 n
R 2
m.
= E N(N-1)
Σ
(f(Zm) - f(Zm)) Z z Z z
2
Just like the univariate case, we only need N evaluations of regardless of the dimensionality D.
3.2 Antithetic categorical variables
To use CARMS in practice, we describe two ways to generate antithetic categorical variables in this section. Both methods are based on transformations of uniform random variables, which makes them amenable to antithetic copulas, such as the Gaussian or Dirichlet copula Dimitriev and Zhou 2021
3.2.1 Inverse CDF sampling
The inverse transform sampling is commonly used to transform uniform random variables, which are easy to generate, to some other distribution:
I ~ (x(x)
> u Unif(0 1), x = F- (x),
where (x) denotes the CDF of the desired distribution. For categorical variable and some ordering of its probability vector p (P1,...,PP) the inverse CDF approach amounts the following algorithm. Define the left and right boundaries:
Σ Pk = l P.-
Pk:
Then setting z = if E [1,, ry where Unif(0, 1) produces categorical variable because P(z P(u E (y) = r. -r Pj. To obtain N antithetic categorical variables, we sample ~ Cn from copula and set z, in a vectorized manner. Importantly, this sampling approach allows us to easily analytically evaluate the bivariate PMF (proof in the Appendix):
P(z i, =) = (,(r,,r)) (c(r,, Be(4.rry) + (c,,,,)))
where &e denotes the bivariate CDF of the copula. However, probability for any (i,j) pair within a particular ordering. For example, if the ordering is = (0.1, 0.2, 0.7) and we use a bivariate copula (u, u). then the probability of obtaining the pair (1,2) is zero:
not guaranteed a non-zero
P(z = 1, =) P(ue [0,0.11, 1 (0.1,0.3) P(u E [0,0.1], E [0.7, 0.97 = 0.
5

---

Algorithm Antithetic inverse CDF categorical sampling
Input: Number of samples N, probabilities p o()). copula C. Sample u (uI,..., uxx ~u Cn-
Shuffle p according to Eq. | and define the left and right boundaries L, according to Eq.
Set vn: Zn = J. where is such that un E r/l
For i,j E [1,.... C): compute Rij = P.P; P(nn i, Zn = 5) according Eq.
return: (21,-.., ,N), R.
Algorithm 2 Antithetic Gumbel categorical sampling
Input: Number of samples N, probabilities p a()), copula C.
Sample u, (Uel, .... UcN) CN for each category
Set Ve, n Jen = = -Inn- (=-((II((uen))
Convert to categorical Vn: Enc 1 if argmax, Idn else 0.
Approximate the bivariate PMF P( = i x' according to Eq
Set Vi, {1...., C}: Rig P.P,/ P(x i, return: (z1,..., nN), R.
However, are always guaranteed non-zero probability for the pair (1, C) (proof for the Dirichlet copula the appendix). Therefore, we can average over C(C 1)/2 orderings of p, where the ordering oi rearranges p such that categories and j are the first and last category of the vector, i.e. (1) = 1, o. = C. Thus, any pair of categories (i..) has non-zer zero probability in at least one ordering, namely the ordering oij. There many sets of orderings that satisfy this, but one simple way to do this is to translate the original vector i elements to the right, then swap the last element with the jth or more precisely:
C, (h - 1) mod C, otherwise
k=
Tiy (k)
k=
(9)
The bivariate PMF given the set of orderings is:
P(z i = =] = TCC-1
Σ ((z))) = =o((f) (ox(p))
k<1
(()= = (p)).
(10)
O
aEO
If the number of categories C is large, we can approximate the sum by sampling random subset as shown above. Note that we don't need to calculate Eq. 10
of orderings O, |O| < for all possible pairs of categories, just those that were in a sample of N. This results in ( total computation per gradient, and is easily parallelizable. We summarize the inverse CDF sampling method in Algorithm
3.2.2 Gumbel max sampling
It is not strictly necessary to analytically calculate the bivariate PMF If the variance is not too large. Monte Carlo estimation suffices, and we can use the well known Gumbel max trick [Gumbel| 1954
Jang et al., 2017 Maddison et al. 2017]:
(o))
1, d = argmax, 9e O. 9c Gumbel(0, 1), 2d
idd
0, otherwise
where, as before, z denotes the one hot encoding of the categorical variable. This is also equivalent to the following exponential racing [Ross] 2014 Caron and Teh 2012 Zhang and Zhou 2018 Yin
et al. 2019) sampling: set z = j iff argmin, C, where <; ~ Exp(e
Since Gumbel variable can be generated using uniform: - (-((-In(())) u Unif 1),
we use a copula to generate u, = [Mel, ..., W.N). This produces a negative correlation between all
pairs of samples for category c, and we do the same for all categories. If we let Z [# 2N]" as

---

Inverse Dirchict
Inverse Geussion
Gumbel Dinchlet
Gumbel Inv. Dirchlet
Gumbel Goussian
04
0.35
0:21
0.00
0.20
0:21
0.29
0.00
0:59
2,40
0.38
0.0
6-1
0.1
0.52
0.51
0:15
6012
6:07
0:07
4046
0.25
0.12
000
0.1
2.
0.28
0.2
0.20
2:20
0:33
5.
0:16
0.19
0:15
0.00
5:0
500
3.07
8:00
0.08
0.00
0.00
0.07
0:0
0.28
0.1
0:00
0.07
001
008
0:12
00:0
00:0
3.00
1:00
3007
-02
0.00
0.08
0.05
0.03
3.03
0.30
0.33
2.00
0.00
0.20
0.15
0.06
0.00
0.20
0.15
0.06
0.03
0.25
0.15
0/22
0.52
0.85
0.13
0.03
Figure I: For an antithetic pair of one-hot encoded categorical variables 2, z' we show the correlation between each pair and zj. which are marginally Bemoulli variables with correlation Pij corr zy). The correlation matrices shown use different categorical sampling methods different copulas, estimated using N = 10 Monte Carlo samples.
before, a simple empirical estimate of the bivariate PMF computed using all pairs and only matrix operations is:
P( - i, ='=' N(- z" (1NxN - IN) Z.
(11)
We summarize the antithetic Gumbel sampling method in Algorithm 2 and in Fig. 1 we show an example of bivariate PMF produced by either the Gumbel or inverse CDF method using different copulas. The results are similar using either a Dirichlet or Gaussian copula, with larger differences between the categorical sampling method. For CARMS to be unbiased we need good estimate of the ratio ((=))p(( 2/)/p(ppppp //((z, 2,), which in turn requires good estimate of the bivariate probability mass function (PMF) the denominator. If we draw N samples to do this empirically, there is small probability that some pair (i.j) does not occur at all, so in experiments we clip the maximum value of the ratio 10 avoid infimities.
4 Experimental results
In this section, we first illustrate the variance reduction that CARMS offers toy example. We also optimize a categorical variational autoencoder (VAE) (Kinmma and Welling 2014 Rezende ee al. 2014), and a stochastic network for structured output prediction, which are standard tasks [Jang ctt al. 2017] for categorical variables, done on three different benchmark datasets. Each experiment uses both antithetic categorical approaches for CARMS: inverse CDF sampling and the Gumbel max trick, denoted as CARMS-I and CARMS-G, respectively, and we use the Dirichlet copula for
both. We compare
approach to three state-of-the-art unbiased estimators: LOORF/VarGrad Kool
et al., 2019a| Richter et al. 2020), the unordered set estimator (UNORD) Kool et al. 2020 and
ARSM Yin et al. 2019 The code for all experiments is freely available
4.1 Toy example
We first showease the variance reduction over other methods in simple toy example, where take the gradient with respect to the logits & of:
D
(()) = E[f( ... zD)]. za Cat(a f #D)
Σ d c zdc-
d=1em1
For simplicity, let the number of categories, dimensions, and samples be D N= = The probabilities are randomly sampled from a Dirichlet distribution: aΦ)) Dir(Ic a), but are identical for all methods for given a. We vary the entropy the probabilities from high (a = 1) to low (a 1000). In high entropy setting, there is little difference between the estimators, as the variance itself is very high, but differences emerge as we increase a and lower the entropy. The combined log variance of the gradient of each logit, for different methods and different a, is shown in Fig 2
7

---

Hijn+-2.27
Hint
8T
CAPACA
CARUSS
LOOns
CARTUSA
CARISS
Unono
LOOPT
CARME
CARIIGO
INDED
LOOSE
CARMOL
CACUSG
UNDEO
LOORT
Figure 2: Log variance of the gradient of different estimators with respect to the logits on toy problem. Columns correspond to different entropy levels of the logits.
4.2 Categorical variational autoencoder
For this task, we follow the experimental setting from ARMS (Dimitriev and Zhou 2021 except we use categorical instead of binary latent variables, and maximize the ELBO:
ELBO( = E p(alz)p(z)
N
p((z)zn)) (p(())
In
,1,...,2N NN ~ Caty (a()))
L (z((a)))
2 (z.
1=
where Caty (o) denotes N-variate categorical distribution with identical marginals. The number of categories is C' E (3. 5, 10} with D = |200/C] /C] latent variables, respectively, to make the total computational effort similar. In the binary case C 2, CARMS reduces to ARMS, for which a thorough comparison has already been produced. The task is training a categorical VAE using either linear nonlinear encoder/decoder pair three different datasets: Dynamic(ally binarized) MNIST [LeCun et al. 2010), Fashion MNIST [Xiao et al. 2017]. and Omniglot [Lake et al. [2015). All datasets are freely available under the MIT license, and do not contain any personally identifiable information or offensive content. For a fair comparison, all methods use the same learning rate, optimizer, model architecture, and number of samples. Since ARSM uses a variable number of function evaluations per step, we use one sample per step, for which ARSM uses around twice as many evaluations as the other methods. The results are combined from five independent runs for each experimental configuration.
The VAE consists of a stochastic layer with 2000/C] units, each of which is C-way categorical variable. For the nonlinear case, there are additionally two layers of 200 units with LeakyReI.U (Mass] al.. 2013) activations. The prior logits are optimized using SGD with a learning rate of 10 whereas the encoder and decoder are optimized using Adam (Kingma and Ba 2015) with learning rate of 10-4, following Yin al. 2019 The optimization is run for 10 steps, with batch size of 50, from which the global dataset mean is subtracted. The models are trained on an Intel Xeon Platinum 8280 2.7GHz CPU, and an individual run takes approximately 16 hours on one core the machine, with total carbon emissions estimated to be 28.19 kg of CO2 [Lacoste et al.) 2019]. An exception UNORD for 10 samples, which was significantly more heavy computation. Although the paper states that step can be performed in (22)), ), the provided code requires O(C!) evaluations per step, forcing us to use N samples at most.
In Fig we plot the training and test log likelihood using 100 samples, and gradient variance w.r.t the logits over time, for nonlinear network on dynamic MNIST. Similar plots for other datasets and network types can be found in the Appendix. Also shown in Table i is the final training log likelihood using 100 samples for all three datasets, network types, categories, and gradient estimators. The corresponding table with the final test log likelihood can be found in the Appendix. In general, both versions of CARMS perform comparably, and result in slightly higher log likelihood than other methods. Because Gumbel CARMS uses an empirical estimate of the debiasing ratio, it has higher variance and slightly worse performance. When limiting the number of function evaluations, there is gap between ARSM (which used on average 2C' evaluations per step, out of maximum of C(C 1) per step) compared to the other methods, which used C evaluations. This is possibly due to the continuous reparameterization that ARSM uses, which adds variance.
4.3 Structured prediction with stochastic categorical networks
We also compare all methods the standard benchmark task of predicting the lower half of an image from the upper half, i.e., the conditional distribution p(arq æx)), where v,, and æy denote the upper and lower half of an image, respectively. We use stochastic categorical network to estimate this

---

Table 1: Final training 100 sample log likelihood of VAES using different estimators, where the stochastic layer contains C=3, 5, or 10 categories, with |200/C latent variables and C samples per gradient step, respectively. Results are reported on three datasets: Dynamic MNIST, Fashion MNIST, and Omniglot over runs, with the best performing methods bold.
Categories
CARMS-I
CARMS-G
LOORF
UNORD
ARSM
Dynamic MNIST
Linear
-105.34 0.25
-105.36 0.24
-105.64 0.23
-105.23 0.23
-107.35 0.56
-103.53 0.13
1003335 0.18
100.55 0.15
-103.50 :: 0.13
130.. 0,53
-103222 1 0.05
-103.12 0.06
-103.48
-103.56 + 0.06
106.77 0.58
Nontinr
-94.85 + 0.28
.94.60 + 0.28
-.21
-95.21 0.22
-9.62 = 0.50
-93.05 = 0.14
-9..00 # 0.12 0.12
-92.91
-92.88 + 0.12
0.12
-98.89 = 0.43
10
-92.13 = 0.05
0.10
-2.44
-93.05
0.14 0.14
997.66 0.41
Fashion MNIST
Linear
-245,44 + 0.22
245.699
-215.8
-245.90 0.21
-247.51 +
-242.06 + 0.13
-241.90 0.10
-2.2.17
-242.43 0.12
-244.63 # 0.47
-240.44 + 0.03
-240.52 0.04
-240.91 0.04
0.06 -241.08
-243.29 0.36
Nonliar
-233.13 1 0.16
-233.20 0.16
233.75
-333.34 0.12
-237.93
0.20
-231.72 + 0.09
-231.67
-232.01 0
2332.19
-237.29 0.34
-230.77 + 0.05
2311.66
231.35
-231.74
-235.88 0.17
Linear
-114,53 +: 0.12
-114.73 # 0.15
-114,90 & 0.13
-114,77 # 0.15
-116.34 * 0.42
Omniglot
-114.01 # 0.10
-113,93 0.11
-114.01 = 0.10
-114.00 0.09
-115.72 0.35
10 -114.97 # 0.06
-114.99 = 0.06
-115.17 # 0.05
-115.55 0.08
-117.48 # 0.35
Nonlint
-110.19 + 0.23
-110.16
-110.33 0.27
-110.31
-115.13 0.21
1199114 + 0.09
-009.3 0.13
-09939 0.16
109.55
0.14
-115.07 0.37
-108.66 1 0.08
-108.62 0.07
-108.98 0.06
109.54 0.15
-115.11 = 0.34
categories
calegorte
10 categories
Traning ELBO
-25
-35
-00
129
110
Tost log Ikelhood
-35
120
130
-15
20.66
10'
ARSU
Log valfance
CAR05-1
107
LOOG?
16'
UNORD
06
2004
400%
603K
D00K
1000X
04
D00K
420K
B00K
200%
100KK
200K
400%
500%
300K
1CO0
Steps
Steps
Steps
Figure Training nonlinear categorical VAE with different estimators on Dynamic MNIST using ELBO. Columns correspond to C' E (3. 5, 10) categories with C' samples per gradient step, respectively. Rows correspond the 100 sample training and test log likelihood, and the variance of the gradient with respect to the logits of the encoder network. Results for different datasets and other networks can be found in the Appendix.
distribution, with the objective: Ez-p(zz a.) M EM Inp(x; [mm)] where z denotes stochastic
categorical layer. The encoder/decoder pair each contains one hidden layer with [200/C] latent variables and LeakyReLU activation, with the optimization performed for C' E {3, 10) categories, on all three datasets, with identical settings for each gradient estimator for a fair comparison. We use M = 1 for training, and M 1000 for evaluation on both the train and test set. In Table we show the final training set log likelihood, with the corresponding test log likelihood table in the appendix, though the results are qualitatively similar. The results are similar to the VAE experiment, with the inverse CDF CARMS having a slightly higher log likelihood. However, the differences between estimators are less pronounced, with the unordered set estimator is being mostly on par with CARMS, and ARSM only slightly trailing the other methods, with much smaller gap.

---

Table 2: Final training log likelihood of categorical network for conditional estimation using different gradient estimators, where the stochastic layer contains = 3, 5, or 10 categories, with [200/C] latent variables and C samples per gradient step, respectively. Results are reported on three datasets: Dynamic MNIST, Fashion MNIST, and Omniglot over 5 runs, with the best performing methods in bold.
Categories
CARMS-1
CARMS-G
1.OORF
UNORD
ARSM
Dynamic MNIST
57.98 0.10
0.14
58.19
58.06 0.04
60.22 0.19
57.57 0.05
0.12
57.78 0.08
57.6 + 0.05
59.38 0.13
10
58.17 0.26
58.33 0.13
58.20 0.14
58.18 0.17
59.32 0.11
Fashion
132.83 + 0.08
132,90 0.08
133.10 = 0.05
0.06 0.06 133.06
134.56 # 0,35
MNIST
132.68 +: 0,05
132.81 0.12
132.91 + 0.07
1
134.09 + 0.12
133.32 = 0.15
133.43 0.23
133.54 0.17
134.02 = 0.21 1
Omniglot
3
65.57 = 0.10
66.05 = 0.14
65.92 0.26
65.81 0.09
68.00 = 0.09
5
65.65 0.18
66.16 0.32
65.92 0.23
65.78 = 0.06
67.99 = 0.32
10
66.76 0.31
66.94 a: 0.24
66.87 # 0.07
66.66 + 0.24
68.35 * 0.16
Discussion
We have presented a novel approach for training categorical variables, which extends the ARMS estimator for the binary case. goes beyond i.i.d. samples by using a copula to generate antithetic categorical samples, but preserves unbiasedness by including a multiplicative term. For i.i.d. samples, the form the estimator reduces to LOORE. We showcase usefulness on several datasets, different number of categories and types of deep neural networks. In variational inference tasks, and conditional estimation, CARMS outperforms other state of the estimators. The main limitation of this work is its specificity to categorical (including binary) variables. We hope to extend this, c.g. to Plackett-Luce models for top-k sampling (Kool et al. 2019b] Grover et al., 2019], which has important applications in ranking (Dadaneh et al.) 2020]. There also general limitation that CARMS shares with other state-of-tt oftte-of-the-arr estimators, which is higher complexity than LOORF, a very simple but strong baseline. Future work includes investigating theoretical properties and large scale applications, as well as possible general antithetic gradient estimators, and we plan to investigate adaptive correlations that take into account the properties of f to further reduce the variance.
Potential societal impact This work is focused on better optimization for categorical variables, which includes any network containing stochastic categorical layer. In particular, generative models such as VAEs are widely used, and are sometimes trained on more human centered datasets. These models can sometimes be used to impersonate person's 's face voice. We only used non-human datasets such as MNIST, but with enough knowledge and using the available code, anyone, including bad actors, can train the same model on human content for malicious purposes. However, strongly believe that wide knowledge dissemination and open source code is crucial for reproducibility in all of science.
Acknowledgments
The authors acknowledge the support of NSF 181269,,, the APX 2019 project sponsored by the Office of the Vice President for Research The University of Texas at Austin, the support of gift fund from ByteDance Inc., and the Texas Advanced Computing Center (TACC) for providing HPC resources that have contributed the research results reported within this paper.
References
Y. Bengio, N. Léonard, and A. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.34.. 2013.
Y. Burda, R. B. Grosse, and R. Salakhutdinov. Importance weighted autoencoders. In 4th International Conference on Learning Representations, ICLR, 2016.
Caron and W. Teh. Bayesian nonparametric models for ranked data. In Advances in Neural Information Processing Systems, volume 25, 2012.
10

---

S. Z. Dadaneh, S. Boluki, M. Zhou, and X. Qian. Arsm gradient estimator for supervised learning to rank. In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 3157-3161, IEEE, 2020.
A. Dimitriev and M. Zhou. Arms: Antihetic-eeinfore-mmlii------ International Conference on Machine Learning, 2021.
gradient for binary variables.
Z. Dong, A. Mnih, and G. Tucker. Disarm: An antithetic gradient estimator for binary latent variables. In Advances in Neural Information Processing Systems 33, 2020.
Z. Dong, A. Mnih, and G. Tucker. Coupled gradient estimators for discrete latent variables. Third Symposium on Advances Approximate Bayesian Inference, 2021.
M. C. Fu. Gradient estimation. Handbooks in operations research and management science, 13: 575-616, 2006.
W. Glynn. Likelihood ratio gradient estimation for stochastic systems. Communications of the ACM, 33(10):75-84, 1990.
W. Grathwohl, D. Choi, Y. Wu, G. Roeder, and D. Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In 6th International Conference on Learning Representations, ICLR, 2018.
A. Grover, E. Wang, A. Zweig, and S. Ermon. Stochastic optimization of sorting networks via continuous relaxations. arXiv preprint arXiv:1933.08850 2019.
S. Gu, S. Levine, Sutskever, and A. Mnih. Muprop: Unbiased backpropagation for stochastic neural networks. In 4th International Conference on Learning Representations, ICLR, 2016.
E. Gumbel. Statistical theory of extreme values and some practical applications: a series of lectures, volume 33. US Govermment Printing Office, 1954.
E. Jang, S. Gu, and B. Poole. Categorical reparameteriization with gumbel-softmax. In 5th Interna- tional Conference on Learning Representations, ICLR. OpenReview.ned 2017.
M. Jordan, Z. Ghahramani, T. S. Jaakkola, and L. K. Saul. An introduction to variational methods for graphical models. In Learning in graphical models, pages 105-161. Springer, 1998.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR, 2015.
D. P. Kingma and M. Welling. Auto-encoding variational bayes. In 2nd International Conference on Learning Representations, ICLR, 2014.
W. Kool, H. van Hoof, and M. Welling. Buy 4 REINFORCE samples, get a baseline for free! In Workshop, Deep Reinforcement Learning Meets Structured Prediction, ICLR, 2019a.
W. Kool, H. Van Hoof, and M. Welling, Stochastic beams and where to find them: The gumbel-top- trick for sampling sequences without replacement. In International Conference on Machine Learning, pages 3499-3508. PMLR, 2019b.
W. Kool, H. Hoof. and M. Welling. Estimating gradients for discrete random variables by sampling without replacement. In International Conference on Learning Representations, 2020.
A. Kucukelbir, D. Tran, R. Ranganath, A. Gelman, and D. M. Blei. Automatic differentiation variational inference. The Journal of Machine Learning Research, 18(1):30-477, 2017.
A. Lacoste, A. Luccioni, V. Schmidt, and T. Dandres. Quantifying the carbon emissions of machine learning. arXiv preprint arXiv: 2019.
B. M. Lake, R. Salakhutdinov, and J. B. Tenenbaum. Human-level concept learning through proba- bilistic program induction. Science, 350(62666:1332 1332-1111 1338, 2015.
Y. LeCun, C. Cortes, and C. Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http:///www..cc...ooo/ee/e////ee
2, 2010.

---

W. Lee and J. Y. Ahn. On the multidimensional extension of countermonotoniictyy and its applications. Insurance: Mathematics and Economics, 56:68-79, 2014.
R. Liu, J. Regier, N. Tripuraneni, M. Jordan, and J. Mcauliffe. Rao-blackwellized stochastic gradients for discrete distributions. In International Conference Machine Learning, pages 4023-4031 PMLR, 2019.
G. Lorberbom, S. Jaakkola, A. Gane, and T. Hazan. Direct optimization through arg max for discrete variational auto-encoder, In Advances in Neural Information Processing Systems 32, pages 6200-6211, 2019.
A. L. Maas, A. Y. Hannun, and A. Y. Ng. Rectifier nonlincarities improve neural network acoustic models. In ICML Workshop Deep Learning for Audio, Speech and Language Processing. Citeseer, 2013.
C. J. Maddison, A. Mnih, and Y. W. Teh. The concrete distribution: A continuous relaxation of discrete random variables. In 5th International Conference Learning Representations, ICL.R OpenReview.net, 2017.
A. J. McNeil and Nešlehová. Multivariate archimedean copulas, d-monotone functions and II Ilnorm symmetric distributions. The Annals of Statistics, 37/5B):3059-3097, 2009.
A. Mnih and K. Gregor. Neural variational inference and learning in belief networks. In International Conference on Machine Learning, pages 1791-1799. PMLR, 2014.
A. Mnih and D. J. Rezende. Variational inference for monte carlo objectives. In Proceedings of the 33nd International Conference on Machine Learning, ICML., volume 48, pages 2188-2196 JMLR.org org, 2016.
S. Mohamed, M. Rosca, M. Figurnov, and A. Mnih. Monte carlo gradient estimation in machine learning. J. Mach. Learn. Res., 21:132:1-132:62, 2020.
A. B. Owen. Monte carlo theory, methods and examples, 2013.
W. Paisley, D. M. Blei, and M. I. Jordan. Variational bayesian inference with stochastic search. In Proceedings of the 29th International Conference Machine Learning, ICML, 2012.
M. B. Paulus, C. J. Maddison, and A. Krause. Rao-blackwellizingg the traight-thhouglh gumbel- softmax gradient estimator. In International Conference on Learning Representations. 2021.
A. Ramesh, M. Pavlov, G. Goh, S. Gray, C. Voss, A. Radford, M. Chen, and I. Sutskever. Zero-shot text-to-image generation. arXiv preprint 1009, 2021.
R. Ranganath, S. Gerrish, and D. Blei. Black box variational inference. In Artificial intelligence and staristics, pages 814-822. PMLR, 2014.
D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning, pages 1278 1286- PMLR, 2014.
I. Richter, A. Boustati, N. Nüsken, F. Ruiz, and D. Akyildiz. Vargrad: A low-variance gradient estimator for variational inference. In Advances in Neural Information Processing Systems, volume 33, pages 13481-13492, 2020.
M. Rosca, M. Figurnov, S. Mohamed, and A. Mnih. Measure-valued derivatives for approximate bayesian inference. 4th workshop Bayesian Deep Learning, NeurIPS, 2019.
M. Ross. Introduction probability models. Academic press, 2014.
Ruiz, M. K. Titsias, and D. M. Blei. The generalized reparameterization gradient. In Advances in Neural Information Processing Systems 29, pages 460-468, 2016.
T. Salimans and D. A. Knowles. On using control variates with stochastic approximation for variational bayes and connection to stochastic linear regression. arXiv preprint arXiv: 1022, 2014.
12

---

M. Titsias and M.
Local expectation gradients for black box variational inference. In Advances neural information processing systems, pages 2620-2628. Citeseer, 2015.
K. Trivedi and D. M. Zimmer. Copula modeling: an introduction for practitioners. 2007.
G. Tucker, A. Mnih, C. Maddison, D. Lawson, and J. Sohl-Dickstein. REBAR: low-variance. unbiased gradient estimates for discrete latent variable models. In Advances in Neural Information Processing Systems 30, pages 2627-2636, 2017.
R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning, Machine learning, 8(3-4):229-256, 1992.
H. Xiao, Rasul, and Vollgraf. Fashion-mnis novel image dataset for benchmarking machine learning algorithms, axxiv.org/ubs/1778.077774 2017.
M. Yin, Y. Yue, and M. Zhou. ARSM: augment-riinffrreswwa---m
estimator for gradient
backpropagation through categorical variables. Proceedings of the 36th International Conference Machine Learning, volume 97, pages 7095-7104. PMLR, 2019.
M. Yin, N. Ho, B. Yan, X. Qian, and M. Zhou. Probabilistic Best Subset Selection by Gradient-Band Based Optimization. arXiv e-prints, 2020.
Q. Zhang and M. Zhou, Nonparametric bayesian lomax delegate racing for survival analysis with competing risks. arXiv preprint 1rx.v:1870008564, 2018.
13


Reviews:
[{"review_id": "fnrGm4CNME", "reviewer": "Reviewer_eSUj", "summary": "The paper proposes a categorical version of ARMS, namely CARMS.\nBy utilizing REINFORCE together with copula sampling, CARMS remains an unbiased estimator and non-duplicated samples reduce the variance of gradients.\nThe authors verify the proposed model through toy example, VAE, and the structured output prediction task.\n", "questions": "", "limitations": "", "rating": 5, "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:\nI do not think the natural extension from the previous idea harms the originality or the novelty of the paper \"in general\".\nHowever, in this case, since CARMS is an extension of ARMS over the number of categories, I think the paper also needs to focus on the obstacle which follows from the natural extension which is discussed in the significance part.\nTherefore, I will suspend the judgment on the originality until the authors' response.\n\n\nQuality and Clarity:\nThe paper is grounded by theorems and the proofs are provided in the appendix.\nIt will be helpful for readers if the authors add a minimum introduction on the copula in the background section.\nAlso, the paper is written in a clear manner.\n\n\nSignificance:\n- While the binary version (ARMS) may not suffer from the complexity, it can be a huge drawback for the proposed CARMS when the number of categories grows. However, there is no specific discussion about the complexity, nor the report on the running time in the experimental section. ARMS is a good estimator for the binary variable, however, CARMS, which is the natural extension of ARMS, seems that it needs to break through the complexity issue. Also, the performance gap is too marginal even though the proposed theorems are grounded. If the complexity really matters, why not use the comparable baselines? (which are simple and powerful) Or, CARMS needs to find a suitable application that CARMS outperforms other baselines with meaningful performance gaps.\n- There are some missing baselines: RELAX, REBAR, etc.\n- As a minor comment, the code is not available for reproducibility.\n\n\nQuestions:\n- What does ECRAMS stand for?\n- How does the experimental result change for the toy example if we sample the probabilities from Dir(\\alpha) \\alpha << 1? I'm wondering about the behavior of the proposed and baseline methods when the probability is extreme.\n- Similarly, how does the performance change for the non-trivial or extreme prior in the VAE case?\n- The complexity seems problematic. What is the exact computational complexity, and how does the complexity changes from ARMS to CARMS? (if there are the same number of latent dimensions) Also, could you compare the running times of the experiments? How does the running time change as the number of categories grows, for examples, C = 100, 200?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors pointed out the limitation, the complexity, of their work in one line in the conclusion.\nHowever, my thought is that the complexity is somewhat crucial to the proposed method as the number of categories grows so that it should be dealt with in a separate subsection.\nMeanwhile, the work seems that it does not have any potential negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "ZHkBJtif3WK", "reviewer": "Reviewer_qwuP", "summary": "**Summary**\n- The paper introduces a multi-sample gradient estimator for categorical variables. \n- The estimator is based on the score-function and uses correlated samples to achieve a variance reduction. \n- The estimator extends several recently proposed estimators naturally. \n- The estimator yields marginal improvements over baselines in the experimental evaluations. \n\n**Edit after Author Response**\n- Thank you for replying to the points I raised. I think the changes you do improve the draft. However, because I find the experimental improvements fairly marginal, I have kept my score unchanged. I believe the paper is above the acceptance threshold.", "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": "**Strengths**\n- Originality. The contribution is a natural extension of existing work. \n- Quality. To me, the contribution appears technically sound. The derivations and arguments are clear.  \n- Clarity. The paper is well-written. I enjoyed reading it and was able to follow the line of argument and the derivations. \n- Significance. Gradient estimator for categorical distributions has wide-spread applications and draws interest from the community.\n\n**Weaknesses**\n- Quality/ Significance The improvements demonstrated in the experimental evaluation are marginal.\n\n**Decision**\n- Weak Accept. I think the paper is above the acceptance threshold, but I wished the experimental results were more convincing. \n\n**Suggestions**\n\n*Are duplicate samples really a significant issue?*\n- Based on the paper, I suppose that the benefit of using this estimator instead of LOORF are most significant, when the duplicate sampling issue is acute. Based on the experiments presented, I am not convinced that duplicate samples are a significant issue in most practical applications. \n- I would recommend to the authors to seek applications where the duplicate sampling issue is more acute to clearly demonstrate the utility of their estimator. Speculative: Maybe, as models converge, the entropy of the latent distributions shrinks. If this is the case, one could consider an experiment, where the estimators are used for fine-tuning a model which has been trained for significantly more iterations than the models presented in 4.2 (10^6 iterations).  \n\n*I am concerned about the statistical significance of the improvements.*\n- The improvements appear very marginal, and I am not sure what the bolding scheme is indicating. Is this based on the standard errors, shout it actually be based on confidence intervals? I am unsure. \n- Rather than bolding, would it be possible to carry out a hypothesis test (t-test) to test whether the means of the proposed method are significantly different from the baselines? \n\n*Isn't it better to tune hyper-parameters exhaustively for each method separately?*\n- I believe it is better to tune hyper-parameters individually for each method rather than choosing a fixed set of hyper-parameters on which all methods are compared against (which is what I understand, the paper did). \n-It may be acceptable in this case though, because all methods are fairly similar anyway, and the optimal hyper-parameters might be similar/ the same as a result. \n\n*What is ecarms in the plots?*\n- The experimental set-up promised Carms-I and Carms-G. Maybe, I missed it, though. \n\n*There is more related work on Rao-Blackwellization for gradient estimation for discrete distributions.*\n- The authors survey related work on Rao-Blackwellization to improve gradient estimation, e.g. Titsias and Lazaro-Gredilla (2015) and Dong et al. (2021).  Rao-Blackwellization is increasngly being used for improving gradient estimators for discrete distributions and I think it's worth highlighting also other work in this area, e.g. Liu et al. (2018), Kool et al. (2020) or Paulus et al. (2021). \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": "QSTd1OcGt-N", "reviewer": "Reviewer_m9re", "summary": "The paper proposed a novel gradient estimator for categorical stochastic hidden variables, named Categorical-Antithetic REINFORCE Multi-Sample (CARMS), which is a categorical extension of ARMS (Dimitriev and Zhou, 2021). CARMS reduces the variance of LOORF (Kool et at., 2019) by the means of importance weighted sampling. Using copula, where Dirichlet copula was employed in the experiments, CARMS generates the antithetic categorical samples to reduce variance. The paper compares two versions of the estimator, using inverse CDF sampling and Gumbel max sampling, against UNORD, LOORF and ARSM for categorical VAE and structured prediction tasks on Dynamic MNIST, Fashion MNIST and Omniglot. They find slightly better results compared to other estimators on both tasks.", "questions": "", "limitations": "", "rating": 7, "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: \nThere are a variety of approaches proposed in the past few years for unbiased gradient estimators over binary stochastic hidden variables. The paper introduces a natural extension for one of previously proposed estimators, ARMS, to categorical variables. Background and related works are adequately cited. \n### Quality: \nThe authors did a great job in derivation of the estimators, design and implementation of the experiments. By closely checking the details, the methods are well-grounded. \n### Clarity: \nThe paper, especially the method section, is clearly written and well explained. \n### Significance: \nThe method proposed can be directly applied to interpretable models and modeling natural systems with discrete variables, e.g. sparse computing.  \n\n## Typos and potential places to improve: \n* Ln 95, which both -> both of which\n* Ln 148, adding a definition of Z^{-d}\n* Eq 7, for l_i^o, the upper limit for the summation should be i-1, instead of j-1\n* Figure 1, how many samples are used for estimating the correlation matrix? Also, how to interpret the numbers on two-axis of the matrices, as z[i] and z[j] should be categorical variables?\n* Ln 198, “For CARMS-G, we clip the empirical ratio values to avoid numerical instabilities.” It is a bit unclear to me: 1. What is “the empirical ratio values”? Is it the sampled value for the ratio tensor R? 2. How did you clip it? E.g. clipping range, etc. I didn’t find the relevant information in either the main text or the appendix.\n* In Figure 2 and 3, there are labels of “CARMS” and “ECARMS”, while in other places of the paper, the proposed methods are mentioned as CARMS-I and CARMS-G. Which CARMS-? Is “CARMS” referring to, and what does “ECARMS” stand for? \n\n## Questions:\n* In Ln 199 of the Experiment results section, it mentioned that the Dirichlet copula is used for the experiments? I am wondering whether there are results by trying Gaussian copula. To my limited knowledge, I find the Dirichlet copula introduces a limited dependence structure, while the Gaussian copula seems more flexible. Is there any insight into which copula to choose?\n* In Ln 234 of the Categorical VAE section, it was mentioned that “an individual run took 5-8 hours” and it was on CPUs. This is surprisingly fast. Based on the description follows, it seems this holds for all estimators in Table 2 other than UNORD with C=10. Does this speed number also hold for ARSM? Also, in Line 245, it mentioned the number of evaluations is limited to 2C for ARSM, but in the original ARSM paper, it requires C(C-1)/2 evaluations for the swapping matrix. How does it (using 2C evaluations rather than the C(C-1)/2) impact the performance of ARSM as a baseline?\n* The experiments show that CARMS-I/G “result in slightly higher log likelihood than other methods”. Is there any insight or result for the case when the number of categories is growing larger? It is common to have a model with C >> 10. I am curious, when C grows larger, shall I expect to see a larger gap between LOORF and CARMS or smaller?\n* Eq 5 introduces a very general form of couplings of categorical variables p(z, z’). In Figure 1, the author compares correlation matrices of 5 different couplings. It seems inverse CDF sampling and Gumbel max sampling favor quite different sets of joint samples. Is there any insight on what kind of correlation structure is better? \n* Based on the results, it seems LOORF is almost as good as the estimators using coupled samples. Does it mean that using the independent samples is good enough? Or putting the question in another way, in what situation, shall we expect a big improvement by using coupled samples?\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "KnRNyDzBsOn", "reviewer": "Reviewer_1mCY", "summary": "This paper tackles the problem of differentiating expectations of discrete random variables. Specifically, this paper proposes a generalization of the recent ARMS approach to more than two categories. The CARMS gradient estimator uses multiple samples from a categorical distribution. Unlike e.g. LOORF, CARMS uses samples that are dependent (typically negatively correlated), hoping to further reduce estimator variance. The authors prove how to construct an unbiased gradient from these dependent samples using importance sampling. The authors propose several joint distributions (copulas) to actually generate such dependent samples. Experimental results are provided on several problems such as a categorical latent VAE. ", "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: This work constitutes a natural extension of ARMS. To the best of my knowledge, this exact approach has not appeared before in the literature. This paper does a very good job of summarizing recent connected work in the field (Sections 1 and 2).\n\nQuality: The core concept (Theorem 1) is well built up to and well explained. It’s clear where the estimator comes from and why it’s unbiased (from the preceding few lines). The approach and maths follows naturally from ARMS. Theorem 2 adds justification for the CARMS approach, the variance can be provably lower. \n\nOne thing that could make the paper better is a more thorough discussion of how to generate antithetic categorical variables (Sec 3.2), because this is more involved than in the 2-category case.\n- I would like more details on the random reordering and the statement “Unfortunately, keeping to one specific ordering does not allow for all possible pairs to have non-zero probabilities”. It’s not clear to me why we cannot have distributions on [0,1]^2 which have positive density across the entire square, have uniform marginals and are not independent. Such a pair of variables, when discretized as per (8) would surely assign positive probability to every pair. Going further, can’t we transform the u, u’ distribution under a reordering by suitably shuffling the joint density between the different sub-rectangles of [0,1]^2? (Maybe this is something obvious that I missed, but it would be nice to have an explanation).\n- How restrictive / how general is the inverse CDF method with uniforms? Is it possible that there exist more “natural” joint distributions for N dependent samples of a categorical?\n- Is it clear why the Gumbel max approach seems to give more strongly negatively correlated samples than the Inverse CDF method? Or is this an artifact of the estimation of the joint distribution of the Gumbel method?\n\nFor both ARMS and CARMS, the benefit of antithetic sampling decreases as we add more samples. Does this decay happen more slowly for CARMS? For example, with a uniform categorical on C items we can make a coupling with N=C samples (e.g. by cycling one sample) that guarantees z_i != z_j. Discussion on this point might further add to the paper.\n\nFor the experiments, I think it’s good to consider a fair range of baselines, which the authors appear to have done. The results are not strong improvements over existing methods, but they do offer marginal improvements in most cases. The presentation of Fig 3 could be improved because it is impossible to visually disentangle most of the methods.\n\nThe computational complexity is mentioned in the discussion, but it could be worth explicitly quantifying it, e.g. in Sec 3.1 to mathematically explain the benefit of only using N evaluations of f we could write out the total cost of a gradient step of CARMS and compare to LOORF. This could also aid a theoretical comparison with UNORD and ARSM, as the intro suggests both are costlier than CARMS.\n\nClarity: the paper is mostly clear. Particularly earlier sections. I have made a few suggestions / queries about Sections 3.1 and 3.2 which would also benefit clarity. Overall, I find the exposition to be clear and helpful.\n\nSignificance: this paper touches on an interesting and relevant area of ML research: differentiating through categorical random variables. It forms a natural step in the logical progression of previous ideas (ARMS) and would therefore be a welcome and valuable component of the literature in this area. The fact that CARMS gets (slightly) better empirical results that previous methods also indicates that it can transfer to the practical setting well. On the other hand, the paper is natural next step with modest empirical improvements, rather than a seminal work.\n\nExplanation of score\n----------------------\nI was torn between Accept and Weak Accept. Whilst I enjoyed reading the paper, the paper makes a modest step from ARMS, and has modest empirical results. It is great to have a more technically focused paper, but with that in mind, I think a bit more detail in Secs 3.1 and 3.2 would step things up to Accept for me.\n\nTypos\n-------\n135: errant full stop\n\n(7) : l_i^{(o)} should sum to i-1 not j-1\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have adequately addressed limitations and societal impact. For more details, see the main review.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The submission proposes a low-variance gradient estimator for discrete random variables, extending recent work on binary to categorical. Whilst the reviewers raised concerns around the marginal improvement over existing methods, all agreed that the submission is clearly written and is a theoretically grounded, novel extension of ARMS. Overall, this is a good contribution to NeurIPS and should be accepted.
Decision: Accept (Poster)