What is Gumbel-Softmax?. A differentiable approximation to... | by Wanshun Wong | Towards Data Science

Very often in deep learning we want to sample discrete data. For example,
text generation by generative adversarial network (GAN)
variational autoencoder with discrete latent variable
deep reinforcement learning with discrete action space
However, the sampling process of discrete data from a categorical distribution is not differentiable, meaning that backpropagation will not work. The Gumbel-Softmax distribution is a continuous distribution that approximates samples from a categorical distribution and also works with backpropagation.
Let Z be a categorical variable with categorical distribution Categorical(π1, ..., πx), where πi are the class probabilities to be learned by our neural network. Assume our discrete data are encoded as one-hot vectors. The most common way of sampling Z is given by
Z = onehot(max{i | π1 + ... + πi−1 ≤ U})
where i = 1, ..., x is the class index, and U ~ Uniform(0, 1). This sampling formula is not differentiable because of the max function. In order to obtain a differentiable approximation, we apply the following:
The Gumbel-Max trick
The Gumbel-Max trick
The Gumbel-Max trick provides a different formula for sampling Z
Z = onehot(argmaxi{Gi + log(πi)})
where Gi ~ Gumbel(0,1) are i.i.d. samples drawn from the standard Gumbel distribution. This is a “reparameterization trick”, refactoring the sampling of Z into a deterministic function of the parameters and some independent noise with a fixed distribution.
The refactoring of the sampling process does not make it differentiable. In our case the non-differentiability comes from the argmax function. However, once we have a differentiable approximation of argmax, we can do backpropagation easily as it will simply be about computing the gradient w.r.t. parameters of a deterministic function. On the other hand, without the reparameterization trick, we would have to compute the gradient w.r.t. parameters of a distribution, which is a harder problem.
2. Using softmax as a differentiable approximation
We use softmax as a differentiable approximation to argmax. The sample vectors y are now given by
yi = exp((Gi + log(πi)) / τ) / Σj exp((Gj + log(πj)) / τ)
for every i = 1, ..., x. τ is the temperature parameter that controls how closely the new samples approximate discrete, one-hot vectors. As τ → 0, the softmax computation smoothly approaches the argmax, and the sample vectors approach one-hot; as τ → ∞, the sample vectors become uniform.
The distribution with the above sampling formula is called the Gumbel-Softmax distribution. Note that continuous vectors are used during training, but the sample vectors are discretized to one-hot vectors during evaluation.
Whenever we have a stochastic neural network with discrete variables, we can use Gumbel-Softmax distributions to approximate the sampling process of discrete data. The network can then be trained using backpropagation, where the performance of the network will depend on the choice of the temperature parameter τ.
There are cases in which we will want to sample discrete data during training:
We are constrained to discrete values because real-valued continuous approximations are not allowed. For example, in deep reinforcement learning where the action space is discrete.
Using continuous approximations during training and using one-hot vectors during evaluation have vastly different dynamics. This leads to a massive drop in model performance during evaluation.
We can discretize our Gumbel-Softmax sample back with argmax in the forward pass, which becomes the same as a sample from the original categorical distribution. This ensures the training and the evaluation dynamics are the same. However, we still use the Gumbel-Softmax sample in the backward pass to approximate the gradients, so that backpropagation would still work. This technique is known as Straight Through Gumbel-Softmax in literature.
Tensorflow: tfp.distributions.RelaxedOneHotCategorical
PyTorch: torch.nn.functional.gumbel_softmax
Q: How do we choose the temperature parameter τ?
A: We can have a fixed value for τ. The optimal value is determined by trial and error and is quite often ≦ 1. We could also use an annealing schedule or make it a learned parameter.
Q: How come we do not need to use Gumbel-Softmax for classification tasks, which involves discrete class labels?
A: It is because usually there is no sampling process involved in solving classification problems. Instead, we only require our models to output probabilities of feature vectors belonging to different classes. Simply applying the softmax function is sufficient for this purpose.
The Gumbel-Softmax distribution was independently discovered by [2] and [3], where it is called the concrete distribution in [3]. Both papers are excellent references, especially about the theoretical aspects of the distribution and also about the reparameterization trick.[5] is a blog post written by the first author of [2]. It is a nice tutorial with an interactive widget about sampling from Gumbel-Softmax distributions.For an application of Gumbel-Softmax distribution in GAN, we refer to [4], which gives a big picture of how Gumbel-Softmax solves the problem of training GAN with discrete data.[1] provides a use case where Straight Through Gumbel-Softmax is needed instead of the ordinary Gumbel-Softmax.
The Gumbel-Softmax distribution was independently discovered by [2] and [3], where it is called the concrete distribution in [3]. Both papers are excellent references, especially about the theoretical aspects of the distribution and also about the reparameterization trick.
[5] is a blog post written by the first author of [2]. It is a nice tutorial with an interactive widget about sampling from Gumbel-Softmax distributions.
For an application of Gumbel-Softmax distribution in GAN, we refer to [4], which gives a big picture of how Gumbel-Softmax solves the problem of training GAN with discrete data.
[1] provides a use case where Straight Through Gumbel-Softmax is needed instead of the ordinary Gumbel-Softmax.
S. Havrylov, and I. Titov. Emergence of Language with Multi-agent Games: Learning to Communicate with Sequences of Symbols (2017), NIPS 2017.E. Jang, S. Gu, and B. Poole. Categorical Reparameterization with Gumbel-Softmax (2017), ICLR 2017.C. J. Maddison, A. Mnih, and Y. W. Teh. The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables (2017), ICLR 2017.W. Nie, N. Narodytska, and A. Patel. RelGAN: Relational Generative Adversarial Networks for Text Generation (2019). ICLR 2019.Tutorial: Categorical Variational Autoencoders using Gumbel-Softmax, https://blog.evjang.com/2016/11/tutorial-categorical-variational.html
S. Havrylov, and I. Titov. Emergence of Language with Multi-agent Games: Learning to Communicate with Sequences of Symbols (2017), NIPS 2017.
E. Jang, S. Gu, and B. Poole. Categorical Reparameterization with Gumbel-Softmax (2017), ICLR 2017.
C. J. Maddison, A. Mnih, and Y. W. Teh. The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables (2017), ICLR 2017.
W. Nie, N. Narodytska, and A. Patel. RelGAN: Relational Generative Adversarial Networks for Text Generation (2019). ICLR 2019.
Tutorial: Categorical Variational Autoencoders using Gumbel-Softmax, https://blog.evjang.com/2016/11/tutorial-categorical-variational.html