Gumbel-Softmax trick vs Softmax with temperature

Question

From what I understand, the Gumbel-Softmax trick is a technique that enables us to sample discrete random variables, in a way that is differentiable (and therefore suited for end-to-end deep learning).

Many papers and articles describe it as a way of selecting instances in the input (i.e. 'pointers') without using the non-differentiable argmax-function. The thing that confuses me is that this effect can be achieved without randomness by just using Softmax with temperature:

Softmax with temperature $$y_𝑖=\frac{exp(\frac{𝑧_𝑖}{\tau})}{\sum_{𝑗}exp(\frac{𝑧_𝑗}{\tau})}$$

Gumbel-Softmax $$y_𝑖=\frac{exp(\frac{log(\pi_𝑖)+g_i}{\tau})}{\sum_{𝑗}exp(\frac{log(\pi_j)+g_j}{\tau})}$$

My question

From a practical and theoretical perspective, when is it beneficial to incorporate Gumbel noise into a neural network, as opposed to just using Softmax with temperature?

A couple of observations:

1. When the temperature is low, both Softmax with temperature and the Gumbel-Softmax functions will approximate a one-hot vector. However, before convergence, the Gumbel-Softmax may more suddenly 'change' its decision because of the noise.
2. When the temperature is higher, the Gumbel noise will get a larger significance and the distribution will become more uniform. Why is this desired?

My best guess is that the introduction of the Gumbel noise enforces stronger exploration before convergence, but I can't recall reading any papers that use this as a motivation to bring in the extra randomness.

Does anyone have any experience or insights on this? Maybe I've completely missed the key point of Gumbel-Softmax :)

Answer 1

For the softmax function, no matter what is the temperature, it is not the exact one-hot vector. If you could accept a soft version, it is good. However, if you choose the argmax to be the one, it is non-differentiable. One alternative way to back-propagate the gradients is by using the Straight Through Estimator (STE)[1] trick, and directly back-propagate the gradients [2], the gradient is an inaccurate approximation.

The advantage of Gumbel Softmax [3] is it samples one-hot according to the current learned distribution of \pi, it is one-hot and it is differentiable and the probability of sampled one-hot vector is according to \pi.

For your send question: at the beginning, the distribution \pi does not have any prior knowledge, so we want to sample one-hot vector by uniform (at this stage, the noise matters), and the distribution will gradually converge to the desired distribution (slightly sharper). As you training for longer epochs, prior knowledge of the distribution is learned enough, gradually decrease the temperature \tau and make \pi converge to a discrete distribution. As you gradually decrease the temperature \tau, the effect of noise is smaller.

PS: The sentence is incorrect:

When the temperature is low, both Softmax with temperature and the Gumbel-Softmax functions will approximate a one-hot vector.

Gumbel-softmax could sample a one-hot vector rather than an approximation. You could read the PyTorch code at [4].

[1] Binaryconnect: Training deep neural networks with binary weights during propagations

[2] LegoNet: Efficient Convolutional Neural Networks with Lego Filters

[3] Categorical Reparameterization with Gumbel-Softmax

[4] https://github.com/pytorch/pytorch/blob/15f9fe1d92a5d1e86278ae25f92dd9677b4956dc/torch/nn/functional.py#L1237