Title: CATEGORICAL REPARAMETERIZATION WITH GUMBEL-SOFTMAX

Abstract: Categorical variables are a natural choice for representing discrete structure in the world. However, stochastic neural networks rarely use categorical latent variables due to the inability to backpropagate through samples. In this work, we present an efficient gradient estimator that replaces the non-differentiable sample from a categorical distribution with a differentiable sample from a novel Gumbel-Softmax distribution. This distribution has the essential property that it can be smoothly annealed into a categorical distribution. We show that our Gumbel-Softmax estimator outperforms state-of-the-art gradient estimators on structured output prediction and unsupervised generative modeling tasks with categorical latent variables, and enables large speedups on semi-supervised classification.

1. INTRODUCTION:

Stochastic neural networks with discrete random variables are a powerful technique for representing distributions encountered in unsupervised learning, language modeling, attention mechanisms, and reinforcement learning domains. For example, discrete variables have been used to learn probabilistic latent representations that correspond to distinct semantic classes (Kingma et al., 2014) , image regions (Xu et al., 2015) , and memory locations (Graves et al., 2014; Graves et al., 2016) . Discrete representations are often more interpretable (Chen et al., 2016) and more computationally efficient (Rae et al., 2016) than their continuous analogues.

However, stochastic networks with discrete variables are difficult to train because the backpropagation algorithm -while permitting efficient computation of parameter gradients -cannot be applied to non-differentiable layers. Prior work on stochastic gradient estimation has traditionally focused on either score function estimators augmented with Monte Carlo variance reduction techniques (Paisley et al., 2012; Mnih & Gregor, 2014; Gu et al., 2016; Gregor et al., 2013) , or biased path derivative estimators for Bernoulli variables (Bengio et al., 2013) . However, no existing gradient estimator has been formulated specifically for categorical variables. The contributions of this work are threefold:

1. We introduce Gumbel-Softmax, a continuous distribution on the simplex that can approximate categorical samples, and whose parameter gradients can be easily computed via the reparameterization trick.

2. We show experimentally that Gumbel-Softmax outperforms all single-sample gradient estimators on both Bernoulli variables and categorical variables.

3. We show that this estimator can be used to efficiently train semi-supervised models (e.g. Kingma et al. (2014) ) without costly marginalization over unobserved categorical latent variables.

The practical outcome of this paper is a simple, differentiable approximate sampling mechanism for categorical variables that can be integrated into neural networks and trained using standard backpropagation.

2. THE GUMBEL-SOFTMAX DISTRIBUTION:

We begin by defining the Gumbel-Softmax distribution, a continuous distribution over the simplex that can approximate samples from a categorical distribution. Let z be a categorical variable with class probabilities π 1 , π 2 , ...π k . For the remainder of this paper we assume categorical samples are encoded as k-dimensional one-hot vectors lying on the corners of the (k -1)-dimensional simplex, ∆ k-1 . This allows us to define quantities such as the element-wise mean E p [z] = [π 1 , ..., π k ] of these vectors.

The Gumbel-Max trick (Gumbel, 1954; Maddison et al., 2014) provides a simple and efficient way to draw samples z from a categorical distribution with class probabilities π:

EQUATION

where g 1 ...g k are i.i.d samples drawn from Gumbel(0, 1) 1 . We use the softmax function as a continuous, differentiable approximation to arg max, and generate k-dimensional sample vectors y ∈ ∆ k-1 where

EQUATION

The density of the Gumbel-Softmax distribution (derived in Appendix B) is:

EQUATION

This distribution was independently discovered by Maddison et al. (2016) , where it is referred to as the concrete distribution. As the softmax temperature τ approaches 0, samples from the Gumbel-Softmax distribution become one-hot and the Gumbel-Softmax distribution becomes identical to the categorical distribution p(z). the expected value of a Gumbel-Softmax random variable approaches the expected value of a categorical random variable with the same logits. As the temperature increases (τ = 1.0, τ = 10.0), the expected value converges to a uniform distribution over the categories. (b) Samples from Gumbel-Softmax distributions are identical to samples from a categorical distribution as τ → 0. At higher temperatures, Gumbel-Softmax samples are no longer one-hot, and become uniform as τ → ∞.

2.1. GUMBEL-SOFTMAX ESTIMATOR:

The Gumbel-Softmax distribution is smooth for τ > 0, and therefore has a well-defined gradient ∂y /∂π with respect to the parameters π. Thus, by replacing categorical samples with Gumbel-Softmax samples we can use backpropagation to compute gradients (see Section 3.1). We denote this procedure of replacing non-differentiable categorical samples with a differentiable approximation during training as the Gumbel-Softmax estimator.

While Gumbel-Softmax samples are differentiable, they are not identical to samples from the corresponding categorical distribution for non-zero temperature. For learning, there is a tradeoff between small temperatures, where samples are close to one-hot but the variance of the gradients is large, and large temperatures, where samples are smooth but the variance of the gradients is small (Figure 1 ). In practice, we start at a high temperature and anneal to a small but non-zero temperature.

In our experiments, we find that the softmax temperature τ can be annealed according to a variety of schedules and still perform well. If τ is a learned parameter (rather than annealed via a fixed schedule), this scheme can be interpreted as entropy regularization (Szegedy et al., 2015; Pereyra et al., 2016) , where the Gumbel-Softmax distribution can adaptively adjust the "confidence" of proposed samples during the training process.

2.2. STRAIGHT-THROUGH GUMBEL-SOFTMAX ESTIMATOR:

Continuous relaxations of one-hot vectors are suitable for problems such as learning hidden representations and sequence modeling. For scenarios in which we are constrained to sampling discrete values (e.g. from a discrete action space for reinforcement learning, or quantized compression), we discretize y using arg max but use our continuous approximation in the backward pass by approximating ∇ θ z ≈ ∇ θ y. We call this the Straight-Through (ST) Gumbel Estimator, as it is reminiscent of the biased path derivative estimator described in Bengio et al. (2013) . ST Gumbel-Softmax allows samples to be sparse even when the temperature τ is high.

3. RELATED WORK:

In this section we review existing stochastic gradient estimation techniques for discrete variables (illustrated in Figure 2 ). Consider a stochastic computation graph (Schulman et al., 2015) with discrete random variable z whose distribution depends on parameter θ, and cost function f (z).

The objective is to minimize the expected cost L(θ) = E z∼p θ (z) [f (z)] via gradient descent, which requires us to estimate ∇ θ E z∼p θ (z) [f (z)].

3.1. PATH DERIVATIVE GRADIENT ESTIMATORS:

For distributions that are reparameterizable, we can compute the sample z as a deterministic function g of the parameters θ and an independent random variable , so that z = g(θ, ). The path-wise gradients from f to θ can then be computed without encountering any stochastic nodes:

EQUATION

For example, the normal distribution z ∼ N (µ, σ) can be re-written as µ + σ • N (0, 1), making it trivial to compute ∂z /∂µ and ∂z /∂σ. This reparameterization trick is commonly applied to training variational autooencoders with continuous latent variables using backpropagation (Kingma & Welling, 2013; Rezende et al., 2014b) . As shown in Figure 2 , we exploit such a trick in the construction of the Gumbel-Softmax estimator.

Biased path derivative estimators can be utilized even when z is not reparameterizable. In general, we can approximate ∇ θ z ≈ ∇ θ m(θ), where m is a differentiable proxy for the stochastic sample.

For Bernoulli variables with mean parameter θ, the Straight-Through (ST) estimator (Bengio et al., 2013) approximates m = µ θ (z), implying ∇ θ m = 1. For k = 2 (Bernoulli), ST Gumbel-Softmax is similar to the slope-annealed Straight-Through estimator proposed by Chung et al. (2016) , but uses a softmax instead of a hard sigmoid to determine the slope. Rolfe ( 2016) considers an alternative approach where each binary latent variable parameterizes a continuous mixture model. Reparameterization gradients are obtained by backpropagating through the continuous variables and marginalizing out the binary variables.

One limitation of the ST estimator is that backpropagating with respect to the sample-independent mean may cause discrepancies between the forward and backward pass, leading to higher variance. Gumbel-Softmax avoids this problem because each sample y is a differentiable proxy of the corresponding discrete sample z.

3.2. SCORE FUNCTION-BASED GRADIENT ESTIMATORS:

The score function estimator (SF, also referred to as REINFORCE (Williams, 1992) and likelihood ratio estimator (Glynn, 1990 )) uses the identity ∇ θ log p θ (z) = p θ (z)∇ θ log p θ (z) to derive the following unbiased estimator:

EQUATION

SF only requires that p θ (z) is continuous in θ, and does not require backpropagating through f or the sample z. However, SF suffers from high variance and is consequently slow to converge. In particular, the variance of SF scales linearly with the number of dimensions of the sample vector (Rezende et al., 2014a) , making it especially challenging to use for categorical distributions.

The variance of a score function estimator can be reduced by subtracting a control variate b(z) from the learning signal f , and adding back its analytical expectation

µ b = E z [b(z)∇ θ log p θ (z)]

to keep the estimator unbiased:

∇ θ E z [f (z)] = E z [f (z)∇ θ log p θ (z) + (b(z)∇ θ log p θ (z) -b(z)∇ θ log p θ (z))] (6) = E z [(f (z) -b(z))∇ θ log p θ (z)] + µ b (7)

We briefly summarize recent stochastic gradient estimators that utilize control variates. We direct the reader to Gu et al. (2016) for further detail on these techniques.

• NVIL (Mnih & Gregor, 2014) uses two baselines: (1) a moving average f of f to center the learning signal, and (2) an input-dependent baseline computed by a 1-layer neural network fitted to f -f (a control variate for the centered learning signal itself). Finally, variance normalization divides the learning signal by max(1, σ f ), where σ 2 f is a moving average of Var[f ].

• DARN (Gregor et al., 2013) 

uses b = f (z) + f (z)(z -z)

, where the baseline corresponds to the first-order Taylor approximation of f (z) from f (z). z is chosen to be 1 /2 for Bernoulli variables, which makes the estimator biased for non-quadratic f , since it ignores the correction term µ b in the estimator expression. • MuProp (Gu et al., 2016 ) also models the baseline as a first-order Taylor expansion:

b = f (z) + f (z)(z -z) and µ b = f (z)∇ θ E z [z].

To overcome backpropagation through discrete sampling, a mean-field approximation f M F (µ θ (z)) is used in place of f (z) to compute the baseline and derive the relevant gradients. • VIMCO (Mnih & Rezende, 2016 ) is a gradient estimator for multi-sample objectives that uses the mean of other samples b = 1 /m j =i f (z j ) to construct a baseline for each sample z i ∈ z 1:m . We exclude VIMCO from our experiments because we are comparing estimators for single-sample objectives, although Gumbel-Softmax can be easily extended to multisample objectives.

3.3. SEMI-SUPERVISED GENERATIVE MODELS:

Semi-supervised learning considers the problem of learning from both labeled data (x, y) ∼ D L and unlabeled data x ∼ D U , where x are observations (i.e. images) and y are corresponding labels (e.g. semantic class). For semi-supervised classification, Kingma et al. (2014) propose a variational autoencoder (VAE) whose latent state is the joint distribution over a Gaussian "style" variable z and a categorical "semantic class" variable y (Figure 6 , Appendix). The VAE objective trains a discriminative network q φ (y|x), inference network q φ (z|x, y), and generative network p θ (x|y, z) end-to-end by maximizing a variational lower bound on the log-likelihood of the observation under the generative model. For labeled data, the class y is observed, so inference is only done on z ∼ q(z|x, y). The variational lower bound on labeled data is given by:

EQUATION

For unlabeled data, difficulties arise because the categorical distribution is not reparameterizable. Kingma et al. (2014) approach this by marginalizing out y over all classes, so that for unlabeled data, inference is still on q φ (z|x, y) for each y. The lower bound on unlabeled data is:

EQUATION

The full maximization objective is:

EQUATION

)

where α is the scalar trade-off between the generative and discriminative objectives.

One limitation of this approach is that marginalization over all k class values becomes prohibitively expensive for models with a large number of classes. If D, I, G are the computational cost of sampling from q φ (y|x), q φ (z|x, y), and 

4. EXPERIMENTAL RESULTS:

In our first set of experiments, we compare Gumbel-Softmax and ST Gumbel-Softmax to other stochastic gradient estimators: Score-Function (SF), DARN, MuProp, Straight-Through (ST), and Slope-Annealed ST. Each estimator is evaluated on two tasks: (1) structured output prediction and

(2) variational training of generative models. We use the MNIST dataset with fixed binarization for training and evaluation, which is common practice for evaluating stochastic gradient estimators (Salakhutdinov & Murray, 2008; Larochelle & Murray, 2011) .

Learning rates are chosen from {3e-5, 1e-5, 3e-4, 1e-4, 3e-3, 1e-3}; we select the best learning rate for each estimator using the MNIST validation set, and report performance on the test set. Samples drawn from the Gumbel-Softmax distribution are continuous during training, but are discretized to one-hot vectors during evaluation. We also found that variance normalization was necessary to obtain competitive performance for SF, DARN, and MuProp. We used sigmoid activation functions for binary (Bernoulli) neural networks and softmax activations for categorical variables. Models were trained using stochastic gradient descent with momentum 0.9.

4.1. STRUCTURED OUTPUT PREDICTION WITH STOCHASTIC BINARY NETWORKS:

The objective of structured output prediction is to predict the lower half of a 28 × 28 MNIST digit given the top half of the image (14×28). This is a common benchmark for training stochastic binary networks (SBN) (Raiko et al., 2014; Gu et al., 2016; Mnih & Rezende, 2016) . The minimization objective for this conditional generative model is an importance-sampled estimate of the likelihood objective, E h∼p θ (hi|xupper) 1 m m i=1 log p θ (x lower |h i ) , where m = 1 is used for training and m = 1000 is used for evaluation.

We trained a SBN with two hidden layers of 200 units each. This corresponds to either 200 Bernoulli variables (denoted as 392-200-200-392) or 20 categorical variables (each with 10 classes) with binarized activations (denoted as 392-(20 × 10)-(20 × 10)-392).

As shown in Figure 3 , ST Gumbel-Softmax is on par with the other estimators for Bernoulli variables and outperforms on categorical variables. Meanwhile, Gumbel-Softmax outperforms other estimators on both Bernoulli and Categorical variables. We found that it was not necessary to anneal the softmax temperature for this task, and used a fixed τ = 1. 

4.2. GENERATIVE MODELING WITH VARIATIONAL AUTOENCODERS:

We train variational autoencoders (Kingma & Welling, 2013) , where the objective is to learn a generative model of binary MNIST images. In our experiments, we modeled the latent variable as a single hidden layer with 200 Bernoulli variables or 20 categorical variables (20×10). We use a learned categorical prior rather than a Gumbel-Softmax prior in the training objective. Thus, the minimization objective during training is no longer a variational bound if the samples are not discrete. In practice, we find that optimizing this objective in combination with temperature annealing still minimizes actual variational bounds on validation and test sets. Like the structured output prediction task, we use a multi-sample bound for evaluation with m = 1000.

The temperature is annealed using the schedule τ = max(0.5, exp(-rt)) of the global training step t, where τ is updated every N steps. N ∈ {500, 1000} and r ∈ {1e-5, 1e-4} are hyperparameters for which we select the best-performing estimator on the validation set and report test performance.

As shown in Figure 4 , ST Gumbel-Softmax outperforms other estimators for Categorical variables, and Gumbel-Softmax drastically outperforms other estimators in both Bernoulli and Categorical variables. We apply the Gumbel-Softmax estimator to semi-supervised classification on the binary MNIST dataset. We compare the original marginalization-based inference approach (Kingma et al., 2014) to single-sample inference with Gumbel-Softmax and ST Gumbel-Softmax.

We trained on a dataset consisting of 100 labeled examples (distributed evenly among each of the 10 classes) and 50,000 unlabeled examples, with dynamic binarization of the unlabeled examples for each minibatch. The discriminative model q φ (y|x) and inference model q φ (z|x, y) are each implemented as 3-layer convolutional neural networks with ReLU activation functions. The generative model p θ (x|y, z) is a 4-layer convolutional-transpose network with ReLU activations. Experimental details are provided in Appendix A.

Estimators were trained and evaluated against several values of α = {0.1, 0.2, 0.3, 0.8, 1.0} and the best unlabeled classification results for test sets were selected for each estimator and reported Published as a conference paper at ICLR 2017 in Table 2 . We used an annealing schedule of τ = max(0.5, exp(-3e-5 • t)), updated every 2000 steps.

In Kingma et al. (2014) , inference over the latent state is done by marginalizing out y and using the reparameterization trick for sampling from q φ (z|x, y). However, this approach has a computational cost that scales linearly with the number of classes. Gumbel-Softmax allows us to backpropagate directly through single samples from the joint q φ (y, z|x), achieving drastic speedups in training without compromising generative or classification performance. (Table 2 , Figure 5 ). In Figure 5 

5. DISCUSSION:

The primary contribution of this work is the reparameterizable Gumbel-Softmax distribution, whose corresponding estimator affords low-variance path derivative gradients for the categorical distribution. We show that Gumbel-Softmax and Straight-Through Gumbel-Softmax are effective on structured output prediction and variational autoencoder tasks, outperforming existing stochastic gradient estimators for both Bernoulli and categorical latent variables. Finally, Gumbel-Softmax enables dramatic speedups in inference over discrete latent variables. 

B DERIVING THE DENSITY OF THE GUMBEL-SOFTMAX DISTRIBUTION:

Here we derive the probability density function of the Gumbel-Softmax distribution with probabilities π 1 , ..., π k and temperature τ . We first define the logits x i = log π i , and Gumbel samples Figure 7 : Network architecture for (a) classification q φ (y|x) (b) inference q φ (z|x, y), and (c) generative p θ (x|y, z) models. The output of these networks parameterize Categorical, Gaussian, and Bernoulli distributions which we sample from. g 1 , ..., g k , where g i ∼ Gumbel(0, 1). A sample from the Gumbel-Softmax can then be computed as:

EQUATION

B.1 CENTERED GUMBEL DENSITY

The mapping from the Gumbel samples g to the Gumbel-Softmax sample y is not invertible as the normalization of the softmax operation removes one degree of freedom. To compensate for this, we define an equivalent sampling process that subtracts off the last element, (x k + g k )/τ before the softmax:

EQUATION

To derive the density of this equivalent sampling process, we first derive the density for the "centered" multivariate Gumbel density corresponding to:

EQUATION

where g i ∼ Gumbel(0, 1). Note the probability density of a Gumbel distribution with scale parameter β = 1 and mean µ at z is: f (z, µ) = e µ-z-e µ-z . We can now compute the density of this distribution by marginalizing out the last Gumbel sample, g k :

p(u 1 , ..., u k-1 ) = ∞ -∞ dg k p(u 1 , ..., u k |g k )p(g k ) = ∞ -∞ dg k p(g k ) k-1 i=1 p(u i |g k ) = ∞ -∞ dg k f (g k , 0) k-1 i=1 f (x k + g k , x i -u i ) = ∞ -∞

dg k e -g k -e -g k k-1 i=1 e xi-ui-x k -g k -e x i -u i -x k -g k We perform a change of variables with v = e -g k , so dv = -e -g k dg k and dg k = -dv e g k = dv/v, and define u k = 0 to simplify notation: p(u 1 , ..., u k,-1 ) = δ(u k = 0)

∞ 0 dv 1 v ve x k -v k-1 i=1

ve xi-ui-x k -ve x i -u i -x k (15)

EQUATION

B.2 TRANSFORMING TO A GUMBEL-SOFTMAX Given samples u 1 , ..., u k,-1 from the centered Gumbel distribution, we can apply a deterministic transformation h to yield the first k -1 coordinates of the sample from the Gumbel-Softmax:

EQUATION

Note that the final coordinate probability, y k , is fixed given the first k -1 as k i=1 y i = 1:

EQUATION

We can thus compute the probability of a sample from the Gumbel-Softmax using the change of variables formula on only the first k -1 variables:

p(y 1:k ) = p h -1 (y 1:k-1 ) ∂h -1 (y 1:k-1 ) ∂y 1:k-1 (21)

So to compute the probability of the Gumbel-Softmax we need two more pieces: the inverse of h and its Jacobian determinant. The inverse of h is:

EQUATION

The determinant of the Jacobian can then be computed:

EQUATION

)

We can then plug into the change of variables formula (Eq. 21) using the density of the centered Gumbel (Eq.15), the inverse of h (Eq. 22) and its Jacobian determinant (Eq. 

:

The Gumbel(0, 1) distribution can be sampled using inverse transform sampling by drawing u ∼ Uniform(0, 1) and computing g = -log(-log(u)).

Figures and Tables:

Figure1: The Gumbel-Softmax distribution interpolates between discrete one-hot-encoded categorical distributions and continuous categorical densities. (a) For low temperatures (τ = 0.1, τ = 0.5), the expected value of a Gumbel-Softmax random variable approaches the expected value of a categorical random variable with the same logits. As the temperature increases (τ = 1.0, τ = 10.0), the expected value converges to a uniform distribution over the categories. (b) Samples from Gumbel-Softmax distributions are identical to samples from a categorical distribution as τ → 0. At higher temperatures, Gumbel-Softmax samples are no longer one-hot, and become uniform as τ → ∞.

Figure 2: Gradient estimation in stochastic computation graphs. (1) ∇ θ f (x) can be computed via backpropagation if x(θ) is deterministic and differentiable. (2) The presence of stochastic node z precludes backpropagation as the sampler function does not have a well-defined gradient. (3) The score function estimator and its variants (NVIL, DARN, MuProp, VIMCO) obtain an unbiased estimate of ∇ θ f (x) by backpropagating along a surrogate loss f log p θ (z), where f = f (x) -b and b is a baseline for variance reduction. (4) The Straight-Through estimator, developed primarily for Bernoulli variables, approximates ∇ θ z ≈ 1. (5) Gumbel-Softmax is a path derivative estimator for a continuous distribution y that approximates z. Reparameterization allows gradients to flow from f (y) to θ. y can be annealed to one-hot categorical variables over the course of training.

p θ (x|y, z) respectively, then training the unsupervised objective requires O(D + k(I + G)) for each forward/backward step. In contrast, Gumbel-Softmax allows us to backpropagate through y ∼ q φ (y|x) for single sample gradient estimation, and achieves a cost of O(D + I + G) per training step. Experimental comparisons in training speed are shown in Figure 5.

Figure 3: Test loss (negative log-likelihood) on the structured output prediction task with binarized MNIST using a stochastic binary network with (a) Bernoulli latent variables (392-200-200-392) and (b) categorical latent variables (392-(20 × 10)-(20 × 10)-392).

Figure 4: Test loss (negative variational lower bound) on binarized MNIST VAE with (a) Bernoulli latent variables (784 -200 -784) and (b) categorical latent variables (784 -(20 × 10) -200).

, we show how Gumbel-Softmax versus marginalization scales with the number of categorical classes. For these experiments, we use MNIST images with randomly generated labels. Training the model with the Gumbel-Softmax estimator is 2× as fast for 10 classes and 9.9× as fast for 100 classes.

Figure 5: Gumbel-Softmax allows us to backpropagate through samples from the posterior q φ (y|x), providing a scalable method for semi-supervised learning for tasks with a large number of classes. (a) Comparison of training speed (steps/sec) between Gumbel-Softmax and marginalization (Kingma et al., 2014) on a semi-supervised VAE. Evaluations were performed on a GTX Titan X R GPU. (b) Visualization of MNIST analogies generated by varying style variable z across each row and class variable y across each

Figures 6 and 7 describe the architecture used in our experiments for semi-supervised classification (Section 4.3).

Figure 6: Semi-supervised generative model proposed by Kingma et al. (2014). (a) Generative model p θ (x|y, z) synthesizes images from latent Gaussian "style" variable z and categorical class variable y. (b) Inference model φ (y, z|x) samples latent state y, z given x. Gaussian z can be differentiated with respect to its parameters because it is reparameterizable. In previous work, when y is not observed, training the VAE objective requires marginalizing over all values of y. (c) Gumbel-Softmax reparameterizes y so that backpropagation is also possible through y without encountering stochastic nodes.

The Gumbel-Softmax estimator outperforms other estimators on Bernoulli and Categorical latent variables. For the structured output prediction (SBN) task, numbers correspond to negative log-likelihoods (nats) of input images (lower is better). For the VAE task, numbers correspond to negative variational lower bounds (nats) on the log-likelihood (lower is better).

Table content: <table><tr><td/><td>SF</td><td colspan="2">DARN MuProp</td><td>ST</td><td colspan="3">Annealed ST Gumbel-S. ST Gumbel-S.</td></tr><tr><td colspan="2">SBN (Bern.) 72.0</td><td>59.7</td><td>58.9</td><td>58.9</td><td>58.7</td><td>58.5</td><td>59.3</td></tr><tr><td>SBN (Cat.)</td><td>73.1</td><td>67.9</td><td>63.0</td><td>61.8</td><td>61.1</td><td>59.0</td><td>59.7</td></tr><tr><td colspan="3">VAE (Bern.) 112.2 110.9</td><td>109.7</td><td>116.0</td><td>111.5</td><td>105.0</td><td>111.5</td></tr><tr><td colspan="3">VAE (Cat.) 110.6 128.8</td><td>107.0</td><td>110.9</td><td>107.8</td><td>101.5</td><td>107.8</td></tr><tr><td colspan="6">4.3 GENERATIVE SEMI-SUPERVISED CLASSIFICATION</td><td/><td/></tr></table>

Marginalizing over y and single-sample variational inference perform equally well when applied to image classification on the binarized MNIST dataset(Larochelle & Murray, 2011). We report variational lower bounds and image classification accuracy for unlabeled data in the test set.

Table content: <table><tr><td/><td>ELBO Accuracy</td></tr><tr><td>Marginalization</td><td>-106.8 92.6%</td></tr><tr><td>Gumbel</td><td>-109.6 92.4%</td></tr><tr><td colspan="2">ST Gumbel-Softmax -110.7 93.6%</td></tr></table>

Peer review (iclr2017):
The authors propose a method for reparameterization gradients with categorical distributions. This is done by using the Gumbel-Softmax distribution, a smoothened version of the Gumbel-Max trick for sampling from a multinomial. The paper is well-written and clear. The application to the semi-supervised model in Kingma et al. (2014) makes sense for large classes, as well as its application to general stochastic computation graphs (Schulman et al., 2015). One disconcerting point is that (from my understanding at least), this does not actually perform variational inference for discrete latent variable models. Rather, it changes the probability model itself and performs approximate inference on the modified (continuous relaxed) version of the model. This is fine in practice given that it*s all approximate inference, but unlike previous variational inference advances either in more expressive approximations or faster computation (as noted by the different gradient estimators they compare to), the probability model is fundamentally changed. Two critical points seem key: the sensitivity of the temperature, and whether this applies for non-one hot encodings of the categorical distribution (and thus sufficiently scale to high dimensions). Comments by the authors on this are welcome. There is a related work by Rolfe (2016) on discrete VAEs, who also consider a continuous relaxed approach. This is worth citing and comparing to (or at least mentioning) in the paper. References Rolfe, J. T. (2016). Discrete Variational Autoencoders. arXiv.org.