从熵不变性看Attention的Scale操作

原文链接 | 2021-12-21 | 苏剑林

当前Transformer架构用的最多的注意力机制，全称为“Scaled Dot-Product Attention”，其中“Scaled”是因为在\(Q,K\)转置相乘之后还要除以一个\(\sqrt{d}\)再做Softmax（下面均不失一般性地假设\(Q,K,V\in\mathbb{R}^{n\times d}\)）：

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

在《浅谈Transformer的初始化、参数化与标准化》中，我们已经初步解释了除以\(\sqrt{d}\)的缘由。而在这篇文章中，笔者将从“熵不变性”的角度来理解这个缩放操作，并且得到****一个新的缩放因子****。在MLM的实验显示，新的缩放因子具有****更好的长度外推性能****。

熵不变性#

我们将一般的Scaled Dot-Product Attention改写成

\[\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}\]

其中\(\lambda\)是缩放因子，它跟\(\boldsymbol{q}_i,\boldsymbol{k}_j\)无关，但原则上可以跟长度\(n\)、维度\(d\)等参数有关，目前主流的就是\(\lambda=1/\sqrt{d}\)。

本文提出一个观点：

为了使得模型结果能够更好地泛化到未知长度，Attention机制的设计应该使得\(a_{i,j}\)尽量具备熵不变性。

怎么理解这句话呢？首先，泛化到未知长度，指的是预测长度和训练不一致时也能有不错的效果，比如\(n=64\)训练然后外推到\(n=128,256\)测试。我们知道，使用RoPE之类的相对位置编码的模型，对长度具有比较好的外推性，但我们依然可以通过更好的设计来增强这种外推性，比如熵不变性就是其中之一。

具体来说，\(a_{i,j}\)可以视为\(i\)为条件、\(j\)为随机变量的条件分布，它的熵为

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

熵不变性是指，\(\mathcal{H}_i\)应该对长度\(n\)不敏感。更具体一点，就是如果在已有的token基础上，再补充几个token，那么新算出来各个\(a_{i,j}\)自然也会有所改变，但我们希望\(\mathcal{H}_i\)不要有太大改变。

为什么希望熵不变呢？我们知道，熵是不确定性的度量（参考《“熵”不起：从熵、最大熵原理到最大熵模型（一）》），换个角度想，我们可以将不确定性视为注意力的“聚焦程度”：如果熵为0，那么注意力将聚焦到某一个token上，如果熵为\(\log n\)，那么注意力均匀分布到所有token上。我们希望熵不变，是希望引入新的token后，已有的token依旧能同样地聚焦到原来的token上，而不希望新token的引入过多地“分摊”了原有的注意力，导致求和结果显著发生变化。

新的因子#

根据熵不变性以及一些合理的假设，我们可以得到一个新的缩放因子，从而得到一种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}\]

这里的\(\kappa\)是一个跟\(n,d\)都无关的超参数，详细推导过程我们下一节再介绍。为了称呼上的方便，这里将式\(\eqref{eq:std}\)描述的常规Scaled Dot-Product Attention称为“Attention-O”（Original），而式\(\eqref{eq:ei}\)以及下面的式\(\eqref{eq:ei2}\)描述的变体称为“Attention-E”（Entropy Invariance）。

可能有读者对引入了一个新参数感到不满意，其实这个不难解决。我们知道当前主流的预训练长度就是512，所以我们假设主流的参数都是为\(n=512\)调试好的，所以当\(n=512\)的时候，上式应退化为普通的Scaled Dot-Product Attention，即\(\frac{\kappa \log 512}{d}=\frac{1}{\sqrt{d}}\)，推出\(\kappa = \frac{\sqrt{d}}{\log 512}\)，代入上式整理后得到

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

这就去掉了超参数\(\lambda\)，下面的实验也是用这个版本。

为了验证该改动是否真如预期那样能提高Transformer的外推效果，笔者分别用Attention-O和Attention-E分别训练了一个RoFormer small版本，训练任务为MLM，训练长度为64，然后在不同长度的验证集下比较MLM的准确率，结果如下：

\[\begin{array}{c} \text{Attention的长度外推实验} \\ {\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}\]

从实验结果可以看出，在与训练长度一致\(n=64\)的情况下，Attention-O和Attention-E的效果是很接近的，但是外推到更大的测试长度时，则明显拉开了差距，比如\(n=256\)时Attention-E要比Attention-O高10个百分点以上的准确率，可真不是一星半点了。

推导过程#

这一节我们介绍式\(\eqref{eq:ei}\)的推导过程。事实上，推导过程和假设都跟《最小熵原理（六）：词向量的维度应该怎么选择？》中的几乎是一样的。

首先，我们代入\(a_{i,j}\)的表达式，就可以得到：

\[\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}\]

要注意，我们仅仅是要做一个半定量的估计，以确定适合的\(\lambda\)来抵消部分长度的影响，让熵完全不受长度影响是做不到的。所以，我们可以做一些假设，比如假设\(\boldsymbol{k}_j\)是一个随机变量，那么可以写出

\[\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}\]

将所有求和都用同样的近似代替，我们得到

\[\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}\]

留意到一般情况下\(\boldsymbol{q}_i,\boldsymbol{k}_j\)都是Layer Norm出来之后再接一个Dense层，而Dense层接近正交变换（参考《从几何视角来理解模型参数的初始化策略》），所以我们近似地假设\(\boldsymbol{q}_i,\boldsymbol{k}_j\)都是模长为\(\sqrt{d}\)的向量，所以\(\boldsymbol{q}_i\cdot \boldsymbol{k}_j=d\cos(\boldsymbol{q}_i,\boldsymbol{k}_j)\)；然后进一步假设\(\boldsymbol{k}_j\)均匀地分布在半径为\(\sqrt{d}\)的球面上，那么对\(\boldsymbol{k}_j\)的期望可以转化为对\(\boldsymbol{q}_i,\boldsymbol{k}_j\)夹角的期望，即

\[\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}\]

其中\(\theta\)服从的分布就是球面上任意两个向量之间的夹角分布，我们在《n维空间下两个随机向量的夹角分布》讨论过。接下来可以像《最小熵原理（六）：词向量的维度应该怎么选择？》的“近似估计”一样，用拉普拉斯近似得到

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

因此，为了抵消长度\(n\)的影响，我们让\(\log n - 0.24\lambda d = 0\)，从而得出\(\lambda = \log n / (0.24 d)\)。当然，我们知道这只是估计，所以没必要保留系数\(0.24\)了，倒不如直接引入超参数\(\kappa\)，使得

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

这就是对应式\(\eqref{eq:ei}\)了。

文章总结#

本文从熵不变性的角度重新推导了Scaled Dot-Product Attention中的Scale操作，得到了一个新的缩放因子。初步的实验结果显示，新的缩放因子不改变已有的训练性能，并且对长度外推具有更好的结果。

熵不变性#

新的因子#

推导过程#

相关结果#

文章总结#