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

Let's say you have two states, $X_1$ and $X_2$, and you have a model, $M$, that produces a score $M(X_i)$ for each state (i.e, the logits). Next you can use the logits to compute some distribution

$$P = softmax(\{M(X_1), M(X_2)\})$$

and take the state with the highest probability

$$X=argmax_{X_i}(P)$$

But what if you actually want to sample from $P$ instead of just taking the argmax - and you want the sample operation to be differentiable! This is where the Gumbel Trick comes in - instead of softmax, you compute

\begin{equation} \ X = argmax_{X_i}(\{M(X_i)+Z_i\}) \end{equation}

Where $Zi$ are i.i.d Gumbel(0,1). It turns out that $X$ will be equal to $X_1$ about $P(X_1)$ of the times and to $X_2$ about $P(X_2)$ of the times. In other words, the equation above samples from $P$.

But it still not differentiable because of the argmax operation. So instead of doing that we'll compute the Gumbel-Softmax distribution. Now if the temperature is low enough then the Gumbel-Softmax will produce something very close to one hot vector, where the probability of the predicted label will be 1 and other labels will have a probability of zero. So for example, if the Gumbel-Softmax gave the highest probability to $X_1$, you can do:

\begin{equation} \ X = \sum_{x_i} P_g(X_i)*X_i = 1*X_1 + 0*X_2 = X_1 \end{equation}

Where $P_g$ is the Gumbel-Softmax operation. No argmax is needed! So with this cool trick we can sample from a discrete distribution in a differentiable way.

Answer 2

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?

You don't necessarily need Gumbel-Softmax to obtain "one-hot like" vectors, or the ability to differentiate through an indexing mechanism.

The LSTM architecture and derived variants are examples of this. They model "forget/input" gates using sigmoid outputs, which are deterministic. A "true" gating mechanism would be either 0 or 1, but to make things differentiable, LSTMs relax that constraint to sigmoided outputs. You'll notice that there is no "random" inputs here, and you can still apply the straight-through trick here to make the gates truly discrete (while backpropagating a biased gradient).

Gumbel-Softmax can be used wherever you would consider using a non-stochastic indexing mechanism (it is a more general formulation). But it's especially useful when you want to backpropagate through random samples of discrete variables.

• VAE with a Gumbel-Softmax or Categorical posterior (encoder) distribution. Notably, you cannot simply use a deterministic softmax here because it would turn your VAE into a standard autoencoder. Autoencoders lack a way to generate new samples from the prior.
• Actor-Critic architecture with a Gumbel-softmax or Categorical actor (most policy gradient implementations assume you can re-parameterize the gradients from the critic through the actor without using a score function estimator to estimate the black-box gradient). You cannot simply substitute the deterministic softmax here, because there is a type mismatch: the critic takes as input a action $a \in A$, while the softmax represents the conditional policy distribution $\pi(a|s)$
• The "probabilistic" interpretation of a non-random quantization such as an LSTM would essentially be mode-seeking behavior in fitting a density. You have loss function that takes in categorical decisions $c$, so the expected loss $\mathbb{E}_c[f(c)]$ is minimized by learning some distribution $p(c)$. Quantizing a softmax without sampling the Gumbel noise (e.g. just using a sigmoid or softmax) is akin to choosing the same $c$ every time. For some $f$ this is okay, and for other $f$ this is highly suboptimal (consider the categorical KL divergence as a loss).

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.

This is an interesting idea, but there are many ways to inject "exploration" noise into the set of parameters you use in a function approximator.

Answer 3

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

Answer 4

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.