Variational Autoencoders (VII): VAE on a Sphere (vMF-VAE)

变分自编码器（七）：球面上的VAE（vMF-VAE） 苏剑林 · 2021-05-17 · Machine Learning variational unsupervised VAE generative model

In Variational Autoencoders (V): VAE + BN = A Better VAE, we discussed the KL divergence vanishing phenomenon commonly observed when training VAEs in NLP, and mentioned that BN could be used to give the KL divergence term a positive lower bound, thereby preventing the KL term from vanishing. In fact, as early as 2018, similar ideas had already been proposed: by switching to new prior and posterior distributions in the VAE, one can likewise ensure the KL divergence term has a positive lower bound.

This idea appeared in two closely related 2018 papers, Hyperspherical Variational Auto-Encoders and Spherical Latent Spaces for Stable Variational Autoencoders, both of which use the von Mises–Fisher (vMF) distribution defined on a hypersphere to construct the prior and posterior. In some sense, this distribution is even simpler and more interesting than the Gaussian we usually work with.

KL Divergence Vanishing

As we know, the training objective of a VAE is

\begin{equation}\mathcal{L} = \mathbb{E}_{x\sim \tilde{p}(x)} \Big[\mathbb{E}_{z\sim p(z|x)}\big[-\log q(x|z)\big]+KL\big(p(z|x)\big\Vert q(z)\big)\Big] \end{equation}more

where the first term is the reconstruction term and the second is the KL divergence term. As discussed in Variational Autoencoders (I): So That's What It Is, these two terms are in a sense "adversarial": the presence of the KL divergence term makes it harder for the decoder to make use of the encoded information. If the KL divergence term becomes 0, it means the decoder is making no use of the encoder's information at all.

In NLP, the objects being encoded and reconstructed are sentences, and to ensure good performance, the decoder is generally an autoregressive model. However, autoregressive models are extremely powerful — powerful enough to complete training even without any input at all (degenerating into an unconditional language model). And as we just said, the KL divergence term makes it harder for the decoder to exploit the encoded information, so the decoder simply gives up on using it altogether. This is exactly how KL divergence vanishing arises.

An early and fairly common remedy was to gradually increase the weight of the KL term, in order to nudge the decoder toward using the encoded information. A more popular approach nowadays is to make certain structural changes so that the KL divergence term is guaranteed a positive lower bound outright. Replacing the prior and posterior with vMF distributions is one classic example of this kind of approach.

The vMF Distribution

The vMF distribution is defined on a $d-1$-dimensional hypersphere, with sample space $S^{d-1}=\{x|x\in\mathbb{R}^d, \Vert x\Vert=1\}$, and probability density function

\begin{equation}p(x) = \frac{e^{\langle\xi,x\rangle}}{Z_{d, \Vert\xi\Vert}},\quad Z_{d, \Vert\xi\Vert}=\int_{S^{d-1}}e^{\langle\xi,x\rangle} dS^{d-1}\end{equation}

where $\xi\in\mathbb{R}^d$ is a pre-specified parameter vector. It's not hard to imagine that this is a distribution on $S^{d-1}$ centered at $\xi$; writing the normalization factor as $Z_{d, \Vert\xi\Vert}$ makes it clear that it depends only on the norm of $\xi$, a consequence of isotropy. Because of this property, the vMF distribution is more commonly written by setting $\mu=\xi/\Vert\xi\Vert, \kappa=\Vert\xi\Vert, C_{d,\kappa}=1/Z_{d, \Vert\xi\Vert}$, giving

\begin{equation}p(x) = C_{d,\kappa} e^{\kappa\langle\mu,x\rangle}\end{equation}

Here $\langle\mu,x\rangle$ is precisely the cosine of the angle between $\mu,x$, which shows that the vMF distribution is essentially a distribution defined in terms of cosine similarity. Since cosine similarity is so often used to measure the similarity between two vectors, models built on the vMF distribution tend to align more naturally with this need. When $\kappa=0$, the vMF distribution reduces to the uniform distribution on the sphere.

Looking at the normalization factor $Z_{d, \Vert\xi\Vert}$ in its integral form, it is in fact also the moment-generating function of the vMF distribution, so the various moments of vMF can be expressed in terms of $Z_{d, \Vert\xi\Vert}$. For example, the first moment is

\begin{equation}\mathbb{E}_{x\sim p(x)} [x] = \nabla_{\xi} \log Z_{d, \Vert\xi\Vert}=\frac{d \log Z_{d,\Vert\xi\Vert}}{d\Vert\xi\Vert}\frac{\xi}{\Vert\xi\Vert}\end{equation}

We can see that $\mathbb{E}_{x\sim p(x)} [x]$ points in the same direction as $\xi$. The exact form of $Z_{d, \Vert\xi\Vert}$ can be worked out, but it's fairly complicated, and in most cases we don't actually need to know this normalization factor exactly, so we won't compute it here.

As for the meaning of the parameter $\kappa$, it may be easier to understand by setting $\tau=1/\kappa$, so that $p(x)\sim e^{\langle\mu,x\rangle/\tau}$. Readers familiar with energy-based models will recognize $\tau$ here as the temperature parameter: the smaller $\tau$ is (equivalently, the larger $\kappa$ is), the more concentrated the distribution is around $\mu$, whereas the opposite makes it more spread out (approaching a uniform distribution over the sphere). This is why $\kappa$ is vividly referred to as the "concentration" parameter.

Sampling from vMF

The first difficulty we need to tackle with the vMF distribution is how to actually draw concrete samples from it. This step is especially critical if we want to apply it within a VAE.

Uniform Distribution

The simplest case is $\kappa=0$, i.e., the uniform distribution on the $d-1$-dimensional sphere. Since the standard normal distribution is itself isotropic, with a probability density proportional to $e^{-\Vert x\Vert^2/2}$ that depends only on the norm, we simply need to sample a vector $z$ from a standard normal $d$ and let $x=z/\Vert z\Vert$ to obtain a uniform sample on the sphere.

A Special Direction

Next, for the general case $\kappa > 0$, let us write $x=[x_1,x_2,\cdots,x_d]$, and first consider a special case: $\mu = [1, 0, \cdots, 0]$. In fact, because of isotropy, in most situations we only need to work out this special case, and the general case follows in parallel.

Here the probability density is proportional to $e^{\kappa x_1}$. Let us switch to spherical coordinates:

\begin{equation} \left\{\begin{aligned} x_1 &= \cos\varphi_1\\ x_2 &= \sin\varphi_1 \cos\varphi_2 \\ x_3 &= \sin\varphi_1 \sin\varphi_2 \cos\varphi_3 \\ &\,\,\vdots \\ x_{d-1} &= \sin\varphi_1 \cdots \sin\varphi_{d-2} \cos\varphi_{d-1}\\ x_d &= \sin\varphi_1 \cdots \sin\varphi_{d-2} \sin\varphi_{d-1} \end{aligned}\right. \end{equation}

Then (for the integral transformation on hyperspherical coordinates, please refer directly to Wikipedia)

\begin{equation}\begin{aligned} e^{\kappa x_1}dS^{d-1} =& e^{\kappa\cos\varphi_1}\sin^{d-2}\varphi_1 \sin^{d-3}\varphi_2 \cdots \sin\varphi_{d-2} d\varphi_1 d\varphi_2 \cdots d\varphi_{d-1} \\ =& \left(e^{\kappa\cos\varphi_1}\sin^{d-2}\varphi_1 d\varphi_1\right)\left(\sin^{d-3}\varphi_2 \cdots \sin\varphi_{d-2} d\varphi_2 \cdots d\varphi_{d-1}\right) \\ =& \left(e^{\kappa\cos\varphi_1}\sin^{d-2}\varphi_1 d\varphi_1\right)dS^{d-2} \\ \end{aligned}\end{equation}

This decomposition shows that sampling from this vMF distribution is equivalent to first sampling a scalar $\varphi_1$ from a distribution with density proportional to $e^{\kappa\cos\varphi_1}\sin^{d-2}\varphi_1$, then uniformly sampling a $d-1$-dimensional vector $\varepsilon = [\varepsilon_2,\varepsilon_3,\cdots,\varepsilon_d]$ on the $d-2$-dimensional hypersphere, and combining them as follows to obtain the final sample

\begin{equation}x = [\cos\varphi_1, \varepsilon_2\sin\varphi_1, \varepsilon_3\sin\varphi_1, \cdots, \varepsilon_d\sin\varphi_1]\end{equation}

Setting $w=\cos\phi_1\in[-1,1]$, we then have

\begin{equation}\left|e^{\kappa\cos\varphi_1}\sin^{d-2}\varphi_1 d\varphi_1\right| = \left|e^{\kappa w} (1-w^2)^{(d-3)/2}dw\right|\end{equation}

So our main task is studying how to sample from a distribution whose density is proportional to $e^{\kappa w} (1-w^2)^{(d-3)/2}$.

However, what I find puzzling is that most papers dealing with the vMF distribution adopt the rejection-sampling scheme based on the beta distribution, proposed in the 1994 paper Simulation of the von Mises Fisher distribution — an overall sampling procedure that is fairly complicated. But here we are in 2021, and for sampling a one-dimensional distribution, do we really still need such an inefficient scheme as rejection sampling?

In fact, for any one-dimensional distribution $p(w)$, if we let its cumulative distribution function be $\Phi(w)$, then $w=\Phi^{-1}(\varepsilon),\varepsilon\sim U[0,1]$ is the most convenient and general sampling scheme available. Some readers might object, "The CDF is hard to compute!" or "Its inverse is even harder to compute!" But when implementing sampling in code, we don't actually need to know what $\Phi(w)$ looks like analytically at all — we can just compute it numerically. A reference implementation is as follows:

import numpy as np

def sample_from_pw(size, kappa, dims, epsilon=1e-7):
    x = np.arange(-1 + epsilon, 1, epsilon)
    y = kappa * x + np.log(1 - x**2) * (dims - 3) / 2
    y = np.cumsum(np.exp(y - y.max()))
    y = y / y[-1]
    return np.interp(np.random.random(size), y, x)

In this implementation, the most computationally expensive part is computing the variable y; but once it's computed, it can be cached, and afterward only the last step needs to be executed to complete the sampling, which is extremely fast. However you look at it, this is simpler and more convenient than rejection sampling from a beta distribution. Incidentally, the implementation here also uses a trick: first compute the log values, subtract the maximum, and only then take the exponential — this prevents overflow, so that even if $\kappa$ is in the thousands or tens of thousands, the computation still succeeds.

The General Case

Now that we've implemented sampling from the vMF distribution with $\mu=[1,0,\cdots,0]$, we can decompose the sampling result as

\begin{equation}x = w\times\underbrace{[1,0,\cdots,0]}_{\text{parameter vector}\mu} + \sqrt{1-w^2}\times\underbrace{[0,\varepsilon_2,\cdots,\varepsilon_d]}_{\begin{array}{c}\text{and}\mu\text{orthogonal}d-2\text{dim}\\ \text{uniform sampling on hypersphere}\end{array}}\end{equation}

Again, by isotropy, for a general $\mu$, the sampling result retains exactly the same form:

\begin{equation}\begin{aligned} &x = w\mu + \sqrt{1-w^2}\nu\\ &w\sim e^{\kappa w} (1-w^2)^{(d-3)/2}\\ &\nu\sim \text{and}\mu\text{orthogonal}d-2\text{uniform distribution on hypersphere in dim} \end{aligned}\end{equation}

For sampling $\nu$, the key point is that it must be orthogonal to $\mu$. This isn't hard either: first sample a $d$-dimensional vector $z$ from a standard normal distribution, then keep only the component orthogonal to $\mu$ and normalize:

\begin{equation}\nu = \frac{\varepsilon - \langle \varepsilon,\mu\rangle \mu}{\Vert \varepsilon - \langle \varepsilon,\mu\rangle \mu\Vert},\quad \varepsilon\sim\mathcal{N}(0,1_d)\end{equation}

vMF-VAE

At this point, we've completed the hardest part of this article, and constructing the vMF-VAE is now a straightforward matter. The vMF-VAE chooses the uniform distribution on the sphere ($\kappa=0$) as the prior $q(z)$, and selects the vMF distribution as the posterior:

\begin{equation}p(z|x) = C_{d,\kappa} e^{\kappa\langle\mu(x),z\rangle}\end{equation}

For simplicity, we treat $\kappa$ as a hyperparameter (which can also be thought of as updating this parameter manually rather than through gradient descent). This means the only source of parameters for $p(z|x)$ is $\mu(x)$. We can then compute the KL divergence term:

\begin{equation}\begin{aligned} \int p(z|x) \log\frac{p(z|x)}{q(z)} dz =&\, \int C_{d,\kappa} e^{\kappa\langle\mu(x),z\rangle}\left(\kappa\langle\mu(x),z\rangle + \log C_{d,\kappa} - \log C_{d,0}\right)dz\\ =&\,\kappa\left\langle\mu(x),\mathbb{E}_{z\sim p(z|x)}[z]\right\rangle + \log C_{d,\kappa} - \log C_{d,0} \end{aligned}\end{equation}

As discussed earlier, the mean direction of the vMF distribution aligns with $\mu(x)$, while its magnitude depends only on $d$ and $\kappa$. Substituting this into the equation above, we see that the KL divergence term depends only on $d$ and $\kappa$; once these two parameters are fixed, it becomes a constant (and by the properties of KL divergence, when $\kappa\neq 0$, it is necessarily greater than 0). There is thus no possibility of KL divergence vanishing.

That leaves the reconstruction term. We need to use the "reparameterization" trick to perform sampling while preserving gradients. Since we've already worked out the vMF sampling process above, this isn't difficult to implement either. The overall pipeline is:

\begin{equation}\begin{aligned} &\mathcal{L} = \Vert x - g(z)\Vert^2\\ &z = w\mu(x) + \sqrt{1-w^2}\nu\\ &w\sim e^{\kappa w} (1-w^2)^{(d-3)/2}\\ &\nu=\frac{\varepsilon - \langle \varepsilon,\mu\rangle \mu}{\Vert \varepsilon - \langle \varepsilon,\mu\rangle \mu\Vert}\\ &\varepsilon\sim\mathcal{N}(0,1_d) \end{aligned}\end{equation}

Here the reconstruction loss is taken as MSE for illustration; for sentence reconstruction, cross-entropy would be used instead. Here $\mu(x)$ is the encoder and $g(z)$ is the decoder. Since the KL divergence term is a constant, it has no effect on the optimization, which means that compared to an ordinary autoencoder, the vMF-VAE only adds one somewhat complicated reparameterization operation (plus manually tuning $\kappa$) — considerably simpler than the standard Gaussian-based VAE.

Moreover, from this pipeline we can also see that, aside from "simplicity," there's another important reason for not making $\kappa$ trainable: $\kappa$ is tied to the sampling of $w$, and it's quite difficult to preserve gradients with respect to $w$ through the sampling process of $\kappa$.

Reference Implementation

The main difficulty in implementing the vMF-VAE lies in the reparameterization step, i.e., sampling from the vMF distribution, and the crux of that is sampling $w$. Earlier we already gave a numpy implementation for sampling $w$, but since TensorFlow lacks an analogous function to np.interp, it isn't straightforward to convert this into a pure TF implementation. Of course, if you're using a dynamic-graph framework like PyTorch or TF2, mixing in numpy code directly is fine, but here I'd still like to construct a more general-purpose solution.

This turns out not to be too hard either. Since $w$ is just a one-dimensional variable, and each training step only needs batch_size sampled values, we can simply pre-sample a large enough batch (hundreds of thousands) of $w$ values in advance using numpy functions, store them, and then during training just draw randomly from this pre-sampled pool. A reference implementation is as follows:

def sampling(mu):
    """vMF分布重参数操作
    """
    dims = K.int_shape(mu)[-1]
    # 预先计算一批w
    epsilon = 1e-7
    x = np.arange(-1 + epsilon, 1, epsilon)
    y = kappa * x + np.log(1 - x**2) * (dims - 3) / 2
    y = np.cumsum(np.exp(y - y.max()))
    y = y / y[-1]
    W = K.constant(np.interp(np.random.random(10**6), y, x))
    # 实时采样w
    idxs = K.random_uniform(K.shape(mu[:, :1]), 0, 10**6, dtype='int32')
    w = K.gather(W, idxs)
    # 实时采样z
    eps = K.random_normal(K.shape(mu))
    nu = eps - K.sum(eps * mu, axis=1, keepdims=True) * mu
    nu = K.l2_normalize(nu, axis=-1)
    return w * mu + (1 - w**2)**0.5 * nu

A complete example based on MNIST can be found here:

https://github.com/bojone/vae/blob/master/vae_vmf_keras.py

As for examples of applying vMF-VAE to NLP, we'll share those another time if the opportunity arises. This article has mainly focused on theoretical introduction and a simple demonstration.

Summary

This article introduced a VAE implementation based on the vMF distribution, the main difficulty of which lies in sampling from the vMF distribution. Overall, the vMF distribution is built on cosine similarity as its measure, and its properties align better with our intuition in certain respects. Using it in a VAE makes the KL divergence term a constant, thereby preventing the KL divergence vanishing phenomenon, while also simplifying the VAE's structure.