Gumbel Softmax

Problem

In DALLE¹, they wanted to maximize $\ln{p_{\theta}, \psi (x,y)}$ and by ELBO, they derive the lower bound to be:

$\ln{p_{\theta}, \psi (x,y)} \geq \mathbb{E}_{z \sim q_{\phi}(z|x)}[\ln{p_{\theta}(x|y, z)}] - \beta D_{KL}(q_{\phi}(y,z|x)||p_{\psi}(y,z))$

Here $z$ is a sample from the categorical distribution $q_{\phi}(z|x)$, which is a discrete variable. This is analogous to a VAE, where $z$ is a latent, but in a VAE the latent space is continuos and in DALLE's case its discrete. The problem is that the sampling operation is not differentiable, so we cannot backpropagate through it. To solve this, we can use the same reparametrization trick used in VAEs, so they use the Gumbel Softmax trick. This defines our problem: We want to sample from a $k$ dimensional categorical distribution with unnormalized probabilities $[{\pi}_1,...,{\pi}_2]$ and allow gradients to flow through.

Gumbel Softmax

The Gumbel Softmax heavily utilizes the results from (Maddison et al²), where the authors present the Concrete distribution: $X \in \triangle^{k-1}$, where $\triangle^{k-1}$ is the $k-1$ simplex, which is the set of all $k$ dimensional vectors with non-negative entries that sum to 1. The Concrete distribution is a continuous relaxation of the categorical distribution. Intuitively, we want to sample from the vertices of this simplex based on our categorical distribution. The authors define the Concrete distribution as:

where $g_i \sim \mathrm{Gumbel(0,1)}$ and temperature $\tau > 0$. Based on this, the joint distribution over the simplex is: