Title: Categorical Reparameterization with Gumbel-Softmax
Authors: Eric Jang
Shixiang Gu
Ben Poole
Venue: ICLR 2017
Keywords: Deep learning
Semi-Supervised Learning
Optimization
Structured prediction

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.

TL;DR: Simple, differentiable sampling mechanism for categorical variables that can be trained in neural nets via standard backprop.

Review discussion:
How does this relate to dirichlet distribution?
Gumbel-Softmax and Dirichlet are both continuous distributions on the simplex;
thus, samples from both distributions can be interpreted as class (categorical) probabilities.
One distinction is that samples from Gumbel-Softmax can be sampled via a straightforward reparameterization trick, while the sampling from a Dirichlet (via stick-breaking or urn drawing) makes it harder to exploit the reparameterization trick for estimating gradients (though recent work https://arxiv.org/abs/1610.05683 enables this via rejection sampling, and a reparameterizable Dirichlet approximation might be possible using https://arxiv.org/abs/1505.05770).
The paper combines Gumbel distribution with the popular softmax function to obtain a continuous distribution on the simplex that can approximate categorical samples. It is not surprising that Gumbel softmax outperforms other single sample gradient estimators.
However, I am curious about how Gumbel compares with Dirichlet experimentally.
The computational efficiency of the estimator when training semi-supervised models is nice.
However, the advantage will be greater when the number of classes are huge, which doesn't seem to be the case in a simple dataset like MNIST. I am wondering why the experiments are not done on a richer dataset.
The presentation of the paper is neat and clean. The experiments settings are clearly explained and the analysis appears to be complete.
The only concern I have is the novelty of this work. I consider this work as a nice but may be incremental (relatively small) contribution to our community.
Thank you for the comments.
Unlike previous approaches to learning with categorical latent variables that rely on complex score function estimators, Gumbel-Softmax is easy to implement in any modern deep learning framework.
Furthermore, the idea of developing continuous relaxations for reparameterization gradients with categorical variables is novel, and represents a large step forward in our ability to develop estimators for discrete variables.
The purpose of the semi-supervised experiments was to demonstrate that stochastic inference using Gumbel-Softmax results in tremendous speedups of semi-supervised architectures without compromising accuracy.
Evaluating on a simple dataset such as MNIST makes the analysis and interpretation of results simpler; the feasibility of the technique comes from stochastic inference with Gumbel-Softmax, and we did not have to tweak the inference, discriminative, or generative components of the model to reproduce reasonable semi-supervised results.
As for scaling up to large class sizes, the interior of the simplex (where the continuous relaxations lie) does indeed get smaller as the number of classes go up, resulting in vanishing gradients. We ran some experiments and found that unbiased estimators like MuProp converge faster than Gumbel-Softmax when the dimensionality of the sample is large (K>800).
However, Gumbel-Softmax remains useful for differentiating through many kinds of lower-dimensional categorical samples (mixture models, discrete action spaces, character sequences).
In practice, large class sizes are not an issue anyway, because categorical samples with large number of classes K can be encoded in base M << K rather than a single one-hot K-vector.
For instance, stochastic binary layer with log2(K) units.
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.
Thank you for the comments. Please see our response to Review 2 for discussion of the temperature parameter.
Our proposed straight-through Gumbel-Softmax estimator can be used to produce samples that are exactly one-hot and can be re-encoded as a K-ary categorical variable.
We have updated our paper with a reference and comparison to the Discrete VAE paper by Rolfe (2016) (Section 3.1). They consider a substantially different model from other prior work; they use a discrete latent space augmented with auxiliary continuous variables, a hierarchical posterior, and a RBM prior. Because the models are so different, it’s difficult to compare quantitative results directly.
This paper introduces a continuous relaxation of categorical distribution,  namely the the Gumbel-Softmax distribution, such that generative models with categorical random variables can be trained using reparameterization (path-derivative) gradients. The method is shown to improve upon other methods in terms of the achieved log-likelihoods of the resulting models. The main contribution, namely the method itself, is simple yet nontrivial and worth publishing, and seems effective in experiments. The paper is well-written, and I applaud the details provided in the appendix. The main application seems to be semi-supervised situations where you really want categorical variables.
- P1: "differentiable sampling mechanism for softmax". "sampling" => "approximate sampling", since it's technically sampling from the Gumbal-softmax.
- P3: "backpropagtion"
- Section 4.1: Interesting experiments.
- It would be interesting to report whether there is any discrepancy between the relaxed and non-relaxed models in terms of log-likelihood. Currently, only the likelihoods under the non-relaxed models are reported.
- It is slightly discouraging that the temperature (a nuisance parameter) is used differently across experiments. It would be nice to give more details on whether you were succesful in learning the temperature, instead of annealing it; it would be interesting if that hyper-parameter could be eliminated.
Thank you for the comments. We agree with the reviewer that Gumbel-Softmax is useful in semi-supervised situations, and also find it to be useful in learning discrete latent variable models in the purely unsupervised setting. The learned discrete latent spaces often yield more interpretable and semantically meaningful latent spaces than a corresponding continuous latent variable model.
Additionally, sampling from Gumbel-Softmax can be used to select discrete “actions” (and then differentiated to obtain policy gradients for Reinforcement Learning), or used in sequential sampling of discrete samples corresponding to character tokens (language models), or sampling from mixture models in a re-parameterizable fashion.
We have made the correction from “sampling” => “approximate sampling” (P1) and the “backpropagtion” typo (P3).
We made a typo in describing the architecture of the VAE - our experiments use a learned categorical prior (consistent with prior work from Gu et al. 2016), not a uniform categorical prior, and we have made a correction to this as well.
For our experiments, fixed temperatures between 0.5 and 1.0 yield good results for both structured output prediction and VAE tasks. The concurrent submission by Maddison et al. 2016 (https://openreview.net/forum?id=S1jE5L5gl) uses a fixed temperature of 2/3 for all experiments. We found that annealing the temperature yielded slightly better results that also converged faster for VAEs. The takeaway is that annealing improves performance, but is not absolutely critical to yield good results.
We ran some follow-up experiments to determine whether the temperature parameter can indeed be learned. Surprisingly, learning the temperature causes the temperature parameter to increase and saturate, rather than decrease and saturate. We suspect that when posterior inference does not do a good job, i.e. sampling q(z|x) chooses a “bad sample”, the models can improve the reconstruction locally (and thus the autoencoding term) simply by increasing the temperature (i.e. smoothing the discrete sample to a more uniform one). Optimizing the model parameters with this higher temperature results in a lower training loss, but because Gumbel-Softmax samples are no longer constrained to be sparse, we do poorly on the original discrete objective. We still observe overfitting behavior on the validation and test sets when training using Straight-Through Gumbel Softmax with learning temperature.

Meta-review:
role: PC
data: comment: This paper proposes a neat general method for relaxing models with discrete softmax choices into closely-related models with continuous random variables. The method is designed to work well with the reparameterization trick used in stochastic variational inference. This work is likely to have wide impact.
 
 Related submissions at ICLR:
 "The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables" by Maddison et al. contains the same core idea. "Discrete variational autoencoders", by Rolfe, contains an alternative relaxation for autoencoders with discrete latents, which I personally find harder to follow.

Decision: Accept (Poster)