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. 915
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. 915