Categorical Reparameterization with Gumbel-Softmax
This addresses a key bottleneck for researchers and practitioners using categorical variables in deep learning, enabling more efficient training of models with discrete structure.
The paper tackles the problem of backpropagating through categorical latent variables in stochastic neural networks by introducing a differentiable Gumbel-Softmax distribution, which outperforms state-of-the-art gradient estimators on tasks like structured output prediction and unsupervised generative modeling, enabling large speedups in semi-supervised classification.
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.