I've been also looking for the answer of this question, and I give my different view of Gumbel softmax just because I think this is a good question.
From a general point of view:
We use softmax normally because we need a so-called score, or a distribution $\pi_1 .. \pi_n$ for representing n probabilities of categorical variable with size n;
We use Gumbel-softmax to sample an one hot sample [0..1..0] from this distribution.
In more concrete examples:
Normally in networks for NLP(which categorize outputs into different word tokens), softmax is used to calculate the distribution of different, say, 5000 word choices at current text position. the cross entropy loss, gives a measure about the difference between the softmax predicted distribution and the real word distribution; For Gumbel-softmax, it is normally used to generate a sample one-hot vector for the constructing the following network, like in some VAE-based models. That's why the temperature factor $\tau$ is a must for Gumbel-softmax, and in most cases not needed in softmax.
At this point, we know the difference of their uses cases, now comes the confusing part: the two formulas are so similar, why they are doing the different things?
The first key factor: What makes the difference is the $g_i$ term in the Gumbel-softmax formula. It represents a point sampled from the distribution $Gumbel(0, 1)$. the adding term $log(\pi_i)$ and the scaling term $\tau$ is just used to reparameterize it to $Gumbel(log(\pi_i), \tau)$ (reparameterization trick is for making it differentiable). This Gumbel distribution is the key distribution for sampling a categorical variable in Gumbel-max method(hard and not differentiable since there is a argmax), and this is also where the name Gumbel comes from in Gumbel-softmax. Every time we use Gumbel-softmax, we need to randomly sample from $Gumbel(0,1)$ and do the reparameterization trick, this is the most different part from softmax.
I would rather name it soft-Gumbel-max, to indicate it is motivated to make a soft version of Gumbel-max, instead of just intending to add a Gumbel term to softmax.
The second most significant difference: is the use of $\tau$. In most neural networks, softmax is not coupled with this term. Since we normally need a distribution, not a nearly one-hot vector. What's more, in some cases, like beam search, we need to get the second or third most probable choices to explore global optimal searches. In softmax, $\tau$ is normally added with some domain-specific knowledges to make the distribution steeper. Smaller $\tau$ not always means better, we need to tune it to best fit the model.
For Gumbel-softmax however, in the later stage of the model training, we need the Gumbel-softmax as close to one-hot vector as possible. That's why we need to anneal it smaller and smaller during training.(the $\tau$ is not too small at the beginning of training for this makes the training more stable).
In some implementation like torch.nn.functional.gumbel_softmax
, it uses the straight through trick hard - (detached soft) + soft
to maintain the output value a one-hot vector as in hard Gumbel-max but with well-defined gradients as in soft differentiable Gumbel-softmax.
The most part I disagree with the previous answer is that, Gumbel-softmax will not give you the exact one-hot vector. It is still an estimation of one-hot vector. Even for straight-through, the forward pass is Gumbel-max and only the backward pass is Gumbel-softmax.
That's almost all how I understand Gumbel-softmax, Gumbel-max and softmax. Please comment if there is anything unclear or incorrect.
MORE If you still confuse
After writing this answer, I found that, although the use cases of softmax and Gumbel-softmax are different, but we can still force to apply softmax in the place where Gumbel-softmax is applied without getting any arithmetic problem, and vice versa. Because both of them are soft and not exact, and both of them are differentiable. To make clearer why this is a problem. let's make two very gross examples of misuse.
Suppose there is a network needing a sample but not a distribution of a categorical variable, e.g. a network need a categorical variable representing whose speech I wanna generate: Obama or Trump. Say 70% of the chance the training data is from Trump and 30% of time is from Obama, Then 70% of time the variable should be [0 1] and 30 % of time should be [1 0]. If we use softmax here instead, the variable will always be about [0.3 0.7], thus it becomes a constant, so it can be ignored by the network. To mitigate the training loss penalty, the resulting network is very probable to produce the hybrid sound a bit like Obama and a bit more like Trump. What's more [0.7 0.3] during training is never used during predicting, this is also a problem about out-of-domain prediction.
Another example is for predicting the next word after "you" in a sentence. suppose there are only two options: "are" and "have".
If we use Gumbel-softmax instead of softmax for representing the probability of the two words. The cross entropy loss is not close to infinity(-log0) only when the training process make the correct sampling. The training process will easily get gradient explosion in this case.