JL Lemma and Entropy-Invariant Attention

从JL引理看熵不变性Attention 苏剑林 · 2023-04-10 · Machine Learning entropy attention

In 《From Entropy Invariance to the Scale Operation in Attention》] and 《A Quick Derivation of Entropy-Invariant Softmax》] I proposed entropy-invariant Softmax, which is simply about multiplying the attention matrix by an extra factor of $\log n$ before Softmax; in theory this helps improve length extrapolation, where $n$ is the sequence length. This factor $\log n$ reminded me of the JL lemma (the Johnson-Lindenstrauss lemma]), because the JL lemma tells us that encoding $n$ vectors only requires a dimensionality of $\mathcal{O}(\log n)$ — both involve a $\log n$, so is there some connection between the two?

Entropy Invariance

As we know, entropy is a measure of uncertainty; in the context of the attention mechanism, we use it as a measure of "how concentrated the attention is." What we mean by entropy invariance is that regardless of the sequence length $n$, attention should stay focused on a few key tokens rather than becoming too dispersed. To this end, we proposed the following entropy-invariant Attention form:

\begin{equation}Attention(Q,K,V) = softmax\left(\frac{\log_{512} n}{\sqrt{d}}QK^{\top}\right)V\label{eq:core}\end{equation}more

Here $Q,K\in\mathbb{R}^{n\times d}$. Compared to conventional attention, the scale factor has an extra term $\log_{512} n$, where the base is taken as 512, on the assumption that all our hyperparameters (such as $d$) have been tuned for a training length of 512. Of course, even if your planned pretraining length isn't 512, you can just use 512 as the base without much loss — the results won't be significantly affected.

The rationale behind this form is quite intuitive: as $n$ increases, it means more tokens are sharing the attention, causing it to become less concentrated. So we multiply by a factor that increases monotonically with $n$; after softmax, this is effectively equivalent to raising the original probabilities to some power. Since probabilities are all less than 1, smaller probabilities become even smaller after exponentiation, which re-concentrates the attention. As for why this factor takes a logarithmic form, that requires looking at the derivation in the earlier articles.

The JL Lemma

The JL lemma, short for the "Johnson-Lindenstrauss lemma," is an important result about vector embeddings. Simply put, it tells us that "to fit $n$ vectors, you only need a $\mathcal{O}(\log n)$-dimensional space" (here $\log$ has no explicit base written, and it defaults to the natural logarithm base $e$). For a detailed introduction, see 《The Astonishing Johnson-Lindenstrauss Lemma: Theory》].

Interestingly, even before I knew about the JL lemma, I had derived the same — and even more specific — result in 《The Minimum Entropy Principle (VI): How Should We Choose the Dimensionality of Word Embeddings?》]: to embed $n$ word vectors, roughly $8\log n$ dimensions suffice. This estimate is quite close to the dimensions used in practice — for instance, when $n$ equals 100,000, $8\log n$ works out to about 92, and the word embedding dimensions we commonly use are also on the order of a hundred or two hundred.

Additionally, the JL lemma can also be used to explain the multi-head nature of the attention mechanism. If we substitute in $n=512$, then $8\log n\approx 50$, which is quite close to the projection dimension commonly used for Q and K in attention (i.e., key_size, which is 64 in BERT — see here]). This tells us that if the sequence length is 512, then the dimensionality used for computing Q and K in attention only needs to be on the order of 50 — there's no need to use the full hidden_size (768 for BERT base). The dimensions saved can instead be used for multi-head attention.

For further related discussion, see 《An Analysis of the Usability of the Dimensionality Formula "n > 8.33 log N"》] and 《The Astonishing Johnson-Lindenstrauss Lemma: Applications》].

Connecting the Two

Now we can try to connect the JL lemma with entropy-invariant attention.

Let's denote the key_size of Q and K as $d$. Then the JL lemma tells us that the optimal choice for $d$ should be $d_n=\lambda \log n$, where $\lambda$ is a proportionality constant whose exact value doesn't matter. In other words, ideally $d$ should vary with $n$, but clearly such a design isn't easy to implement and isn't conducive to parallelizing computation, so in practice we can only use a fixed $d$.

Suppose we've chosen a fixed $d$, and suppose this $d$ was designed for a training length of 512. Then we can derive $d = \lambda \log 512$, that is, $\lambda = \frac{d}{\log 512}$, and

\begin{equation}d_n = \frac{d}{\log 512}\log n=d\log_{512} n\end{equation}

For $n\neq 512$, ideally we should use a projection dimension of $d_n$, but in practice we use $d$ dimensions. According to the definition of the inner product $\langle q,k\rangle = \sum\limits_{i=1}^d q_i k_i$, the number of summation terms equals exactly the number of dimensions $d$ — that is, ideally there should be a sum of $d_n$ terms, but in practice it becomes a sum of $d$ terms. Intuitively, if the contribution of each term is roughly comparable, then multiplying the result by $\frac{d_n}{d}$ should bring it closer to the ideal case of summing $d_n$ terms. So we conclude that we should multiply $\langle q,k\rangle$ by the factor

\begin{equation}\frac{d_n}{d} = \log_{512} n\end{equation}

to compensate for the gap between the actual and ideal situations. And multiplying conventional Scaled-Dot Attention by $\log_{512} n$ gives exactly entropy-invariant attention, i.e., equation $\eqref{eq:core}$.

In this way, we've connected the JL lemma with entropy-invariant attention. Note that this is only an intuitive, qualitative understanding — it's hard to make this rigorous in a quantitative sense, and in fact there's no real need to, since the JL lemma itself is, for the most part, only a qualitative result.

Summary

This article has built a simple connection between the JL lemma and entropy-invariant attention.