没研究过RoPE的小伙伴,可以先看一下:一文讲透ROPE
什么是长度外推,有什么问题?
假设我们训练模型的时候,最大的序列长度是4096,但是在推理的时候,来了一个8192的长度序列,这就是要求模型具备处理超出训练长度的能力,这就是长度外推。
前文说过,RoPE通过在在源数据中添加绝对位置,使得内积后只和相对距离有关,理论上直接输入8192的长度,公式照样能算。但是实际上,如果直接这么做,模型的推理结果将会非常糟糕。
为什么会这样?简单来说就是外推会导致没见过的角度出现
未见过的相位角
也就是我们常说的Out-of-Distribution(OOD,分布外)问题。
在4096长度的训练中,模型见过最大的相对距离是4096,意味着对某个维度来说,最多只见过的角度。但是在推理时,却要转
,模型能不懵逼吗。其实这个问题对高频区(即转的快的维度区域)和低频区的影响是完全不一样的:
- 对高频区来说,在训练的时候,在4096长度内早就已经转了好多圈了,0-360度内每个角度都已经见过,所以高频区不怕外推。例如对第0维来说:
(默认Base=10000),即对第0维来说,相距一个token就转了57.2角度,转的飞快,那么在4096的长度内已经转了650圈,阅历相当丰富,完全不怕外推。
- 但是对低频区来说,外推却是非常的致命。对一个128维的多头注意力来说,最后一个维度
,即
,每过1个token,才转0.006611度,在4096长度最多才转了27度,而在8192长度token下最大却转了54度,这就是模型懵逼的原因了。
那么有什么办法来解决这个问题呢?
方案一:线性位置内插(Position Interpolation, PI)
核心思想:压缩
既然训练时只见过0~4096的长度,现在要处理8192,那么就等比例把0~8192压缩到0~4096.
PI是怎么做的
只需要引入一个缩放因子即可(scale factor,s)
还是用训练时的L=4096,而推理的最大长度为8192,那么缩放因子s=8192/4086=2。
在ROPE中,对位置m的token的旋转角度是,现在变成了
,把原本外推8192的角度变成了4096的了,保证了低频旋转角度的绝对安全。
PI有什么问题
PI 看起来很完美,但它有一个致命的副作用,导致它不能直接拿来就用(Zero-shot),必须收集长文本数据对模型进行微调(Fine-tuning)才能恢复效果。
- 对于低频区:这种压缩是救命的,因为它把超出的角度拉回了安全区
- 对高频区:这种压缩是毁灭性的。原本的高频区(转得飞快的秒针),用来精准区分相邻的词汇。词元 A 在位置 1,词元 B 在位置 2,它们之间的夹角原本是
,现在除以s=2后,他们之间的夹角变成了
,直接砍半。模型变得“近视”了,它看远处的宏观结构看得很清楚,但看紧挨着的两个词却觉得极其模糊,难以分辨它们到底是谁在前、谁在后。这种高频特征的粘连,导致如果没有后续的微调让模型重新适应这种“近视眼”,它输出文本的质量依然会大幅下降。简单来说,就是原本在高频区相邻词的角度是相差很大的(例如57度),在这么大的角度上直接砍半会导致得到的角度和原先的角度相差太大,模型适应不了,因此必须微调。
方案二:NTK-Aware
NTK-Aware是怎么解决PI高频区的问题的
既然PI会导致高频区崩溃,那么能不能通过只修改ROPE的底数(Base,原始默认底数为10000),使得第0维不压缩,最后一维变为PI的1/s的安全压缩,从第0到底d/2-1维是一种缓慢变化的过程?
首先对于第0维,不管Base变成多少,,也就是在高频区做到几乎不压缩。
那么现在的问题变成求解,使得在最后一维(d/2 - 1),有以下方程关系:
简化后:
继续简化:
最终得到为:
我们来看一下和Base的比例关系:
NTK-Aware的缺陷
NTK-Aware当然不是完美的,它有什么问题呢?从上述的分析过程,我们得知NTK-Aware除了在第0维(i=0)处完全不压缩,其他维都进行了压缩。 我们不禁有一个思考:真的需要对除了i=0外的维度都进行压缩吗?假如在训练集长度L内,且维度在n内,即,这些维度在训练集中都已经转过了360,那么这些维度对0-360的角度全部都已经见过了,即使是推理时的长度比训练集要大,但是三角函数是周期性的,所以哪怕推理时转480度都没有问题。
哪些维度是不需要压缩的?
上面说过,我们认为在训练长度L内,哪些能完整转过360度的维度是不需要压缩的,那么具体是哪些维度i呢? 现在我们来计算一下:
对每一维度都对应一个旋转角度
,那么波长
就代表维度i转了360度(2Π弧度)所需要的token距离,那么:
那么当时,这些维度i在训练集长度L内,都是已经转过360的,也就是说完全不需要压缩,经过简单的变换可以得到:
对训练长度L=4096,多头注意力维度d=128,Base=10000,计算,即45以内的维度其实都是不需要压缩的,但是NTK-Aware为了追求平滑的曲线,反而导致好心办坏事了。但是NTK-Aware却把这些维度都强行进行压缩,这会带来混乱问题:
- 训练时的记忆:模型在这个中频维度上已经学得非常精细了。它知道,“如果两个词在这个维度的夹角是90度,说明它们相隔 50 个词(比如刚好跨越一个段落)”。
- NTK 压缩后的错乱:现在你为了追求“平滑过渡”,给这个维度强行打了个 0.7 倍的折。原本相隔 50 个词的夹角变成了63度,模型看到63度会怎么想?它会翻阅训练记忆,得出结论:“哦,这俩词相隔 35 个词。”
- 注意力发散:角度被压缩了,说明相似度变小了,q和k在空间中的分布更加密集了,计算的注意力得分时整体的差异就变小了,这就会导致原本高度集中的注意力变成一盘散沙,模型抓不住重点,注意力变得涣散了。
方案三:YaRN(Yet another RoPE for NTK)一统江湖
既然NTK-Aware的问题是把一些本不应该压缩的维度压缩了,YaRN的解决办法也很简单,把特征维度划分成三个部分:
分段压缩
- 高频区:不压缩,纯外推
- 低频区:完全PI插值,即压缩1/s
- 中频区(过渡区):使用一个平滑的斜坡函数在这高频区和低频区之间做线性过渡,避免相邻维度间产生剧烈的突变
上述已经有波长公式:
YaRN引入两个超参和
(通常设定
),在波长空间中画出两条线来把特征维度划分成高中低频三个部分:
- 下边界(高频边界):
,在训练长度L内,至少转了
圈(如1圈),这是高频区
- 上边界(低频边界):
,在训练长度L内,连
圈(如
圈)都转不完,这是低频区
那么波长介于就是中频过度区
现在根据波长公式计算出高中低频的维度i,根据上述波长公式解出i的公式维:
那么:
则:
到此,我们已经确定了不压缩的高频区的维度是哪些,完全PI内插(压缩1/s)的低频区维度是哪些,当然中频区的维度是哪些也是知道的,但是还不知道中频区是如何从高频区的压缩系数为1过度到低频区压缩系数为1/s的,也就是要找到关于中频区的压缩系数的函数:
YaRN定义了一个斜坡函数:
有了插值权重,中频区每个维度的最终频率采用凸组合进行加权融合:
由于,上述公式简化为:
那么最终得到中频区的压缩系数为:
为什么要设计这样的斜坡函数呢?能不能不要这个过渡区,也就是除了高频区不压缩之外(压缩系数等于1),剩下的全部压缩1/s,这样很符合我们的直观:高频区全部都是至少转了360度的,剩下的维度全部压缩1/s,多安全!但是这会导致一个问题:
会导致频谱相位不连续。相邻的两个维度和
之间,原本在空间中承担着连续的正交表征功能,却因为缩放比例从 1.0瞬间骤降到1/s,导致它们的相位变化率产生了巨大的阶跃断层。
在自注意力机制中,这种断层会在点乘矩阵中引入高频噪声,破坏特征向量的高维正交性,引发类似信号处理中的吉布斯现象(震荡)
通过斜坡函数的缓冲,缩放系数
变成了一条平滑连接的折线:
这条斜坡既保证了高频区的“零扭曲”,又保证了低频区的“防爆 OOD”,同时消除了中频区的频谱断层噪声。
到此YaRN解决了NTK-Aware对高频区也进行压缩的问题,但是注意力涣散问题还没得到解决
注意力温度缩放
self-attention 的注意力公式为:
YaRN的做法很简单,不改变注意力公式的主体结构,而是直接在生成ROPE向量时,把cos和sin项乘以一个固定的标量mscale(mscale>1)。
当QK矩阵融入了mscale之后,其内积被放大了
倍:
将此公式改写为统计学中常见的带温度参数的 Softmax,形式如下:
YaRN 论文中给出了mscale随压缩系数s变化的最佳经验公式:
为什么只是把mscale乘以ROPE的旋转矩阵中的cos和sin就能放大注意力呢,推导过程如下:
对原始Q矩阵中某一组二维向量:
对其旋转的旋转矩阵为:
旋转后的q向量为:
YaRN将cos和sin修改为:
带入旋转矩阵,得到新的旋转矩阵为:
那么将mscale乘以cos和sin后,新的旋转q向量为:
同理在n处新的旋转k向量为:
那么和
向量的内积为:
那么从二维q和k向量推广到Q和K矩阵有:
矩阵乘法:
那么logits为:
在机器学习中,温度 Scaling(Temperature Scaling)的标准数学形式是将除以一个温度参数 T:
那么容易看出:
最后推导得出最终的注意力公式为:
其中,
mscale公式的由来
我们不禁要思考,mscale为什么要写成,而不是其他的,是YaRN团队拍脑袋想出来的吗?当然不是。
大模型计算某个关键词的概率为:
,Z为所有词的总和
分母Z其实大多数都是背景噪音,当外推文本s倍时,新加入的词也是和原来差不多的背景噪音,那么新的分母,此时关键词新的概率变小了:
在深度学习中,对概率的分析通常放到对数空间中来进行分析,对两边取对数:
经过简单的简化:
也就是说文本长度拉伸s倍之后,对某个关键词的概率在对数空间中的差值恰好约等于,即损失了
.
YaRN通过将cos和sin乘以mscale,使得Q和K的点积结果放大了
倍。
假设没放大之前:
- 关键词得分
- 背景噪声得分
- 两者的差为:
那么施加mscale之后:
- 关键词得分
- 背景噪声得分
- 两者的差为:
添加mscale是为了带来扩大关键词和背景噪声之间的差距的收益,带来的差距收益为:
那么现在我们的目标很明确了:让添加mscale之后带来的收益恰好能够抵消掉文本拉伸s倍后损失掉的,即:
那么:
当x<<1时,有泰勒展开:
令,则:
可能大家有疑惑:为什么可以认为是一个<<1的值?当s=2,
,
,
,
,经过充分学习
一般在
之间,他们的比值也确实<<1,即使是后面令a=0.1,他们的比值不再<<1,经过学习也能吸收掉这个误差。
令常数项,则:
优美的公式:
- 1.0时边界,不外推时,不做处理
是文本拉伸s倍后,softmax分母也扩增s倍,导致关键词在对数空间中损失
,是在真实模型(如 LLaMA)上的平均得分差拟合值,是一个经验值
全流程可总结为:
到此我们已经把PI,NTK-Aware,YaRN完整的学习一遍了,如果有错误的地方,欢迎大家指出,我们一起交流学习!