In the previous article "The Amazing Johnson-Lindenstrauss Lemma: Theory", we introduced the theoretical derivation of the Johnson-Lindenstrauss Lemma (JL Lemma) in detail. In this post, we will focus on its applications.
As a conclusion inherently related to dimensionality reduction, the most fundamental use of the JL Lemma is, naturally, as a dimensionality reduction method. However, beyond this direct application, many seemingly unrelated algorithms—such as Locality-Sensitive Hashing (LSH) and Randomized SVD—essentially rely on the JL Lemma. Furthermore, for machine learning models, the JL Lemma often provides a theoretical explanation for our choice of dimensions.
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, \dots, v_N \in \mathbb{R}^m, if you want to reduce them to n dimensions, you only need to sample an n \times m matrix A from \mathcal{N}(0, 1/n). Then Av_1, Av_2, \dots, Av_N are the results after dimensionality reduction.
There is no doubt that this approach is simple and fast. The question that naturally follows for readers is: How does it perform compared to dimensionality reduction methods like PCA or t-SNE?
In fact, just as "existence is reasonable," the fact that more complex methods like PCA and t-SNE have not been phased out indicates they must have advantages over random projection. Indeed, the random projection of the JL Lemma provides only a very basic dimensionality reduction method, showing that even with such a simple approach, the reduced dimension only needs to be \mathcal{O}(\log N). It serves more as a theoretical proof.
Therefore, if one truly seeks high precision in dimensionality reduction, specialized methods like PCA or t-SNE will almost certainly perform better than random projection in most cases. Moreover, as mentioned in the previous article, the JL Lemma provides a very sufficient condition. The bounds it yields, such as n > \frac{24\log N}{\varepsilon^2} or even n > \frac{16\log N}{\varepsilon^2}, are extremely conservative. For instance, if we take \varepsilon=0.1, we get n > 1600\log N, which has little practical value. By using more precise dimensionality reduction methods like PCA or t-SNE, this requirement can be relaxed, achieving better results at much lower dimensions.
Locality-Sensitive Hashing
Locality-Sensitive Hashing (LSH) is a scheme for approximate nearest neighbor search under a certain metric. Generally, we rarely associate LSH with the JL Lemma, but I believe that the choice of hash functions in LSH is actually closely related to the JL Lemma. Simply put, LSH is an algorithm that binarizes vectors such that the binarized vectors approximately preserve the metric. A common scheme is to use random projection to (approximately) preserve the cosine similarity.
Specifically, according to the JL Lemma, if we sample an n \times m matrix A from \mathcal{N}(0, 1/n), 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 is not the entirety of LSH. We notice that after projection by A, the distribution of positive and negative signs in Av_i, Av_j is relatively uniform. Thus, we further approximate: \cos(v_i, v_j) \approx \cos(Av_i, Av_j) \approx \cos(\text{sign}(Av_i), \text{sign}(Av_j)) That is, we binarize each element to \pm 1 based on its sign. This achieves vector binarization while keeping the cosine value approximately unchanged. Once we have binarized vectors, we can build indices or buckets to speed up retrieval, which I won’t detail here.
In short, a key step in the LSH process is random projection, which itself is closely related to the JL Lemma. Of course, binarization usually sacrifices precision significantly. Depending on the actual scenario, we do not always "reduce dimensions"; that is, n is not always smaller than m. Sometimes we might even choose n > m. For related discussions, readers can refer to my previous post "How does a binarized word vector model relate to fruit flies?".
Randomized Decomposition
Matrix decomposition is a powerful tool for solving many machine learning problems, and Singular Value Decomposition (SVD) is one of the typical methods. However, when the matrix is large, the cost of computing an exact SVD is quite high. In practical scenarios, while the matrix to be decomposed is large, it is often low-rank, making an exact SVD unnecessary. This is where "Randomized SVD" comes into play.
Let the matrix to be decomposed be M \in \mathbb{R}^{m \times n}, where m and n are both large. According to the JL Lemma, we can choose a relatively small k < \min(m, n) such that an n \times k matrix Q sampled from \mathcal{N}(0, 1/k) still satisfies QQ^{\top} \approx I (an approximately orthogonal matrix) with high precision, so that M \approx MQQ^{\top}. Thus, we can perform SVD only on the m \times k matrix B = MQ. If we obtain MQ = B = U_B \Sigma_B V_B^{\top}, then: M \approx MQQ^{\top} = U_B \Sigma_B V_B^{\top} Q^{\top} = U_B \Sigma_B (QV_B)^{\top} This gives us an approximate SVD of the original matrix M. Note that the aforementioned Q is only an approximately orthogonal matrix; we can make it strictly orthogonal through QR decomposition (or Gram-Schmidt orthogonalization), which is a minor detail. Throughout this process, the JL Lemma tells us that k can be chosen small enough so that performing SVD on B = MQ is low-cost, yet the overall precision remains acceptable.
Word Vector Dimensions
We say that the intuitive understanding of the JL Lemma is: "To fit N vectors, you only need an \mathcal{O}(\log N)-dimensional space." Returning to the problem of choosing word vector dimensions, this means that if the vocabulary size is N, then a word vector dimension of \mathcal{O}(\log N) is sufficient.
What is truly astonishing is that in my previous article "The Principle of Minimum Entropy (VI): How to Choose the Dimension of Word Vectors?", I calculated a dimension selection formula for the Skip-Gram word vector model: n > 8.33 \log N The result is identical to the \mathcal{O}(\log N) given by the JL Lemma! The above formula was estimated based on the idea of entropy and has almost no intersection with the starting point of the JL Lemma, yet it reached the same conclusion of \log N through a different path.
Furthermore, it’s not just the main term \log N. We also see that based on the entropy estimation, we even calculated the coefficient 8.33 in front of \log N. Previous experimental experience has shown that the result 8.33 \log N is quite consistent with empirical findings; while it may not be optimal, it is at least in the right ballpark. Does this conversely suggest that we can use entropy to more accurately estimate the coefficient in front of \log N for specific problems?
Multi-Head Attention
Regarding the Attention mechanism, common interview questions include "Why multi-head?" and "What is the difference between a single-head attention with head_size 768 and a 12-head attention with head_size 64?" In other words, for an Attention model like BERT, why reduce the head_size to 64 before calculating the dot product? Is 64 really enough?
This question essentially asks whether the Attention mechanism is sufficient to fit any probability pattern. Specifically, the formula for Attention is: a_{i,j} = \frac{e^{\langle q_i, k_j\rangle}}{\sum\limits_{j=1}^L e^{\langle q_i, k_j\rangle}} where q_i, k_j \in \mathbb{R}^d. The question of "is it enough" refers to whether the defined a_{i,j} can well-approximate any given probability matrix p_{i,j}.
Does the definition of a_{i,j} look familiar to any readers? If we set aside the context of Attention and treat q_i, k_j as two "word vectors," then the definition of a_{i,j} is exactly the same as the Skip-Gram model! That is to say, looking purely at the calculation formula of the Attention matrix, it is essentially the same as the Skip-Gram model. Therefore, the choice of head_size in Attention is essentially the same as the choice of dimension for word vectors.
Let’s re-trace the process. We want to answer "how much head_size is enough," which becomes the question of "can a_{i,j} approximate any probability matrix p_{i,j}." That is, for a given p_{i,j}, can we find a set of q_1, \dots, q_L, k_1, \dots, k_L \in \mathbb{R}^d such that a_{i,j} is sufficiently close to p_{i,j}? This problem is mathematically equivalent to the dimension selection of Skip-Gram word vectors.
Therefore, the result of word vector dimension selection can be applied to the choice of head_size for Attention, except the vocabulary size becomes the sequence length L. That is, d > 8.33 \log L. For a common pre-training length of L=512, the calculation yields approximately 52. This is again very surprising, as it is indeed very close to the common head_size of 64! So, 64 is truly enough; increasing it further won’t yield significant improvements. It is better to use the extra computation to increase the number of heads.
(Note: For related discussions, one can also refer to the paper "On the Expressive Power of Self-Attention Matrices".)
Summary
This article mainly introduced several direct or indirect applications of the Johnson-Lindenstrauss Lemma (JL Lemma). As we have seen, from dimensionality reduction and hashing methods to word vector dimensions and Attention head sizes, all are more or less related to the JL Lemma. This further demonstrates the wide range of applicability of the JL Lemma.
Reprinting: Please include the original address of this article: https://kexue.fm/archives/8706
For more detailed reprinting matters, please refer to: "Scientific Space FAQ"