CARMS: Categorical-Antithetic-REINFORCE Multi-Sample Gradient Estimator
Year: 2021

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.1

heading: 1 Introduction
text: When optimizing an expectation based objective of the form Ez∼qϕ(z)[f(z)], we sometimes require the gradients with respect to the parameters ϕ of the distribution. This is 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 et al., 2017, Maddison et al., 2017]. 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 f to be the differentiable. In this paper, we focus on the case when z is a high dimensional categorical variable with logits ϕ. An common example of this form is 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 a large number of categorical variables, they require Monte Carlo gradients with respect to the parameters of the stochastic distribution, but have excellent performance, as a categorical VAE has in practice achieved state-of-the-art zero-shot image generation [Ramesh et al., 2021].
Our main contribution is a novel unbiased and low-variance gradient estimator for categorical variables. The Categorical-Antithetic-REINFORCE-Multi-Sample (CARMS) estimator uses a copula to generate any number of antithetic (mutually negatively correlated) [Owen, 2013] categorical samples, and constructs an unbiased estimator by combining them into a 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
1https://github.com/alekdimi/carms
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
reduces to it, while for independent samples, CARMS reduces to the leave-one-out-REINFORCE (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 [Lorberbom et al., 2019], as well as the concurrently developed equivalent Gumbel-Softmax (GS) [Jang et al., 2017] or Concrete [Maddison et al., 2017] 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 [Grathwohl et al., 2018], which uses a free-form neural network.
In practice, REINFORCE is almost always augmented by baselines, e.g., in variational inference [Mnih and Gregor, 2014, Ranganath et al., 2014, Paisley et al., 2012, Ruiz et al., 2016, Kucukelbir et al., 2017]. MuProp [Gu et al., 2016] uses a first order mean-field Taylor approximation, but requires f to be differentiable. Other estimators apply Rao-Blackwellization [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 et al., 2021]. Besides REINFORCE and reparameterization type gradients, there is also the measure valued gradient [Rosca et al., 2019] and finite differences [Fu, 2006], but are less common in practice. A comprehensive review can be found in Mohamed et al. [2020].
The more recent approaches for categorical variables are unbiased and REINFORCE based, building on the idea of using multiple samples to construct a baseline for variance reduction. One such baseline is LOORF [Kool et al., 2019a], originally introduced in Salimans and Knowles [2014] and also known as VarGrad [Richter et al., 2020], where its theoretical properties are further analyzed. VIMCO [Mnih and Rezende, 2016] has a similar form, but is specific to the importance weighted multi sample bound [Burda et al., 2016]. A different approach is ARSM [Yin et al., 2019], which reparameterizes the gradient with a Dirichlet distribution and uses swaps to obtain 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.

heading: 2 Background
text: Let z = (z1, ...,zD) denote D independent categorical variables, where zd is a one hot encoded categorical sample from zd ∼ qϕd(zd) = Cat(σ(ϕd)), with σ(x)i = exi/ ∑ j e
xj being the softmax function. Let also ı̂ be the basis vector with its ith coordinate set to one, and otherwise zero, such that P (z = ı̂) = σ(ϕ)i, 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):
L(ϕ) = Ez∼qϕ(z)[f(z)], qϕ(z) = D∏
d=1
qϕd(zd).
Although the score function gradient contains two terms: ∇ϕL(ϕ) = Ez∼qϕ(z) [ ∇ϕf(z) + f(z)∇ϕ ln qϕ(z) ] 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.

heading: 2.1 LOORF
text: Leave-one-out-REINFORCE (LOORF) [Salimans and Knowles, 2014, Kool et al., 2019a], also known as VarGrad [Richter et al., 2020], is a general score function based gradient estimator that uses N i.i.d. samples to construct a baseline for variance reduction, and is 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 ∼ qϕ(z) and evaluate f(z). Given N samples z1, ...,zN iid∼ Cat(σ(ϕ)), it has the form:
gLOORF = 1
N−1 N∑ n=1 ( f(zn)− f̄(z) ) ∇ϕ ln qϕ(zn) = 1N−1 N∑ n=1 ( f(zn)− f̄(z) )( zn − σ(ϕn) ) , (1)
where f̄(z) = 1N ∑N n=1 f(zn). It is commonly used due to its simplicity and strong performance.

heading: 2.2 ARMS
text: The Antithetic-REINFORCE-MultiSample (ARMS) [Dimitriev and Zhou, 2021] estimator is a recent work that uses any number N of mutually negatively correlated samples. However, although unbiased, it 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 use N = 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 = (u1, ..., uN ) ∼ CN , ∀i : ui ∼ 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 Φ(ui,uj) 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 bi = 1ui<p, ∀i. For a set of N antithetic D-dimensional Bernoulli variables b1, ..., bN with probabilities σ(ϕ), the N -sample unbiased estimator has the following simple form:
gARMS = 1
N − 1 N∑ n=1
( f(bn)− 1
n N∑ m=1 f(bm)
) bn − σ(ϕ)
1− ρ , (2)
with ρ = (ρ1, ..., ρD) and ρd = corr(bdi , b d j ), which is simple to compute given the bivariate CDF of the copula.

heading: 3 CARMS
text: The CARMS estimator has a 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 be 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 N samples. Next, we generalize CARMS to any number of categorical variables. Lastly, in Section 3.2, we 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 z, z′ ∼ Cat(σ(ϕ)), LOORF has the following simple form:
gLOORF(z, z ′) =
1
2
( f(z)− f(z′) ) (z − z′), (3)
which, importantly, is unbiased [Kool et al., 2019a, Dimitriev and Zhou, 2021]. This means that using a simple importance weight preserves its unbiasedness, for an arbitrary bivariate categorical distribution (z, z′) ∼ Cat ( σ(ϕ), σ(ϕ) ) :
E [ gLOORF(z, z ′) P (z = ı̂)P (z′ = ȷ̂)
P (z = ı̂, z′ = ȷ̂) ] = ∑ i,j P (z = ı̂, z′ = ȷ̂) P (z = ı̂)P (z′ = ȷ̂) P (z = ı̂, z′ = ȷ̂) gLOORF(̂ı, ȷ̂)
= ∑ i,j P (z = ı̂)P (z′ = ȷ̂)gLOORF(̂ı, ȷ̂) = E z,z′ iid∼ Cat(σ(ϕ)) [ gLOORF(z, z ′) ] = ∇ϕL(ϕ).
We summarize the derivation of the two sample version of CARMS, which we denote as the Categorical-Antithetic-REINFORCE-Two-Sample (CARTS) estimator in the following theorem.
Theorem 1 Let (z, z′) be a sample from an arbitrary bivariate Categorical distribution with marginal distributions z, z′ ∼ Cat(σ(ϕ)), and a known bivariate PMF. An unbiased estimator of ∇ϕE[f(z)] is:
gCARTS(z, z ′) =
1
2
( f(z)− f(z′) ) (z − z′)zTRz′, Rij =
σ(ϕ)iσ(ϕ)j P (z = ı̂, z′ = ȷ̂) . (4)
The intuition behind using negatively correlated variables is that we want to avoid the case when z = z′, because the sample is then “wasted.” We formalize this intuition below, and defer the proof to the Appendix.
Theorem 2 Let (z, z′) ∼ Cat(σ(ϕ), σ(ϕ)). If the bivariate PMF satisfies: ∀i ̸= j : P (z = ı̂, z′ = ȷ̂) ≥ P (z = ı̂)P (z′ = ȷ̂),
with strict inequality for at least one pair, then Var[gCARTS(z, z′)] < Var[gLOORF(z, z′)].
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 -sample LOORF is equivalent to averaging 2-sample LOORF over all ( N 2 ) pairs:
gLOORF(z1, ..,zN ) = 1
N N∑ n=1
( f(zn)− 1
N N∑ m=1 f(zm)
) (zn − σ(ϕ))
= 1N(N−1) ∑ n ̸=m 1 2 ( f(zn)− f(zm) )( zn − zm ) = 1N(N−1) ∑ n ̸=m gLOORF(zn, zm).
With the above identity it easily follows that given N antithetic categorical samples Z = [z1, ...,zN ]T , applying CARTS to all pairs results in an unbiased estimator, due to linearity of expectations:
E [gCARMS(Z,R)] = E [ 1 N(N−1) ∑ n ̸=m gCARTS(zn, zm,R) ] = E [ 1 N(N−1) ∑ n ̸=m gLOORF(zn, zm) ] = E [gLOORF(Z)] = ∇ϕL(ϕ).
We summarize CARMS in the next theorem and also rewrite it in a 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 = [z1, ...,zN ]T be a sample from an arbitrary N -variate Categorical distribution with identical marginal and bivariate distributions, such that zi ∼ Cat(σ(ϕ)), and Rij = σ(ϕ)iσ(ϕ)j/P (zn = ı̂, zm = ȷ̂). An unbiased estimator of ∇ϕE[f(z)] is:
gCARMS(Z) = 1
N(N−1) ∑ n ̸=m 1 2 ( f(zn)− f(zm) ) (zn − zm)zTnRz′m
Lemma 4 Let f(Z) = [f(z1), ..., f(zN )]T , ◦ denote the Hadamard product, 1N×N a matrix of ones, and IN the identity matrix. Define O = 1N−1 (1NxN − IN )◦ (ZRZ T ), and D = diag(O1N ),
to be a diagonal matrix. The CARMS estimator can equivalently be written in the following form:
gCARMS(Z) = 1
N f(Z)T (D −O)
( Z − 1Nσ(ϕ)T ) . (5)
Furthermore, for independent samples, ZRZT = 1N×N and LOORF has the form:
gLOORF(Z) = 1
N f(Z)T
( IN×N − 1N−1 (1NxN − IN ) ) ( Z − 1Nσ(ϕ)T ) (6)

heading: 3.1 Multivariate CARMS
text: 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 are required. For D dimensions, we can sample Zd ∼ Cat(σ(ϕd)) independently and combine them into a D ×N × C tensor Z, with the corresponding D × C × C importance ratio tensor R. Below, we use superscripts and subscripts to index the first and second dimension of the tensors, with Z−d denoting excluding the dth row of Z . Focusing on the dth dimension, we have:
∇ϕdE [f(Z)] = EZ−d [ ∇ϕdEZd [ f(Z−d,Zd) ]] = EZ−d EZd  1
N(N−1) ∑ n̸=m 1 2 ( f(Z−dn ,Z d n)− f(Z−dm ,Zdm) ) (Zdn −Zdm)Zdn TRdZdm  = E  1 N(N−1)
∑ n ̸=m 1 2 (f(Zn)− f(Zm)) ( Zdn −Zdm ) Zdn TRdZdm  . Just like the univariate case, we only need N evaluations of f regardless of the dimensionality D.

heading: 3.2 Antithetic categorical variables
text: 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].

heading: 3.2.1 Inverse CDF sampling
text: The inverse transform sampling is commonly used to transform uniform random variables, which are easy to generate, to some other distribution:
x ∼ Fx(x) ⇐⇒ u ∼ Unif(0, 1), x = F−1x (u), where Fx(x) denotes the CDF of the desired distribution. For a categorical variable and some ordering of its probability vector p = (p1, ..., pC) the inverse CDF approach amounts to the following algorithm. Define the left and right boundaries:
li = i−1∑ k=1 pk, ri = i∑ k=1 pk = li + pi. (7)
Then setting z = ȷ̂ if u ∈ [lj , rj ], where u ∼ Unif(0, 1) produces a categorical variable because P (z = ȷ̂) = P (u ∈ [lj , rj ]) = rj− lj = pj . To obtain N antithetic categorical variables, we sample u ∼ CN from a copula and set zn in a vectorized manner. Importantly, this sampling approach allows us to easily analytically evaluate the bivariate PMF (proof in the Appendix):
P (z = ı̂, z′ = ȷ̂) = ΦC(ri, rj)− ΦC(ri, lj)− ΦC(li, rj) + ΦC(li, lj), (8) where ΦC denotes the bivariate CDF of the copula. However, we are not guaranteed a non-zero probability for any (i, j) pair within a particular ordering. For example, if the ordering is p = (0.1, 0.2, 0.7) and we use a bivariate copula (u, 1− u), then the probability of obtaining the pair (1,2) is zero:
P (z = 1, z′ = 2) = P (u ∈ [0, 0.1], 1− u ∈ [0.1, 0.3]) = P (u ∈ [0, 0.1], u ∈ [0.7, 0.9]) = 0.
Algorithm 1 Antithetic inverse CDF categorical sampling Input: Number of samples N , probabilities p = σ(ϕ), copula C. Sample u = (u1, ..., uN ) ∼ CN . Shuffle p according to Eq. 9, and define the left and right boundaries l, r according to Eq. 7. Set ∀n: zn = ȷ̂, where j is such that un ∈ [lj , rj ]. For i, j ∈ {1, ..., C}: compute Rij = pipj/P (zn = ı̂, zn′ = ȷ̂) according to Eq. 10. return: (z1, ...,zN ), R.
Algorithm 2 Antithetic Gumbel categorical sampling Input: Number of samples N , probabilities p = σ(ϕ), copula C. Sample uc = (uc1, ..., ucN )
iid∼ CN for each category c. Set ∀c, n : gcn = − ln(− ln(ucn)). Convert to categorical ∀n: znc = 1 if c = argmaxd gdn else 0. Approximate the bivariate PMF P (z = ı̂, z′ = ȷ̂) according to Eq 11. Set ∀i, j ∈ {1, ..., C}: Rij = pipj/P (z = ı̂, z′ = ȷ̂). return: (z1, ...,zN ), R.
However, we are always guaranteed a non-zero probability for the pair (1, C) (proof for the Dirichlet copula in the appendix). Therefore, we can average over C(C − 1)/2 orderings of p, where the ordering σij rearranges p such that categories i and j are the first and last category of the vector, i.e. σij(i) = 1, σij(j) = C. Thus, any pair of categories (i, j) has non-zero probability in at least one ordering, namely the ordering σij . There are many sets of orderings that satisfy this, but one simple way to do this is to translate the original vector p i elements to the right, then swap the last element with the jth, or more precisely:
σij(k) =  C, k = j j, k = C
(k − i+ 1) mod C, otherwise . (9)
The bivariate PMF given the set of orderings is:
P (z = ı̂, z′ = ȷ̂) = 2C(C−1) ∑ k<l P ( z = σkl(̂ı), z ′ = σkl(ȷ̂) | σkl(p) )
≈ 1 |O| ∑ o∈O P ( z = σo(̂ı), z ′ = σo(ȷ̂) | σo(p) ) . (10)
If the number of categories C is large, we can approximate the sum by sampling a random subset of orderings O, |O| ≪ C(C − 1)/2, as shown above. Note that we don’t need to calculate Eq. 10 for all possible pairs of categories, just those that were in a sample of N . This results in O(N2|O|) total computation per gradient, and is easily parallelizable. We summarize the inverse CDF sampling method in Algorithm 1.

heading: 3.2.2 Gumbel max sampling
text: It is not strictly necessary to analytically calculate the bivariate PMF. If the variance is not too large, a Monte Carlo estimation suffices, and we can use the well known Gumbel max trick [Gumbel, 1954, Jang et al., 2017, Maddison et al., 2017]:
z ∼ Cat(σ(θ)) ⇐⇒ gc iid∼ Gumbel(0, 1), zd = { 1, d = argmaxc gc + θc 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 = ȷ̂ iff j = argminiϵi, where ϵi ∼ Exp(eϕi). Since a Gumbel variable can be generated using a uniform: g = − ln(− ln(u)), u ∼ Unif(0, 1), we use a copula to generate uc = (uc1, ..., ucN ). 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 = [z1, ..,zN ]T as
before, a simple empirical estimate of the bivariate PMF computed using all pairs and only matrix operations is:
P (z = ı̂, z′ = ȷ̂) ≈ 1N(N−1)Z T (1NxN − IN )Z. (11)
We summarize the antithetic Gumbel sampling method in Algorithm 2, and in Fig. 1 we show an example of a 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 a good estimate of the ratio p(zi)p(zj)/p(zi, zj), which in turn requires a good estimate of the bivariate probability mass function (PMF) in the denominator. If we draw N samples to do this empirically, there is a small probability that some pair (i, j) does not occur at all, so in experiments we clip the maximum value of the ratio to 10 to avoid infinities.

heading: 4 Experimental results
text: In this section, we first illustrate the variance reduction that CARMS offers on a toy example. We also optimize a categorical variational autoencoder (VAE) [Kingma and Welling, 2014, Rezende et al., 2014], and a stochastic network for structured output prediction, which are standard tasks [Jang et 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 our 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 available2.

heading: 4.1 Toy example
text: We first showcase the variance reduction over other methods in a simple toy example, where we take the gradient with respect to the logits ϕ of:
L(ϕ) = E[f(z1, . . . ,zD)], zd ∼ Cat(σ(ϕd)), f(z1, . . . ,zD) = D∑
d=1 C∑ c=1 d · c · zdc.
For simplicity, let the number of categories, dimensions, and samples be C = D = N = 3. The probabilities are randomly sampled from a Dirichlet distribution: σ(ϕ) ∼ Dir(1C · α), but are identical for all methods for a given α. We vary the entropy of the probabilities from high (α = 1) to low (α = 1000). In a high entropy setting, there is little difference between the estimators, as the variance itself is very high, but differences emerge as we increase α and lower the entropy. The combined log variance of the gradient of each logit, for different methods and different α, is shown in Fig 2.
2https://github.com/alekdimi/carms

heading: 4.2 Categorical variational autoencoder
text: 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 [ ln
p(x|z)p(z) qϕ(z|x)
] ≈ N∑ n=1 ln p(x|zn)p(zn) qϕ(zn|x) , z1, . . . , zN ∼ CatN (σ(ϕ)),
where CatN (σ(ϕ)) denotes an N -variate categorical distribution with identical marginals. The number of categories is C ∈ {3, 5, 10} with D = ⌊200/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 a linear or nonlinear encoder/decoder pair on 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 ⌊200/C⌋ units, each of which is a C-way categorical variable. For the nonlinear case, there are additionally two layers of 200 units with LeakyReLU [Maas et al., 2013] activations. The prior logits are optimized using SGD with a learning rate of 10−2, whereas the encoder and decoder are optimized using Adam [Kingma and Ba, 2015] with a learning rate of 10−4, following Yin et al. [2019]. The optimization is run for 106 steps, with a 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 of the machine, with a total carbon emissions estimated to be 28.19 kg of CO2 [Lacoste et al., 2019]. An exception is UNORD for 10 samples, which was significantly more heavy in computation. Although the paper states that a step can be performed in O(2C), the provided code requires O(C!) evaluations per step, forcing us to use N = 8 samples at most.
In Fig 3, we plot the training and test log likelihood using 100 samples, and gradient variance w.r.t the logits over time, for a nonlinear network on dynamic MNIST. Similar plots for other datasets and network types can be found in the Appendix. Also shown in Table 1 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 a gap between ARSM (which used on average 2C evaluations per step, out of a maximum of C(C − 1)/2 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.

heading: 4.3 Structured prediction with stochastic categorical networks
text: We also compare all methods on the standard benchmark task of predicting the lower half of an image from the upper half, i.e., the conditional distribution p(xl |xu), where xu and xl denote the upper and lower half of an image, respectively. We use a stochastic categorical network to estimate this
distribution, with the objective: Ez∼p(zm |xu) [ 1 M ∑M m=1 ln p(xl | zm) ] , where z denotes a stochastic categorical layer. The encoder/decoder pair each contains one hidden layer with ⌊200/C⌋ latent variables and a LeakyReLU activation, with the optimization performed for C ∈ {3, 5, 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 2, 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 a much smaller gap.

heading: 5 Discussion
text: We have presented a novel approach for training categorical variables, which extends the ARMS estimator for the binary case. It 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 of the estimator reduces to LOORF. We showcase its usefulness on several datasets, a different number of categories and types of deep neural networks. In variational inference tasks, and conditional estimation, CARMS outperforms other state of the art estimators. The main limitation of this work is its specificity to categorical (including binary) variables. We hope to extend this, e.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 is also a general limitation that CARMS shares with other state-of-the-art 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 a 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 a person’s face or 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, we strongly believe that wide knowledge dissemination and open source code is crucial for reproducibility in all of science.

heading: Acknowledgments
text: The authors acknowledge the support of NSF IIS-1812699, the APX 2019 project sponsored by the Office of the Vice President for Research at The University of Texas at Austin, the support of a gift fund from ByteDance Inc., and the Texas Advanced Computing Center (TACC) for providing HPC resources that have contributed to the research results reported within this paper.

Review summaries:
**Summary** - The paper introduces a multi-sample gradient estimator for categorical variables.  - The estimator is based on the score-function and uses correlated samples to achieve a variance reduction.  - The estimator extends several recently proposed estimators naturally.  - The estimator yields marginal improvements over baselines in the experimental evaluations.   **Edit after Author Response** - 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.

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.

The paper proposes a categorical version of ARMS, namely CARMS. By utilizing REINFORCE together with copula sampling, CARMS remains an unbiased estimator and non-duplicated samples reduce the variance of gradients. The authors verify the proposed model through toy example, VAE, and the structured output prediction task. 

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.