从熵不变性看Attention的Scale操作
By 苏剑林 | 2021-12-21 | 115322位读者 |当前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://www.kexue.fm/archives/8823
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (Dec. 21, 2021). 《从熵不变性看Attention的Scale操作 》[Blog post]. Retrieved from https://www.kexue.fm/archives/8823
@online{kexuefm-8823,
title={从熵不变性看Attention的Scale操作},
author={苏剑林},
year={2021},
month={Dec},
url={\url{https://www.kexue.fm/archives/8823}},
}
March 20th, 2024
很棒的思路!我在长度外推模型的推理侧简单做了logn调整的尝试,loss曲线被明显压低了,文本越长loss的压制作用越显著。但同时,我发现经过logn调整后的模型,在大海捞针的实验中的性能不升反降。经过分析,初步认为是做了logn调整后其实是让模型更加注意“该注意”的序列信息,而所谓“该注意”的信息通常距离当前预测位置较近,这样一来,其实反而让模型更加“不注意”远程的信息了。
你指的是直接长度外推做更长文本的大海捞针任务吗?具体是什么长度外推方法呢?
基础模型:约1B大小,4096窗口预训练,rotary theta = 10000;
长度外推模型:基于基础模型做长度外推微调,128K窗口,rotary theta = 5000000;
【400K长序列上各区间loss实验】
在40万窗口内,不同序列长度区间上的loss评测结果显示,“长度外推模型”虽然loss没有“崩溃”,但随着序列长度轻微上扬的趋势(2.2逐渐上扬到2.5的样子);
而经过logn调整(时间原因,目前仅应用在了推理时)的“长度外推模型”模型,loss曲线更加平缓,(2.2逐渐上扬到2.3的样子);
这说明logn的调整的确有效~起码是让模型在大窗口下“更会说话”了;
【64k窗口内大海捞针实验】
“长度外推模型”在该实验中表现“尚可”,而经过logn调整,性能有所下降。更具体的分析还没来得及做。
顺便提一句,我觉得大海捞针实验评测基础模型上不够合理,捞不到“针”不一定是没看到“针”,而也可能是因为模型根本没有在遵循指令或者模型在用已有的知识经验作答而忽略了文章信息。
如果我没理解错,你这里的“基于基础模型做长度外推微调”,是改了theta之后,用128k的语料去微调了一定步骤?
你说的logn“仅应用在了推理时”,是指在推理阶段加入$\max(1, \log_{\text{128k}} n)$作为scale因子吗?如果是,那么它在64k时应该不会生效?
看到$max(1,\log_{128k}n)$有一个疑问,取小于1的值会有问题吗?
希望有大佬能解答一下,非常感谢。
如果你训练没有用logn scale,只在推理时长度外推用,那么肯定要保证训练长度内(这里是128k)的形式完全不变啊,训练长度内你没scale,就相当于scale是1,并不是小于1不行。如果你训练阶段就加入这个logn scale,那么自然是允许小于1的
August 3rd, 2024
调低温度是否也是相同效果.都是对logits scale up,都会让整体分布尖锐
本文就是调低温度的策略啊,本文的scale的倒数不就是温度的意义了吗