From Entropy Invariance: A Look at the Scale Operation in Attention

从熵不变性看Attention的Scale操作 苏剑林 · 2021-12-21 · Machine Learning probability entropy attention

The attention mechanism most widely used in current Transformer architectures is called "Scaled Dot-Product Attention," where "Scaled" refers to the fact that after multiplying $Q,K$ by its transpose, the result is divided by a $\sqrt{d}$ before the softmax is applied (below we assume $Q,K,V\in\mathbb{R}^{n\times d}$ without loss of generality):

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

In Some Thoughts on the Initialization, Parameterization, and Normalization of Transformers, we already gave a preliminary explanation for dividing by $\sqrt{d}$. In this article, I'll instead approach this scaling operation from the perspective of "entropy invariance," and arrive at a new scaling factor. Experiments on MLM show that this new scaling factor achieves better length-extrapolation performance. more

Entropy Invariance

Let's rewrite the general form of Scaled Dot-Product Attention as

\begin{equation}\boldsymbol{o}_i = \sum_{j=1}^n a_{i,j}\boldsymbol{v}_j,\quad a_{i,j}=\frac{e^{\lambda \boldsymbol{q}_i\cdot \boldsymbol{k}_j}}{\sum\limits_{j=1}^n e^{\lambda \boldsymbol{q}_i\cdot \boldsymbol{k}_j}}\end{equation}

where $\lambda$ is the scaling factor. It doesn't depend on $\boldsymbol{q}_i,\boldsymbol{k}_j$, but in principle it could depend on parameters such as the length $n$ or the dimension $d$. The current mainstream choice is $\lambda=1/\sqrt{d}$.

This article proposes the following viewpoint:

In order for the model to generalize better to unseen lengths, the design of the attention mechanism should aim to make $a_{i,j}$ as close to entropy-invariant as possible.

How should we understand this statement? First, "generalizing to unseen lengths" means the model performs well even when the test length differs from the training length — for instance, training at length $n=64$ and then extrapolating to test length $n=128,256$. We know that models using relative positional encodings such as RoPE already extrapolate reasonably well to different lengths, but we can still push this extrapolation ability further through better design choices — entropy invariance being one of them.

Concretely, $a_{i,j}$ can be viewed as a conditional distribution over the random variable $j$, conditioned on $i$. Its entropy is

\begin{equation}\mathcal{H}_i = -\sum_{j=1}^n a_{i,j}\log a_{i,j}\end{equation}

Entropy invariance means that $\mathcal{H}_i$ should be insensitive to the length $n$. More specifically, if we append a few more tokens to an existing sequence, the newly computed $a_{i,j}$ values will naturally change somewhat, but we would like $\mathcal{H}_i$ itself not to change by much.

Why would we want the entropy to stay invariant? Recall that entropy measures uncertainty (see Making Sense of "Entropy": From Entropy and the Maximum Entropy Principle to Maximum Entropy Models (Part 1)). Put another way, we can think of this uncertainty as the "degree of focus" of attention: if the entropy is 0, attention concentrates entirely on a single token; if the entropy is $\log n$, attention spreads uniformly across all tokens. What we want, by keeping the entropy invariant, is that after new tokens are introduced, the existing tokens can still focus on the same original tokens to the same degree — we don't want the newly introduced tokens to "siphon off" too much of the original attention and thereby significantly change the weighted sum.

A New Factor

Based on entropy invariance together with a few reasonable assumptions, we can derive a new scaling factor, giving rise to a variant of Scaled Dot-Product Attention:

\begin{equation}Attention(Q,K,V) = softmax\left(\frac{\kappa \log n}{d}QK^{\top}\right)V\label{eq:ei}\end{equation}

Here $\kappa$ is a hyperparameter unrelated to $n,d$; we'll walk through the detailed derivation in the next section. For convenience, we'll refer to the standard Scaled Dot-Product Attention described by equation $\eqref{eq:std}$ as "Attention-O" (Original), and the variant described by equations $\eqref{eq:ei}$ and $\eqref{eq:ei2}$ below as "Attention-E" (Entropy Invariance).

Some readers might be unhappy about introducing a new hyperparameter, but this is easy to resolve. Since the current mainstream pretraining length is 512, we can assume that most existing hyperparameters have already been tuned for $n=512$. So when $n=512$, the formula above should reduce to the ordinary Scaled Dot-Product Attention, i.e., $\frac{\kappa \log 512}{d}=\frac{1}{\sqrt{d}}$, which gives us $\kappa = \frac{\sqrt{d}}{\log 512}$. Substituting this back and simplifying, we obtain

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

This eliminates the hyperparameter $\lambda$, and it's this version we use in the experiments below.

To verify whether this modification actually improves the Transformer's extrapolation performance as expected, I trained a small RoFormer using Attention-O and Attention-E respectively, on an MLM task with a training length of 64, and then compared MLM accuracy on validation sets of different lengths. The results are as follows:

$$\begin{array}{c} \text{attention length extrapolation experiment} \\ {\begin{array}{c|ccccc} \hline & n=64 & n=128 & n=256 & n=512 & 1024 \\ \hline \text{Attention-O} & 43.27 & 36.53 & 23.02 & 15.12 & 11.54\\ \text{Attention-E} & 43.11 & 41.17 & 34.04 & 20.15 & 13.58\\ \hline \end{array}$$}

\end{array}

The experimental results show that when the test length matches the training length $n=64$, Attention-O and Attention-E perform very similarly. However, when extrapolating to larger test lengths, the gap widens considerably — for example, at $n=256$, Attention-E outperforms Attention-O by more than 10 percentage points in accuracy, which is no small margin.

Derivation

In this section we walk through the derivation of equation $\eqref{eq:ei}$. In fact, the derivation and the assumptions involved are nearly identical to those in The Minimum Entropy Principle (Part 6): How Should We Choose the Dimensionality of Word Embeddings?.

First, substituting the expression for $a_{i,j}$, we get:

\begin{equation}\mathcal{H}_i = -\sum_{j=1}^n a_{i,j}\log a_{i,j}=\log \sum_{j=1}^n e^{\lambda \boldsymbol{q}_i\cdot \boldsymbol{k}_j} - \frac{\sum\limits_{j=1}^n e^{\lambda \boldsymbol{q}_i\cdot \boldsymbol{k}_j}(\lambda \boldsymbol{q}_i\cdot \boldsymbol{k}_j)}{\sum\limits_{j=1}^n e^{\lambda \boldsymbol{q}_i\cdot \boldsymbol{k}_j}}\end{equation}

Note that we're only aiming for a semi-quantitative estimate, in order to determine a suitable $\lambda$ that offsets part of the effect of length — making the entropy completely independent of length is not actually achievable. So we can make some assumptions; for instance, assuming $\boldsymbol{k}_j$ is a random variable, we can write

\begin{equation}\sum_{j=1}^n e^{\lambda \boldsymbol{q}_i\cdot \boldsymbol{k}_j} = n\times \frac{1}{n}\sum_{j=1}^n e^{\lambda \boldsymbol{q}_i\cdot \boldsymbol{k}_j}\approx n\,\mathbb{E}_j[e^{\lambda \boldsymbol{q}_i\cdot \boldsymbol{k}_j}]\end{equation}

Replacing all the sums with the same approximation, we obtain

\begin{equation}\mathcal{H}_i \approx \log n + \log \mathbb{E}_j[e^{\lambda \boldsymbol{q}_i\cdot \boldsymbol{k}_j}] - \frac{\lambda\,\mathbb{E}_j[e^{\lambda \boldsymbol{q}_i\cdot \boldsymbol{k}_j}(\boldsymbol{q}_i\cdot \boldsymbol{k}_j)]}{\mathbb{E}_j[e^{\lambda \boldsymbol{q}_i\cdot \boldsymbol{k}_j}]} \end{equation}

Note that in general $\boldsymbol{q}_i,\boldsymbol{k}_j$ is produced by applying a Layer Norm followed by a Dense layer, and a Dense layer is approximately an orthogonal transformation (see Understanding Parameter Initialization Strategies from a Geometric Perspective). So we can approximately assume that all instances of $\boldsymbol{q}_i,\boldsymbol{k}_j$ are vectors of the same norm $\sqrt{d}$, giving us $\boldsymbol{q}_i\cdot \boldsymbol{k}_j=d\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)$. Further assuming that $\boldsymbol{k}_j$ is uniformly distributed on a sphere of radius $\sqrt{d}$, the expectation over $\boldsymbol{k}_j$ can be converted into an expectation over the angle between $\boldsymbol{q}_i,\boldsymbol{k}_j$, i.e.,

\begin{equation}\mathcal{H}_i \approx \log n + \log \mathbb{E}_{\theta}[e^{\lambda d \cos\theta}] - \frac{\lambda d\,\mathbb{E}_{\theta}[e^{\lambda d \cos\theta}\cos\theta]}{\mathbb{E}_{\theta}[e^{\lambda d \cos\theta}]} \end{equation}

where the distribution followed by $\theta$ is exactly the distribution of the angle between two random vectors on a sphere, which we discussed in The Distribution of the Angle Between Two Random Vectors in n-Dimensional Space. From here, following the same approach as the "Approximate Estimation" section of The Minimum Entropy Principle (Part 6): How Should We Choose the Dimensionality of Word Embeddings?, we can apply a Laplace approximation to obtain

\begin{equation}\mathcal{H}_i \approx \log n - 0.24\lambda d + \mathcal{O}(1) \end{equation}

Therefore, in order to offset the effect of length $n$, we set $\log n - 0.24\lambda d = 0$, from which it follows that $\lambda = \log n / (0.24 d)$. Of course, since this is just an estimate, there's no need to keep the coefficient $0.24$ exactly — instead, we might as well introduce a hyperparameter $\kappa$, giving

\begin{equation}\lambda = \frac{\kappa\log n}{d}\end{equation}

which is exactly equation $\eqref{eq:ei}$.

Related Results

While reading through ACL 2022 submissions, I came across a paper titled Overcoming a Theoretical Limitation of Self-Attention, which arrives at a similar result (Equation 1 in Section 4.3 of the paper):

\begin{equation}Attention(Q,K,V) = softmax\left(\frac{\log n}{\sqrt{d}}QK^{\top}\right)V\end{equation}

However, that paper doesn't offer much theoretical analysis; it simply constructs two special test cases to evaluate attention's performance, finds that multiplying the scaling factor by $\log n$ helps with length generalization, and proposes it on that basis.

That said, if we follow the default convention that $\log$ uses the natural logarithm, the formula above is clearly problematic: when $n$ is large, the scaling factor becomes too large, causing severe vanishing gradients. It's just that this paper only ran experiments on machine translation, where the sequences tested were all around length $n=20$, so the vanishing-gradient issue never showed up.

Summary

This article re-derived the scale operation in Scaled Dot-Product Attention from the perspective of entropy invariance, arriving at a new scaling factor. Preliminary experimental results show that this new scaling factor doesn't degrade existing training performance, while achieving better length-extrapolation results.