Understanding the Scale Operation in Attention from the Perspective of Entropy Invariance

Entropy Invariance

We can rewrite the general Scaled Dot-Product Attention as: \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}} where \lambda is the scaling factor. It is independent of \boldsymbol{q}_i and \boldsymbol{k}_j, but in principle, it can be related to parameters such as the length n and the dimension d. Currently, the mainstream choice is \lambda=1/\sqrt{d}.

This article proposes the following viewpoint:

In order to make the model generalize better to unknown lengths, the design of the Attention mechanism should ensure that a_{i,j} possesses entropy invariance as much as possible.

How should we understand this statement? First, generalizing to unknown lengths refers to the ability to maintain good performance when the sequence length during prediction is inconsistent with the training length (e.g., training on n=64 and extrapolating to n=128, 256 during testing). We know that models using relative position encodings like RoPE have relatively good length extrapolation, but we can still enhance this extrapolation through better design, and entropy invariance is one such approach.

Specifically, a_{i,j} can be viewed as a conditional distribution where i is the condition and j is the random variable. Its entropy is: \mathcal{H}_i = -\sum_{j=1}^n a_{i,j}\log a_{i,j} Entropy invariance means that \mathcal{H}_i should be insensitive to the length n. More specifically, if we add a few more tokens to the existing ones, the newly calculated a_{i,j} values will naturally change, but we hope that \mathcal{H}_i does not change significantly.

Why do we want entropy invariance? We know that entropy is a measure of uncertainty (refer to "Entropy is Affordable: From Entropy and the Maximum Entropy Principle to Maximum Entropy Models (I)"). From another perspective, we can view uncertainty as the "degree of focus" of attention: if the entropy is 0, the attention is focused on a single token; if the entropy is \log n, the attention is uniformly distributed across all tokens. We want entropy invariance because we hope that after introducing new tokens, the existing tokens can still focus on the original tokens in the same way, rather than having the new tokens "dilute" the original attention too much, which would cause the summation result to change significantly.

The New Factor

Based on entropy invariance and some reasonable assumptions, we can derive a new scaling factor, resulting in a variant of Scaled Dot-Product Attention: \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{\kappa \log n}{d}QK^{\top}\right)V \label{eq:ei} Here, \kappa is a hyperparameter independent of n and d. The detailed derivation process will be introduced in the next section. For convenience, we refer to the conventional Scaled Dot-Product Attention described in Eq. [eq:std] as "Attention-O" (Original), and the variant described in Eq. [eq:ei] and Eq. [eq:ei2] below as "Attention-E" (Entropy Invariance).

Some readers might be dissatisfied with the introduction of a new parameter. This is actually easy to resolve. We know that the current mainstream pre-training length is 512, so we assume that the mainstream parameters are tuned for n=512. Therefore, when n=512, the above formula should degenerate into the standard Scaled Dot-Product Attention, i.e., \frac{\kappa \log 512}{d}=\frac{1}{\sqrt{d}}, which implies \kappa = \frac{\sqrt{d}}{\log 512}. Substituting this back and simplifying, we get: \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{\log_{512} n}{\sqrt{d}}QK^{\top}\right)V \label{eq:ei2} This removes the hyperparameter \lambda. The following experiments also use this version.

To verify whether this modification truly improves the extrapolation effect of the Transformer as expected, the author trained two versions of RoFormer small using Attention-O and Attention-E respectively. The training task was MLM, the training length was 64, and the MLM accuracy was compared across validation sets of different lengths. The results are as follows:

From the experimental results, it can be seen that when the test length is the same as the training length (n=64), the performance of Attention-O and Attention-E is very close. However, when extrapolating to larger test lengths, a significant gap emerges. For example, at n=256, Attention-E is more than 10 percentage points higher in accuracy than Attention-O, which is certainly not a small difference.

Derivation Process

In this section, we introduce the derivation of Eq. [eq:ei]. In fact, the derivation process and assumptions are almost identical to those in "The Principle of Minimum Entropy (VI): How to Choose the Dimension of Word Vectors?".

First, by substituting the expression for a_{i,j}, we get: \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}} Note that we only want to make a semi-quantitative estimate to determine a suitable \lambda to offset part of the length’s influence; making entropy completely independent of length is impossible. Therefore, we can make some assumptions, such as assuming \boldsymbol{k}_j is a random variable, then we can write: \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}] Using the same approximation for all sums, we obtain: \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}]} Noting that in general, \boldsymbol{q}_i, \boldsymbol{k}_j are passed through a Dense layer after Layer Norm, and Dense layers are close to orthogonal transformations (refer to "Understanding Model Parameter Initialization Strategies from a Geometric Perspective"), we approximately assume that \boldsymbol{q}_i, \boldsymbol{k}_j are vectors with a norm of \sqrt{d}, so \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 transformed into an expectation over the angle between \boldsymbol{q}_i and \boldsymbol{k}_j: \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}]} where the distribution of \theta is the distribution of the angle between any two random vectors on the sphere, which we discussed in "Distribution of the Angle Between Two Random Vectors in n-Dimensional Space". Next, similar to the "Approximate Estimation" in "The Principle of Minimum Entropy (VI)", we can use the Laplace approximation to get: \mathcal{H}_i \approx \log n - 0.24\lambda d + \mathcal{O}(1) Therefore, to offset the influence of length n, we set \log n - 0.24\lambda d = 0, which leads to \lambda = \log n / (0.24 d). Of course, since this is just an estimate, there is no need to keep the coefficient 0.24. Instead, we introduce the hyperparameter \kappa, such that: \lambda = \frac{\kappa\log n}{d} This corresponds to Eq. [eq:ei].

Related Results

While reading submissions for ACL 2022, I found a paper titled "Overcoming a Theoretical Limitation of Self-Attention", which provides a similar result (Equation 1 in Section 4.3 of the paper): \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{\log n}{\sqrt{d}}QK^{\top}\right)V However, that paper does not provide deep theoretical analysis; it merely constructs two special cases to test Attention performance and finds that multiplying the scaling factor by \log n helps with length generalization.

However, it can be seen that if the default convention for \log is the natural logarithm, the above formula is clearly unreasonable because when n is large, the scaling factor becomes too large, leading to severe gradient vanishing. It’s just that the paper only conducted experiments on machine translation with sequences around n=20, so the gradient vanishing problem did not manifest.

Conclusion

This article re-derives the Scale operation in Scaled Dot-Product Attention from the perspective of entropy invariance and obtains a new scaling factor. Preliminary experimental results show that the new scaling factor does not change existing training performance and yields better results for length extrapolation.

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

For more details on reprinting, please refer to: "Scientific Space FAQ"