从熵不变性看Attention的Scale操作
By 苏剑林 | 2021-12-21 | 113755位读者 |当前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}$了。
相关结果 #
在阅读ACL2022的投稿论文时,发现上面有一篇《Overcoming a Theoretical Limitation of Self-Attention》,给出了相近的结果(论文4.3节的公式1):
\begin{equation}Attention(Q,K,V) = softmax\left(\frac{\log n}{\sqrt{d}}QK^{\top}\right)V\end{equation}
不过,该论文并没有太深刻的理论分析,只是构建了两个特殊的case来测试Attention的性能,测试发现往缩放因子乘上$\log n$有助于泛化长度,所以就提出来了。
然而可以看出,如果按照默认约定$\log$用自然对数的话,那么上式很明显是不合理的,因为当$n$较大时,缩放因子过大,会导致严重的梯度消失。只不过该论文只是在机器翻译上做实验,测得都是$n=20$级别的序列,所以就没有显示出梯度消失问题。
文章总结 #
本文从熵不变性的角度重新推导了Scaled Dot-Product Attention中的Scale操作,得到了一个新的缩放因子。初步的实验结果显示,新的缩放因子不改变已有的训练性能,并且对长度外推具有更好的结果。
转载到请包括本文地址:https://kexue.fm/archives/8823
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (Dec. 21, 2021). 《从熵不变性看Attention的Scale操作 》[Blog post]. Retrieved from https://kexue.fm/archives/8823
@online{kexuefm-8823,
title={从熵不变性看Attention的Scale操作},
author={苏剑林},
year={2021},
month={Dec},
url={\url{https://kexue.fm/archives/8823}},
}
March 26th, 2023
训练的数据必须是固定的吗?比如可以是在(0,50)随机分布的吗?然后使用了熵不变性,在一定程度上使得模型对于长度大于50的数据的预测效果还不错。
可以随机分布啊,随机分布更好呀。
March 27th, 2023
请问苏神:
如果训练集和测试集的数据长度都是变化的,也就是m和n不是固定的,那么这时候如何来确定一个高效的尺度因子呢?
从零训练的话,直接固定scale因子为$\frac{\log_{512}n}{\sqrt{d}}$就行了,$n$是当前样本长度。
当前样本的长度是变化的,那么我把n固定成平均长度可以吗?
如果在训练开始之时就引入这个scale,那么是直接$\log
_{512} n$,$n$是当前样本长度。
那如果我是在sft阶段呢,还是设置成512,还是max_position_embedding值?
sft阶段如果有更长的文本,可以考虑$\text{clip}(\log_{512} n, 1, \infty)$,如果sft也只是512,那么就没必要加$\log n$了。
April 11th, 2023
想请教一下,原文中提到“因为当n较大时,缩放因子过大,会导致严重的梯度消失”,这里缩放因子过大为什么会导致梯度消失呢?个人理解这里应该是对QKV求导,那么因子会使得梯度变大,而不是变小
举个简单的例子:$\sigma(ax)=\frac{1}{1+e^{-ax}}$,你觉得$a\to\infty$时,梯度是变大还是变小?
苏神您好,上面说从熵不变的角度乘完logn以后能让注意力更集中, 但是从温度来看,乘完logn再softmax注意力不应该是更分散吗,我这里是哪里理解不对了吗
可以证明乘上的数字越大(或者说温度参数越小),熵越小,即注意力更集中。
April 19th, 2023
式子 (2) 中的分子分母裏的公式很像是 heat kernel。你提出的式子 (11) 說不定能解決 「heat kernel 被應用在 graph signal processing 中一直無法確定 diffusion time 取何值纔是最佳」的問題。
这个问题没了解(捂脸)
July 16th, 2023
如果将点积注意力换成加性注意力,请问应该怎么推导scale因子呢?
很好的问题,但结果可能很复杂,我暂时无法分析出来。
感谢回复,期待您的更新!
September 21st, 2023
苏神您这篇文章的相关思路有发表过吗
没
October 25th, 2023
[...]对于base=500在16K长度上续训得结果,由于其困惑度曲线足够平坦,笔者认为这在某种程度上已经打破 LM-Infinite 和 熵不变性 等研究中提出的,RoPE自注意力分布熵增导致外推效果变差的“诅咒”。笔者得到了一个大胆的猜想,当base足够小并且训练长度足够长,$bm{q}_t,bm{k}_s$的每个维度在训练阶段就已经感知到了位置编码的周期性变化,由此可以胜任近乎无限的外推。进一步使用[...]
November 24th, 2023
太厉害了~
December 6th, 2023
式4到式5是不是公式漏了点东西,看上去有些令人困惑
两个公式之间不是还有说明文字?
December 28th, 2023
是不是给softmax加个温度系数也可以达到相同目的。
现在不就是讨论温度系数怎么设置嘛
数值上看,增加了这个“温度”,是让n较大的序列更加去注意那些a_{i,j}数值高的token,防止被增加的token分散注意力。