The Astonishing Johnson-Lindenstrauss Lemma: Applications
In the previous post, The Astonishing Johnson-Lindenstrauss Lemma: Theory], we went through a fairly detailed derivation of the Johnson-Lindenstrauss lemma (JL lemma). In this post, let's turn to its applications.
Since the JL lemma is, at its core, a result about dimensionality reduction, the most natural application is simply to use it as a dimensionality-reduction method. But beyond this direct use, many seemingly unrelated algorithms — such as locality-sensitive hashing (LSH) and randomized SVD — also rely on the JL lemma at their core. Moreover, for machine learning models, the JL lemma often gives us a theoretical justification for our choice of dimension.
A Tool for Dimensionality Reduction
The JL lemma provides a very simple and direct "random projection" approach to dimensionality reduction:
Given $N$ vectors $v_1,v_2,\cdots,v_N\in\mathbb{R}^m$, if we want to reduce them to $n$ dimensions, we only need to sample an $n\times m$ matrix $A$ from $\mathcal{N}(0,1/n)$, and then $Av_1,Av_2,\cdots,Av_N$ is the reduced-dimension result.
There's no doubt this approach is simple and fast, and the reader's natural follow-up question is: how does it compare to methods like PCA and t-SNE?
Actually, as the saying goes, "existence justifies itself" — the fact that more sophisticated methods like PCA and t-SNE haven't been rendered obsolete tells us they must have advantages over random projection. Indeed, the random projection given by the JL lemma is only a very basic dimensionality-reduction method; the fact that even such a simple method only requires the reduced dimension to be $\mathcal{O}(\log N)$ is mainly meant as a theoretical proof, rather than a practical recommendation.
So if we're really after reduction accuracy, in most cases dedicated methods like PCA and t-SNE will outperform random projection. And as we mentioned in the previous post, the JL lemma is a very sufficient condition — the bound $n > \frac{24\log N}{\varepsilon^2}$, or even $n > \frac{16\log N}{\varepsilon^2}$, that it gives is a rather loose sufficient bound. For instance, taking $\varepsilon=0.1$ gives $n > 1600\log N$, which is basically of no practical value. Switching to more precise dimensionality-reduction methods like PCA or t-SNE lets us relax this requirement, achieving better results at smaller dimensions.
Locality-Sensitive Hashing
Locality-sensitive hashing (LSH) is a scheme for approximately finding nearest neighbors under some metric. We rarely think to connect LSH with the JL lemma, but in the author's view, the choice of hash function in LSH is in fact closely tied to the JL lemma. Put simply, LSH is an algorithm that binarizes vectors, such that the binarized vectors approximately preserve the metric. One common scheme uses random projections to (approximately) preserve cosine similarity.
Specifically, by the JL lemma, we sample an $n\times m$ matrix $A$ from $\mathcal{N}(0,1/n)$, and then for any $v_i,v_j\in\mathbb{R}^m$, we have $\cos(v_i,v_j)\approx \cos(Av_i, Av_j)$. Of course, random projection alone isn't the whole story of LSH — we also note that, after projecting with $A$, the positive/negative signs of the entries of $Av_i,Av_j$ are fairly evenly distributed, so we go a step further and approximate
\begin{equation}\cos(v_i,v_j)\approx \cos(Av_i, Av_j)\approx \cos(\text{sign}(Av_i), \text{sign}(Av_j))\end{equation}
That is, we binarize each entry according to its sign as $\pm 1$. This gives us binarized vectors, while approximately preserving cosine similarity. Once we have binarized vectors, we can build indexes, bucket them, and so on to speed up retrieval — details we won't go into here.
In short, the key step in LSH is also random projection, and this step is itself closely tied to the JL lemma. Of course, binarization typically sacrifices a fair amount of precision, so depending on the actual scenario, we're not always "reducing" dimension — that is, $n$ isn't always smaller than $m$; sometimes we might even choose $n > m$. Interested readers can refer to the author's earlier post How Does a Binarized Word Embedding Model End Up Related to Fruit Flies?] for further discussion.
Randomized Decomposition
Matrix factorization is a powerful tool for solving many machine learning problems, and singular value decomposition (SVD) is one of its most typical methods. However, when the matrix is large, computing an exact SVD is quite costly. In practice, the matrix to be decomposed, though large, is often also low-rank, in which case computing an exact SVD isn't necessary. This is where "randomized SVD" comes in handy.
Let the matrix to be decomposed be $M\in\mathbb{R}^{m\times n}$, where $m,n$ are both fairly large. By the JL lemma, we can choose a fairly small $k < \min(m,n)$ such that a matrix $n\times k$ sampled from $\mathcal{N}(0,1/k)$, namely $Q$, still satisfies $QQ^{\top}\approx I$ (approximate orthogonality) to reasonably high precision, so that $M\approx MQQ^{\top}$. This lets us perform SVD only on the $m\times k$ matrix $B=MQ$, obtaining $MQ=B=U_B\Sigma_B V_B^{\top}$, so that
\begin{equation}M\approx MQQ^{\top} = U_B\Sigma_B V_B^{\top}Q^{\top} = U_B \Sigma_B (QV_B)^{\top}\end{equation}
gives us an approximate SVD of the original matrix $M$. Note that the $Q$ above is still only approximately orthogonal; we can use QR decomposition (or Gram-Schmidt orthogonalization) to make it strictly orthogonal — a minor detail. Throughout this process, what the JL lemma tells us is that $k$ can be chosen fairly small, so that performing SVD on $B=MQ$ is relatively cheap, while overall accuracy remains reasonably good.
Word Embedding Dimension
We said the intuitive takeaway of the JL lemma is "fitting $N$ vectors only requires a $\mathcal{O}(\log N)$-dimensional space." Returning to the question of word embedding dimension: if the vocabulary size is $N$, then a word embedding dimension of $\mathcal{O}(\log N)$ should suffice.
Remarkably, in an earlier post, The Minimum Entropy Principle (VI): How Should We Choose the Dimension of Word Embeddings?], we derived a dimension-selection formula for the Skip-Gram word embedding model:
\begin{equation}n > 8.33\log N\end{equation}
The result matches exactly the $\mathcal{O}(\log N)$ given by the JL lemma! That formula was derived from an entropy-based line of reasoning, which shares almost no common ground with the starting point of the JL lemma, yet the two arrive at the same conclusion $\log N$ by completely different routes.
Moreover, it's not just the main term $\log N$ — with the entropy-based estimate, we were also able to compute the coefficient $8.33$ in front of $\log N$, and past experimental experience has shown that this result $8.33\log N$ is actually quite consistent with practice; while it might not be exactly optimal, it's at least in the right ballpark. Doesn't this suggest that we could use entropy to more precisely estimate the coefficient in front of $\log N$ for a given problem?
Multi-Head Attention
A common interview question about the attention mechanism is: "Why use multiple heads?" or "What's the difference between single-head attention with head_size=768 and 12-head attention with head_size=64?" In other words: why does an attention model like BERT first reduce head_size down to 64 before taking the inner product? Is 64 really enough?
Fundamentally, this question is about whether the attention mechanism is expressive enough to fit any probability pattern. Specifically, the attention computation is given by:
\begin{equation}a_{i,j} = \frac{e^{\langle q_i, k_j\rangle}}{\sum\limits_{j=1}^L e^{\langle q_i, k_j\rangle}}\end{equation}
where $q_i,k_j\in\mathbb{R}^{d}$. The question of whether it's "enough" amounts to asking: for any given probability matrix $p_{i,j}$, can the $a_{i,j}$ defined above approximate it well?
Looking at the definition of $a_{i,j}$, does it look familiar to any readers? If we set aside the attention context and simply regard $q_i,k_j$ as two sets of "word embeddings," then the definition of $a_{i,j}$ is exactly the same as in the Skip-Gram model! In other words, purely in terms of its computation formula, the attention matrix is essentially the same object as the Skip-Gram model, so the choice of head_size in attention is, in essence, the same problem as choosing a word embedding dimension.
Let's walk through this reasoning again. The question we're trying to answer is "how large does head_size need to be," which becomes the question "can $a_{i,j}$ approximate an arbitrary probability matrix $p_{i,j}$?" That is, given $p_{i,j}$, can we find a set of $q_1,\cdots,q_L,k_1,\cdots,k_L\in\mathbb{R}^d$ such that $a_{i,j}$ is sufficiently close to $p_{i,j}$? This is mathematically equivalent to the dimension-selection problem for Skip-Gram word embeddings.
Therefore, the results on word embedding dimension selection can also be applied to attention head_size selection — the only change is that the vocabulary size becomes the sequence length, i.e. $d > 8.33\log L$. For a common pretraining length of $L=512$, plugging this in gives roughly 52 — again, remarkably close to the commonly used head_size = 64! So 64 really is enough; going any larger won't give a noticeable improvement, and it would be better to spend that extra compute on adding more heads instead.
(Note: for related discussion, see also On the Expressive Power of Self-Attention Matrices].)
Summary, Once Again
This post has introduced several direct and indirect applications of the Johnson-Lindenstrauss lemma (JL lemma). As we've seen, from dimensionality reduction and hashing methods to word embedding dimensions and attention head sizes, there is, to varying degrees, a connection to the JL lemma — further demonstrating just how broadly applicable it is.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.