Year: 2018

Abstract:
Reparameterization of variational auto-encoders is an effective method for reducing the variance of their gradient estimates. However, when the latent variables are discrete, a reparameterization is problematic due to discontinuities in the discrete space. In this work, we extend the direct loss minimization technique to discrete variational auto-encoders. We first reparameterize a discrete random variable using the argmax function of the Gumbel-Max perturbation model. We then use direct optimization to propagate gradients through the non-differentiable argmax using two perturbed argmax operations.

heading: 1 INTRODUCTION
text: Models with discrete latent variables drive extensive research in machine learning applications, such as language classification and generation (Yogatama et al., 2016; Hu et al., 2017; Shen et al., 2018), molecular synthesis (Kusner et al., 2017), or game solving (Mena et al., 2018). Compared to their continuous counterparts, discrete latent variable models can decrease the computational complexity of inference calculations, for instance, by discarding alternatives in hard attention models (Lawson et al., 2017), they can improve interpretability by illustrating which terms contributed to the solution (Mordatch & Abbeel, 2017; Yogatama et al., 2016), and they can facilitate the encoding of inductive biases in the learning process, such as images consisting of a small number of objects (Eslami et al., 2016) or tasks requiring intermediate alignments (Mena et al., 2018). Finally, in some cases, discrete latent variables are natural choices, including when modeling datasets with discrete classes (Rolfe, 2016; Jang et al., 2016; Maddison et al., 2016).
Models involving discrete latent variables are hard to train, with the key issue being to estimate the gradients of the resulting non-differentiable objectives. While one can use unbiased estimators, such as REINFORCE (Williams, 1992), their variance is typically high (Paisley et al., 2012; Mnih & Gregor, 2014; Titsias, 2015; Gu et al., 2015; Mnih & Rezende, 2016; Tucker et al., 2017). In variational auto-encoders (VAEs) with continuous latent variables, the reparameterization trick provides a successful alternative (Kingma & Welling, 2013; Rezende et al., 2014). However, it cannot be directly applied to non-differentiable objectives.
Recent work (Maddison et al., 2016; Jang et al., 2016) uses a relaxation of the discrete VAE objective, where latent variables follow a Gumbel-Softmax distribution. The Gumbel-Softmax approach is motivated by the connection of the Gumbel-Max perturbation model and the Gibbs distribution (Papandreou & Yuille, 2011; Tarlow et al., 2012; Hazan et al., 2013; Maddison et al., 2014). The Gumbel-Softmax distribution relaxes the non-differentiable Gumbel-Max trick for sampling from the categorical distribution by replacing the argmax operation with a softmax operation. The reformulation results in a continuous objective function, which allows the use of the reparameterization trick. (Mena et al., 2018) provide an extension of this approach for discrete latent structures, namely distributions over latent matchings.
Our work proposes minimizing the non-differentiable objective, by extending the direct loss minimization technique to generative models (McAllester et al., 2010; Song et al., 2016). Since categorical variables are represented by the Gibbs distribution, we start from the argmax formulation of the Gibbs distribution. We subsequently derive an optimization method that can propagate (biased) gradients through reparameterized argmax. The gradient of the resulting expectation is estimated by the difference of gradients of two argmax-perturbations.
We begin by introducing the notation, the VAE formulation and the equivalence between the Gibbs distribution and the Gumbel-Max perturbation model in Section 3. In Section 4.1, we use these components to reformulate the discrete VAE objective, with respect to the argmax prediction operation. We subsequently state and prove the main result that allows us to differentiate through the argmax function in Section 4.2. In Section 5, we extend this result to mixed discrete-continuous VAEs and to semi-supervised VAE objectives. Finally, we demonstrate the effectiveness of our approach on image generation.

heading: 2 RELATED WORK
text: Variational inference has been extensively studied in machine learning, see Blei et al. (2017) for a review paper. In our work, we consider variational Bayes bounds with a discrete latent space. Many approaches to optimizing the variational Bayes objective, that are based on samples from the distribution, can be seen as applications of the REINFORCE gradient estimator (Williams, 1992). These estimators are unbiased, but without a carefully chosen baseline, their variance tends to be too high for the estimator to be useful and considerable work has gone into finding effective baselines (Paisley et al., 2012). Other methods use various techniques to reduce the estimator variance (Ranganath et al., 2014; Mnih & Gregor, 2014; Gu et al., 2015).
Reparameterization is an effective method to reduce the gradient estimate variance in generative learning. Kingma & Welling (2013) have shown its effectiveness in auto-encoding variational Bayes (also called variational auto-encoders, VAEs) for continuous latent spaces. Rezende et al. (2014) demonstrated its effectiveness in deep latent models. The success of these works led to reparameterization approaches in discrete latent spaces. Rolfe (2016) and Vahdat et al. (2018) represent the marginal distribution of each binary latent variable as a continuous variable in the unit interval. This reparameterization allows the backpropogation of gradients through the continuous representation. These works are restricted to binary random variables, and as a by-product, it encourages highdimensional representations for which inference is exponential in the dimension size. In contrast, our work reparameterizes the discrete Gibbs latent model, using a Gumbel-Max perturbation model and directly propagates gradients through the reparameterized objective.
Maddison et al. (2016) and Jang et al. (2016) recently introduced a novel distribution, the Concrete distribution or the Gumbel-Softmax, that continuously relaxes discrete random variables. Replacing every discrete random variable in a model with a Concrete random variable, results in a continuous model, where the reparameterization trick is applicable. These works are close to ours, with a few notable differences. They use the Gumbel-Softmax function and their model is smooth and reparameterization may use the chain rule to propagate gradients. Similar to our setting, the GumbelSoftmax operation results in a biased estimate of the gradient. Different from our setting, the softmax operation relaxes the variational Bayes objective and results in a non-tight representation. Our work uses the Gumbel-Max perturbation model, which is an equivalent representation of the Gibbs distribution. With that, we do not relax the variational Bayes objective, while our argmax prediction remains non-differentiable and we cannot naively use the chain rule to propagate gradients. Instead, we develop a direct optimization method to propagate gradients through argmax operation using the difference of gradients of two max-perturbations. Our gradient estimate is biased, except for its limit.
Differentiating through argmax prediction was previously done in discriminative learning, in the context of direct loss minimization (McAllester et al., 2010; Song et al., 2016). Unfortunately, direct loss minimization cannot be applied to generative learning, since it does not have a posterior distribution around its argmax prediction. We apply the Gumbel-Max perturbation model to transform the argmax prediction to the Gibbs distribution. This also allows us to overcome the “general position” assumption in (McAllester et al., 2010; Song et al., 2016) using our “prediction generating function”.

heading: 3 BACKGROUND
text: To model the data generating distribution, we consider samples S = {x1, ..., xm} originating from some unknown underlying distribution. We explain the generation process of a parameterized model pθ(x), by minimizing its log-loss when marginalizing over its hidden variables z. Using variational
Bayes, we upper bound the log-loss of an observed data point∑ x∈S − log pθ(x) ≤ ∑ x∈S −Ez∼qφ log pθ(x|z) + ∑ x∈S KL(qφ(z|x)||pθ(z)) (1)
Typically, the model distribution pθ(x|z) follows the Gibbs distribution law pθ(x|z) = e−θ(x,z). When considering a discrete latent space, i.e., z ∈ {1, ..., k}, the approximated posterior distribution follows the Gibbs distribution law qφ(z|x) ∝ eφ(x,z). The challenge in generative learning is to reparameterize and optimize
−Ez∼qφ log pθ(x|z) = k∑ z=1 eφ(x,z)∑ ẑ e φ(x,ẑ) θ(x, z) (2)
In our work, we reparameterize the variational bound using the equivalence between Gibbs models and Gumbel-Max perturbation models.
Gumbel-Max perturbation models allow an alternative representation of Gibbs distributions qφ(z|x) ∝ eφ(x,z) that is based on the extreme value statistics of Gumbel-distributed random variables. Let γ be a random function that associates random variables γ(z) for each z = 1, ..., k. When the random perturbations follow the zero mean Gumbel distribution law, whose probability density function is g(γ) = ∏k z=1 e
−(γ(z)+c+e−(γ(z)+c)) for the Euler constant c ≈ 0.57, we obtain the following identity between Gibbs models and Gumbel-Max perturbation models1 (cf. Kotz & Nadarajah (2000))
eφ(x,z)∑ ẑ e φ(x,ẑ) = Pγ∼g[zφ+γ = z], where zφ+γ def = arg max ẑ=1,...,k {φ(x, ẑ) + γ(ẑ)} (3)
For completeness, a proof for this statement appears in Appendix A.

heading: 4 REPARAMETERIZATION AND DIRECT OPTIMIZATION
text: In the following section, we reformulate the discrete VAE objective using the Gumbel-Max perturbation model. We then derive an optimization method that directly propagates gradients through the reparameterized argmax function.

heading: 4.1 VAES WITH GUMBEL-MAX PERTURBATION MODELS
text: Perturbation models allow an alternative representation of Gibbs distributions qφ(z|x) ∝ eφ(x,z). Using the Gumbel-Max perturbation model in Equation (3), the negative log likelihood in Equation (2) takes the form
−Ez∼qφ log pθ(x|z) = k∑ z=1 Pγ∼g[zφ+γ = z]θ(x, z) = k∑ z=1 Pγ∼g[zφ+γ = z]θ(x, zφ+γ) (4)
= k∑ z=1 Eγ∼g[1zφ+γ=zθ(x, zφ+γ)] = Eγ∼g[θ(x, zφ+γ)] (5)
These quantities result from applying the law of total expectation, while realizing the probability events for zφ+γ : The last equality in Equation (4) holds since we restrict to the space zφ+γ = z. The last equality in Equation (5) holds by the linearity of the expectation, i.e., ∑k z=1 Eγ∼g[1zφ+γ=zθ(x, zφ+γ)] = Eγ∼g[ ∑k z=1 1zφ+γ=zθ(x, z
φ+γ)] and the fact that∑k z=1 1zφ+γ=z = 1.
The gradient of the decoder log probability θ(x, zφ+γ) with respect to its parameters is derived by the chain rule. The main challenge is to evaluate the gradient of Eγ∼g[θ(x, zφ+γ)] with respect to the encoder parameters, since the chain rule does not propagate through the argmax function zφ+γ .
1The set argmaxẑ=1,...,k{φ(x, ẑ) + γ(ẑ)} is the set of all maximal arguments, and does not always consist of a single element. However, since the Gumbel distribution is continuous, the γ for which their set argmaxẑ=1,...,k{φ(x, ẑ) + γ(ẑ)} consists more than a single element has a measure of zero. For notational convenience, when we consider integrals (or probability distributions), we ignore measure zero sets.
4.2 DIRECT OPTIMIZATION THROUGH argmax
Our main result is presented in Theorem 1 and shows how to compute the gradient of the reparameterized discrete VAE, i.e., Eγ∼g[θ(x, zφ+γ)] with respect to the encoder parameters v. In the following we omit γ ∼ g for brevity. To make the encoder parameters v explicit in our notation, we denote the encoder function by φv(x, z).
Instrumental to our approach, is a novel “prediction generating function”.
G(v, ) = Eγ [max ẑ { θ(x, ẑ) + φv(x, ẑ) + γ(ẑ)}] (6)
The proof of Theorem 1 is composed from three steps:
1. We prove that G(v, ) is a smooth function of v, . Therefore, the Hessian of G(v, ) exists and it is symmetric, namely
∂v∂ G(v, ) = ∂ ∂vG(v, ). (7)
2. We show that encoder gradient is apparent in the Hessian:
∂v∂ G(v, 0) = ∇vEγ [θ(x, zφv+γ)]. (8)
3. We derive our update rule as the complement representation of the Hessian:
∂ ∂vG(v, 0) = lim →0
1 ( Eγ [∇vφv(x, z θ+φv+γ)−∇vφv(x, zφv+γ)] ) (9)
The reparameterized gradient computation appears in Equation (10) and it is formally derived in Theorem 1.
Theorem 1. Assume φv(x, z) is a smooth function of v. Then
∇vEγ [θ(x, zφv+γ)] = lim →0
1 ( Eγ [∇vφv(x, z θ+φv+γ)−∇vφv(x, zφv+γ)] ) (10)
Proof. First, we prove that G(v, ) is a smooth function. Recall, g(γ) is the zero mean Gumbel probability density function. Applying a change of variable γ̂(z) = θ(x, ẑ) + φv(x, ẑ) + γ(ẑ), we obtain
G(v, ) = ∫ ∞ −∞ g(γ)max ẑ { θ(x, ẑ) + φv(x, ẑ) + γ(ẑ)}dγ = ∫ ∞ −∞ g(γ̂ − θ − φv)max ẑ {γ̂(ẑ)}dγ̂.
Since g(γ̂ − θ− φv) is a smooth function of and φv(x, z) and φv(x, z) is a smooth function of v, we conclude that G(v, ) is a smooth function of v, . Therefore, the Hessian of G(v, ) exists and symmetric, i.e., ∂v∂ G(v, ) = ∂ ∂vG(v, ). We thus proved Equation (7).
To prove Equations (8) and (9) we differentiate under the integral, both with respect to and with respect to v. We are able to differentiate under the integral, since g(γ̂− θ−φv) is a smooth function of and v and its gradient is bounded by an integrable function (cf. Folland (1999), Theorem 2.27).
We turn to prove Equation (8). We begin by noting that maxẑ{ θ(x, ẑ) + φv(x, ẑ) + γ(ẑ)} is a maximum over linear function of , thus by Danskin Theorem (cf. Bertsekas et al. (2003), Proposition 4.5.1) holds ∂ (maxẑ{ θ(x, ẑ) + φv(x, ẑ) + γ(ẑ)}) = θ(x, z θ+φv+γ). By differentiating under the integral, ∂ G(v, ) = Eγ [∂ (θ(x, z θ+φv+γ) + φv(x, z θ+φv+γ) + γ(z θ+φv+γ))] = Eγ [θ(x, z θ+φv+γ)]. We obtain Equation (8) by differentiating under the integral, now with respect to v, and setting = 0.
Finally, we turn to prove Equation (9). By differentiating under the integral ∂vG(v, ) = Eγ [∇vφv(x, z θ+φv+γ)]. Equation (9) is attained by taking the derivative with respect to = 0 on both sides.
The theorem follows by combining Equation (7) when = 0, i.e., ∂v∂ G(v, 0) = ∂ ∂vG(v, 0) with the equalities in Equations (8) and (9).
The gradient estimate in Theorem 1 is unbiased in the limit → 0. However, for small epsilon the gradient is either zero, when z θ+φ+γ = zφ+γ , or very large, since the gradients difference is multiplied by 1/ . For large we often obtain a moderate non-zero gradient. In practice we use ≥ 0.1 which means that the gradient estimate is biased. This bias-variance tradeoff is evaluated in Figure 1, on an encoder with three layers of sizes (784, 300, 10), a discrete latent structure of size 10 and a matching decoder. In this evaluation we sampled 200 gradients for a fixed to evaluate the bias and standard deviation for each coordinate of the gradient. We report in the graph these quantities.
The above theorem closely relates to the direct loss minimization technique (cf. McAllester et al. (2010); Song et al. (2016)), which, in our setting, can be used to compute the gradient of Exθ(x, zφ). The direct loss minimization predicts a single zφ for a given x and, therefore, cannot generate a posterior distribution on all z = 1, ..., k, i.e., it lacks a generative model that exists in Gumbel-Max perturbation models.

heading: 5 EXTENSIONS

heading: 5.1 SEMI SUPERVISED GENERATIVE MODELS
text: The main advantage in our framework is that learning a discrete VAE using Gumbel-Max reparameterization is intimately related to predicting a discrete latent label through the argmax operation zφ+γ . Therefore, semi-supervised VAEs are naturally integrated into our reparameterization framework using any loss function. Formally, assume that a subset of the data is labeled, i.e., S1 = {(x1, z1), ..., (xm1 , zm1)}. In semi-supervised learning, we may add to the learning objective the loss function `(z, zφ+γ), for any (x, z) ∈ S1, to better control the prediction of the latent space. The semi-supervised discrete VAEs objective function is∑
x∈S Eγ [θ(x, zφ+γ)] + ∑ (x,z)∈S1 Eγ [`(z, zφ+γ)] + ∑ x∈S KL(qφ(z|x)||pθ(z)) (11)
The supervised component is explicitly handled by Theorem 1 and optimizing a semi-supervised discrete VAEs is straight forward in our framework. Our supervised component is intimately related to direct loss minimization (McAllester et al., 2010; Song et al., 2016), which, in our setting, minimizes ∑ (x,z)∈S1 `(z, z
φ). Compared to direct loss minimization, our work adds random perturbation γ to the encoder and thus overcomes the “general position” assumption of direct loss minimization. This addition allows us to introduce the “prediction generating function”, which greatly simplifies our proof. In addition, the added random perturbation γ allows us to use a generative model to prediction, namely, we can randomly generate different explanations zφ+γ while the direct loss minimization allows a single explanation in the form of zφ.
Semi-supervised generative models require the discrete space to have a semantic meaning. For example, in generating images of digits, the image may be represented both by the discrete digit class (e.g., 0, 1, ...) and the continuous style (e.g., bold, tilted,...). Following (Kingma et al., 2014; Jang et al., 2016) we use a mixture of discrete and continuous latent variables in our semi-supervised setting, to capture both the discrete class information and the continuous class style. The network architecture consists of an initial encoder φ to process the input, which is fed into two separate encoders φd and φc that encode the discrete and continuous latent spaces respectively. The mixed discrete-continuous latent space consists of the matrix diag(zφd+γ) · zc, i.e, if zφd+γ = i then this matrix is all zero, except for the i-th row, which consists of the independent Gaussian random variables zc. An illustration of our mixed discrete-continuous VAE architecture appears in Figure 1. The gradient of the continuous component of the encoder is computed by reparameterization trick. The gradient of the discrete component is computed according to Theorem 1.

heading: 5.2 DISCRETE PRODUCT SPACES
text: Discrete VAEs are being applied also to discrete product spaces (Maddison et al., 2016; Jang et al., 2016). A single discrete random variable z ∈ {1, .., k} may not represent the variability of the generative process. Therefore, one may represent the latent space with a discrete product space z = (z1, ..., zn), where zi ∈ {1, ..., k}. We note that Theorem 1 holds without any change, if we match the structures of θ(x, z), φ(x, z) and γ(z), i.e., we have an independent Gumbel random variable for each z = (z1, ..., zn). However, this may be computationally inefficient, since the number of Gumbel random variables in this case is exponential in n. The gradient in Theorem 1 requires two argmax-perturbations: zφ+γ that requires the encoder output and z θ+φ+γ that requires both the encoder and the decoder log probability. Similarly to Maddison et al. (2016) and Jang et al. (2016) we use a conditionally independent encoder, i.e., φ(x, z) = ∑n i=1 φi(x, zi) for which we can compute zφ+γ efficiently. In contrast to their approach, we use an approximation for the decoder log probability in order approximate z θ+φ+γ efficiently.
In discrete product spaces we use low-dimensional random perturbation γ(z) = ∑n i=1 γi(zi). Thus instead of exponential number of random variables we use a linear number of random variables (linear in k and n). To compute zφ+γ efficiently, we first note that the φ(x, z) decomposes according to its dimensions, i.e., φ(x, z) = ∑n i=1 φi(x, zi). Thus the perturb-max argument zφ+γ = (zφ+γ1 , ..., z φ+γ n ) also decomposes, i.e., z φ+γ i = argmaxẑ=1,...,k{φi(x, ẑ)+γi(ẑ)}. Therefore, our approach is exact for the argmax-perturbation zφ+γ
Our approximation is not necessarily exact for the argmax-perturbation z θ+φ+γ since the decoder does not decompose according to its dimensions. To be able to compute z θ+φ+γ efficiently, we use the fact we can compute zφ+γ efficiently and apply a low dimensional approximation for the decoder log probability θ̃(x, z) = ∑n i=1 θ̃i(zi), where θ̃i(zi) = θ(z γ1+φ1 , ..., zi, ..., z γn+φn). With this in mind, we approximate z θ+φ+γ using z θ̃+φ+γ , and its coordinates are computed efficiently by z θ̃+φ+γi = argmaxẑ=1,...,k{ θ̃i(x, ẑi) + φi(x, ẑ) + γi(ẑ)}.
Lastly, the KL-divergence in Equation (1) can utilize the decomposed encoder φ(x, z) =∑n i=1 φi(x, zi). In this case q(z|x) ∝ eφ(x,z) and q(zi|x) ∝ eφi(x,zi), therefore q(z|x) =∏n i=1 q(z|x) and KL(qφ(z|x)||pθ(z)) = ∑n i=1KL(qφi(zi|x)||pθ(zi)), where pθ(zi) is the marginal probability of pθ(z) with respect to its i−th entry.

heading: 6 EXPERIMENTS
text: We start by comparing our method to the state-of-the-art (Maddison et al., 2016; Jang et al., 2016; Tucker et al., 2017; Grathwohl et al., 2017). We also investigate our approximation for discrete product spaces. We conclude with a set of experiments that demonstrate the effectiveness of our approach in semi-supervised learning of VAEs and the importance of weak supervision in image generation.

heading: 6.1 DISCRETE VAES
text: We begin our experiments by comparing the test loss bound of our direct optimization biased gradient estimator with the biased gradient estimator Gumbel-Softmax (GSM) of Maddison et al. (2016) and Jang et al. (2016), and the unbiased estimators REBAR (Tucker et al., 2017) and RELAX (Grathwohl et al., 2017). We performed these experiments using the binarized MNIST dataset (Salakhutdinov & Murray, 2008), and the standard 50,000/10,000 split into training/testing sets. Following Jang et al. (2016) we set our learning rate to 1e − 3 and the annealing rate to 1e − 5 and we used their annealing schedule every 1000 step, setting the minimal to be 0.1. When considering a latent space with n random variables z = (z1, ..., zn) where zi ∈ {1, ..., k}, our encoder has a (784, n×k) linear layer and the decoder has a (n × k, 784) linear layer. We applied REBAR and RELAX only to discrete binary random variables2. The results appear in Figure 2. For the 20 × 2 and 20 × 10 architectures, we applied our approximation that is described in Section 5.2.
Next, we explore our discrete product space approximation. We compare a latent space with n random variables z = (z1, ..., zn) with zi ∈ {1, ..., k} with a latent space that has a single random variable z ∈ {1, ..., kn}. When n = 3, k = 2, a discrete VAE with a single categorial random variable that has 23 possible values achieved a test log-loss bound of 168.57, while 3 binary variables achieved test log-loss bound of 177.9. When n = 6, k = 2, a single categorial random variable achieved a test log-loss bound of 147.59, while 6 binary variables achieved test log-loss bound of 161.09. One can see that there is a gap when using our approximation that increases with n.
The main advantage of our framework is that it seamlessly integrates semi-supervised learning. For these experiments, we used a mixed continuous discrete architecture, where the architecture of the encoder consists of a network φ which has a (784, 400) linear layer, followed by a ReLU, and a (400, 200) linear layer. The output of this later is fed both to a discrete encoder φd and a continuous encoder φc. The discrete latent space consists of zd ∈ {1, ..., 10} and its encoder φd consists of a (200, 100) linear layer, a ReLU, and a (100, 10) linear layer. The continuous latent space considers k = 10, c = 20, and its encoder φc consists of a (200, 100) linear layer, a (100, 66) linear layer followed by ReLU and dropout and a (66, 40) linear layer to estimate the mean and variance of c−dimensional Gaussian random variables z1, ..., zk.
2For REBAR and RELAX we used the code in https://github.com/duvenaud/relax.
Following Kingma et al. (2014), we conducted a quantitive experiment with weak supervision on MNIST with 50/100/300/600 labeled examples out of the 50, 000 training examples. For labeled examples, we set the perturbed label z θ+φ+γ+` to be the true label. This is equivalent to using the indicator function over the space of correct predictions. A comparison of our method with Gumbel-Softmax appears in Figure 3. We can see that direct loss minimization combined with VAE improves the performance even with weak supervision, e.g., with only 50 examples direct loss minimization with VAE achieves better accuracy than semi-supervised Gumbel-Softmax VAE. We note that we cannot compare the objective function of both methods, as our objective considers direct loss minimization, however, we are able to compare the test log-loss bound, see Figure 3.
Supervision in generative models also greatly helps to control discrete semantics within images. We learn to generate images using k = 8 discrete attributes of the CelebA dataset (cf. Liu et al. (2015)) while using our semi-supervised VAE. For this task, we use convolutional layers for both the encoder and the decoder, except the last two layers of the continuous latent model which are linear layers that share parameters over the 8 possible representations of the image. In Figure 4, we show generated images with discrete semantics turned on/off (with/without glasses, with/without smile, woman/man).

heading: 7 DISCUSSION AND FUTURE WORK
text: In this work, we use the Gumbel-Max trick to reparameterize discrete VAEs using the argmax prediction operator and show how to propagate gradients through the non-differentiable argmax function. We show that this approach compares favorably to state-of-the-art methods, and extend it to semi-supervised learning and image attribute generation.
These results can be taken in a number of different directions. Our gradient estimation is practically biased, while REINFORCE is an unbiased estimator. Our methods may benefit from the REBAR/RELAX framework, which directs biased gradients towards the unbiased gradient (Tucker et al., 2017). There are also open problems when fitting this approach to structured latent spaces (Mena et al., 2018; Jin et al., 2018), as well as estimating its KL-divergence (Roeder et al., 2017). There are also optimization-related questions that arise from our work: the interplay of and the learning rate is unexplored and might be correlated. The number of stochastic gradient steps, interleaving Gumbel perturbation with batch samples, might also benefit from a rigorous investigation.
The direct optimization approach we present is general and may be applied beyond VAEs, including reinforcement learning and attention models. Further investigation in this direction is required.

heading: A GUMBEL-MAX PERTURBATION MODEL AND THE GIBBS DISTRIBUTION
text: Theorem 2. Gumbel & Lieblein (1954); Luce (1959); McFadden (1973) Let γ be a random function that associates random variable γ(z) for each z = 1, ..., k whose distribution follows the zero mean Gumbel distribution law, i.e., its probability density function is g(t) = e−(t+c+e
−(t+c)) for the Euler constant c ≈ 0.57. Then
eφ(x,z)∑ ẑ e φ(x,ẑ) = Pγ∼g[z = zφ+γ ], where zφ+γ def = arg max ẑ=1,...,k {φ(x, ẑ) + γ(ẑ)} (12)
Proof. Let G(t) = e−e −(t+c) be the Gumbel cumulative distribution function. Then
Pγ∼g[z = zφ+γ ] = Pγ∼g[z ∈ arg max ẑ=1,..,k {φ(x, ẑ) + γ(ẑ)}] (13)
= ∫ g(t− φ(x, z)) ∏ ẑ 6=z G(t− φ(x, ẑ))dt (14)
Since g(t) = e−(t+c)G(t) it holds that∫ g(t− φ(z)) ∏ ẑ 6=z G(t− φ(ẑ))dt = ∫ e−(t−φ(x,z)+c)G(t− φ(x, z)) ∏ ẑ 6=z G(t− φ(x, ẑ))dt
= eφ(x,z)
Z (15)
where 1Z = ∫ e−(t+c) ∏k ẑ=1G(t − φ(ẑ))dt is independent of z. Since Pγ∼g[z = zφ+γ ] is a distri-
bution then Z must equal to ∑k ẑ=1 e φ(x,ẑ).

Review summaries:
This paper proposes combining the Gumbel-max trick and "direct loss optimization" for variance reduction in VAEs with discrete latent variables. This is a natural combination (in hindsight), since the Gumbel-max trick turns sampling into non-differentiable optimization, and direct loss optimization provides a way to optimize the expected value of a non-differentiable loss. The paper is well-written for the most part and is backed by good experimental results. However it like some of the mathematical details and some of the exposition could be greatly improved.

The authors propose a method to apply the reparametrization trick when the random variables of interest are discrete. Their technique is based on a formulation of the objective function in terms of Gumbel-Max operators. They propose a derivation of the gradient in terms of an auxiliary variable \epsilon, such that the resulting gradient estimate is biased but the bias is reduced as \epsilon approaches zero, at the cost of increasing variance. Experiments are performed with VAE including discrete latent variable models. The authors show how their method converges faster than other baselines formed by estimators of the gradient given by the REBAR, RELAX and Gumbel-soft-max methods. In experiments with semi-supervised VAEs, their method outperforms the Gumbel softmax method in terms of accuracy and objective function.