In "Variational Autoencoders (V): VAE + BN = Better VAE", we discussed the common phenomenon of KL divergence vanishing when training VAEs in NLP, and mentioned using Batch Normalization (BN) to ensure a positive lower bound for the KL divergence term, thereby preventing it from vanishing. In fact, as early as 2018, work based on similar ideas was proposed. These approaches involve using new prior and posterior distributions in the VAE to ensure the KL divergence term has a positive lower bound.
This idea appeared in two similar papers in 2018: "Hyperspherical Variational Auto-Encoders" and "Spherical Latent Spaces for Stable Variational Autoencoders". Both use the von Mises–Fisher (vMF) distribution defined on a hypersphere to construct the prior and posterior distributions. In some ways, this distribution is even simpler and more interesting than the Gaussian distribution we commonly use.
KL Divergence Vanishing
We know that the training objective of a VAE is: \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] where the first term is the reconstruction term and the second is the KL divergence term. In "Variational Autoencoders (I): What is it all about?", we mentioned that these two terms are in a sense "adversarial." The existence of the KL divergence term increases the difficulty for the decoder to utilize the encoded information. If the KL divergence term is 0, it means the decoder is not using any information from the encoder at all.
In NLP, the input and the object of reconstruction are sentences. To ensure performance, the decoder is generally an autoregressive model. However, autoregressive models are very powerful—so powerful that they can complete training even without input (degenerating into an unconditional language model). As mentioned, since the KL divergence term makes it harder for the decoder to use encoded information, the decoder simply ignores it, leading to the KL divergence vanishing phenomenon.
An early common solution was to gradually increase the weight of the KL term to guide the decoder to use encoded information. A more popular current solution is to make certain modifications so that the KL divergence term directly has a positive lower bound. Replacing the prior and posterior distributions with the vMF distribution is a classic example of this approach.
vMF Distribution
The vMF distribution is defined on a (d-1)-dimensional hypersphere. Its sample space is S^{d-1}=\{x|x\in\mathbb{R}^d, \Vert x\Vert=1\}, and its probability density function is: 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} where \xi\in\mathbb{R}^d is a pre-specified parameter vector. It is easy to imagine that this is a distribution on S^{d-1} centered at \xi. Writing the normalization factor as Z_{d, \Vert\xi\Vert} implies it only depends on the norm of \xi, which is due to isotropy. Because of this property, a more common notation for the vMF distribution is to set \mu=\xi/\Vert\xi\Vert, \kappa=\Vert\xi\Vert, C_{d,\kappa}=1/Z_{d, \Vert\xi\Vert}, resulting in: p(x) = C_{d,\kappa} e^{\kappa\langle\mu,x\rangle} In this case, \langle\mu,x\rangle is the cosine of the angle between \mu and x. Thus, the vMF distribution is essentially a distribution based on cosine similarity. Since we often use cosine values to measure the similarity between two vectors, models based on the vMF distribution usually satisfy this requirement better. When \kappa=0, the vMF distribution is a uniform distribution on the sphere.
From the integral form of the normalization factor Z_{d, \Vert\xi\Vert}, it is actually the moment-generating function of the vMF distribution. Consequently, the moments of the vMF distribution can also be expressed through Z_{d, \Vert\xi\Vert}. For example, the first moment is: \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} It can be seen that the direction of \mathbb{E}_{x\sim p(x)} [x] is consistent with \xi. The exact form of Z_{d, \Vert\xi\Vert} can be calculated, but it is quite complex, and in many cases, we do not need to know the exact normalization factor, so we will not calculate it here.
As for the meaning of the parameter \kappa, it might be easier to understand if we set \tau=1/\kappa. Then p(x)\sim e^{\langle\mu,x\rangle/\tau}. Students familiar with energy-based models know that \tau is the temperature parameter. If \tau is smaller (\kappa is larger), the distribution is more concentrated around \mu; otherwise, it is more dispersed (closer to a uniform distribution on the sphere). Therefore, \kappa is also vividly called the "concentration" parameter.
Sampling from vMF
For the vMF distribution, the first major challenge is how to sample specific instances from it. This step is crucial if we want to apply it to VAEs.
Uniform Distribution
The simplest case is \kappa=0, which is the uniform distribution on the (d-1)-dimensional sphere. Since the standard normal distribution is isotropic and its probability density is proportional to e^{-\Vert x\Vert^2/2} (depending only on the norm), we only need to sample a z from a d-dimensional standard normal distribution and then let x=z/\Vert z\Vert to obtain a uniform sample on the sphere.
Special Direction
Next, for the case where \kappa > 0, let x=[x_1,x_2,\cdots,x_d]. First, consider a special case: \mu = [1, 0, \cdots, 0]. In fact, due to isotropy, we often only need to consider this special case and then generalize it to the general case in parallel.
In this case, the probability density is proportional to e^{\kappa x_1}. We then transform to spherical coordinates: \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. Then (for the integral transformation of hyperspherical coordinates, please refer directly to Wikipedia): \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} This decomposition shows that sampling from this vMF distribution is equivalent to first sampling a \varphi_1 from a distribution with probability density proportional to e^{\kappa\cos\varphi_1}\sin^{d-2}\varphi_1, and then uniformly sampling a (d-1)-dimensional vector \varepsilon = [\varepsilon_2,\varepsilon_3,\cdots,\varepsilon_d] from a (d-2)-dimensional hypersphere, combining them into the final sample as follows: x = [\cos\varphi_1, \varepsilon_2\sin\varphi_1, \varepsilon_3\sin\varphi_1, \cdots, \varepsilon_d\sin\varphi_1] Let w=\cos\varphi_1\in[-1,1], then: \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| So we mainly study sampling from a distribution with probability density proportional to e^{\kappa w} (1-w^2)^{(d-3)/2}.
However, what I do not understand is that most papers involving the vMF distribution use the rejection sampling scheme based on the Beta distribution proposed in the 1994 paper "Simulation of the von Mises Fisher distribution". The entire sampling process is quite complex. But it is 2021 now; do we still need such an inefficient scheme as rejection sampling for a one-dimensional distribution?
In fact, for any one-dimensional distribution p(w), let its cumulative distribution function be \Phi(w). Then w=\Phi^{-1}(\varepsilon), \varepsilon\sim U[0,1] is the most convenient and universal sampling scheme. Some readers might protest that "the cumulative distribution function is hard to calculate" or "its inverse function is even harder to calculate." But when implementing sampling in code, we don’t need to know what \Phi(w) looks like; we can just use numerical computation directly. 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 intensive part is
the calculation of variable y. Once calculated, it can be
cached, and subsequent sampling only requires executing the last step,
which is very fast. This is much simpler and more convenient than
rejection sampling from a Beta distribution. Incidentally, a trick is
used here: calculate the logarithm first, subtract the maximum value,
and then calculate the exponential to prevent overflow. Even if \kappa is in the thousands, it can be
calculated successfully.
General Case
Now that we have implemented sampling from a vMF distribution with \mu=[1,0,\cdots,0], we can decompose the sampling result as: 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{uniform sample from } d-2\text{ dim}\\ \text{hypersphere orthogonal to } \mu\end{array}} Similarly, due to isotropy, for a general \mu, the sampling result still has the same form: \begin{aligned} &x = w\mu + \sqrt{1-w^2}\nu\\ &w\sim e^{\kappa w} (1-w^2)^{(d-3)/2}\\ &\nu\sim \text{uniform distribution on } d-2\text{ dim hypersphere orthogonal to } \mu \end{aligned} To sample \nu, the key is orthogonality to \mu. This is also easy to implement: first sample a d-dimensional vector \varepsilon from a standard normal distribution, then keep the component orthogonal to \mu and normalize it: \nu = \frac{\varepsilon - \langle \varepsilon,\mu\rangle \mu}{\Vert \varepsilon - \langle \varepsilon,\mu\rangle \mu\Vert},\quad \varepsilon\sim\mathcal{N}(0,I_d)
vMF-VAE
At this point, we have completed the most difficult part of this article. Building the vMF-VAE is now straightforward. vMF-VAE chooses the uniform distribution on the sphere (\kappa=0) as the prior distribution q(z) and the vMF distribution as the posterior distribution: p(z|x) = C_{d,\kappa} e^{\kappa\langle\mu(x),z\rangle} For simplicity, we set \kappa as a hyperparameter (which can also be understood as updating this parameter manually rather than through gradient descent). Consequently, the only source of parameters for p(z|x) is \mu(x). We can then calculate the KL divergence term: \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} As discussed earlier, the direction of the mean of the vMF distribution is consistent with \mu(x), and the norm depends only on d and \kappa. Substituting this into the equation, we see that the KL divergence term depends only on d and \kappa. Once these two parameters are selected, it becomes a constant (by the properties of KL divergence, it must be greater than 0 when \kappa \neq 0), and the KL divergence vanishing phenomenon will never occur.
Now only the reconstruction term remains. We need to use "Reparameterization" to complete the sampling while preserving gradients. Since we have already studied the sampling process of vMF, it is not difficult to implement. The overall process is: \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,I_d) \end{aligned} Here, the reconstruction loss is shown as MSE as an example; if it were 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 and does not affect optimization, vMF-VAE is just like a regular autoencoder but with a slightly more complex reparameterization operation (and manual adjustment of \kappa). Compared to the standard VAE based on the Gaussian distribution, it is significantly simplified.
Furthermore, from this process, we can see that besides "simplicity," another main reason for not making \kappa trainable is that \kappa is involved in the sampling of w, and it is difficult to preserve the gradient of \kappa during the sampling of w.
Reference Implementation
The difficulty of implementing vMF-VAE lies in the reparameterization
part, specifically sampling from the vMF distribution, where the key is
sampling w. We previously provided a
NumPy implementation for sampling w,
but there is no function similar to np.interp in
TensorFlow, making it difficult to convert to a pure TF implementation.
Of course, for dynamic graph frameworks like PyTorch or TF2, mixing
NumPy code is fine, but here I want to construct a more general
solution.
It’s actually not hard. Since w is
just a one-dimensional variable and each training step only requires
batch_size samples, we can pre-sample a large enough number
(hundreds of thousands) of w values
using NumPy and store them. During training, we can simply draw randomly
from this batch of pre-sampled results. A reference implementation is as
follows:
def sampling(mu):
"""vMF distribution reparameterization
"""
dims = K.int_shape(mu)[-1]
# Pre-calculate a batch of 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))
# Real-time sampling of w
idxs = K.random_uniform(K.shape(mu[:, :1]), 0, 10**6, dtype='int32')
w = K.gather(W, idxs)
# Real-time sampling of 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 * nuA complete example based on MNIST can be found at:
As for examples of vMF-VAE applied to NLP, we will share those when the opportunity arises. This article mainly focuses on theoretical introduction and simple demonstrations.
Summary
This article introduced the implementation of VAE based on the vMF distribution. The main difficulty lies in sampling from the vMF distribution. Overall, the vMF distribution is built on the cosine similarity metric, and its properties in certain aspects are more consistent with our intuitive understanding. Using it in a VAE allows the KL divergence term to be a constant, thereby preventing the KL divergence vanishing phenomenon and simplifying the VAE structure.
When reprinting, please include the original address of this article: https://kexue.fm/archives/8404
For more detailed reprinting matters, please refer to: "Scientific Space FAQ"