The Amazing Johnson-Lindenstrauss Lemma: Theory Edition

Logarithmic Dimensions

The JL Lemma can be expressed very colloquially as:

Colloquial JL Lemma: To fit N vectors, one only needs \mathcal{O}(\log N) dimensions.

Specifically, the JL Lemma states that regardless of the original dimensionality of these N vectors, we can reduce them to \mathcal{O}(\log N) dimensions while keeping the relative distance error within a certain range. One can imagine that this is a very strong, counter-intuitive, and practical conclusion. For instance, if we are performing vector retrieval, the original vector dimensions might be very large, making the cost of a full-scale search prohibitive. The JL Lemma tells us that we can transform them into \mathcal{O}(\log N) dimensions with nearly unchanged retrieval performance. This is simply a "gift from the heavens"!

Readers might wonder: for such a powerful conclusion, would the corresponding dimensionality reduction method be particularly complex? The answer is quite the opposite: the dimensionality reduction process only involves random linear projection! Some have even remarked that the JL Lemma is a conclusion that is "easier to prove than to understand." That is to say, proving it mathematically is not particularly difficult, but intuitively understanding this counter-intuitive result is not so easy.

Coincidentally, we have previously introduced two counter-intuitive results: in the article "Distribution of the Angle Between Two Random Vectors in n-dimensional Space", we introduced the idea that "any two vectors in high-dimensional space are almost always orthogonal," which clearly differs significantly from results in two- or three-dimensional space. In the article "Understanding Model Parameter Initialization Strategies from a Geometric Perspective", this result was further upgraded to "an n \times n matrix sampled from \mathcal{N}(0, 1/n) is almost an orthogonal matrix," which deviates significantly from our usual understanding that "orthogonality is very strict (requiring the transpose to equal the inverse)."

In fact, these two conclusions are not only correct but also directly related to the JL Lemma. It can be said that the JL Lemma is a refinement and application of them. Therefore, we first need to use more quantitative language to characterize these conclusions—for example, if they are "almost orthogonal," what exactly is the probability of orthogonality, and if they are "approximately orthogonal," how large is the error?

Probability Inequalities

To this end, we need some knowledge of probability, most importantly "Markov’s Inequality":

Markov’s Inequality: If x is a non-negative random variable and a > 0, then P(x \geq a) \leq \frac{\mathbb{E}[x]}{a}

Note that this inequality does not impose other special restrictions on the distribution of x, only requiring that the sample space of the random variable is non-negative (or equivalently, the probability of negative x is zero). The proof is quite simple: \mathbb{E}[x] = \int_0^{\infty} x p(x) dx \geq \int_a^{\infty} x p(x) dx \geq \int_a^{\infty} a p(x) dx = a P(x \geq a)

Markov’s inequality requires the random variable to be non-negative, but the random variables we usually deal with are not necessarily non-negative, so we usually need to transform them. For example, x - \mathbb{E}[x] is not non-negative, but |x - \mathbb{E}[x]| is. Thus, using Markov’s inequality, we have: P(|x - \mathbb{E}[x]| \geq a) = P((x - \mathbb{E}[x])^2 \geq a^2) \leq \frac{\mathbb{E}[(x - \mathbb{E}[x])^2]}{a^2} = \frac{\mathbb{V}ar[x]}{a^2} This is "Chebyshev’s Inequality."

Another classic technique is called the "Cramér-Chernoff method," which is the primary method we will use later. It uses the exponential function to make the random variable non-negative: for any \lambda > 0, we have x \geq a \quad \Leftrightarrow \quad \lambda x \geq \lambda a \quad \Leftrightarrow \quad e^{\lambda x} \geq e^{\lambda a} So, using Markov’s inequality, we have P(x \geq a) = P(e^{\lambda x} \geq e^{\lambda a}) \leq e^{-\lambda a} \mathbb{E}[e^{\lambda x}] The leftmost side is independent of \lambda, but the rightmost side contains \lambda, and this inequality holds for any \lambda > 0. Therefore, in theory, we can find the \lambda that minimizes the rightmost side to obtain the highest estimation accuracy: P(x \geq a) \leq \min_{\lambda > 0} e^{-\lambda a} \mathbb{E}[e^{\lambda x}]

The Lemma of the Lemma

Now, we can introduce the following result, which is the lemma of the JL Lemma; one might even say it is the theoretical basis for all conclusions in this article:

Unit Norm Lemma: Let u \in \mathbb{R}^n be a vector sampled independently and identically from \mathcal{N}(0, 1/n), and let \varepsilon \in (0, 1) be a given constant. Then we have P(|\Vert u\Vert^2 - 1| \geq \varepsilon) \leq 2\exp\left(-\frac{\varepsilon^2 n}{8}\right)

This lemma tells us that when n is sufficiently large, the probability that the norm of u deviates significantly from 1 is very small (for a given \varepsilon, it will decrease to 0 exponentially with n). Thus, an n-dimensional vector sampled from \mathcal{N}(0, 1/n) will be very close to a unit vector.

Its proof uses the "Cramér-Chernoff method": first, |\Vert u\Vert^2 - 1| \geq \varepsilon means \Vert u\Vert^2 - 1 \geq \varepsilon or 1 - \Vert u\Vert^2 \geq \varepsilon. We need to derive these separately. Without loss of generality, let’s first derive the probability of \Vert u\Vert^2 - 1 \geq \varepsilon. According to the Cramér-Chernoff method: P(\Vert u\Vert^2 - 1 \geq \varepsilon) \leq \min_{\lambda > 0} e^{-\lambda \varepsilon} \mathbb{E}\big[e^{\lambda (\Vert u\Vert^2 - 1)}\big] = \min_{\lambda > 0} e^{-\lambda (\varepsilon + 1)} \mathbb{E}\big[e^{\lambda \Vert u\Vert^2}\big] Writing u in component form (u_1, u_2, \dots, u_n), where each component is independent and follows \mathcal{N}(0, 1/n), we have \mathbb{E}\big[e^{\lambda \Vert u\Vert^2}\big] = \mathbb{E}\big[e^{\lambda \sum_i u_i^2}\big] = \mathbb{E}\big[\prod_i e^{\lambda u_i^2}\big] = \prod_i \mathbb{E}\big[ e^{\lambda u_i^2}\big] Since \mathbb{E}\big[ e^{\lambda u_i^2}\big] = \int_{-\infty}^{\infty} \sqrt{\frac{n}{2\pi}} e^{-n u_i^2/2} e^{\lambda u_i^2} du_i = \sqrt{n/(n-2\lambda)}, we have P(\Vert u\Vert^2 - 1 \geq \varepsilon) \leq \min_{\lambda > 0} e^{-\lambda (\varepsilon + 1)} \left(\frac{n}{n-2\lambda}\right)^{n/2} The minimum value on the right side is reached at \lambda = \frac{n\varepsilon}{2(1+\varepsilon)} (the derivation is left to the reader). Substituting this back, we get P(\Vert u\Vert^2 - 1 \geq \varepsilon) \leq e^{n(\log(1+\varepsilon) - \varepsilon)/2} \leq e^{-n\varepsilon^2/8} The proof that \log(1+\varepsilon) - \varepsilon \leq -\varepsilon^2/4 is also left to the reader. Similarly, we can derive the probability for 1 - \Vert u\Vert^2 \geq \varepsilon: P(1 - \Vert u\Vert^2 \geq \varepsilon) \leq e^{n(\log(1-\varepsilon) + \varepsilon)/2} \leq e^{-n\varepsilon^2/8} Since it can be proven that \log(1-\varepsilon) + \varepsilon \leq \log(1+\varepsilon) - \varepsilon, the above equation uses the same inequality. Adding the two results, we get P(|\Vert u\Vert^2 - 1| \geq \varepsilon) \leq 2e^{-n\varepsilon^2/8}. Q.E.D.

Starting from the "Unit Norm Lemma," we can prove the "Orthogonality Lemma":

Orthogonality Lemma: Let u, v \in \mathbb{R}^n be two vectors sampled independently and identically from \mathcal{N}(0, 1/n), and let \varepsilon \in (0, 1) be a given constant. Then we have P(|\langle u, v\rangle| \geq \varepsilon) \leq 4\exp\left(-\frac{\varepsilon^2 n}{8}\right)

This lemma tells us that when n is sufficiently large, the probability that the inner product of u and v deviates significantly from 0 is very small. Combined with the "Unit Norm Lemma," we obtain the conclusion that "an n \times n matrix sampled from \mathcal{N}(0, 1/n) is almost an orthogonal matrix."

With the "Unit Norm Lemma" as a foundation, the proof is not difficult. We know that if u, v \sim \mathcal{N}(0, 1/n), then \frac{u \pm v}{\sqrt{2}} \sim \mathcal{N}(0, 1/n). Thus, according to the proof of the "Unit Norm Lemma," we have P\left(\left\Vert \frac{u+v}{\sqrt{2}}\right\Vert^2 - 1 \geq \varepsilon\right) \leq e^{-n\varepsilon^2/8}, \quad P\left(1 - \left\Vert \frac{u-v}{\sqrt{2}}\right\Vert^2 \geq \varepsilon\right) \leq e^{-n\varepsilon^2/8} Note that adding \left\Vert \frac{u+v}{\sqrt{2}}\right\Vert^2 - 1 \geq \varepsilon and 1 - \left\Vert \frac{u-v}{\sqrt{2}}\right\Vert^2 \geq \varepsilon leads to \langle u, v\rangle \geq \varepsilon, so P(\langle u, v\rangle \geq \varepsilon) \leq P\left(\left\Vert \frac{u+v}{\sqrt{2}}\right\Vert^2 - 1 \geq \varepsilon\right) + P\left(1 - \left\Vert \frac{u-v}{\sqrt{2}}\right\Vert^2 \geq \varepsilon\right) \leq 2e^{-n\varepsilon^2/8} Similarly, it can be proven that P(-\langle u, v\rangle \geq \varepsilon) \leq 2e^{-n\varepsilon^2/8}. Combining the two yields the "Orthogonality Lemma."

The Proof Process

Now we can proceed to prove the JL Lemma. Below is its mathematical formulation:

Mathematical JL Lemma: Given N vectors v_1, v_2, \dots, v_N \in \mathbb{R}^m and n > \frac{24\log N}{\varepsilon^2}, let the random matrix A \in \mathbb{R}^{n \times m} be sampled independently and identically from \mathcal{N}(0, 1/n), and let \varepsilon \in (0, 1) be a given constant. Then with a probability of at least \frac{N-1}{N}, for all i \neq j, the following holds: (1-\varepsilon)\Vert v_i - v_j\Vert^2 \leq \Vert Av_i - A v_j\Vert^2 \leq (1+\varepsilon)\Vert v_i - v_j\Vert^2 \label{eq:bound}

The lemma tells us that regardless of the original dimensionality m, we only need a dimensionality of n > \frac{24\log N}{\varepsilon^2} to accommodate N vectors such that their relative distance deviations do not exceed \varepsilon. Furthermore, the JL Lemma provides the dimensionality reduction method: simply sample a random n \times m matrix A from \mathcal{N}(0, 1/n), and the transformation v \to Av has a probability of \frac{N-1}{N} to achieve the goal. It is truly simple and practical!

The proof process is a direct application of the "Unit Norm Lemma." First, if u \in \mathbb{R}^m is a given unit vector and A \in \mathbb{R}^{n \times m} is sampled independently from \mathcal{N}(0, 1/n), then each component of Au independently follows \mathcal{N}(0, 1/n). The proof is not hard: by definition, each component (Au)_i = \sum_j A_{i,j} u_j. Since A_{i,j} are independent, (Au)_i are clearly independent. Since A_{i,j} \sim \mathcal{N}(0, 1/n), the sum of normal random variables is still a normal distribution. Thus, (Au)_i follows a normal distribution with mean \sum_j u_j \times 0 = 0 and variance \sum_j u_j^2 \times \frac{1}{n} = \frac{1}{n}.

In short, Au is equivalent to an n-dimensional vector sampled independently from \mathcal{N}(0, 1/n). Now, substituting u = \frac{v_i - v_j}{\Vert v_i - v_j\Vert} and using the "Unit Norm Lemma," we get P\left(\left|\left\Vert \frac{A(v_i - v_j)}{\Vert v_i - v_j\Vert}\right\Vert^2 - 1\right| \geq \varepsilon\right) \leq 2\exp\left(-\frac{\varepsilon^2 n}{8}\right) This result holds for any i \neq j. By iterating through all combinations of i \neq j, the probability that at least one term is \geq \varepsilon does not exceed P\left(\exists (i,j):\,\left|\left\Vert \frac{A(v_i - v_j)}{\Vert v_i - v_j\Vert}\right\Vert^2 - 1\right| \geq \varepsilon\right) \leq 2 \binom{N}{2} \exp\left(-\frac{\varepsilon^2 n}{8}\right) Conversely, the probability that for all i \neq j, \left|\left\Vert \frac{A(v_i - v_j)}{\Vert v_i - v_j\Vert}\right\Vert^2 - 1\right| \leq \varepsilon (equivalent to [eq:bound]) holds is at least 1 - 2 \binom{N}{2} \exp\left(-\frac{\varepsilon^2 n}{8}\right) = 1 - N(N-1)\exp\left(-\frac{\varepsilon^2 n}{8}\right) Substituting n > \frac{24\log N}{\varepsilon^2}, we get 1 - N(N-1)\exp\left(-\frac{\varepsilon^2 n}{8}\right) \geq 1 - N(N-1)N^{-3} \geq 1 - N^{-1} The proof is now complete.

The JL Lemma above maintains the approximate Euclidean distance. Often, we use inner products (such as cosine similarity) for retrieval instead of Euclidean distance. For this, we have:

Inner Product Version of JL Lemma: Given N unit vectors v_1, v_2, \dots, v_N \in \mathbb{R}^m and n > \frac{24\log N}{\varepsilon^2}, let the random matrix A \in \mathbb{R}^{n \times m} be sampled independently from \mathcal{N}(0, 1/n), and let \varepsilon \in (0, 1) be a given constant. Then with a probability of at least \frac{N-2}{N}, for all i \neq j, the following holds: \left|\langle Av_i, Av_j\rangle - \langle v_i, v_j\rangle\right| \leq \varepsilon

The proof is simple, mimicking the proof of the "Orthogonality Lemma." According to the proof of the JL Lemma, under the same conditions, there is at least a \frac{N-2}{N} probability that for all i \neq j: \begin{aligned} (1-\varepsilon)\Vert v_i - v_j\Vert^2 \leq \Vert Av_i - A v_j\Vert^2 \leq (1+\varepsilon)\Vert v_i - v_j\Vert^2 \\ (1-\varepsilon)\Vert v_i + v_j\Vert^2 \leq \Vert Av_i + A v_j\Vert^2 \leq (1+\varepsilon)\Vert v_i + v_j\Vert^2 \end{aligned} Multiplying the first by -1 gives -(1+\varepsilon)\Vert v_i - v_j\Vert^2 \leq -\Vert Av_i - A v_j\Vert^2 \leq -(1-\varepsilon)\Vert v_i - v_j\Vert^2. Adding this to the second equation yields: 4\langle v_i, v_j\rangle - 2\varepsilon(\Vert v_i\Vert^2 + \Vert v_j\Vert^2) \leq 4\langle Av_i, Av_j\rangle \leq 4\langle v_i, v_j\rangle + 2\varepsilon(\Vert v_i\Vert^2 + \Vert v_j\Vert^2) Since v_i, v_j are unit vectors, this is equivalent to \left|\langle Av_i, Av_j\rangle - \langle v_i, v_j\rangle\right| \leq \varepsilon.

Extreme Sufficiency

Students who have personally derived the proof of the JL Lemma should notice that the reason \log N appears in the conclusion is essentially because the probability term 2\exp\left(-\frac{\varepsilon^2 n}{8}\right) in the "Unit Norm Lemma" is exponentially decaying. We can relax this decay rate to polynomial decay, which results in \log N.

In general, the JL Lemma tells us that to fit N vectors with an error \varepsilon, we only need a space of \mathcal{O}\left(\frac{\log N}{\varepsilon^2}\right) dimensions. The specific constant in front of \frac{\log N}{\varepsilon^2} is not very important. This is because, in fact, the JL Lemma is a very sufficient condition; in practice, conditions are often much looser. For example, in the proof of the JL Lemma, if we change the condition to n > \frac{16\log N}{\varepsilon^2}, the probability that equation [eq:bound] holds is at least: 1 - N(N-1)\exp\left(-\frac{\varepsilon^2 n}{8}\right) \geq 1 - N(N-1)N^{-2} = 1/N Note that although 1/N is small, it is still greater than 0. Thus, there still exists an A such that [eq:bound] holds; it’s just that the cost of finding A is higher (the probability of success each time is only 1/N). If we only care about existence, this is sufficient.

Moreover, the JL Lemma only considers dimensionality reduction under random linear projection and already yields n > \frac{16\log N}{\varepsilon^2}. If more refined dimensionality reduction methods are used, such as SVD-based reduction, it is possible to obtain better results (smaller coefficients). If non-linear dimensionality reduction methods are also considered, the results can be even more optimal. Therefore, there is no need to worry too much about the constant in front of \frac{\log N}{\varepsilon^2}. We only need to know the order of magnitude \mathcal{O}\left(\frac{\log N}{\varepsilon^2}\right). If we actually need to use it, we usually determine the constant based on the actual situation rather than just calling upon the theoretical result.

To Be Continued

In this article, we introduced the Johnson-Lindenstrauss Lemma (JL Lemma), an important and wonderful conclusion about dimensionality reduction and a key manifestation of the unusual nature of high-dimensional spaces. It tells us that "only \mathcal{O}(\log N) dimensions are needed to fit N vectors," allowing retrieval problems in high-dimensional spaces to be reduced to \mathcal{O}(\log N) dimensions.

This article mainly discussed the theoretical proof details of the JL Lemma. In the next article, we will attempt to apply it to understand some machine learning problems. Stay tuned!

Original Address: https://kexue.fm/archives/8679