Generalized Gumbel-Softmax Gradient Estimator for Generic Discrete Random Variables

1 INTRODUCTION . Stochastic computational graphs , including deep generative models such as variational autoencoders , are widely used for representation learning . Optimizing the network parameters through gradient methods requires an estimation of the gradient values , but the stochasticity requires the computation of expectation , which differentiates this problem from the deterministic gradient of ordinary neural networks . There are two common ways of obtaining the gradients : score function-based methods and reparameterization methods . The score function-based estimators tend to result in unbiased gradients with high variances , while the reparameterization estimators seem to result in biased gradients with low variances ( Xu et al. , 2019 ) . Hence , the core technique of the score-function based estimators becomes reducing the variances of gradients to achieve stable and fast optimizations . Meanwhile , utilizing the reparameterization estimators requires the differentiable non-centered parameterization ( Kingma & Welling , 2014a ) of random variables . If we focus on the reparameterization estimators , one of the most popular examples is the reparameterization in the Gaussian variational autoencoder ( VAE ) ( Kingma & Welling , 2014b ) , which has an exact reparameterization form . Other VAEs with explicit priors suggest the reparameterization tricks with approximations ( Nalisnick & Smyth , 2017 ; Joo et al. , 2020 ) . For continuous random variables , it is feasible to estimate gradients with automatic differentiation by utilizing a transport equation ( Jankowiak & Obermeyer , 2018 ) or an implicit reparameterization ( Figurnov et al. , 2018 ) . However , these methods are not applicable to discrete random variables , due to the non-differentiability . Recently , some discrete random variables , such as Bernoulli or categorical random variables , have been well-explored in terms of the reparameterization method by overcoming such difficulty through a continuous relaxation ( Jang et al. , 2017 ; Maddison et al. , 2017 ) . However , other discrete distributions have not been explored from the learning perspective in the deep generative modeling community , for example , Poisson , binomial , multinomial , geometric , negative binomial distributions , etc . Prior works on graphical models , such as Ranganath et al . ( 2015 ; 2016 ) , utilized Poisson latent variables for the latent counting . Another line of work ( Wu et al. , 2020 ) utilized the Gaussian approximation on the Poisson latent variable to count the number of words , which can be a poor approximation if the rate parameter is small . In this sense , study on the stochastic gradient estimator for discrete distributions is required in the deep generative modeling community , which broadens the choice of prior assumptions and the utilization of various distributions . This paper proposes a generalized version of the Gumbel-Softmax reparameterization trick , which can be utilized to generic discrete random variables through continuous relaxation , namely Generalized Gumbel-Softmax ( GENGS ) . The key ideas of GENGS are ( 1 ) a conversion of the sampling process to one-hot categorical selection process ; ( 2 ) a reversion of the selected category in the one-hot form to the original sample value ; and ( 3 ) a relaxation of the categorical selection process into the continuous form . To implement these steps , GENGS first truncates discrete random variables to approximate the distribution with the finite number of possible outcomes . Afterward , GENGS utilizes the Gumbel-Softmax trick together with a special form of a linear transformation . Our main theorem supports that the proposed GENGS is applicable to general discrete random variables , other than the Bernoulli and the categorical . The GENGS experiments show the efficacy with synthetic examples and VAEs , as well as the usability in topic model application . 2 PRELIMINARY : REPARAMETERIZATION TRICK & GUMBEL-SOFTMAX . 2.1 BACKPROPAGATION THROUGH STOCHASTIC NODES WITH REPARAMETERIZATION TRICK . Suppose we have a stochastic node , or a latent variable , z ∼ p ( z|θ ) , where the distribution depends on parameter θ . The goal is optimizing the loss function L ( θ , η ) = Ez∼p ( z|θ ) [ fη ( z ) ] , where fη is a continuous and differentiable function with respect to η , i.e. , neural networks . To optimize the loss function in terms of θ through the gradient methods , we need to find ∇θL ( θ , η ) = ∇θEz∼p ( z|θ ) [ fη ( z ) ] , which can not be directly computed with its original form . To compute ∇θL ( θ , η ) , the reparameterization trick alternatively introduces an auxiliary variable ∼ p ( ) , which takes over all randomness of the latent variable z , so the sampled value z can be re-written as z = g ( θ , ) , with a deterministic and differentiable function g in terms of θ . Figure 1 ( a ) illustrates the reparameterization trick : the shaded nodes indicate random nodes , and the dotted lines denote sampling processes . Here , the gradient of the loss function with respect to θ is derived as ∇θL = ∇θEz∼p ( z|θ ) [ fη ( z ) ] = E ∼p ( ) [ ∇θfη ( g ( θ , ) ) ] = E ∼p ( ) [ ∇gfη ( g ( θ , ) ) ∇θg ( θ , ) ] ( 1 ) where the last term of Equation 1 is now achievable . A condition on enabling the reparameterization trick is the assumption of the continuity of the random variable z , so the distribution of z is limited to a class of continuous distributions . To utilize the differentiable reparameterization trick on discrete random variables , continuous relaxation can be applied : for example , a relaxation from the categorical distribution to the Gumbel-Softmax distribution , described in the next subsection . 2.2 REPARAMETERIZATION TRICK ON CATEGORICAL RANDOM VARIABLE . A Gumbel-Max trick ( Gumbel , 1948 ) is a procedure for sampling a one-hot categorical value from the Gumbel distribution , instead of direct sampling from a categorical distribution . This implies that the categorical random variable X ∼ Categorical ( π ) , where π lies on the ( n− 1 ) -dimensional simplex ∆n−1 , can be reparameterized by the Gumbel-Max trick : ( 1 ) sample uj ∼ Uniform ( 0 , 1 ) to generate a gumbel sample gj = − log ( − log uj ) for each j = 1 , · · · , n ; and ( 2 ) compute k = argmaxnj=1 [ log πj + gj ] , where π is a categorical parameter . This procedure generates a one-hot sample x such that xj = 0 for j 6= k and xk = 1 with P ( Xk = 1 ) = πk . We denote GM ( π ) to be the distribution whose samples are generated by the Gumbel-Max trick . A Gumbel-Softmax trick ( Jang et al. , 2017 ; Maddison et al. , 2017 ) is an alternative to the GumbelMax trick that continuously relaxes a categorical random variable . The Gumbel-Softmax utilizes the softmax with a temperature τ > 0 , instead of the argmax in the sampling process , which enables ( 1 ) relaxing the discreteness of the categorical random variable to the one-hot-like form x ( τ ) = softmax ( log π+g τ ) in the continuous domain ; and ( 2 ) approximating the Gumbel-Max by taking τ small enough . Lately , the Gumbel-Softmax estimator has been widely used to reparameterize categorical random variables , such as RelaxedOneHotCategorical in TensorFlow ( Abadi et al. , 2016 ) . We denote GS ( π , τ ) to be the distribution generated by the Gumbel-Softmax trick . 3 PROCESS OF GENGS REPARAMETERIZATION . This section discusses the process of GENGS to help understand the concept with minimal theoretical details , and Section 4 provides the theoretical background of GENGS . The three steps of GENGS are the following : ( 1 ) approximate a discrete distribution by truncating the distribution ; ( 2 ) reparameterize the truncated distribution with the Gumbel-Max trick and the linear transformation T , which will be introduced below ; and ( 3 ) relax the discreteness by replacing the Gumbel-Max trick in Step 2 with the Gumbel-Softmax trick . Figure 1 ( b ) illustrates the full steps of the GENGS trick . Step 1 . Truncate the discrete distribution to finitize the number of possible outcomes . Suppose X ∼ Poisson ( 100 ) , which has a mode near at x = 100 , and near-zero probabilities at x < 50 and x > 150 . The key idea of the first step is ignoring the outcomes of near-zero probabilities at certain levels ( ex . x < 50 and x > 150 ) and only focusing on the probable samples of meaningful probabilities ( ex . 50 ≤ x ≤ 150 ) , i.e. , truncating the distribution , which finitizes the support of the distribution . Now , suppose we have a discrete random variable X ∼ D ( λ ) , and its truncated random variable Z ∼ TD ( λ , R ) , where R denotes the truncation range that needs to be pre-defined . Proposition 3 in Section 4 provides theoretical reason that Z approximates X . Since we finitized the support by the truncation , we may assume Z has a support C = { c0 , · · · , cn−1 } of n possible outcomes and its corresponding constant outcome vector c = ( c0 , · · · , cn−1 ) . Note that the ordering of ck is not significant , and Appendix E provides examples of the setting on c. Step 2 . Divide sampling process of Z into two-fold : select a one-hot category of Z , and revert the selected one-hot category into the original value . For example , if the sampled value of Z is c2 ∈ C , we will first focus on the one-hot category class vector one hot ( c2 ) = ( 0 , 0 , 1 , 0 , · · · , 0 ) , rather than the sampled value c2 . Such a one-hot categorical selection process is possible by utilizing the categorical selection w ∼ Categorical ( π ) or its reparameterized version , the Gumbel-Max trick GM ( π ) . Here , the categorical parameter π = ( π0 , · · · , πn−1 ) can be directly calculated by the explicit probability mass funciton ( PMF ) of the distribution , i.e. , πk = P ( Z = ck ) . However , the PMF of the truncated distribution requires a modification from the PMF of the original distribution , which is determined by how we define Z from X . See Definition 1 , 2 , and Appendix A for detailed configuration of π . Suppose we now have a one-hot categorical sample w from the categorical parameter π . Afterward , we revert the one-hot selected categorical vector w = ( w0 , · · · , wn−1 ) into the original sample value with a linear transformation T ( w ) = ∑ k wkck = ∑ k w c. Proposition 4 shows the validity of the alternative sampling process in Section 4 . Step 3 . Relax the one-hot categorical selection into the continuous form by utilizing the GumbelSoftmax trick . Up to now , the sole shortage of the reparameterization trick is the differentiability due to the one-hot categorical selecting Gumbel-Max process . Then , as in Section 2.2 , the process can be continuously relaxed with the Gumbel-Softmax trick GS ( π , τ ) for some temperature τ . Theorem 5 in Section 4 shows that the alternative sampling process still holds under the continuous relaxation .

Summary: The paper presents a generalization of the Gumbel-Softmax gradient estimator. The original Gumbel-Softmax is usually applied to Bernoulli and categorical random variables. The method proposed in the paper attempts to extend it applicability to other discrete distributions, such as Poisson, multinomial, geometric, among others. The main ideas of the approach are: (1) Random variables that may take countably infinite values are truncated, (2) The sampling process of the random variable is converted to a one-hot scheme (where Gumbel-Softmax relaxation is applied), (3) ``One-hot'' samples are reverted to the original sample space.