Just Relax It! Leveraging relaxation for discrete variables optimization

We assume that you are already familiar with the concept of a Variational Autoencoder (VAE), at least at the level that it consists of two parts: an encoder and a decoder, and tries to approximate the distribution $p(\mathbf{x})$ of data $\mathbf{x}$ using latent variables $\mathbf{z}$. If this is not the case, then we recommend that you to read a blog-post by Lilian Weng.

The original VAE (Kingma & Welling, 2014) has two main parts:

1. Encoder $q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})$: A neural network $g_{\boldsymbol{\phi}}(\mathbf{x})$ that outputs parameters of the latent Gaussian distribution.
2. Decoder $p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z})$: A neural network $f_{\boldsymbol{\theta}}(\mathbf{z})$ that outputs parameters of the sample distribution (usually Gaussian or Bernoulli).

Training a VAE involves maximizing the ELBO (evidence lower bound) with respect to the encoder and decoder parameters:

$$ \mathcal{L}_{\boldsymbol{\phi}, \boldsymbol{\theta}}(\mathbf{x}) = \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})} \log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z}) - KL(q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) | p(\mathbf{z})) \to \max_{\boldsymbol{\phi}, \boldsymbol{\theta}}. $$

During the M-step, we derive an unbiased estimator for the gradient $\nabla_{\boldsymbol{\theta}}\mathcal{L}_{\boldsymbol{\phi}, \boldsymbol{\theta}}(\mathbf{x})$:

$$ \begin{aligned} \nabla_{\boldsymbol{\theta}}\mathcal{L}_{\boldsymbol{\phi}, \boldsymbol{\theta}}(\mathbf{x}) &= \textcolor{blue}{\nabla_{\boldsymbol{\theta}}} \int q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) \log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z}) d\mathbf{z} \\ &= \int q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) \textcolor{blue}{\nabla_{\boldsymbol{\theta}}} \log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z}) d\mathbf{z} \\ &\approx \nabla_{\boldsymbol{\theta}} \log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z}^*), \quad \mathbf{z}^* \sim q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}), \end{aligned} $$

where the last approximation uses Monte-Carlo sampling.

However, during the E-step, getting an unbiased estimator for the gradient $\nabla_{\boldsymbol{\phi}}\mathcal{L}_{\boldsymbol{\phi}, \boldsymbol{\theta}}(\mathbf{x})$ is tricky because the density function $q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})$ depends on $\boldsymbol{\phi}$. This makes Monte-Carlo estimation impossible:

$$ \begin{aligned} \nabla_{\boldsymbol{\phi}}\mathcal{L}_{\boldsymbol{\phi}, \boldsymbol{\theta}}(\mathbf{x}) &= \textcolor{blue}{\nabla_{\boldsymbol{\phi}}} \int q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) \log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z}) d\mathbf{z} - \nabla_{\boldsymbol{\phi}} KL(q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) | p(\mathbf{z})) \\ &\textcolor{red}{\neq} \int q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) \textcolor{blue}{\nabla_{\boldsymbol{\phi}}} \log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z}) d\mathbf{z} - \nabla_{\boldsymbol{\phi}} KL(q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) | p(\mathbf{z})), \\ \end{aligned} $$

This is where the reparameterization trick comes in. Assuming $q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})$ to be Gaussian, we reparameterize the encoder’s outputs:

$$ \begin{aligned} \nabla_{\boldsymbol{\phi}} \int \textcolor{blue}{q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})} \log p_{\boldsymbol{\theta}}(\mathbf{x}|\textcolor{olive}{\mathbf{z}}) d\mathbf{z} &= \int \textcolor{blue}{p(\boldsymbol{\epsilon})} \nabla_{\boldsymbol{\phi}} \log p_{\boldsymbol{\theta}}(\mathbf{x}|\textcolor{olive}{\mathbf{g}_{\boldsymbol{\phi}}(\mathbf{x}, \boldsymbol{\epsilon})}) d\boldsymbol{\epsilon} \\ &\approx \nabla_{\boldsymbol{\phi}} \log p_{\boldsymbol{\theta}}(\mathbf{x}|\textcolor{olive}{\boldsymbol{\sigma}_{\boldsymbol{\phi}}(\mathbf{x})} \odot \textcolor{blue}{\boldsymbol{\epsilon}^*} + \textcolor{olive}{\boldsymbol{\mu}_{\boldsymbol{\phi}}(\mathbf{x})}), \quad \textcolor{blue}{\boldsymbol{\epsilon}^*} \sim \mathcal{N}(0, \mathbf{I}), \end{aligned} $$

so we move the randomness from $\mathbf{z} \sim q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})$ to $\boldsymbol{\epsilon} \sim p(\boldsymbol{\epsilon})$ and use a deterministic transform $\mathbf{z} = \mathbf{g}_{\boldsymbol{\phi}}(\mathbf{x}, \boldsymbol{\epsilon})$ to get an unbiased gradient. Moreover, the normal assumptions for $q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})$ and $p(\mathbf{z})$ allow us to compute $KL$ analytically and thus calculate the gradient $\nabla_{\boldsymbol{\phi}}KL(q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x}) | p(\mathbf{z}))$.

This example highlights the reparameterization trick, crucial for getting unbiased gradient estimates in continuous latent space VAEs. However, discrete representations $\mathbf{z}$ are often more natural for modalities like texts or graphs, leading us to discrete VAE latents. Therefore:

• Our encoder should output a discrete distribution.
• We need an analogue of the reparameterization trick for discrete distributions.
• Our decoder should input discrete random variables.

The classical solution for discrete variable reparameterization is Gumbel-Softmax (Jang et al. 2017) or Concrete distribution (Maddison et al. 2017). This is one distribution proposed in parallel by two research groups.

Gumbel Distribution

$$ g \sim \mathrm{Gumbel}(0, 1) \quad \Leftrightarrow \quad g = -\log( - \log u), \quad u \sim \mathrm{Uniform}[0, 1] $$

Theorem (Gumbel-Max Trick)

Let $g_k \sim \mathrm{Gumbel}(0, 1)$ for $k = 1, \ldots, K$. Then a discrete random variable

$$ c = \arg\max_k [\log \pi_k + g_k] $$

has a categorical distribution $c \sim \mathrm{Categorical}(\boldsymbol{\pi})$.

• We can sample from the discrete distribution using Gumbel-Max reparameterization.
• Parameters and random variable sampling are separated (reparameterization trick).
• Problem: The $\arg\max$ operation is non-differentiable.

To overcome the above problem, the authors propose to consider one-hot Categorical distribution, that is simply replace random variable $c \in \{1, \ldots, K\}$ with a random vector $\mathbf{c} \in \{0, 1\}^{K}$, such that $$ c = k \quad \Leftrightarrow \quad \mathbf{c}_k = 1. $$

Gumbel-Softmax Relaxation

$$ \hat{\mathbf{c}} = \mathrm{softmax}\left( \frac{\log \boldsymbol{\pi} + \mathbf{g}}{\tau} \right) $$

Further, they relax $\arg\max$ operation into $\mathrm{softmax}$, introducing a temperature parameter $\tau > 0$. As temperature $\tau \to \infty$, Gumbel-Softmax distribution $\mathrm{GS}(\boldsymbol{\pi}, \tau)$ becomes more uniformly distributed (consider the analogy with melting ice). In contrast, as temperature limits to zero, that is $\tau \to 0$, the distribution becomes more and more similar with initial one — defined with $\arg\max$.

The key idea is that $\mathrm{softmax}$ is differential function, which allows us to do backward pass in neural networks! Such technique is called relaxation.

This method relaxes Bernoulli random variable $c \sim \mathrm{Be}(\pi)$. The idea is to clip a Gaussian random variable onto $(0, 1)$ and tune the mean $\mu$ and scale $\sigma$ instead of $\pi$: $$ \begin{aligned} \epsilon &\sim \mathcal{N}(0, 1),\\ z &= \sigma \cdot \epsilon + \mu,\\ \hat{c} &= \min (1, \max (0, z)), \end{aligned} $$ where $\mu$ is trainable and $\sigma$ is fixed during training.

This method relaxes Bernoulli random variable $c \sim \mathrm{Be}(\pi)$. It is inspired by Relaxed Bernoulli, but uses another random variable to clip onto $(0, 1)$ — Gumbel-Softmax. Moreover, they additionally “stretch” it from $(0, 1)$ to the wider interval $(a, b)$, where $a < 0$ and $b > 0$ as follows: $$ \begin{aligned} z &\sim \mathrm{GS}(\pi, \tau),\\ \tilde{z} &= (b - a) \cdot z + a,\\ \hat{c} &= \min (1, \max (0, \tilde{z})), \end{aligned} $$ This method applies hard-sigmoid technique to make two delta peaks at zero and one.

This method relaxes Bernoulli random variable $c \sim \mathrm{Be}(\pi)$. In order to achieve gradient flow through distribution parameter $\pi$, this method replaces the gradient w.r.t. $c$ with the gradient of a continuous relaxation $\pi$. Other words, it can be formulated using torch notation as follows: $$ \hat{c} = \pi + (c - \pi)\verb|.detach()|. $$ It means that on the forward pass $\hat{c} = c$, but on the backward pass $\hat{c} = \pi$.

This method relaxes Bernoulli random variable $c \sim \mathrm{Be}(\pi)$ by the following: $$ \begin{aligned} p &= \sigma(a), \\ b &\sim \mathrm{Be}(\sqrt{p}), \\ \hat{c} &= b \sqrt{p}, \end{aligned} $$ where $a$ is a parameter of this distribution.

This method relaxes multivariate Bernoulli random variable $\mathbf{c} \sim \mathrm{MultiBe}(\boldsymbol{\pi}, \mathbf{R})$.

Multivariate Bernoulli distribution assumes that each component $c_{k}$ can be correlated with others. It can be constructed via Gaussian copula $C_{\mathbf{R}}$ (Nelsen, 2007): $$ C_{\mathbf{R}}(u_1, \ldots, u_p) = \Phi_{\mathbf{R}}(\Phi^{-1}(u_1), \ldots, \Phi^{-1}(u_p)), $$ where $\Phi_{\mathbf{R}}$ is the joint CDF of a multivariate Gaussian distribution with correlation matrix $\mathbf{R}$, and $\Phi^{-1}$ is the inverse CDF of the standard univariate Gaussian distribution. Each component is defined as follows: $$ c_{k} = \begin{cases} 1, \quad u_k \leq \pi_k,\\ 0, \quad u_k > \pi_k. \end{cases} $$

Relaxation is applied as follows: $$ \hat{\mathbf{c}} = \sigma \left( \frac{1}{\tau} \left( \log \frac{\mathbf{u}}{1 - \mathbf{u}} + \log \frac{\boldsymbol{\pi}}{1 - \boldsymbol{\pi}} \right) \right), $$ where $\sigma(x) = (1 + \exp(-x))^{-1}$ is a sigmoid function, and $\tau$ is a temperature hyperparameter.

This method relaxes Categorical random variable $\mathbf{c} \sim \mathrm{Cat}(\boldsymbol{\pi})$. The idea is to transform Gaussian noise $\boldsymbol{\epsilon}$ through invertible transformation $g(\cdot, \tau)$ with temperature $\tau$ onto the simplex: $$ \begin{aligned} \boldsymbol{\epsilon} &\sim \mathcal{N}(0, \mathbf{I}),\\ \mathbf{z} &= \boldsymbol{\sigma} \odot \boldsymbol{\epsilon} + \boldsymbol{\mu},\\ \hat{\mathbf{c}} &= g(\mathbf{z}, \tau) = \mathrm{softmax}_{++}(\mathbf{z/\tau}), \end{aligned} $$ where the latter function remains invertible via introduction $\delta > 0$ as follows: $$ \mathrm{softmax}_{++}(\mathbf{z/\tau}) = \frac{\exp(\mathbf{z}/\tau)}{\sum_{k=1}^{K} \exp(z_k/\tau) + \delta}. $$

This method relaxes Categorical random variable $\mathbf{c} \sim \mathrm{Cat}(\boldsymbol{\pi})$. However, it solves quite different problem. Suppose we want to get $K$ samples without replacement (i.e., not repeating) according to the Categorical distribution with probabilities $\boldsymbol{\pi}$. Similar to the Gumbel-Max method, let $g_k \sim \mathrm{Gumbel}(0, 1)$ for $k = 1, \ldots, K$, then the Gumbel-Max TOP-K Theorem says, that the values of the form

$$ c_1, \ldots , c_K = \mathrm{arg}\underset{k}{\mathrm{top}\text{-}\mathrm{K}} [ \log\pi_k + g_k] $$

have the $\mathrm{Cat}(\boldsymbol{\pi})$ distribution without replacement.

This approach has all the same pros and cons as the classical Gumbel-Max trick, however, they can be fixed with the Gumbel-Softmax relaxation using a simple loop:

$$ \begin{aligned} &\text{for } k = 1, \dots, K \\ &\quad 1. \quad c_k \sim \mathrm{GS}(\boldsymbol{\pi}, \tau) \\ &\quad 2. \quad \pi_k = -\infty \end{aligned} $$

This method is not a relaxation technique, but is quite useful in terms of dicrete variables approximation. This is an approach of approximating Dirichlet distribution with Logistic-Normal, and vice versa.

In particular, the analytic map from the Dirichlet distribution parameter $\boldsymbol{\alpha} \in \mathbb{R}_{+}^{K}$ to the parameters of the Gaussian $\boldsymbol{\mu} \in \mathbb{R}^{K}$ and symmetric positive definite $\boldsymbol{\Sigma} \in \mathbb{R}^{K \times K}$ is given by

$$ \begin{aligned} \mu_i &= \log \alpha_i - \frac{1}{K} \sum_{k=1}^{K} \log \alpha_k,\\ \Sigma_{ij} &= \delta_{ij} \frac{1}{\alpha_i} - \frac{1}{K} \left( \frac{1}{\alpha_i} + \frac{1}{\alpha_j} - \frac{1}{K} \sum_{k=1}^{K} \frac{1}{\alpha_k} \right), \end{aligned} $$

and the pseudo-inverse of this one, which maps the Gaussian parameters to those of the Dirichlet as

$$ \alpha_k = \frac{1}{\Sigma_{kk}} \left( 1 - \frac{2}{K} + \frac{e^{\mu_k}}{K^2} \sum_{l=1}^{K} e^{-\mu_l} \right). $$

And this is what is called Laplace Bridge between Dirichlet and Logistic-Normal distributions.