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

1 Introduction . 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 . 2 Background . 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 . 2.1 LOORF . 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 . 2.2 ARMS . 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 . 3 CARMS . 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 )

Summary: **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.