The Principle of Minimum Entropy (VI): How Should We Choose the Dimensionality of Word Embeddings?

最小熵原理（六）：词向量的维度应该怎么选择？ 苏剑林 · 2020-08-20 · Machine Learning entropy word vector minimum entropy

With the development of NLP, word embedding models like Word2Vec and GloVe are gradually being replaced by Transformer-based models such as BERT. But classics are classics — word embedding models still shine in plenty of scenarios, and there remain aspects of them worth studying. In this post, I want to address a question that anyone building a word embedding model might wonder about: roughly how large should the embedding dimension be?

Let me state the conclusion first. My estimate is

\begin{equation}n > 8.33\log N\label{eq:final}\end{equation}

or, more simply, just remember $n > 8\log N$, where $N$ is the vocabulary size, $n$ is the word embedding dimension, and $\log$ is the natural logarithm. When $n$ exceeds this threshold, it means the model has enough capacity to accommodate these $N$ words (of course, the larger $n$ is, the greater the overfitting risk). So, when $N=100000$, the resulting $n$ is about 96, meaning that for a word embedding model with a vocabulary of 100,000 words, a dimension of 96 is already sufficient; if we want to accommodate 5 million words, then $n$ comes out to about 128. more

Background

The reason this question came to mind is that yesterday I came across the paper Word2vec Skip-gram Dimensionality Selection via Sequential Normalized Maximum Likelihood on Arxiv. Unfortunately, I didn't find the answer I was looking for in that paper. While searching around, I also found other papers studying the same problem, such as On the Dimensionality of Word Embedding, but the answer still wasn't what I wanted.

Why do I say this? Well, it's obvious that the most authoritative answer to this question should come from repeated experiments to determine the optimal dimension, so we can't expect theoretical analysis to give a very precise answer. The word embedding dimensions we typically use in practice are things like 64, 100, 128, 256, 300, and so on — the difference in performance between these dimensions isn't actually that large. So all I want here is to derive, in the simplest and most intuitive way possible, the order of magnitude of the dimension a typical word embedding model needs — say, tens or hundreds — without resorting to overly complicated analysis.

Since I couldn't find a satisfying existing result, I decided to analyze this from the perspective of the principle of minimum entropy, and arrived at an answer close to what I had in mind.

Analysis

What we're analyzing here is a word embedding model based on the Skip-Gram idea. Most word embedding models are actually variants of it. As for CBOW-type models, based on past experiments, their performance is actually quite close to Skip-Gram (especially with larger amounts of data), so we can consider the analysis result for Skip-Gram to be broadly applicable.

Minimizing entropy

Our starting point is information entropy. As we know, entropy is a measure of uncertainty (see Entropy Ain't So Simple: From Entropy and the Maximum Entropy Principle to the Maximum Entropy Model (I)). Language itself carries a certain amount of uncertainty, and when we encode words as vectors, the result of that encoding should be equal to, or even smaller than, this uncertainty, in order to guarantee that the encoding is effective and fully preserves the information of the original language. So, we want to eliminate uncertainty — that is, we want to minimize entropy.

Note that since word embeddings here are based on the Skip-Gram model, what we need to compute is not the average entropy of individual words, but the average entropy of the whole Skip-Gram model. Suppose the frequency of the word pair $(w_i, w_j)$ is $\tilde{p}(w_i, w_j)$; then we can estimate its entropy as

\begin{equation}\tilde{H}=-\sum_{i, j} \tilde{p}(w_i, w_j)\log \tilde{p}(w_i, w_j)\end{equation}

Different word embedding training objectives differ somewhat — some fit the joint probability $p(w_i, w_j)$, others fit the conditional probability $p(w_j|w_i)$ — but the difference isn't significant. As mentioned before, this post is only aiming for a rough order-of-magnitude estimate. So here we uniformly assume that the word embedding model takes the form

\begin{equation}p(w_i, w_j) = \frac{e^{\langle\boldsymbol{u}_i, \boldsymbol{v}_j\rangle}}{Z},\quad Z = \sum_{i,j}e^{\langle\boldsymbol{u}_i, \boldsymbol{v}_j\rangle}\end{equation}

where $\boldsymbol{u},\boldsymbol{v}$ represents two distinct sets of word vectors (center-word vectors and context-word vectors), $\boldsymbol{u}_i$ represents the word $w_i$, and $\boldsymbol{v}_j$ represents the word $w_j$. In that case, the information entropy is

\begin{equation}H=-\sum_{i, j} p(w_i, w_j)\log p(w_i, w_j)=\log Z-\frac{1}{Z}\sum_{i, j} e^{\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle} \langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle\end{equation}

Sampling approximation

To approximate the above expression, we replace the summations with sampling approximations, for instance

\begin{equation}Z = \sum_{i, j} e^{\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle} = N^2\times \frac{1}{N^2}\sum_{i, j} e^{\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle}\approx N^2\mathbb{E}_{\boldsymbol{u}_i,\boldsymbol{v}_j}\left[e^{\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle}\right]\end{equation}

where $N$ is the vocabulary size. Similarly,

\begin{equation}\sum_{i, j} e^{\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle} \langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle\approx N^2\mathbb{E}_{\boldsymbol{u}_i,\boldsymbol{v}_j}\left[e^{\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle} \langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle\right]\end{equation}

so we get the approximation

\begin{equation}H\approx\log N^2 + \log \mathbb{E}_{\boldsymbol{u}_i,\boldsymbol{v}_j}\left[e^{\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle}\right]-\frac{\mathbb{E}_{\boldsymbol{u}_i,\boldsymbol{v}_j}\left[e^{\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle} \langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle\right]}{\mathbb{E}_{\boldsymbol{u}_i,\boldsymbol{v}_j}\left[e^{\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle}\right]}\end{equation}

Distributional assumption

Looking at existing word embedding models, we can observe that the values along each dimension are both positive and negative, and their absolute values tend to be fairly uniform in magnitude. Let's therefore assume that the absolute value of each component is roughly 1, so the norm of each word vector is approximately $\sqrt{n}$ (where $n$ is the word embedding dimension — the very quantity we want to estimate; if this approximation seems too crude, feel free to adjust it). We further assume that all word vectors are uniformly distributed on a $n$-dimensional hypersphere of radius $\sqrt{n}$, so that $\langle \boldsymbol{u}_i, \boldsymbol{v}_j\rangle=n\cos\theta$, where $\theta$ is the angle between them. Hence

\begin{equation}H\approx\log N^2 + \log \mathbb{E}_{\theta}\left[e^{n\cos\theta}\right]-\frac{\mathbb{E}_{\theta}\left[e^{n\cos\theta} n\cos\theta\right]}{\mathbb{E}_{\theta}\left[e^{n\cos\theta}\right]}\label{eq:H-theta}\end{equation}

Here $\theta$ is exactly the angle between two arbitrary vectors in $n$-dimensional space, and we already worked out its distribution in The Distribution of the Angle Between Two Random Vectors in n-Dimensional Space:

\begin{equation}p_n(\theta) = \frac{\Gamma\left(\frac{n}{2}\right)}{\Gamma\left(\frac{n-1}{2}\right)\sqrt{\pi}}\sin^{n-2} \label{eq:jqfb}\theta\end{equation}

Once the probability density function is fixed, for given $N$ and $n$, the approximation $\eqref{eq:H-theta}$ can be computed numerically without any trouble, and from $\tilde{H} > H$ we can then solve for the corresponding $n$.

Comparing the results

First, let's numerically compute some values of $h_n=\log \mathbb{E}_{\theta}\left[e^{n\cos\theta}\right]-\frac{\mathbb{E}_{\theta}\left[e^{n\cos\theta} n\cos\theta\right]}{\mathbb{E}_{\theta}\left[e^{n\cos\theta}\right]}$:

$$\begin{array}{c|cccccc} \hline n & 32 & 64 & 96 & 128 & 256 & 512\\ \hline h_n & -7.77471 & -15.4734 & -23.1726 & -30.8718 & -61.6692 & -123.264\\ \hline \end{array}$$

So, for example, when $n=64,N=100000$, we get $H\approx \log 100000^2 - 15.4734 = 7.55245$. Readers might find this odd — when $n=128,N=100000$, doesn't $H$ become negative? How can a discrete entropy be negative? In fact, this happens because, in the derivation above, we mixed sampling approximation with exact integration. When the dimensionality of the space $n$ is large enough, even sampling hundreds of thousands of points may not accurately estimate certain statistics — so this sampling-approximation step is what introduces the error.

However, this actually gives us another way to pin down $n$: once $H < 0$ occurs, it means that $N$ samples are no longer enough to properly estimate the statistic in question — in other words, at that point, the $n$-dimensional space has more than enough room to accommodate $N$ samples. So we can simply use $H < 0$ to determine a boundary, without needing to estimate $\tilde{H}$ itself. (Or, seen from another angle: $\tilde{H}$ must always be greater than 0, so $H < 0$ is a sufficient condition for $H < \tilde{H}$.)

Finally, we observe that $h_n$ is roughly linear in $n$, i.e., $h_n/n\approx -0.24$, so $H\approx\log N^2 -0.24n$; setting this to be less than 0 lets us solve for the formula $\eqref{eq:final}$.

An approximate estimate

We can also derive the result $h_n/n\approx -0.24$ theoretically. By the definition of $h_n$, we have

\begin{equation}h_n = \log\frac{\int_0^{\pi} \sin^{n-2}\theta\, e^{n\cos\theta}d\theta}{\int_0^{\pi} \sin^{n-2}\theta d\theta} - n\frac{\int_0^{\pi} \sin^{n-2}\theta\, e^{n\cos\theta}\cos\theta d\theta}{\int_0^{\pi} \sin^{n-2}\theta\,e^{n\cos\theta} d\theta}\end{equation}

So this becomes a matter of finding the asymptotic estimates of a few definite integrals. We'll mainly rely on the idea of Laplace's approximation, assuming $n$ is large enough that we can ignore the distinction between $n$ and $n-2$, giving

\begin{equation}\log\left[\sin^{n-2}\theta\, e^{n\cos\theta}\right]=(n-2)\log \sin\theta + n\cos\theta \approx n (\log \sin\theta + \cos\theta)\end{equation}

We can compute that $\log\sin\theta + \cos\theta$ attains its maximum over $[0,\pi]$ at $\arctan\sqrt{\frac{\sqrt{5}+1}{2}}\approx 0.904557$, so expanding to second order there gives

\begin{equation}\log \sin\theta + \cos\theta\approx 0.377428 -1.11803 (\theta - 0.904557)^2\end{equation}

from which

\begin{equation}\begin{aligned} \int_0^{\pi} \sin^{n-2}\theta\, e^{n\cos\theta}d\theta\approx&\, \int_{-\infty}^{\infty} e^{n[0.377428 -1.11803 (\theta - 0.904557)^2]}d\theta\\ \approx&\, \frac{1.67629}{\sqrt{n}}e^{0.377428n} \end{aligned}\end{equation}

Also note that when $n$ is large enough, the integrand differs appreciably from 0 only near $\theta=0.904557$ (or, equivalently, treat the limit of the normal distribution as an approximation to the Dirac delta function), so we get

\begin{equation}\begin{aligned} \frac{\int_0^{\pi} \sin^{n-2}\theta\, e^{n\cos\theta}\cos\theta d\theta}{\int_0^{\pi} \sin^{n-2}\theta\, e^{n\cos\theta} d\theta}\approx&\, \frac{\int_{-\infty}^{\infty} e^{n[0.377428 -1.11803 (\theta - 0.904557)^2]}\cos\theta d\theta}{\int_{-\infty}^{\infty} e^{n[0.377428 -1.11803 (\theta - 0.904557)^2]}d\theta} \\ \approx&\, \int_{-\infty}^{\infty} \delta(\theta - 0.904557) \cos\theta d\theta\\ =&\,\cos 0.904557 \approx 0.618034 \end{aligned}\end{equation}

Using Laplace's approximation once more, we obtain $\int_0^{\pi} \sin^{n-2}\theta d\theta\approx\frac{2.50663}{\sqrt{n}}$. Putting all these results together, we arrive at

\begin{equation}h_n \approx \log \frac{\frac{1.67629}{\sqrt{n}}e^{0.377428n}}{\frac{2.50663}{\sqrt{n}}}-0.618034n\approx -0.240606 n\end{equation}

which gives us the result $h_n/n\approx -0.24$. The derivation above used numerical values throughout; if we instead keep everything in closed form, the analytic solution for this coefficient works out to $\frac{1}{2}\log\frac{\sqrt{5}+1}{2}$.

Summary

In this post, starting from the principle of minimum entropy, we analyzed the problem of choosing the dimensionality of word embeddings, and ultimately arrived at an approximate estimation formula. The computed results show that this formula is consistent with the empirical wisdom accumulated from past model training practice.