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